1 files changed, 70 insertions, 53 deletions

diff --git a/functions.c b/functions.c
index ac5c517..0b4f085 100644
--- a/functions.c
+++ b/functions.c
@@ -101,7 +101,7 @@ fndyad dyadfunctiondefs[] = {
 	0, /* ⍋ */
 	0, /* ⍒ */
 	fnIndexOf, /* ⍳ */
-	0, /* ⍸ */
+	fnIntervalIndex, /* ⍸ */
 	fnMembership, /* ∊ */
 	0, /* ⍷ */
 	fnUnion, /* ∪ */
@@ -120,6 +120,7 @@ fndyad dyadfunctiondefs[] = {
 
 vlong gcd_int(vlong, vlong);
 double gcd_float(double, double);
+Array *indexOfHelper(Array *, Array *, int);
 
 /* Runner function */
 Array *
@@ -1479,59 +1480,13 @@ 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);
-	}
+	return indexOfHelper(left, right, 0);
+}
 
-	for(i = 0; i < n; i++)
-		freearray(lefts[i]);
-	free(lefts);
-	freearray(rightshape);
-	freearray(right);
-	return result;
+Array *
+fnIntervalIndex(Array *left, Array *right)
+{
+	return indexOfHelper(left, right, 1);
 }
 
 Array *
@@ -1785,4 +1740,66 @@ gcd_float(double a, double b)
 		return a;
 	else
 		return gcd_float(b, fmod(a,b));
+}
+
+Array *
+indexOfHelper(Array *left, Array *right, int interval)
+{
+	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(interval && comparearray(lefts[j], target, 1) > 0){
+				result->intdata[i] = j-1+io;
+				break;
+			}else if(!interval && comparearray(lefts[j], target, 1) == 0){
+				result->intdata[i] = j+io;
+				break;
+			}
+		}
+		if(j == n)
+			result->intdata[i] = n + io - (interval ? 1 : 0);
+		freearray(target);
+	}
+
+	for(i = 0; i < n; i++)
+		freearray(lefts[i]);
+	free(lefts);
+	freearray(rightshape);
+	freearray(right);
+	return result;
 }
\ No newline at end of file