Implement drop in terms of take

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 *