Fix array return types for tuples of arrays

compiler/passes/normalize.cpp

@@ -71,6 +71,7 @@ static bool        isArrayFormal(ArgSymbol* arg);
 static Expr*       arrayTypeEltTypeExprOrNull(Expr* expr);
 
 static bool        returnsArray(FnSymbol* fn);
+static bool        returnsArrayOrTupleOfArrays(FnSymbol* fn);
 static void        makeExportWrapper(FnSymbol* fn);
 
 static void        fixupArrayFormals(FnSymbol* fn);
@@ -2202,26 +2203,42 @@ static bool isVoidReturn(CallExpr* call) {
   return retval;
 }
 
-static bool hasGenericArrayReturn(FnSymbol* fn) {
-  if (returnsArray(fn)) {
-    BlockStmt* typeExpr = fn->retExprType;
+static bool isGenericArray(CallExpr* call) {
+  if (!call->isNamed("chpl__buildArrayRuntimeType")) return false;
 
-    // returnsArray ensured this was a call to "chpl__buildArrayRuntimeType"
-    CallExpr* call = toCallExpr(typeExpr->body.tail);
-    int nArgs = call->numActuals();
-    Expr* domExpr = call->get(1);
-    Expr* eltExpr = nArgs == 2 ? call->get(2) : NULL;
-    bool noDom = (isSymExpr(domExpr) && toSymExpr(domExpr)->symbol() == gNil);
-
-    if (noDom || eltExpr == NULL) {
-      // Either the domain is not provided explicitly as part of the return
-      // type, or the element type is not provided, or both
+  int nArgs = call->numActuals();
+  Expr* domExpr = call->get(1);
+  Expr* eltExpr = nArgs == 2 ? call->get(2) : nullptr;
+  bool noDom = (isSymExpr(domExpr) && toSymExpr(domExpr)->symbol() == gNil);
 
-      return true;
+  // Either the domain is not provided explicitly as part of the return
+  // type, or the element type is not provided, or both
+  return noDom || eltExpr == nullptr;
+}
+static bool containsGenericArray(CallExpr* call) {
+  if (call->isNamed("_build_tuple")) {
+    for_actuals(arg, call) {
+      if (CallExpr* argCall = toCallExpr(arg)) {
+        if (containsGenericArray(argCall))
+          return true;
+      }
     }
+    return false;
+  } else {
+    return isGenericArray(call);
   }
+}
+
+static bool hasGenericArrayReturn(FnSymbol* fn) {
+  if (!returnsArrayOrTupleOfArrays(fn)) return false;
+  BlockStmt* typeExpr = fn->retExprType;
+
+  // returnsArrayOrTupleOfArrays ensured this was a call to
+  // "chpl__buildArrayRuntimeType" or a call to "_build_tuple" with
+  // "chpl__buildArrayRuntimeType" as an argument
+  CallExpr* call = toCallExpr(typeExpr->body.tail);
+  return containsGenericArray(call);
 
-  return false;
 }
 
 //
@@ -2265,29 +2282,109 @@ static void insertElementTypeCheck(Expr* declaredRet, Expr* actualRet,
   retVar->insertBefore(checkEltType);
 }
 
+// Validates the actual return type is an array of some kind
+static void insertGenericArrayCheck(Expr* actualRet, CallExpr* retVar) {
+  CallExpr* checkGenericArray = new CallExpr("chpl__checkGenericArrayReturn",
+                                             actualRet->copy());
+  retVar->insertBefore(checkGenericArray);
+}
+
+static void modifyPartiallyGenericArrayReturnSimple(FnSymbol* fn,
+                                                    VarSymbol* retval,
+                                                    CallExpr* ret,
+                                                    Expr* retExpr);
+static Expr* modifyPartiallyGenericArrayReturnRecurse(FnSymbol* fn,
+                                                      CallExpr* ret,
+                                                      Expr* typeExpr,
+                                                      Expr* retExpr);
 static void modifyPartiallyGenericArrayReturn(FnSymbol* fn,
                                               VarSymbol* retval,
                                               CallExpr* ret,
                                               Expr* retExpr) {
   BlockStmt* typeExpr = fn->retExprType;
+  if (toCallExpr(typeExpr->body.tail)->isNamed("chpl__buildArrayRuntimeType")) {
+    modifyPartiallyGenericArrayReturnSimple(fn, retval, ret, retExpr);
+    return;
+  }
+  auto newRetExpr =
+    modifyPartiallyGenericArrayReturnRecurse(fn, ret, typeExpr->body.tail, retExpr);
+  ret->insertBefore(new CallExpr(PRIM_MOVE, retval, newRetExpr));
+}
+static Expr* modifyPartiallyGenericArrayReturnRecurse(FnSymbol* fn,
+                                                      CallExpr* ret,
+                                                      Expr* typeExpr,
+                                                      Expr* retExpr) {
+  auto typeCall = toCallExpr(typeExpr);
+  auto retCall = toCallExpr(retExpr);
+  if (typeCall && typeCall->isNamed("_build_tuple")) {
+    if (!(retCall && retCall->isNamed("_build_tuple"))) {
+      USR_FATAL(fn, "return type is a tuple, but return value is not");
+    }
+    int nTypeArgs = typeCall->numActuals();
+    int nRetArgs = retCall->numActuals();
+    if (nTypeArgs != nRetArgs) {
+      USR_FATAL(fn, "return type is a tuple of size %d, but return value is a tuple of size %d", nTypeArgs, nRetArgs);
+    }
+    for (int i = 1; i <= nTypeArgs; i++) {
+      Expr* typeArg = typeCall->get(i);
+      Expr* retArg = retCall->get(i);
+      retCall->get(i)->replace(new SymExpr(gNil)); // dummy replacement
+      auto newRetArg =
+        modifyPartiallyGenericArrayReturnRecurse(fn, ret, typeArg, retArg);
+      retCall->get(i)->replace(newRetArg);
+    }
+  } else if (typeCall && isGenericArray(typeCall)) {
+    int nArgs = typeCall->numActuals();
+    Expr* domExpr = typeCall->get(1);
+    Expr* retEltExpr = nArgs == 2 ? typeCall->get(2) : nullptr;
+    bool noDom = (isSymExpr(domExpr) && toSymExpr(domExpr)->symbol() == gNil);
+    if (!noDom || retEltExpr != nullptr) {
+      prepareRetExpr(retExpr, ret);
+    }
+    if (!noDom) {
+      // Add checks against the declared domain
+      insertDomainCheck(retExpr, ret, domExpr);
+    }
+    if (retEltExpr != nullptr) {
+      insertElementTypeCheck(retEltExpr, retExpr, ret);
+    }
+    if (noDom && retEltExpr == nullptr) {
+    insertGenericArrayCheck(retExpr, ret);
+  }
+  }
+  return retExpr;
+}
+
+
+
+
+static void modifyPartiallyGenericArrayReturnSimple(FnSymbol* fn,
+                                                    VarSymbol* retval,
+                                                    CallExpr* ret,
+                                                    Expr* retExpr) {
+  BlockStmt* typeExpr = fn->retExprType;
 
   CallExpr* call = toCallExpr(typeExpr->body.tail);
   int nArgs = call->numActuals();
   Expr* domExpr = call->get(1);
-  Expr* retEltExpr = nArgs == 2 ? call->get(2) : NULL;
+  Expr* retEltExpr = nArgs == 2 ? call->get(2) : nullptr;
   bool noDom = (isSymExpr(domExpr) && toSymExpr(domExpr)->symbol() == gNil);
 
-  if (!noDom) {
+  if (!noDom || retEltExpr != nullptr) {
     prepareRetExpr(retExpr, ret);
+  }
+  if (!noDom) {
     // Add checks against the declared domain
     insertDomainCheck(retExpr, ret, domExpr);
   }
-
-  if (retEltExpr != NULL) {
-    prepareRetExpr(retExpr, ret);
+  if (retEltExpr != nullptr) {
     insertElementTypeCheck(retEltExpr, retExpr, ret);
   }
 
+  if (noDom && retEltExpr == nullptr) {
+    insertGenericArrayCheck(retExpr, ret);
+  }
+
   // TODO: Do something about coercion
 
   ret->insertBefore(new CallExpr(PRIM_MOVE, retval, retExpr));
@@ -3916,6 +4013,21 @@ static bool returnsArray(FnSymbol* fn) {
   // If we don't have a declared return type, assume we don't return an array
   return false;
 }
+static bool isTupleOrArray(CallExpr* call) {
+  if (!call) return false;
+  if (call->isNamed("chpl__buildArrayRuntimeType"))
+    return true;
+  else if (call->isNamed("_build_tuple"))
+    for_actuals(actual, call) {
+      if (isTupleOrArray(toCallExpr(actual)))
+        return true;
+    }
+  return false;
+}
+static bool returnsArrayOrTupleOfArrays(FnSymbol* fn) {
+  return fn->retExprType != NULL &&
+         isTupleOrArray(toCallExpr(fn->retExprType->body.tail));
+}
 
 
 /************************************* | **************************************

modules/internal/ChapelArray.chpl

@@ -648,6 +648,11 @@ module ChapelArray {
                     b: string);
     }
   }
+  pragma "last resort"
+  proc chpl__checkRetEltTypeMatch(a, type b) {
+    compilerError("expected an array or iterator but got a value of type ",
+                  a.type:string);
+  }
   proc chpl__checkOutEltTypeMatch(a: [], type b) {
     if (a.eltType != b) {
       compilerError("array element type mismatch in initializing out formal ",
@@ -669,6 +674,14 @@ module ChapelArray {
     }
   }
 
+  proc chpl__checkGenericArrayReturn(a: []) { }
+  proc chpl__checkGenericArrayReturn(a: _iteratorRecord) { }
+  pragma "last resort"
+  proc chpl__checkGenericArrayReturn(a) {
+    compilerError("expected an array or iterator but got a value of type ",
+                  a.type:string);
+  }
+
   //
   // Support for distributions
   //

modules/internal/ChapelDomain.chpl

@@ -351,6 +351,12 @@ module ChapelDomain {
     }
   }
 
+  pragma "last resort"
+  proc chpl__checkDomainsMatch(a, b) {
+    compilerError("expected an array or iterator but got a value of type ",
+                  a.type:string);
+  }
+
   proc chpl_countDomHelp(dom, counts) {
     var ranges = dom.dims();
     for param i in 0..dom.rank-1 do

test/arrays/return/returnArbitraryArray.chpl

@@ -1,8 +1,20 @@
 proc inc(X: [] real): [] real {
   return X + 1;
 }
+proc incByTuple(X: [] real): (real, [] real) {
+  return (X[1], X + 1);
+}
+proc incByTuple2(X: [] real): ([] real, ([] real, [] real)) {
+  return (X + 1, (X + 2, X + 3));
+}
+proc incByTuple3(X: [] real): (real, ([] real, [] real), [] real) {
+  return (X[1], (X + 2, X + 3), X + 4);
+}
 
 var A = [1.0, 2.0, 3.0];
 
 writeln(inc(A));
+writeln(incByTuple(A));
+writeln(incByTuple2(A));
+writeln(incByTuple3(A));
 writeln(A);

test/arrays/return/returnArbitraryArray.good

@@ -1,2 +1,5 @@
 2.0 3.0 4.0
+(2.0, 2.0 3.0 4.0)
+(2.0 3.0 4.0, (3.0 4.0 5.0, 4.0 5.0 6.0))
+(2.0, (3.0 4.0 5.0, 4.0 5.0 6.0), 5.0 6.0 7.0)
 1.0 2.0 3.0

test/arrays/return/returnArbitraryBadDomain2.bad

@@ -0,0 +1,3 @@
+returnArbitraryBadDomain2.chpl:1: In function 'inc':
+returnArbitraryBadDomain2.chpl:2: error: illegal cast from promoted expression yielding int(64) to [domain(1,int(64),one)] int(64)
+  returnArbitraryBadDomain2.chpl:7: called as inc(X: [domain(1,int(64),one)] int(64))

test/arrays/return/returnArbitraryBadDomain2.chpl

@@ -0,0 +1,8 @@
+proc inc(X: [0..4] int): (int, (int, [1..6] int)) {
+  return (X[0], (X[0], X + 3));
+}
+
+var A = [1, 2, 3, 4, 5];
+
+writeln(inc(A));
+writeln(A);

test/arrays/return/returnArbitraryBadDomain2.future

@@ -0,0 +1,2 @@
+bug: cannot return promoted expression in tuple with explicit domain
+#29261

test/arrays/return/returnArbitraryBadDomain2.good

@@ -0,0 +1 @@
+returnArbitraryBadDomain2.chpl:2: error: halt reached - domain mismatch on return

test/arrays/return/returnArbitraryBadDomain2.skipif

@@ -0,0 +1 @@
+COMPOPTS <= --fast

test/arrays/return/returnArbitraryBadEltType2.chpl

@@ -0,0 +1,8 @@
+proc inc(X: [] real): ([] int, [] int) {
+  return (X + 1, X + 2);
+}
+
+var A = [1.0, 2.0, 3.0];
+
+writeln(inc(A));
+writeln(A);

test/arrays/return/returnArbitraryBadEltType2.good

@@ -0,0 +1,3 @@
+returnArbitraryBadEltType2.chpl:1: In function 'inc':
+returnArbitraryBadEltType2.chpl:2: error: array element type mismatch in return from real(64) to int(64)
+  returnArbitraryBadEltType2.chpl:7: called as inc(X: [domain(1,int(64),one)] real(64))

test/arrays/return/returnArbitraryBadEltType3.chpl

@@ -0,0 +1,8 @@
+proc inc(X: [] real): ([] real, ([] int, int)) {
+  return (X + 1, (X + 2, 3));
+}
+
+var A = [1.0, 2.0, 3.0];
+
+writeln(inc(A));
+writeln(A);

test/arrays/return/returnArbitraryBadEltType3.good

@@ -0,0 +1,3 @@
+returnArbitraryBadEltType3.chpl:1: In function 'inc':
+returnArbitraryBadEltType3.chpl:2: error: array element type mismatch in return from real(64) to int(64)
+  returnArbitraryBadEltType3.chpl:7: called as inc(X: [domain(1,int(64),one)] real(64))

test/arrays/return/returnArbitraryMultiDiffRank.chpl

@@ -7,8 +7,21 @@ proc foo(type t, param fiveInts: bool): [] {
     return x;
   }
 }
+proc bar(type t, param fiveInts: bool): ([], [], ([], [])) {
+  if (fiveInts) {
+    var x: [1..5] int;
+    return (x, x, (x, x));
+  } else {
+    var x: [1..3, 1..4] t;
+    return (x, x, (x, x));
+  }
+}
 
 var a = foo(real, true);
 var b = foo(real, false);
 writeln(a);
 writeln(b);
+var c = bar(real, true);
+var d = bar(real, false);
+writeln(c);
+writeln(d);

test/arrays/return/returnArbitraryMultiDiffRank.good

@@ -2,3 +2,13 @@
 0.0 0.0 0.0 0.0
 0.0 0.0 0.0 0.0
 0.0 0.0 0.0 0.0
+(0 0 0 0 0, 0 0 0 0 0, (0 0 0 0 0, 0 0 0 0 0))
+(0.0 0.0 0.0 0.0
+0.0 0.0 0.0 0.0
+0.0 0.0 0.0 0.0, 0.0 0.0 0.0 0.0
+0.0 0.0 0.0 0.0
+0.0 0.0 0.0 0.0, (0.0 0.0 0.0 0.0
+0.0 0.0 0.0 0.0
+0.0 0.0 0.0 0.0, 0.0 0.0 0.0 0.0
+0.0 0.0 0.0 0.0
+0.0 0.0 0.0 0.0))

test/arrays/return/returnArbitraryMultiDiffType.chpl

@@ -7,8 +7,21 @@ proc foo(type t, param fiveInts: bool): [] {
     return x;
   }
 }
+proc bar(type t, param fiveInts: bool): ([], ([], []), []) {
+  if (fiveInts) {
+    var x: [1..5] int;
+    return (x, (x, x), x);
+  } else {
+    var x: [1..3] t;
+    return (x, (x, x), x);
+  }
+}
 
 var a = foo(real, true);
 var b = foo(real, false);
 writeln(a);
 writeln(b);
+var c = bar(real, true);
+var d = bar(real, false);
+writeln(c);
+writeln(d);

test/arrays/return/returnArbitraryMultiDiffType.good

@@ -1,2 +1,4 @@
 0 0 0 0 0
 0.0 0.0 0.0
+(0 0 0 0 0, (0 0 0 0 0, 0 0 0 0 0), 0 0 0 0 0)
+(0.0 0.0 0.0, (0.0 0.0 0.0, 0.0 0.0 0.0), 0.0 0.0 0.0)