From 52cf7dfde01ca70e867607c2df7de66263c51025 Mon Sep 17 00:00:00 2001
From: metagn <metagngn@gmail.com>
Date: Fri, 18 Oct 2024 08:37:05 +0300
Subject: [PATCH] shallow fold prevention for `addr`, `nkHiddenAddr` (#24322)

fixes #24305, refs #23807

Since #23014 `nkHiddenAddr` is produced to fast assign array elements in
iterators. However the array access inside this `nkHiddenAddr` can get
folded at compile time, generating invalid code. In #23807, compile time
folding of regular `addr` expressions was changed to be prevented in
`transf` but `nkHiddenAddr` was not updated alongside it.

The method for preventing folding in `addr` in #23807 was also faulty,
it should only trigger on the immediate child node of the address rather
than all nodes nested inside it. This caused a regression as outlined in
[this
comment](https://github.com/nim-lang/Nim/pull/24322#issuecomment-2419560182).

To fix both issues, `addr` and `nkHiddenAddr` now both shallowly prevent
constant folding for their immediate children.
---
 compiler/transf.nim        | 22 +++++++++-------------
 tests/iter/tfoldedaddr.nim | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 41 insertions(+), 13 deletions(-)
 create mode 100644 tests/iter/tfoldedaddr.nim

diff --git a/compiler/transf.nim b/compiler/transf.nim
index cfeff18cac..7ee9683743 100644
--- a/compiler/transf.nim
+++ b/compiler/transf.nim
@@ -56,7 +56,6 @@ type
     contSyms, breakSyms: seq[PSym]  # to transform 'continue' and 'break'
     deferDetected, tooEarly: bool
     isIntroducingNewLocalVars: bool  # true if we are in `introducingNewLocalVars` (don't transform yields)
-    inAddr: bool
     flags: TransformFlags
     graph: ModuleGraph
     idgen: IdGenerator
@@ -103,12 +102,12 @@ proc newTemp(c: PTransf, typ: PType, info: TLineInfo): PNode =
   else:
     result = newSymNode(r)
 
-proc transform(c: PTransf, n: PNode): PNode
+proc transform(c: PTransf, n: PNode, noConstFold = false): PNode
 
-proc transformSons(c: PTransf, n: PNode): PNode =
+proc transformSons(c: PTransf, n: PNode, noConstFold = false): PNode =
   result = newTransNode(n)
   for i in 0..<n.len:
-    result[i] = transform(c, n[i])
+    result[i] = transform(c, n[i], noConstFold)
 
 proc newAsgnStmt(c: PTransf, kind: TNodeKind, le: PNode, ri: PNode; isFirstWrite: bool): PNode =
   result = newTransNode(kind, ri.info, 2)
@@ -481,8 +480,8 @@ proc transformYield(c: PTransf, n: PNode): PNode =
     result.add(introduceNewLocalVars(c, c.transCon.forLoopBody))
     c.isIntroducingNewLocalVars = false
 
-proc transformAddrDeref(c: PTransf, n: PNode, kinds: TNodeKinds): PNode =
-  result = transformSons(c, n)
+proc transformAddrDeref(c: PTransf, n: PNode, kinds: TNodeKinds, isAddr = false): PNode =
+  result = transformSons(c, n, noConstFold = isAddr)
   # inlining of 'var openarray' iterators; bug #19977
   if n.typ.kind != tyOpenArray and (c.graph.config.backend == backendCpp or sfCompileToCpp in c.module.flags): return
   var n = result
@@ -1010,7 +1009,7 @@ proc transformDerefBlock(c: PTransf, n: PNode): PNode =
     result[i] = e0[i]
   result[e0.len-1] = newTreeIT(nkHiddenDeref, n.info, n.typ, e0[e0.len-1])
 
-proc transform(c: PTransf, n: PNode): PNode =
+proc transform(c: PTransf, n: PNode, noConstFold = false): PNode =
   when false:
     var oldDeferAnchor: PNode
     if n.kind in {nkElifBranch, nkOfBranch, nkExceptBranch, nkElifExpr,
@@ -1077,12 +1076,9 @@ proc transform(c: PTransf, n: PNode): PNode =
   of nkCallKinds:
     result = transformCall(c, n)
   of nkHiddenAddr:
-    result = transformAddrDeref(c, n, {nkHiddenDeref})
+    result = transformAddrDeref(c, n, {nkHiddenDeref}, isAddr = true)
   of nkAddr:
-    let oldInAddr = c.inAddr
-    c.inAddr = true
-    result = transformAddrDeref(c, n, {nkDerefExpr, nkHiddenDeref})
-    c.inAddr = oldInAddr
+    result = transformAddrDeref(c, n, {nkDerefExpr, nkHiddenDeref}, isAddr = true)
   of nkDerefExpr:
     result = transformAddrDeref(c, n, {nkAddr, nkHiddenAddr})
   of nkHiddenDeref:
@@ -1162,7 +1158,7 @@ proc transform(c: PTransf, n: PNode): PNode =
   let exprIsPointerCast = n.kind in {nkCast, nkConv, nkHiddenStdConv} and
                           n.typ != nil and
                           n.typ.kind == tyPointer
-  if not exprIsPointerCast and not c.inAddr:
+  if not exprIsPointerCast and not noConstFold:
     var cnst = getConstExpr(c.module, result, c.idgen, c.graph)
     # we inline constants if they are not complex constants:
     if cnst != nil and not dontInlineConstant(n, cnst):
diff --git a/tests/iter/tfoldedaddr.nim b/tests/iter/tfoldedaddr.nim
new file mode 100644
index 0000000000..cbf4ea5fa7
--- /dev/null
+++ b/tests/iter/tfoldedaddr.nim
@@ -0,0 +1,32 @@
+discard """
+  output: '''
+23
+23
+23
+23
+23
+23
+'''
+"""
+
+block: # issue #24305
+  iterator demo(a: openArray[int]): int =
+    for k in countUp(a[0], 19):
+      yield 23
+
+  for k in demo(@[17]):
+    echo k
+
+block: # issue #24305 with array
+  iterator demo(a: array[1, int]): int =
+    for k in countUp(a[0], 19):
+      yield 23
+
+  for k in demo([17]):
+    echo k
+
+block: # related regression
+  proc main =
+    let a = [0, 1, 2]
+    let x = addr a[low(a)]
+  main()