improvements for first class iterators

3 files changed, 35 insertions, 6 deletions

diff --git a/compiler/ccgstmts.nim b/compiler/ccgstmts.nim
index 9c9b61088..5e3504cd9 100755
--- a/compiler/ccgstmts.nim
+++ b/compiler/ccgstmts.nim
@@ -115,7 +115,12 @@ proc genGotoState(p: BProc, n: PNode) =
 proc genBreakState(p: BProc, n: PNode) =
   var a: TLoc
   initLocExpr(p, n.sons[0], a)
-  lineF(p, cpsStmts, "if (($1) < 0) break;$n", [rdLoc(a)])
+  if n.sons[0].kind == nkClosure:
+    # XXX this produces quite inefficient code!
+    # the environment is guaranteed to contain the 'state' field at offset 0:
+    lineF(p, cpsStmts, "if ((((NI*) $1.ClEnv)[0]) < 0) break;$n", [rdLoc(a)])
+  else:
+    lineF(p, cpsStmts, "if (($1) < 0) break;$n", [rdLoc(a)])
 
 proc genSingleVar(p: BProc, a: PNode) =
   var v = a.sons[0].sym
diff --git a/compiler/lambdalifting.nim b/compiler/lambdalifting.nim
index ada14edf2..9d7e0e961 100644
--- a/compiler/lambdalifting.nim
+++ b/compiler/lambdalifting.nim
@@ -693,7 +693,7 @@ proc liftIterator*(iter: PSym, body: PNode): PNode =
   result.add(stateAsgnStmt)
 
 proc liftForLoop*(body: PNode): PNode =
-  # BIG problem ahead: the iterator could be invoked indirectly, but then
+  # problem ahead: the iterator could be invoked indirectly, but then
   # we don't know what environment to create here: 
   # 
   # iterator count(): int =
@@ -750,7 +750,10 @@ proc liftForLoop*(body: PNode): PNode =
   # gather vars in a tuple:
   var v2 = newNodeI(nkLetSection, body.info)
   var vpart = newNodeI(if L == 3: nkIdentDefs else: nkVarTuple, body.info)
-  for i in 0 .. L-3: addSon(vpart, body[i])
+  for i in 0 .. L-3: 
+    assert body[i].kind == nkSym
+    body[i].sym.kind = skLet
+    addSon(vpart, body[i])
 
   addSon(vpart, ast.emptyNode) # no explicit type
   if not env.isnil:
@@ -760,7 +763,10 @@ proc liftForLoop*(body: PNode): PNode =
 
   loopBody.sons[0] = v2
   var bs = newNodeI(nkBreakState, body.info)
-  bs.addSon(indirectAccess(env, 
-      newSym(skField, getIdent(":state"), env, env.info), body.info))
+  if not env.isNil:
+    bs.addSon(indirectAccess(env, 
+      newSym(skField, getIdent":state", env, env.info), body.info))
+  else:
+    bs.addSon(call.sons[0])
   loopBody.sons[1] = bs
   loopBody.sons[2] = body[L-1]
diff --git a/tests/run/titer8.nim b/tests/run/titer8.nim
index ae82c7d3c..d77159353 100644
--- a/tests/run/titer8.nim
+++ b/tests/run/titer8.nim
@@ -1,6 +1,11 @@
 discard """
   output: '''tada
-ta da'''
+1
+2
+3
+ta da1
+2
+3'''
 """
 # Test first class iterator:
 
@@ -21,13 +26,26 @@ iterator tokenize2(s: string, seps: set[char] = Whitespace): tuple[
         yield (substr(s, i, j-1), false)
     i = j
 
+iterator count3(): int {.closure.} =
+  yield 1
+  yield 2
+  yield 3
+
 for word, isSep in tokenize2("ta da", whiteSpace):
   if not isSep:
     stdout.write(word)
 echo ""
 
 proc inProc() =
+  for c in count3():
+    echo c
+  
   for word, isSep in tokenize2("ta da", whiteSpace):
     stdout.write(word)
 
+  for c in count3():
+    echo c
+
+
 inProc()
+