From 69faab3cf637416308da8e800e4ed52e09271fc5 Mon Sep 17 00:00:00 2001
From: Bung <crc32@qq.com>
Date: Sun, 27 Aug 2023 20:29:24 +0800
Subject: [PATCH] =?UTF-8?q?fix=20#22548;environment=20misses=20for=20type?=
 =?UTF-8?q?=20reference=20in=20iterator=20access=20n=E2=80=A6=20(#22559)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* fix #22548;environment misses for type reference in iterator access nested in closure

* fix #21737

* Update lambdalifting.nim

* remove containsCallKinds

* simplify

(cherry picked from commit 0b78b7f595ef96a9769e0d59167239c611b9857a)
---
 compiler/lambdalifting.nim | 24 +++++++++++++++++++-----
 tests/iter/t21737.nim      | 22 ++++++++++++++++++++++
 tests/iter/t22548.nim      | 21 +++++++++++++++++++++
 3 files changed, 62 insertions(+), 5 deletions(-)
 create mode 100644 tests/iter/t21737.nim
 create mode 100644 tests/iter/t22548.nim

diff --git a/compiler/lambdalifting.nim b/compiler/lambdalifting.nim
index b806fed8b7..e2c9bca988 100644
--- a/compiler/lambdalifting.nim
+++ b/compiler/lambdalifting.nim
@@ -298,6 +298,7 @@ type
     processed, capturedVars: IntSet
     ownerToType: Table[int, PType]
     somethingToDo: bool
+    inTypeOf: bool
     graph: ModuleGraph
     idgen: IdGenerator
 
@@ -401,6 +402,9 @@ Consider:
 
 """
 
+proc isTypeOf(n: PNode): bool =
+  n.kind == nkSym and n.sym.magic in {mTypeOf, mType}
+
 proc addClosureParam(c: var DetectionPass; fn: PSym; info: TLineInfo) =
   var cp = getEnvParam(fn)
   let owner = if fn.kind == skIterator: fn else: fn.skipGenericOwner
@@ -437,7 +441,8 @@ proc detectCapturedVars(n: PNode; owner: PSym; c: var DetectionPass) =
         c.somethingToDo = true
         addClosureParam(c, owner, n.info)
         if interestingIterVar(s):
-          if not c.capturedVars.containsOrIncl(s.id):
+          if not c.capturedVars.contains(s.id):
+            if not c.inTypeOf: c.capturedVars.incl(s.id)
             let obj = getHiddenParam(c.graph, owner).typ.skipTypes({tyOwned, tyRef, tyPtr})
             #let obj = c.getEnvTypeForOwner(s.owner).skipTypes({tyOwned, tyRef, tyPtr})
 
@@ -463,10 +468,12 @@ proc detectCapturedVars(n: PNode; owner: PSym; c: var DetectionPass) =
       addClosureParam(c, owner, n.info)
       #echo "capturing ", n.info
       # variable 's' is actually captured:
-      if interestingVar(s) and not c.capturedVars.containsOrIncl(s.id):
-        let obj = c.getEnvTypeForOwner(ow, n.info).skipTypes({tyOwned, tyRef, tyPtr})
-        #getHiddenParam(owner).typ.skipTypes({tyOwned, tyRef, tyPtr})
-        discard addField(obj, s, c.graph.cache, c.idgen)
+      if interestingVar(s):
+        if not c.capturedVars.contains(s.id):
+          if not c.inTypeOf: c.capturedVars.incl(s.id)
+          let obj = c.getEnvTypeForOwner(ow, n.info).skipTypes({tyOwned, tyRef, tyPtr})
+          #getHiddenParam(owner).typ.skipTypes({tyOwned, tyRef, tyPtr})
+          discard addField(obj, s, c.graph.cache, c.idgen)
       # create required upFields:
       var w = owner.skipGenericOwner
       if isInnerProc(w) or owner.isIterator:
@@ -498,9 +505,14 @@ proc detectCapturedVars(n: PNode; owner: PSym; c: var DetectionPass) =
       detectCapturedVars(n[namePos], owner, c)
   of nkReturnStmt:
     detectCapturedVars(n[0], owner, c)
+  of nkIdentDefs:
+    detectCapturedVars(n[^1], owner, c)
   else:
+    if n.isCallExpr and n[0].isTypeOf:
+      c.inTypeOf = true
     for i in 0..<n.len:
       detectCapturedVars(n[i], owner, c)
+    c.inTypeOf = false
 
 type
   LiftingPass = object
@@ -780,6 +792,8 @@ proc liftCapturedVars(n: PNode; owner: PSym; d: var DetectionPass;
   of nkTypeOfExpr:
     result = n
   else:
+    if n.isCallExpr and n[0].isTypeOf:
+      return
     if owner.isIterator:
       if nfLL in n.flags:
         # special case 'when nimVm' due to bug #3636:
diff --git a/tests/iter/t21737.nim b/tests/iter/t21737.nim
new file mode 100644
index 0000000000..da06faea71
--- /dev/null
+++ b/tests/iter/t21737.nim
@@ -0,0 +1,22 @@
+discard """
+  action: compile
+"""
+
+template mytoSeq*(iter: untyped): untyped =
+  var result: seq[typeof(iter)]# = @[]
+  for x in iter:
+    result.add(x)
+  result
+
+iterator test(dir:int): int =
+  yield 1234
+
+iterator walkGlobKinds (): int =
+  let dir2 = 123
+  let it = mytoSeq(test(dir2))  
+
+proc main()=
+    let it = iterator(): int=
+      for path in walkGlobKinds():
+          yield path
+main()
diff --git a/tests/iter/t22548.nim b/tests/iter/t22548.nim
new file mode 100644
index 0000000000..b9abb75d03
--- /dev/null
+++ b/tests/iter/t22548.nim
@@ -0,0 +1,21 @@
+discard """
+  action: compile
+"""
+
+type Xxx[T] = object
+
+iterator x(v: string): char =
+  var v2: Xxx[int]
+
+  var y: v2.T
+
+  echo y
+
+proc bbb(vv: string): proc () =
+  proc xxx() =
+    for c in x(vv):
+      echo c
+
+  return xxx
+
+bbb("test")()