1 files changed, 81 insertions, 10 deletions

diff --git a/array.c b/array.c
index 7675144..4208d9f 100644
--- a/array.c
+++ b/array.c
@@ -4,8 +4,11 @@
 
 #include "apl9.h"
 
+Array *inttofloatarray(Array *);
+
 int datasizes[] = {
 	[AtypeInt] = sizeof(vlong),
+	[AtypeFloat] = sizeof(double),
 	[AtypeArray] = sizeof(Array *)
 };
 
@@ -14,7 +17,14 @@ mkscalarint(vlong i)
 {
 	Array *a = allocarray(AtypeInt, 0, 1);
 	a->intdata[0] = i;
+	return a;
+}
 
+Array *
+mkscalarfloat(double f)
+{
+	Array *a = allocarray(AtypeFloat, 0, 1);
+	a->floatdata[0] = f;
 	return a;
 }
 
@@ -39,7 +49,7 @@ simplearray(Array *a)
 int
 simplescalar(Array *a)
 {
-	return simplearray(a) && a->rank == 0;
+	return a->rank == 0 && a->type != AtypeArray;
 }
 
 Array *
@@ -81,6 +91,40 @@ scalarextend(Array *a, Array *b, Array **aa, Array **bb)
 }
 
 Array *
+inttofloatarray(Array *a)
+{
+	Array *b = allocarray(AtypeFloat, a->rank, a->size);
+	for(int i = 0; i < a->rank; i++)
+		b->shape[i] = a->shape[i];
+	for(int i = 0; i < a->size; i++)
+		b->floatdata[i] = a->intdata[i];
+	return b;
+}
+
+int
+commontype(Array *a, Array *b, Array **aa, Array **bb, int forcefloat)
+{
+	/* When A and B are numeric arrays, set aa and bb
+	   to arrays that have compatible types, with the same data.
+	*/
+	if(forcefloat){
+		*aa = a->type == AtypeFloat ? fnSame(a) : inttofloatarray(a);
+		*bb = b->type == AtypeFloat ? fnSame(b) : inttofloatarray(b);
+	}else if(a->type == b->type){
+		*aa = fnSame(a);
+		*bb = fnSame(b);
+	}else if(a->type == AtypeFloat && b->type == AtypeInt){
+		*aa = fnSame(a);
+		*bb = inttofloatarray(b);
+	}else if(a->type == AtypeInt && b->type == AtypeFloat){
+		*aa = inttofloatarray(a);
+		*bb = fnSame(b);
+	}else
+		return 0;
+	return 1;
+}
+
+Array *
 arrayitem(Array *a, int index)
 {
 	Array *res = nil;
@@ -88,6 +132,9 @@ arrayitem(Array *a, int index)
 	case AtypeInt:
 		res = mkscalarint(a->intdata[index]);
 		break;
+	case AtypeFloat:
+		res = mkscalarfloat(a->floatdata[index]);
+		break;
 	case AtypeArray:
 		res = a->arraydata[index];
 		incref(res);
@@ -106,14 +153,38 @@ simplifyarray(Array *a)
 	if(a->type != AtypeArray || a->size == 0)
 		return fnSame(a);
 	int type = a->arraydata[0]->type;
+	int canfloat = type == AtypeFloat || type == AtypeInt;
+	int sametype = 1;
 	int i;
-	for(i = 0; i < a->size; i++)
-		if(a->arraydata[i]->type != type || a->arraydata[i]->rank != 0)
-			return fnSame(a);
-	Array *b = allocarray(type, a->rank, a->size);
-	for(i = 0; i < a->rank; i++)
-		b->shape[i] = a->shape[i];
-	for(i = 0; i < a->size; i++)
-		memcpy(b->rawdata + i * datasizes[type], a->arraydata[i]->rawdata, datasizes[type]);
-	return b;
+
+	for(i = 0; i < a->size; i++){
+		int t = a->arraydata[i]->type;
+		canfloat = canfloat && (t == AtypeFloat || t == AtypeInt);
+		sametype = sametype && (t == type);
+		if(a->arraydata[i]->rank != 0)
+			return fnSame(a); /* cannot be simplified */
+	}
+
+	if(sametype){
+		Array *b = allocarray(type, a->rank, a->size);
+		b->stranded = a->stranded;
+		for(i = 0; i < a->rank; i++)
+			b->shape[i] = a->shape[i];
+		for(i = 0; i < a->size; i++)
+			memcpy(b->rawdata + i * datasizes[type], a->arraydata[i]->rawdata, datasizes[type]);
+		return b;
+	}else if(canfloat){
+		Array *b = allocarray(AtypeFloat, a->rank, a->size);
+		b->stranded = a->stranded;
+		for(i = 0; i < a->rank; i++)
+			b->shape[i] = a->shape[i];
+		for(i = 0; i < a->size; i++){
+			if(a->arraydata[i]->type == AtypeFloat)
+				b->floatdata[i] = a->arraydata[i]->floatdata[0];
+			else
+				b->floatdata[i] = a->arraydata[i]->intdata[0];
+		}
+		return b;
+	}else
+		return fnSame(a);
 }
\ No newline at end of file