Fixes #25202

Author: yglukhov

compiler/closureiters.nim

@@ -166,13 +166,16 @@ type
     tmpResultSym: PSym # Used when we return, but finally has to interfere
     finallyPathSym: PSym
     curExcLevelSym: PSym # Current exception level (because exceptions are stacked)
+    curExcSym: PSym # Current exception
+    externExcSym: PSym # Extern exception: what would getCurrentException() return outside of closure iter
 
     states: seq[State] # The resulting states. Label is int literal.
     finallyPathStack: seq[FinallyTarget] # Stack of split blocks, whiles and finallies
     stateLoopLabel: PSym # Label to break on, when jumping between states.
     tempVarId: int # unique name counter
     hasExceptions: bool # Does closure have yield in try?
     curExcLandingState: PNode
+    curExceptLevel: int
     curFinallyLevel: int
     idgen: IdGenerator
     varStates: Table[ItemId, int] # Used to detect if local variable belongs to multiple states
@@ -242,6 +245,12 @@ proc newFinallyPathAssign(ctx: var Ctx, level: int, label: PNode, info: TLineInf
   let fp = newFinallyPathAccess(ctx, level, info)
   result = newTree(nkAsgn, fp, label)
 
+proc newCurExcAccess(ctx: var Ctx): PNode =
+  if ctx.curExcSym.isNil:
+    let getCurExc = ctx.g.callCodegenProc("getCurrentException")
+    ctx.curExcSym = ctx.newEnvVar(":curExc", getCurExc.typ)
+  ctx.newEnvVarAccess(ctx.curExcSym)
+
 proc newCurExcLevelAccess(ctx: var Ctx): PNode =
   if ctx.curExcLevelSym.isNil:
     ctx.curExcLevelSym = ctx.newEnvVar(":curExcLevel", ctx.g.getSysType(ctx.fn.info, tyInt16))
@@ -284,6 +293,15 @@ proc newTempVar(ctx: var Ctx, typ: PType, parent: PNode, initialValue: PNode = n
   assert(not typ.isNil, "Temp var needs a type")
   parent.add(ctx.newTempVarDef(result, initialValue))
 
+proc newExternExcAccess(ctx: var Ctx): PNode =
+  if ctx.externExcSym == nil:
+    ctx.externExcSym = newSym(skVar, getIdent(ctx.g.cache, ":externExc"), ctx.idgen, ctx.fn, ctx.fn.info)
+    ctx.externExcSym.typ = ctx.curExcSym.typ
+  newSymNode(ctx.externExcSym, ctx.fn.info)
+
+proc newRestoreExternException(ctx: var Ctx): PNode =
+  ctx.g.callCodegenProc("closureIterSetExc", ctx.fn.info, ctx.newExternExcAccess())
+
 proc hasYields(n: PNode): bool =
   # TODO: This is very inefficient. It traverses the node, looking for nkYieldStmt.
   case n.kind
@@ -298,8 +316,15 @@ proc hasYields(n: PNode): bool =
         result = true
         break
 
-proc newNullifyCurExcLevel(ctx: var Ctx, info: TLineInfo, decrement = false): PNode =
-  # :curEcx = 0
+proc newNullifyCurExc(ctx: var Ctx, info: TLineInfo): PNode =
+  # :curEcx = nil
+  let curExc = ctx.newCurExcAccess()
+  curExc.info = info
+  let nilnode = newNodeIT(nkNilLit, info, getSysType(ctx.g, info, tyNil))
+  result = newTree(nkAsgn, curExc, nilnode)
+
+proc newNullifyCurExcLevel(ctx: var Ctx, info: TLineInfo): PNode =
+  # :curEcxLevel = 0
   let curExc = ctx.newCurExcLevelAccess()
   curExc.info = info
   let nilnode = ctx.g.newIntLit(info, 0)
@@ -344,17 +369,20 @@ proc collectExceptState(ctx: var Ctx, n: PNode): PNode {.inline.} =
         else:
           ifBranch = newNodeI(nkElse, c.info)
 
-      ifBranch.add(newTreeI(nkStmtList, c.info, ctx.newChangeCurExcLevel(c.info, -1), c[^1]))
+      ifBranch.add(c[^1])
+      # ifBranch.add(newTreeI(nkStmtList, c.info, ctx.newChangeCurExcLevel(c.info, -1), c[^1]))
       ifStmt.add(ifBranch)
 
   if ifStmt.len != 0:
     result = newTree(nkStmtList, ifStmt)
+    # result = newTree(nkStmtList, ctx.newChangeCurExcLevel(n.info, -1), ifStmt)
   else:
     result = ctx.g.emptyNode
 
-proc addElseToExcept(ctx: var Ctx, n, gotoOut: PNode) =
+proc addElseToExcept(ctx: var Ctx, n, gotoOut: PNode): PNode =
   # We should adjust finallyPath to gotoOut if exception is handled
   # if there is no finally node next to this except, gotoOut must be nil
+  result = n
   if n.kind == nkStmtList:
     if n[0].kind == nkIfStmt and n[0][^1].kind != nkElse:
       # Not all cases are covered, which means exception is not handled
@@ -377,10 +405,12 @@ proc addElseToExcept(ctx: var Ctx, n, gotoOut: PNode) =
     # raised one.
     n.add newTree(nkCall,
       newSymNode(ctx.g.getCompilerProc("popCurrentException")))
+    n.add ctx.newNullifyCurExc(n.info)
     if gotoOut != nil:
       # We have a finally node following this except block, and exception is handled
       # Configure its path to continue normally
       n.add(ctx.newFinallyPathAssign(ctx.curFinallyLevel - 1, gotoOut[0], n.info))
+    
 
 proc getFinallyNode(ctx: var Ctx, n: PNode): PNode =
   result = n[^1]
@@ -837,9 +867,9 @@ proc newEndFinallyNode(ctx: var Ctx, info: TLineInfo): PNode =
 
   let excNilCmp = newTreeIT(nkCall,
                       info, ctx.g.getSysType(info, tyBool),
-                      newSymNode(ctx.g.getSysMagic(info, "==", mEqI), info),
-                      ctx.newCurExcLevelAccess(),
-                      ctx.g.newIntLit(info, 0))
+                      newSymNode(ctx.g.getSysMagic(info, "==", mEqRef), info),
+                      ctx.newCurExcAccess(),
+                      newNodeIT(nkNilLit, info, getSysType(ctx.g, info, tyNil)))
 
   let retStmt =
     block:
@@ -919,6 +949,7 @@ proc transformReturnStmt(ctx: var Ctx, n: PNode): PNode =
 
   # Returns prevent exception propagation
   result.add(ctx.newNullifyCurExcLevel(n.info))
+  result.add(ctx.newNullifyCurExc(n.info))
 
   var finallyChain = newSeq[PNode]()
 
@@ -986,6 +1017,8 @@ proc transformClosureIteratorBody(ctx: var Ctx, n: PNode, gotoOut: PNode): PNode
 
   of nkYieldStmt:
     result = addGotoOut(result, gotoOut)
+    if ctx.curExceptLevel > 0 or ctx.curFinallyLevel > 0:
+      result = newTree(nkStmtList, ctx.newRestoreExternException(), result)
 
   of nkElse, nkElseExpr:
     result[0] = addGotoOut(result[0], gotoOut)
@@ -1055,7 +1088,7 @@ proc transformClosureIteratorBody(ctx: var Ctx, n: PNode, gotoOut: PNode): PNode
     result.add(tryLabel)
     var tryBody = toStmtList(n[0])
 
-    let exceptBody = ctx.collectExceptState(n)
+    var exceptBody = ctx.collectExceptState(n)
     var finallyBody = ctx.getFinallyNode(n)
     var exceptLabel, finallyLabel = ctx.g.emptyNode
 
@@ -1094,28 +1127,27 @@ proc transformClosureIteratorBody(ctx: var Ctx, n: PNode, gotoOut: PNode): PNode
         inc ctx.curFinallyLevel
         ctx.finallyPathStack.add(FinallyTarget(n: n[^1], label: finallyLabel))
 
-      if ctx.transformClosureIteratorBody(tryBody, tryOut) != tryBody:
-        internalError(ctx.g.config, "transformClosureIteratorBody != tryBody")
+      tryBody = ctx.transformClosureIteratorBody(tryBody, tryOut)
 
       if exceptBody.kind != nkEmpty:
+        inc ctx.curExceptLevel
         ctx.curExcLandingState = if finallyBody.kind != nkEmpty: finallyLabel
                                  else: oldExcLandingState
         discard ctx.newState(exceptBody, false, exceptLabel)
 
         let normalOut = if finallyBody.kind != nkEmpty: gotoOut else: nil
-        ctx.addElseToExcept(exceptBody, normalOut)
+        exceptBody = ctx.addElseToExcept(exceptBody, normalOut)
         # echo "EXCEPT: ", renderTree(exceptBody)
-        if ctx.transformClosureIteratorBody(exceptBody, tryOut) != exceptBody:
-          internalError(ctx.g.config, "transformClosureIteratorBody != exceptBody")
+        exceptBody = ctx.transformClosureIteratorBody(exceptBody, tryOut)
+        inc ctx.curExceptLevel
 
       ctx.curExcLandingState = oldExcLandingState
 
       if finallyBody.kind != nkEmpty:
         discard ctx.finallyPathStack.pop()
         discard ctx.newState(finallyBody, false, finallyLabel)
         let finallyExit = newTree(nkGotoState, ctx.newFinallyPathAccess(ctx.curFinallyLevel - 1, finallyBody.info))
-        if ctx.transformClosureIteratorBody(finallyBody, finallyExit) != finallyBody:
-          internalError(ctx.g.config, "transformClosureIteratorBody != finallyBody")
+        finallyBody = ctx.transformClosureIteratorBody(finallyBody, finallyExit)
         dec ctx.curFinallyLevel
 
   of nkGotoState, nkForStmt:
@@ -1209,25 +1241,16 @@ proc newExceptBody(ctx: var Ctx, info: TLineInfo): PNode {.inline.} =
   result = newNodeI(nkStmtList, info)
 
   let intTyp = ctx.g.getSysType(info, tyInt)
-  let boolTyp = ctx.g.getSysType(info, tyBool)
 
   # :state = exceptionTable[:state]
   result.add ctx.newStateAssgn(
     newTreeIT(nkBracketExpr, info, intTyp,
               ctx.createExceptionTable(),
               ctx.newStateAccess()))
 
-  # if :state == 0: raise
-  block:
-    let cond = newTreeIT(nkCall, info, boolTyp,
-      ctx.g.getSysMagic(info, "==", mEqI).newSymNode(),
-      ctx.newStateAccess(),
-      newIntTypeNode(0, intTyp))
+  let getCurExc = ctx.g.callCodegenProc("getCurrentException")
+  result.add newTree(nkFastAsgn, ctx.newCurExcAccess(), getCurExc)
 
-    let raiseStmt = newTree(nkRaiseStmt, ctx.g.emptyNode)
-    let ifBranch = newTree(nkElifBranch, cond, raiseStmt)
-    let ifStmt = newTree(nkIfStmt, ifBranch)
-    result.add(ifStmt)
 
 proc wrapIntoTryExcept(ctx: var Ctx, n: PNode): PNode {.inline.} =
   # Generates code:
@@ -1236,11 +1259,12 @@ proc wrapIntoTryExcept(ctx: var Ctx, n: PNode): PNode {.inline.} =
   #   body
   # except:
   #   :state = exceptionTable[:state]
-  #   if :state == 0:
-  #     raise
-  #   :tmp = getCurrentException()
+  #   :curExc = getCurrentException()
+  # if :state == 0:
+  #   raise
   #
-  # pushCurrentException(:tmp)
+  # pushCurrentException(:curExc)
+  # inc :curExcLevel
 
   let tryBody = newTree(nkStmtList, n)
   let exceptBody = ctx.newExceptBody(ctx.fn.info)
@@ -1251,8 +1275,25 @@ proc wrapIntoTryExcept(ctx: var Ctx, n: PNode): PNode {.inline.} =
   let tempExc = ctx.newTempVar(getCurExc.typ, result)
   result.add newTree(nkTryStmt, tryBody, exceptBranch)
   exceptBody.add ctx.newTempVarAsgn(tempExc, getCurExc)
+  # exceptBody.add newTree(nkFastAsgn, ctx.newCurExcAccess(), getCurExc)
+
+  # if :state == 0: raise
+  block:
+    let boolTyp = ctx.g.getSysType(ctx.fn.info, tyBool)
+    let intTyp = ctx.g.getSysType(ctx.fn.info, tyInt)
+    let cond = newTreeIT(nkCall, ctx.fn.info, boolTyp,
+      ctx.g.getSysMagic(ctx.fn.info, "==", mEqI).newSymNode(),
+      ctx.newStateAccess(),
+      newIntTypeNode(0, intTyp))
+
+    let raiseStmt = newTree(nkRaiseStmt, ctx.newCurExcAccess())
+    let ifBody = newTree(nkStmtList, ctx.newRestoreExternException(), raiseStmt)
+    let ifBranch = newTree(nkElifBranch, cond, ifBody)
+    let ifStmt = newTree(nkIfStmt, ifBranch)
+    result.add(ifStmt)
 
-  result.add newTree(nkCall, newSymNode(ctx.g.getCompilerProc("pushCurrentException")), ctx.newTempVarAccess(tempExc))
+  # result.add newTree(nkCall, newSymNode(ctx.g.getCompilerProc("pushCurrentException")), ctx.newTempVarAccess(tempExc))
+  result.add newTree(nkCall, newSymNode(ctx.g.getCompilerProc("pushCurrentException")), ctx.newCurExcAccess())
   result.add ctx.newChangeCurExcLevel(n.info, 1)
 
 proc wrapIntoStateLoop(ctx: var Ctx, n: PNode): PNode =
@@ -1276,6 +1317,17 @@ proc wrapIntoStateLoop(ctx: var Ctx, n: PNode): PNode =
   blockStmt.add(blockBody)
   loopBody.add(blockStmt)
 
+  if ctx.hasExceptions:
+    let res = newNodeI(nkStmtList, n.info)
+    let getCurExc = ctx.g.callCodegenProc("getCurrentException")
+    discard ctx.newExternExcAccess()
+    res.add(ctx.newTempVarDef(ctx.externExcSym, getCurExc))
+    # let externExc = ctx.newTempVar(getCurExc.typ, res, getCurExc)
+    let setCurExc = ctx.g.callCodegenProc("closureIterSetExc", n.info, ctx.newCurExcAccess())
+    res.add(setCurExc)
+    res.add(result)
+    result = res
+
 proc countStateOccurences(ctx: var Ctx, n: PNode, stateOccurences: var openArray[int]) =
   ## Find all nkGotoState(stateIdx) nodes that do not follow nkYield.
   ## For every such node increment stateOccurences[stateIdx]
@@ -1381,7 +1433,7 @@ proc detectCapturedVars(c: var Ctx, n: PNode, stateIdx: int) =
   case n.kind
   of nkSym:
     let s = n.sym
-    if s.kind in {skResult, skVar, skLet, skForVar, skTemp} and sfGlobal notin s.flags and s.owner == c.fn:
+    if s.kind in {skResult, skVar, skLet, skForVar, skTemp} and sfGlobal notin s.flags and s.owner == c.fn and s != c.externExcSym:
       let vs = c.varStates.getOrDefault(s.itemId, localNotSeen)
       if vs == localNotSeen: # First seing this variable
         c.varStates[s.itemId] = stateIdx
@@ -1458,7 +1510,9 @@ proc transformClosureIterator*(g: ModuleGraph; idgen: IdGenerator; fn: PSym, n:
   # echo "transformed into ", n
 
   discard ctx.newState(n, false, nil)
-  let gotoOut = newTree(nkGotoState, g.newIntLit(n.info, -1))
+
+  let finalState = ctx.newStateLabel()
+  let gotoOut = newTree(nkGotoState, finalState)
 
   var ns = false
   n = ctx.lowerStmtListExprs(n, ns)
@@ -1470,6 +1524,12 @@ proc transformClosureIterator*(g: ModuleGraph; idgen: IdGenerator; fn: PSym, n:
   # Splitting transformation
   discard ctx.transformClosureIteratorBody(n, gotoOut)
 
+  let finalStateBody = newTree(nkStmtList)
+  if ctx.hasExceptions:
+    finalStateBody.add(ctx.newRestoreExternException())
+  finalStateBody.add(newTree(nkGotoState, g.newIntLit(n.info, -1)))
+  discard ctx.newState(finalStateBody, true, finalState)
+
   # Assign state label indexes
   for i in 0 .. ctx.states.high:
     ctx.states[i].label.intVal = i

lib/system.nim

@@ -2329,6 +2329,10 @@ when notJSnotNims and hostOS != "standalone":
     ## Retrieves the current exception; if there is none, `nil` is returned.
     result = currException
 
+  proc setCurrentExceptionForClosureIter*(e: ref Exception) {.compilerRtl, inl, benign.} =
+    ## Retrieves the current exception; if there is none, `nil` is returned.
+    currException = e
+
   proc nimBorrowCurrentException(): ref Exception {.compilerRtl, inl, benign, nodestroy.} =
     # .nodestroy here so that we do not produce a write barrier as the
     # C codegen only uses it in a borrowed way:

lib/system/embedded.nim

@@ -24,6 +24,7 @@ when not gotoBasedExceptions:
   proc popSafePoint {.compilerRtl, inl.} = discard
 proc pushCurrentException(e: ref Exception) {.compilerRtl, inl.} = discard
 proc popCurrentException {.compilerRtl, inl.} = discard
+proc closureIterSetExc(e: ref Exception) {.compilerRtl, inl.} = discard
 
 # some platforms have native support for stack traces:
 const

lib/system/excpt.nim

@@ -159,6 +159,9 @@ proc popCurrentException {.compilerRtl, inl.} =
   currException = currException.up
   #showErrorMessage2 "B"
 
+proc closureIterSetExc(e: ref Exception) {.compilerRtl, inl.} =
+  currException = e
+
 proc popCurrentExceptionEx(id: uint) {.compilerRtl.} =
   discard "only for bootstrapping compatbility"
 

tests/iter/tyieldintry.nim

@@ -752,3 +752,49 @@ block: #25038
       0
 
   test(d)
+
+block: #25202
+  proc p() =
+    iterator p_1073741828(checkpoints: var seq[int]): int {.
+        closure, raises: [].} =
+      var closureSucceeded_1073741827 = true
+      try:
+        try:
+          try:
+            yield 0
+            raise newException(ValueError, "value error")
+          except ValueError:
+            checkpoints.add(1)
+            raise newException(IOError, "io error")
+        finally:
+          yield 2
+      except IOError as exc:
+        closureSucceeded_1073741827 = false
+        checkpoints.add(3)
+      finally:
+        checkpoints.add(4)
+        if closureSucceeded_1073741827:
+          discard
+
+    var internalClosure = p_1073741828
+    var internalClosure2 = p_1073741828
+
+    var checkpoints1 = newSeq[int]()
+    var checkpoints2 = newSeq[int]()
+
+    while true:
+      if not internalClosure.finished():
+        checkpoints1.add internalClosure(checkpoints1)
+        doAssert(getCurrentException() == nil)
+      if not internalClosure2.finished():
+        checkpoints2.add internalClosure2(checkpoints2)
+        doAssert(getCurrentException() == nil)
+      if internalClosure.finished() and internalClosure2.finished():
+        break
+
+    if checkpoints1[^1] == 0: checkpoints1.del(checkpoints1.high)
+    if checkpoints2[^1] == 0: checkpoints2.del(checkpoints2.high)
+    doAssert(checkpoints1 == @[0, 1, 2, 3, 4])
+    doAssert(checkpoints1 == checkpoints2)
+
+  p()