int luaC_array__call(lua_State *L)
{
struct Array *A = lunum_checkarray1(L, 1);
int nind = lua_gettop(L) - 1;
if (nind != A->ndims) {
luaL_error(L, "wrong number of indices (%d) for array of dimension %d",
nind, A->ndims);
return 0;
}
const int Nd = A->ndims;
int *stride = (int*) malloc(A->ndims * sizeof(int));
stride[Nd-1] = 1;
for (int d=Nd-2; d>=0; --d) {
stride[d] = stride[d+1] * A->shape[d+1];
}
int m = 0;
for (int d=0; d<A->ndims; ++d) {
int i = lua_tointeger(L, d+2);
m += i*stride[d];
}
_push_value(L, A->dtype, (char*)A->data + m*array_sizeof(A->dtype));
free(stride);
return 1;
}
static int luaC_lunum_slice(lua_State *L)
{
// The first part of this function extracts a slice of the array 'A' according
// to the convention start:stop:skip. The result is a contiguous array 'B'
// having the same number of dimensions as 'A'.
// ---------------------------------------------------------------------------
size_t Nd0_t, Nd1_t, Nd2_t, Nd3_t;
const Array *A = lunum_checkarray1(L, 1); // the array to resize
size_t *start = (size_t*) lunum_checkarray2(L, 2, ARRAY_TYPE_SIZE_T, &Nd0_t);
size_t *stop = (size_t*) lunum_checkarray2(L, 3, ARRAY_TYPE_SIZE_T, &Nd1_t);
size_t *skip = (size_t*) lunum_checkarray2(L, 4, ARRAY_TYPE_SIZE_T, &Nd2_t);
size_t *squeeze = (size_t*) lunum_checkarray2(L, 5, ARRAY_TYPE_SIZE_T, &Nd3_t);
int Nd0 = (int)Nd0_t, Nd1 = (int)Nd1_t, Nd2 = (int)Nd2_t, Nd3 = (int)Nd3_t;
if (Nd0 != A->ndims || Nd1 != A->ndims || Nd2 != A->ndims || Nd3 != A->ndims) {
return luaL_error(L, "slice has wrong number of dimensions for array");
}
for (int d=0; d<A->ndims; ++d) {
if (start[d] < 0 || stop[d] > A->shape[d]) {
return luaL_error(L, "slice not within array extent");
}
}
Array B = array_new_from_slice(A, start, stop, skip, Nd0);
// The rest of this function deals with squeezing out the size-1 dimensions of
// 'B' which are marked by the 'squeeze' array.
// ---------------------------------------------------------------------------
size_t Nd_new = 0;
for (int d=0; d<Nd0; ++d) Nd_new += !squeeze[d];
// In case we're left with a 0-dimensional (scalar) slice
if (Nd_new == 0) {
_push_value(L, B.dtype, B.data);
return 1;
}
// In case there are any dims to squeeze out
else if (Nd_new != Nd0) {
size_t *shape_new = (size_t*) malloc(Nd_new * sizeof(size_t));
for (int d=0,e=0; d<Nd0; ++d) {
if (B.shape[d] > 1 || !squeeze[d]) {
shape_new[e] = B.shape[d];
++e;
}
}
array_resize_t(&B, shape_new, Nd_new);
free(shape_new);
}
lunum_pusharray1(L, &B);
return 1;
}
/*
lightuserdata pattern
string format "ixrsmb"
integer size
lightuserdata buffer
integer buffer_len
*/
static int
_pattern_unpack(lua_State *L) {
struct pbc_pattern * pat = (struct pbc_pattern *)checkuserdata(L, 1);
if (pat == NULL) {
return luaL_error(L, "unpack pattern is NULL");
}
size_t format_sz = 0;
const char * format = lua_tolstring(L,2,&format_sz);
int size = lua_tointeger(L,3);
struct pbc_slice slice;
if (lua_isstring(L,4)) {
size_t buffer_len = 0;
const char *buffer = luaL_checklstring(L,4,&buffer_len);
slice.buffer = (void *)buffer;
slice.len = buffer_len;
} else {
if (!lua_isuserdata(L,4)) {
return luaL_error(L, "Need a userdata");
}
slice.buffer = lua_touserdata(L,4);
slice.len = luaL_checkinteger(L,5);
}
char * temp = (char *)alloca(size);
int ret = pbc_pattern_unpack(pat, &slice, temp);
if (ret < 0) {
return 0;
}
lua_checkstack(L, format_sz + 3);
int i;
char * ptr = temp;
bool array = false;
for (i=0;i<format_sz;i++) {
char type = format[i];
if (type >= 'a' && type <='z') {
ptr = (char *)_push_value(L,ptr,type);
} else {
array = true;
int n = pbc_array_size((struct _pbc_array *)ptr);
lua_createtable(L,n,0);
int j;
for (j=0;j<n;j++) {
_push_array(L,(struct _pbc_array *)ptr, type, j);
}
ptr += sizeof(pbc_array);
}
}
if (array) {
pbc_pattern_close_arrays(pat, temp);
}
return format_sz;
}
/*
lightuserdata pattern
string format "ixrsmb"
integer size
lightuserdata buffer
integer buffer_len
*/
static int
_pattern_unpack(lua_State *L) {
struct pbc_pattern * pat = lua_touserdata(L, 1);
size_t format_sz = 0;
const char * format = lua_tolstring(L,2,&format_sz);
int size = lua_tointeger(L,3);
struct pbc_slice slice;
if (lua_isstring(L,4)) {
size_t buffer_len = 0;
const char *buffer = lua_tolstring(L,4,&buffer_len);
slice.buffer = (void *)buffer;
slice.len = buffer_len;
} else {
slice.buffer = lua_touserdata(L,4);
slice.len = lua_tointeger(L,5);
}
char temp[size];
int ret = pbc_pattern_unpack(pat, &slice, temp);
if (ret < 0)
return 0;
lua_checkstack(L, format_sz + 3);
int i;
char * ptr = temp;
bool array = false;
for (i=0;i<format_sz;i++) {
char type = format[i];
if (type >= 'a' && type <='z') {
ptr = _push_value(L,ptr,type);
} else {
array = true;
int n = pbc_array_size((void *)ptr);
lua_createtable(L,n,0);
int j;
for (j=0;j<n;j++) {
_push_array(L,(void *)ptr, type, j);
}
}
}
if (array) {
pbc_pattern_close_arrays(pat, temp);
}
return format_sz;
}
static int luaC_array__call(lua_State *L)
{
Array *A = lunum_checkarray1(L, 1);
/* slicing done here to split concerns between indexing and slicing */
if (lua_type(L, 2) == LUA_TTABLE || lua_type(L, 2) == LUA_TSTRING) {
/* make slice */
lua_getglobal(L, "lunum");
lua_getfield(L, -1, "__build_slice");
lua_insert(L, 1);
lua_settop(L, 3);
lua_call(L, 2, 1);
return 1;
}
/* index */
const int nind = lua_gettop(L) - 1;
if (nind != A->ndims) {
return luaL_error(L, "wrong number of indices (%d) for array of dimension %d",
nind, A->ndims);
return 0;
}
int isnum;
size_t m = 0;
for (int d=0; d < nind; ++d) {
const size_t i = lua_tointegerx(L, d+2, &isnum);
if (i >= A->shape[d]) {
return luaL_error(L, "array indexed out of bounds (%d) on dimension %d of size %d",
i, d, A->shape[d]);
} else if (!isnum) {
return luaL_error(L, "non-integer index encountered");
}
m = m * A->shape[d] + i;
}
_push_value(L, A->dtype, (char*)A->data + m*array_sizeof(A->dtype));
return 1;
}
static int luaC_array__index(lua_State *L)
{
Array *A = lunum_checkarray1(L, 1);
// Figure out what is the format of the input index. If it's a number or a
// table of numbers, then pass it along to _get_index. If it's an array of bools,
// then use it as a mask.
// ---------------------------------------------------------------------------
if (lunum_hasmetatable(L, 2, "array")) {
Array *M = lunum_checkarray1(L, 2);
if (M->dtype != ARRAY_TYPE_BOOL) {
return luaL_error(L, "index array must be of type bool");
}
Array B = array_new_from_mask(A, M);
lunum_pusharray1(L, &B);
return 1;
}
/* try to index into array */
int success;
const size_t m = _get_index(L, A, &success);
if (success) {
_push_value(L, A->dtype, (char*)A->data + array_sizeof(A->dtype)*m);
return 1;
}
/* check metatable */
lua_getmetatable(L, 1);
lua_pushvalue(L, 2);
if (lua_gettable(L, -2) != LUA_TNIL) {
return 1;
}
return 0;
}
int luaC_array__index(lua_State *L)
{
struct Array *A = lunum_checkarray1(L, 1);
// Figure out what is the format of the input index. If it's a number or a
// table of numbers, then pass it along to _get_index. If it's a table of
// tables or numbers, then assume it's a slice. If it's an array of bools,
// then use it as a mask.
// ---------------------------------------------------------------------------
if (lunum_hasmetatable(L, 2, "array")) {
struct Array *M = lunum_checkarray1(L, 2);
if (M->dtype != ARRAY_TYPE_BOOL) {
luaL_error(L, "index array must be of type bool");
}
struct Array B = array_new_from_mask(A, M);
lunum_pusharray1(L, &B);
return 1;
}
else if (lua_type(L, 2) == LUA_TTABLE || lua_type(L, 2) == LUA_TSTRING) {
lua_getglobal(L, "lunum");
lua_getfield(L, -1, "__build_slice");
lua_remove(L, -2);
lua_pushvalue(L, 1);
lua_pushvalue(L, 2);
lua_call(L, 2, 1);
return 1;
}
const int m = _get_index(L, A);
_push_value(L, A->dtype, (char*)A->data + array_sizeof(A->dtype)*m);
return 1;
}