Implement partition, but make it work on the first axis instead of the last.

2 files changed, 73 insertions, 2 deletions

diff --git a/apl9.h b/apl9.h
index ace4722..90cdda0 100644
--- a/apl9.h
+++ b/apl9.h
@@ -347,7 +347,9 @@ Array *fnAnd(Array *, Array *);
 Array *fnNand(Array *, Array *);
 Array *fnNor(Array *, Array *);
 Array *fnTake(Array *, Array *);
+Array *fnDrop(Array *, Array *);
 Array *fnPick(Array *, Array *);
+Array *fnPartition(Array *, Array *);
 Array *fnIndex(Array *, Array *);
 Array *fnIndexOf(Array *, Array *);
 Array *fnMembership(Array *, Array *);
diff --git a/functions.c b/functions.c
index a0090a9..7a961a3 100644
--- a/functions.c
+++ b/functions.c
@@ -93,10 +93,10 @@ fndyad dyadfunctiondefs[] = {
 	fnNand, /* ⍲ */
 	fnNor, /* ⍱ */
 	fnTake, /* ↑ */
-	0, /* ↓ */
+	fnDrop, /* ↓ */
 	0, /* ⊂ */
 	fnPick, /* ⊃ */
-	0, /* ⊆ */
+	fnPartition, /* ⊆ */
 	fnIndex, /* ⌷ */
 	0, /* ⍋ */
 	0, /* ⍒ */
@@ -1235,6 +1235,23 @@ 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]);
+	}
+	return res;
+}
+
+Array *
 fnPick(Array *left, Array *right)
 {
 	if(left->rank > 1)
@@ -1265,6 +1282,58 @@ fnPick(Array *left, Array *right)
 }
 
 Array *
+fnPartition(Array *left, Array *right)
+{
+	if(right->rank == 0)
+		throwerror(nil, ERank);
+	if(left->rank > 2)
+		throwerror(nil, ERank);
+	if(left->type != AtypeInt)
+		throwerror(nil, EType);
+	if(left->size != right->shape[0])
+		throwerror(nil, ELength);
+	if(right->shape[0] == 0)
+		return fnSame(right); /* TODO figure out if this is correct */
+
+	Array *ixlist = allocarray(AtypeArray, 1, 0);
+	ixlist->shape[0] = 0;
+
+	int i;
+	int prev = 0;
+	int start = -1;
+	int io = globalIO();
+	int active = 0;
+	for(i = 0; i <= left->size; i++){
+		vlong n = i == left->size ? 0 : left->intdata[i];
+		if(n < 0)
+			throwerror(nil, EDomain);
+		if(active && (n == 0 || n > prev)){
+			int len = i - start;
+			Array *group = allocarray(AtypeInt, 1, len);
+			Array *groupe = fnEnclose(group);
+			group->shape[0] = len;
+			for(int j = 0; j < len; j++)
+				group->intdata[j] = start+j+io;
+			Array *tmp = ixlist;
+			ixlist = fnCatenateFirst(ixlist, groupe);
+			freearray(group);
+			freearray(groupe);
+			freearray(tmp);
+			active = 0;
+		}
+		if(n > prev){
+			start = i;
+			active = 1;
+		}
+		prev = n;
+	}
+
+	Array *result = rundfn(L"↑⍵∘{(⌷∘⍺)¨((⊂⍵)∘⍪¨⍳1↓⍴⍺)}¨⍺", nil, nil, ixlist, right);
+	freearray(ixlist);
+	return result;
+}
+
+Array *
 fnIndex(Array *left, Array *right)
 {
 	int io = globalIO();