diff options
| -rw-r--r-- | functions.c | 52 |
1 files changed, 41 insertions, 11 deletions
diff --git a/functions.c b/functions.c index f5f5fac..ac5c517 100644 --- a/functions.c +++ b/functions.c @@ -1262,18 +1262,48 @@ fnTake(Array *left, Array *right) Array * fnDrop(Array *left, Array *right) { - if(right->rank != 1 || left->rank != 0 || left->type != AtypeInt || left->intdata[0] < 0 || left->intdata[0] > right->size) - throwerror(L"drop only partly implemented", ENotImplemented); - Array *res = allocarray(right->type, 1, right->size - left->intdata[0]); - res->shape[0] = res->size; - for(int i = 0; i < res->size; i++){ - memcpy(res->rawdata + i*datasizes[res->type], - right->rawdata + (i+left->intdata[0])*datasizes[res->type], - datasizes[res->type]); - if(res->type == AtypeArray) - incref(res->arraydata[i]); + int i; + if(left->type != AtypeInt) + throwerror(nil, EType); + if(left->rank > 1) + throwerror(nil, ERank); + + if(right->rank == 0){ + right = duparray(right); + right->rank = left->size; + right->shape = realloc(right->shape, sizeof(int) * right->rank); + for(i = 0; i < right->rank; i++) + right->shape[i] = 1; + }else + right = fnSame(right); + + if(left->size > right->rank) + throwerror(nil, ELength); + else if(left->size == right->rank) + left = fnSame(left); + else{ + Array *old = left; + left = allocarray(AtypeInt, 1, right->rank); + left->shape[0] = left->size; + for(i = 0; i < old->size; i++) + left->intdata[i] = old->intdata[i]; + for(; i < left->size; i++) + left->intdata[i] = 0; } - return res; + for(int i = 0; i < left->size; i++){ + vlong n = left->intdata[i]; + vlong m = right->shape[i]; + if(n > m || (-n) > m) + left->intdata[i] = 0; + else if(n > 0) + left->intdata[i] = -(m-n); + else + left->intdata[i] = m+n; + } + Array *result = fnTake(left, right); + freearray(left); + freearray(right); + return result; } Array * |