fixes iterator codegen regression

2 files changed, 12 insertions, 1 deletions

diff --git a/compiler/lambdalifting.nim b/compiler/lambdalifting.nim
index 3dd116078..20f903da9 100644
--- a/compiler/lambdalifting.nim
+++ b/compiler/lambdalifting.nim
@@ -213,6 +213,8 @@ proc makeClosure*(prc: PSym; env: PNode; info: TLineInfo): PNode =
   if env == nil:
     result.add(newNodeIT(nkNilLit, info, getSysType(tyNil)))
   else:
+    if env.kind == nkClosure:
+      localError(info, "internal error: taking closure of closure")
     result.add(env)
 
 proc interestingIterVar(s: PSym): bool {.inline.} =
@@ -709,6 +711,7 @@ proc liftCapturedVars(n: PNode; owner: PSym; d: DetectionPass;
   of nkClosure:
     if n[1].kind == nkNilLit:
       n.sons[0] = liftCapturedVars(n[0], owner, d, c)
+      #if n.sons[0].kind == nkClosure: result = n.sons[0]
   of nkLambdaKinds, nkIteratorDef:
     if n.typ != nil and n[namePos].kind == nkSym:
       let m = newSymNode(n[namePos].sym)
diff --git a/compiler/transf.nim b/compiler/transf.nim
index 296ea0a0d..50db2c7e8 100644
--- a/compiler/transf.nim
+++ b/compiler/transf.nim
@@ -313,10 +313,18 @@ proc introduceNewLocalVars(c: PTransf, n: PNode): PTransNode =
     result = PTransNode(n)
   of nkVarSection, nkLetSection:
     result = transformVarSection(c, n)
+  of nkClosure:
+    # it can happen that for-loop-inlining produced a fresh
+    # set of variables, including some computed environment
+    # (bug #2604). We need to patch this environment here too:
+    let a = n[1]
+    if a.kind == nkSym:
+      n.sons[1] = transformSymAux(c, a)
+    return PTransNode(n)
   else:
     result = newTransNode(n)
     for i in countup(0, sonsLen(n)-1):
-      result[i] =  introduceNewLocalVars(c, n.sons[i])
+      result[i] = introduceNewLocalVars(c, n.sons[i])
 
 proc transformYield(c: PTransf, n: PNode): PTransNode =
   result = newTransNode(nkStmtList, n.info, 0)