From 1fc3a68205697c2ad33b2f0e1cfcb526e73b0677 Mon Sep 17 00:00:00 2001
From: LemonBoy <LemonBoy@users.noreply.github.com>
Date: Mon, 30 Jul 2018 08:34:15 +0200
Subject: [PATCH] Access implicit `result` trough envP in closures (#8471)

Reuse the existing machinery, most of the changes are only needed to
handle the `result = result` node in nkReturnStmt produced by the
closure iterator transform.

Fixes #338
---
 compiler/lambdalifting.nim | 14 +++++++++++++-
 tests/iter/t338.nim        | 20 ++++++++++++++++++++
 2 files changed, 33 insertions(+), 1 deletion(-)
 create mode 100644 tests/iter/t338.nim

diff --git a/compiler/lambdalifting.nim b/compiler/lambdalifting.nim
index 0a48011504..d8c0461ce8 100644
--- a/compiler/lambdalifting.nim
+++ b/compiler/lambdalifting.nim
@@ -222,7 +222,7 @@ proc interestingIterVar(s: PSym): bool {.inline.} =
   # XXX optimization: Only lift the variable if it lives across
   # yield/return boundaries! This can potentially speed up
   # closure iterators quite a bit.
-  result = s.kind in {skVar, skLet, skTemp, skForVar} and sfGlobal notin s.flags
+  result = s.kind in {skResult, skVar, skLet, skTemp, skForVar} and sfGlobal notin s.flags
 
 template isIterator*(owner: PSym): bool =
   owner.kind == skIterator and owner.typ.callConv == ccClosure
@@ -458,6 +458,10 @@ proc detectCapturedVars(n: PNode; owner: PSym; c: var DetectionPass) =
   of nkLambdaKinds, nkIteratorDef, nkFuncDef:
     if n.typ != nil:
       detectCapturedVars(n[namePos], owner, c)
+  of nkReturnStmt:
+    if n[0].kind in {nkAsgn, nkFastAsgn}:
+      detectCapturedVars(n[0].sons[1], owner, c)
+    else: assert n[0].kind == nkEmpty
   else:
     for i in 0..<n.len:
       detectCapturedVars(n[i], owner, c)
@@ -687,6 +691,13 @@ proc liftCapturedVars(n: PNode; owner: PSym; d: DetectionPass;
     if n.len == 2:
       n.sons[1] = liftCapturedVars(n[1], owner, d, c)
       if n[1].kind == nkClosure: result = n[1]
+  of nkReturnStmt:
+    if n[0].kind in {nkAsgn, nkFastAsgn}:
+      # we have a `result = result` expression produced by the closure
+      # transform, let's not touch the LHS in order to make the lifting pass
+      # correct when `result` is lifted
+      n[0].sons[1] = liftCapturedVars(n[0].sons[1], owner, d, c)
+    else: assert n[0].kind == nkEmpty
   else:
     if owner.isIterator:
       if nfLL in n.flags:
@@ -757,6 +768,7 @@ proc liftLambdas*(g: ModuleGraph; fn: PSym, body: PNode; tooEarly: var bool): PN
     if d.somethingToDo:
       var c = initLiftingPass(fn)
       result = liftCapturedVars(body, fn, d, c)
+      # echo renderTree(result, {renderIds})
       if c.envvars.getOrDefault(fn.id) != nil:
         result = newTree(nkStmtList, rawClosureCreation(fn, d, c), result)
     else:
diff --git a/tests/iter/t338.nim b/tests/iter/t338.nim
new file mode 100644
index 0000000000..dbced8b600
--- /dev/null
+++ b/tests/iter/t338.nim
@@ -0,0 +1,20 @@
+discard """
+  output: '''0
+1
+2
+3
+4
+'''
+"""
+
+proc moo(): iterator (): int =
+  iterator fooGen: int {.closure.} =
+    while true:
+      yield result
+      result.inc
+  return fooGen
+
+var foo = moo()
+
+for i in 0 .. 4:
+  echo foo()