fixes #3636

4 files changed, 31 insertions, 9 deletions

diff --git a/compiler/ast.nim b/compiler/ast.nim
index f34358788..2ce0afcc3 100644
--- a/compiler/ast.nim
+++ b/compiler/ast.nim
@@ -932,7 +932,7 @@ const
     skMacro, skTemplate, skConverter, skEnumField, skLet, skStub, skAlias}
   PersistentNodeFlags*: TNodeFlags = {nfBase2, nfBase8, nfBase16,
                                       nfDotSetter, nfDotField,
-                                      nfIsRef, nfIsCursor}
+                                      nfIsRef, nfIsCursor, nfLL}
   namePos* = 0
   patternPos* = 1    # empty except for term rewriting macros
   genericParamsPos* = 2
diff --git a/compiler/lambdalifting.nim b/compiler/lambdalifting.nim
index 9b83204fe..9c64037a8 100644
--- a/compiler/lambdalifting.nim
+++ b/compiler/lambdalifting.nim
@@ -633,7 +633,6 @@ proc wrapIterBody(n: PNode; owner: PSym): PNode =
                     getStateField(owner), info))
   stateAsgnStmt.add(newIntTypeNode(nkIntLit, -1, getSysType(tyInt)))
   result.add(stateAsgnStmt)
-  result.flags.incl nfLL
 
 proc symToClosure(n: PNode; owner: PSym; d: DetectionPass;
                   c: var LiftingPass): PNode =
@@ -671,6 +670,8 @@ proc liftCapturedVars(n: PNode; owner: PSym; d: DetectionPass;
     let s = n.sym
     if isInnerProc(s):
       if not c.processed.containsOrIncl(s.id):
+        #if s.name.s == "temp":
+        #  echo renderTree(s.getBody, {renderIds})
         let body = wrapIterBody(liftCapturedVars(s.getBody, s, d, c), s)
         if c.envvars.getOrDefault(s.id).isNil:
           s.ast.sons[bodyPos] = body
@@ -696,13 +697,17 @@ proc liftCapturedVars(n: PNode; owner: PSym; d: DetectionPass;
       m.typ = n.typ
       result = liftCapturedVars(m, owner, d, c)
   else:
-    if owner.isIterator and n.kind == nkYieldStmt:
-      result = transformYield(n, owner, d, c)
-    elif owner.isIterator and n.kind == nkReturnStmt:
-      result = transformReturn(n, owner, d, c)
-    else:
-      for i in 0..<n.len:
-        n.sons[i] = liftCapturedVars(n[i], owner, d, c)
+    if owner.isIterator:
+      if n.kind == nkYieldStmt:
+        return transformYield(n, owner, d, c)
+      elif n.kind == nkReturnStmt:
+        return transformReturn(n, owner, d, c)
+      elif nfLL in n.flags:
+        # special case 'when nimVm' due to bug #3636:
+        n.sons[1] = liftCapturedVars(n[1], owner, d, c)
+        return
+    for i in 0..<n.len:
+      n.sons[i] = liftCapturedVars(n[i], owner, d, c)
 
 # ------------------ old stuff -------------------------------------------
 
diff --git a/compiler/semexprs.nim b/compiler/semexprs.nim
index 43c98859c..87d7764a2 100644
--- a/compiler/semexprs.nim
+++ b/compiler/semexprs.nim
@@ -1817,6 +1817,7 @@ proc semWhen(c: PContext, n: PNode, semCheck = true): PNode =
       whenNimvm = lookUp(c, exprNode).magic == mNimvm
     elif exprNode.kind == nkSym:
       whenNimvm = exprNode.sym.magic == mNimvm
+    if whenNimvm: n.flags.incl nfLL
 
   for i in countup(0, sonsLen(n) - 1):
     var it = n.sons[i]
diff --git a/tests/iter/twrap_walkdir.nim b/tests/iter/twrap_walkdir.nim
new file mode 100644
index 000000000..4ac487d8e
--- /dev/null
+++ b/tests/iter/twrap_walkdir.nim
@@ -0,0 +1,16 @@
+
+
+
+import os
+
+# bug #3636
+
+proc fooIt(foo: string): iterator(): (string) =
+  iterator temp(): (string) =
+    for f in walkDirRec(foo): # No problem with walkFiles
+      yield f
+  return temp
+
+let it = fooIt(".")
+for x in it():
+  echo x