Optimize spread iteration inside iteration

civet.dev/reference.md

@@ -1861,6 +1861,13 @@ function flatJoin<T>(list: T[][], sep: T): T[]
       ...sublist
 </Playground>
 
+<Playground>
+flatImage :=
+  for x of [0...nx]
+    ...for y of [0...ny]
+      image.get x, y
+</Playground>
+
 If you don't specify a body, `for` loops list the item being iterated over:
 
 <Playground>

source/parser.hera

@@ -329,6 +329,17 @@ CommaExpression
 # Variation on CommaExpression that allows ... spreads,
 # which we allow in statements for the sake of expressionized iterations
 CommaExpressionSpread
+  _?:ws DotDotDot _?:ws2 IterationActualStatement:iteration ->
+    if (iteration.subtype === "do" || iteration.subtype === "comptime") return $skip
+    if (ws2) {
+      if (ws) {
+        ws = [ws, ws2]
+      } else {
+        ws = ws2
+      }
+    }
+    iteration = { ...iteration, resultsParent: true }
+    return prepend(ws, iteration)
   AssignmentExpressionSpread ( CommaDelimiter AssignmentExpressionSpread )* ->
     if($2.length == 0) return $1
     return $0
@@ -349,6 +360,21 @@ AssignmentExpressionSpread
       expression,
       names: expression.names,
     }
+  _? DotDotDot AssignmentExpression:expression ->
+    return {
+      type: "SpreadElement",
+      children: $0,
+      expression,
+      names: expression.names,
+    }
+  AssignmentExpression:expression ( _? DotDotDot )? ->
+    if (!$2) return $1
+    return {
+      type: "SpreadElement",
+      children: [ ...$2, $1 ],
+      expression,
+      names: expression.names,
+    }
 
 # https://262.ecma-international.org/#prod-Arguments
 Arguments
@@ -4386,12 +4412,7 @@ Statement
   KeywordStatement
   VariableStatement
   IfStatement !ShouldExpressionize -> $1
-  IterationStatement !ShouldExpressionize ->
-    // Leave generator forms for IterationExpression
-    if ($1.generator) return $skip
-    // Leave `for` reductions for IterationExpression
-    if ($1.reduction) return $skip
-    return $1
+  IterationActualStatement
   SwitchStatement !ShouldExpressionize -> $1
   TryStatement !ShouldExpressionize -> $1
   FinallyClause
@@ -4407,6 +4428,14 @@ Statement
 
   # NOTE: no WithStatement
 
+IterationActualStatement
+  IterationStatement !ShouldExpressionize ->
+    // Leave generator forms for IterationExpression
+    if ($1.generator) return $skip
+    // Leave `for` reductions for IterationExpression
+    if ($1.reduction) return $skip
+    return $1
+
 # NOTE: Leave statement with trailing call expressions or pipe operator
 # to be expressionized by
 # ExpressionizedStatementWithTrailingCallExpressions

source/parser/declaration.civet

@@ -156,7 +156,7 @@ function prependStatementExpressionBlock(initializer: Initializer, statement: AS
     ws = exp[0]
     exp = exp[1]!
 
-  return unless exp?.type is "StatementExpression"
+  return unless exp is like {type: "StatementExpression"}, {type: "SpreadElement", expression: {type: "StatementExpression"}}
 
   pre: ASTNode[] := []
   statementExp := exp.statement

source/parser/function.civet

@@ -677,7 +677,24 @@ function wrapIterationReturningResults(
   // to implement `implicitlyReturned`
   return if statement.resultsRef?
 
-  resultsRef := statement.resultsRef = makeRef "results"
+  // Inherit parent's resultsRef if requested
+  if statement.resultsParent
+    { ancestor } := findAncestor statement,
+      .type is "ForStatement" or .type is "IterationStatement"
+      isFunction
+    unless ancestor
+      statement.children.unshift
+        type: "Error"
+        message: "Could not find ancestor of spread iteration"
+      return
+    resultsRef := statement.resultsRef = ancestor.resultsRef
+    iterationDefaultBody statement
+    { block } := statement
+    unless block.empty
+      assignResults block, (node) => [ resultsRef, ".push(", node, ")" ]
+    return
+
+  resultsRef := statement.resultsRef ?= makeRef "results"
 
   declaration := iterationDeclaration statement
   { ancestor, child } := findAncestor statement, .type is "BlockStatement"
@@ -1145,6 +1162,7 @@ function expressionizeIteration(exp: IterationExpression): void
 
     statements =
       . ["", statement]
+
   else
     resultsRef := statement.resultsRef ??= makeRef "results"
 

source/parser/lib.civet

@@ -1002,7 +1002,7 @@ function processAssignments(statements): void
 
     let block?: BlockStatement
     // Extract StatementExpression as block
-    if exp.parent?.type is "BlockStatement" and !$1.-1?.-1?.special// can only prepend to assignments that are children of blocks
+    if blockContainingStatement(exp) and !$1.-1?.-1?.special// can only prepend to assignments that are children of blocks
       block = makeBlockFragment()
       if ref := prependStatementExpressionBlock(
         {type: "Initializer", expression: $2, children: [undefined, undefined, $2]}

source/parser/types.civet

@@ -382,6 +382,7 @@ export type IterationStatement
   generator?: ASTNode
   object?: boolean
   resultsRef: ASTRef?
+  resultsParent?: boolean // inherit resultsRef from parent loop
 
 export type BreakStatement
   type: "BreakStatement"
@@ -425,6 +426,7 @@ export type ForStatement
   hoistDec: ASTNode
   generator?: ASTNode
   resultsRef: ASTRef?
+  resultsParent?: boolean // inherit resultsRef from parent loop
   reduction?: ForReduction
   object?: boolean
 

test/for.civet

@@ -1582,3 +1582,30 @@ describe "for", ->
           results.push(item.pre, ...item.mid, item.post, ... item.tail)
         };const result =results
     """
+
+    testCase """
+      spread iteration
+      ---
+      function flatSquare(arrays)
+        for array of arrays
+          ... for item of array
+            item * item
+      ---
+      function flatSquare(arrays) {
+        const results=[];for (const array of arrays) {
+           for (const item of array) {
+            results.push(item * item)
+          }
+        };return results;
+      }
+    """
+
+    throws """
+      spread iteration without parent
+      ---
+      function flatSquare(arrays)
+        ... for item of array
+          item * item
+      ---
+      ParseErrors: unknown:2:3 Could not find ancestor of spread iteration
+    """