Implement index of (X⍳Y)

1 files changed, 59 insertions, 1 deletions

diff --git a/functions.c b/functions.c
index 997c1b4..6e0c6b6 100644
--- a/functions.c
+++ b/functions.c
@@ -100,7 +100,7 @@ fndyad dyadfunctiondefs[] = {
 	fnIndex, /* ⌷ */
 	0, /* ⍋ */
 	0, /* ⍒ */
-	0, /* ⍳ */
+	fnIndexOf, /* ⍳ */
 	0, /* ⍸ */
 	fnMembership, /* ∊ */
 	0, /* ⍷ */
@@ -1165,6 +1165,64 @@ fnIndex(Array *left, Array *right)
 }
 
 Array *
+fnIndexOf(Array *left, Array *right)
+{
+	if(left->rank < 1)
+		throwerror(nil, ERank);
+	if(right->rank < left->rank-1)
+		throwerror(nil, ERank);
+	int i, j;
+	for(i = 0; i < left->rank-1; i++)
+		if(left->shape[left->rank-1-i] != right->shape[right->rank-1-i])
+			throwerror(nil, EShape);
+
+	int rank = right->rank + 1 - left->rank;
+	int size = 1;
+	int io = globalIO();
+	int n = left->shape[0];
+	for(i = 0; i < rank; i++)
+		size *= right->shape[i];
+
+	Array *result = allocarray(AtypeInt, rank, size);
+	for(i = 0; i < rank; i++)
+		result->shape[i] = right->shape[i];
+
+	/* Reshape Y to allow easy picking of sub-arrays */
+	Array *rightshape = allocarray(AtypeInt, 1, right->rank - rank + 1);
+	rightshape->shape[0] = rightshape->size;
+	rightshape->intdata[0] = size;
+	for(i = 0; i < rightshape->size - 1; i++)
+		rightshape->intdata[rightshape->size-1-i] = right->shape[right->rank-1-i];
+	right = fnReshape(rightshape, right);
+
+	Array **lefts = malloc(sizeof(Array *) * n);
+	Array *index = mkscalarint(io);
+	for(i = 0; i < n; i++){
+		index->intdata[0] = i+io;
+		lefts[i] = fnIndex(index, left);
+	}
+	for(i = 0; i < size; i++){
+		index->intdata[0] = i+io;
+		Array *target = fnIndex(index, right);
+		for(j = 0; j < n; j++)
+			if(comparearray(lefts[j], target, 1) == 0){
+				result->intdata[i] = j+io;
+				break;
+			}
+		if(j == n)
+			result->intdata[i] = n+io;
+		freearray(target);
+	}
+
+	for(i = 0; i < n; i++)
+		freearray(lefts[i]);
+	free(lefts);
+	freearray(rightshape);
+	freearray(right);
+	return result;
+}
+
+Array *
 fnMembership(Array *left, Array *right)
 {
 	/* TODO avoid outer product */