Introduce a new for/assign pattern in loops (#55)

Introduces a new convenience pattern to for loops:

for x := expr

is equivalent to

for x in [expr]

and is handy for "sandwiching" between nested for loops.

This commit also fixes a minor bug in which expressions like

for 3 in list

threw an exception rather than reporting a Syntax error.

lib/core/src/main/antlr/org/cgsuite/lang/parser/Cgsuite.g

@@ -741,6 +741,7 @@ loopAntecedent
 
 forLoopAntecedent
     : (forClause FROM) => forFromLoopAntecedent
+    | (forClause ASSIGN) => forAssignLoopAntecedent
     | forInLoopAntecedent
     ;
 
@@ -749,6 +750,11 @@ forFromLoopAntecedent
       -> ^(LOOP_SPEC[$forClause.tree.getToken()] forClause fromClause toClause? byClause? whileClause? whereClause?)
     ;
 
+forAssignLoopAntecedent
+    : forClause ASSIGN expression
+      -> ^(LOOP_SPEC[$forClause.tree.getToken()] forClause ^(ASSIGN expression))
+    ;
+
 forInLoopAntecedent
     : forClause inClause whileClause? whereClause?
       -> ^(LOOP_SPEC[$forClause.tree.getToken()] forClause inClause whileClause? whereClause?)
@@ -760,7 +766,7 @@ whileLoopAntecedent
     ;
 
 forClause
-    : FOR^ expression
+    : FOR^ IDENTIFIER
     ;
 
 fromClause

lib/core/src/main/scala/org/cgsuite/lang/node/LoopNode.scala

@@ -30,13 +30,14 @@ object LoopNode {
       Vector(LoopNode(
         loopSpecTree,
         loopType,
-        loopSpecTree.children find { _.getType == FOR   } map { t => IdentifierNode(t.head) },
-        loopSpecTree.children find { _.getType == IN    } map { t => EvalNode(t.head) },
-        loopSpecTree.children find { _.getType == FROM  } map { t => EvalNode(t.head) },
-        loopSpecTree.children find { _.getType == TO    } map { t => EvalNode(t.head) },
-        loopSpecTree.children find { _.getType == BY    } map { t => EvalNode(t.head) },
-        loopSpecTree.children find { _.getType == WHILE } map { t => EvalNode(t.head) },
-        loopSpecTree.children find { _.getType == WHERE } map { t => EvalNode(t.head) },
+        loopSpecTree.children find { _.getType == FOR    } map { t => IdentifierNode(t.head) },
+        loopSpecTree.children find { _.getType == IN     } map { t => EvalNode(t.head) },
+        loopSpecTree.children find { _.getType == FROM   } map { t => EvalNode(t.head) },
+        loopSpecTree.children find { _.getType == TO     } map { t => EvalNode(t.head) },
+        loopSpecTree.children find { _.getType == BY     } map { t => EvalNode(t.head) },
+        loopSpecTree.children find { _.getType == WHILE  } map { t => EvalNode(t.head) },
+        loopSpecTree.children find { _.getType == WHERE  } map { t => EvalNode(t.head) },
+        loopSpecTree.children find { _.getType == ASSIGN } map { t => EvalNode(t.head) },
         nextNodes
       ))
     }
@@ -63,10 +64,11 @@ case class LoopNode(
   by      : Option[EvalNode],
   `while` : Option[EvalNode],
   where   : Option[EvalNode],
+  assign  : Option[EvalNode],
   body    : Vector[EvalNode]
 ) extends EvalNode {
 
-  override val children = forId.toVector ++ in ++ from ++ to ++ by ++ `while` ++ where ++ body
+  override val children = forId.toVector ++ in ++ from ++ to ++ by ++ `while` ++ where ++ assign ++ body
 
   private val prepareLoop = Symbol(s"PrepareLoop [${tree.location}]")
   private val loop = Symbol(s"Loop [${tree.location}]")
@@ -129,7 +131,8 @@ case class LoopNode(
       to map { "to " + _.toNodeString },
       by map { "by " + _.toNodeString },
       `while` map { "while " + _.toNodeString },
-      where map { "where " + _.toNodeString }
+      where map { "where " + _.toNodeString },
+      assign map { ":= " + _.toNodeString },
     ).flatten.mkString(" ")
     val bodyStr = body map { loopTypeStr + " " + _.toNodeString } mkString " "
     antecedent + " " + bodyStr + " end"
@@ -155,14 +158,19 @@ case class LoopNode(
 
     var continue = true
 
+    if (assign.isDefined) {
+      domain.localScope(forIndex) = assign.get.evaluate(domain)
+    }
+
     while (continue) {
 
       if (Thread.interrupted())
         throw CalculationCanceledException("Calculation canceled by user.", token = Some(token))
 
       if (iterator == null) {
-        if (forIndex >= 0)
+        if (forIndex >= 0 && counter != null) {
           domain.localScope(forIndex) = counter
+        }
         continue = toVal == null || (checkLeq && Ops.leq(counter, toVal)) || (!checkLeq && Ops.leq(toVal, counter))
       } else {
         if (iterator.hasNext)
@@ -195,6 +203,10 @@ case class LoopNode(
           counter = Ops.Plus(tree, counter, byVal)
       }
 
+      if (assign.isDefined) {
+        continue = false
+      }
+
     }
 
     Profiler.stop(loop)

lib/core/src/test/scala/org/cgsuite/lang/EvalTest.scala

@@ -281,7 +281,7 @@ class EvalTest extends CgscriptSpec {
 
   it should "handle various types of loops correctly" in {
 
-    // (name, initializer, fn, for-snippet, result, optional-sorted-result, sum)
+    // (name, initializer, fn, for-snippet, result, multi-result, optional-sorted-result, sum)
     val loopScenarios = Seq(
       ("for-from-to", "", "x*x", "for x from 1 to 5", "1,4,9,16,25", "1,2,4,5,9,10,16,17,25,26", None, "55"),
       ("for-from-to-by", "", "x*x", "for x from 1 to 10 by 3", "1,16,49,100", "1,2,16,17,49,50,100,101", None, "166"),
@@ -293,7 +293,8 @@ class EvalTest extends CgscriptSpec {
       ("for-in", "", "x*x", "for x in [1,2,3,2]", "1,4,9,4", "1,2,4,5,9,10,4,5", Some("1,4,9"), "18"),
       ("for-in-where", "", "x", "for x in [1,2,3,2] where x % 2 == 1", "1,3", "1,2,3,4", None, "4"),
       ("for-in-while", "", "x", "for x in [1,2,3,2] while x % 2 == 1", "1", "1,2", None, "1"),
-      ("for-in-while-where", "", "x", "for x in [1,2,3,2] while x % 2 == 1 where x != 1", "", "", Some(""), "!!That `Collection` is empty.")
+      ("for-in-while-where", "", "x", "for x in [1,2,3,2] while x % 2 == 1 where x != 1", "", "", Some(""), "!!That `Collection` is empty."),
+      ("for-assign", "", "x*x", "for x := 5", "25", "25,26", None, "25")
     )
 
     val listComprehensionLoops = loopScenarios map { case (name, init, fn, snippet, result, _, _, _) =>
@@ -325,6 +326,15 @@ class EvalTest extends CgscriptSpec {
 
   }
 
+  it should "give the right error messages on malformed loops" in {
+
+    executeTests(Table(
+      header,
+      ("malformed forClause", "for 3 from 1 to 5 yield 6 end", "!!Syntax error.")
+    ))
+
+  }
+
   it should "properly construct and evaluate functions" in {
 
     executeTests(Table(