Fixes nim-lang#25261 (nim-lang#25310)

Returning or yielding from a closureiter must restore "external"
exception, but `popCurrentException` from `blockLeaveActions` was
getting in the way. So now `blockLeaveActions` doesn't emit
`popCurrentException` for returns in closureiters. I'm not a fan of this
"abstraction leakage", but don't see a better solution yet. Any input is
much appreciated.

---------

Co-authored-by: Andreas Rumpf <[email protected]>

Author: yglukhov

compiler/ccgstmts.nim

@@ -226,7 +226,7 @@ proc genState(p: BProc, n: PNode) =
   elif n0.kind == nkStrLit:
     p.s(cpsStmts).addLabel(n0.strVal)
 
-proc blockLeaveActions(p: BProc, howManyTrys, howManyExcepts: int) =
+proc blockLeaveActions(p: BProc, howManyTrys, howManyExcepts: int, isReturnStmt = false) =
   # Called by return and break stmts.
   # Deals with issues faced when jumping out of try/except/finally stmts.
 
@@ -258,7 +258,7 @@ proc blockLeaveActions(p: BProc, howManyTrys, howManyExcepts: int) =
 
   # Pop exceptions that was handled by the
   # except-blocks we are in
-  if noSafePoints notin p.flags:
+  if noSafePoints notin p.flags and not (isReturnStmt and isClosureIterator(p.prc.typ)):
     for i in countdown(howManyExcepts-1, 0):
       p.s(cpsStmts).addCallStmt(cgsymValue(p.module, "popCurrentException"))
 
@@ -557,7 +557,8 @@ proc genReturnStmt(p: BProc, t: PNode) =
   if (t[0].kind != nkEmpty): genStmts(p, t[0])
   blockLeaveActions(p,
     howManyTrys    = p.nestedTryStmts.len,
-    howManyExcepts = p.inExceptBlockLen)
+    howManyExcepts = p.inExceptBlockLen,
+    isReturnStmt = true)
   if (p.finallySafePoints.len > 0) and noSafePoints notin p.flags:
     # If we're in a finally block, and we came here by exception
     # consume it before we return.

compiler/closureiters.nim

@@ -174,12 +174,14 @@ type
     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
     finallyPathLen: PNode # int literal
 
+    nullifyCurExc: PNode # Empty node, if no yields in tries
+    restoreExternExc: PNode # Empty node, id no yields in tries
+
 const
   nkSkip = {nkEmpty..nkNilLit, nkTemplateDef, nkTypeSection, nkStaticStmt,
             nkCommentStmt, nkMixinStmt, nkBindStmt, nkTypeOfExpr} + procDefs
@@ -311,7 +313,7 @@ proc hasYields(n: PNode): bool =
         break
 
 proc newNullifyCurExc(ctx: var Ctx, info: TLineInfo): PNode =
-  # :curEcx = nil
+  # :curExc = nil
   let curExc = ctx.newCurExcAccess()
   curExc.info = info
   let nilnode = newNodeIT(nkNilLit, info, getSysType(ctx.g, info, tyNil))
@@ -862,7 +864,7 @@ proc newEndFinallyNode(ctx: var Ctx, info: TLineInfo): PNode =
   retStmt.flags.incl(nfNoRewrite)
 
   let ifBody = newTree(nkIfStmt,
-                       newTree(nkElifBranch, excNilCmp, retStmt),
+                       newTree(nkElifBranch, excNilCmp, newTree(nkStmtList, ctx.newRestoreExternException(), retStmt)),
                        newTree(nkElse,
                            newTree(nkStmtList,
                                    newTreeI(nkRaiseStmt, info, ctx.g.emptyNode))))
@@ -917,16 +919,15 @@ proc transformBreakStmt(ctx: var Ctx, n: PNode): PNode =
     result = n
 
 proc transformReturnStmt(ctx: var Ctx, n: PNode): PNode =
-  # "Returning" involves jumping along all the cureent finally path.
+  # "Returning" involves jumping along all the current finally path.
   # The last finally should exit to state 0 which is a special case for last exit
   # (either return or propagating exception to the caller).
   # It is eccounted for in newEndFinallyNode.
   result = newNodeI(nkStmtList, n.info)
 
   # Returns prevent exception propagation
-  result.add(ctx.newNullifyCurExc(n.info))
+  result.add(ctx.nullifyCurExc)
 
-  result.add(ctx.newRestoreExternException())
 
   var finallyChain = newSeq[PNode]()
 
@@ -950,6 +951,7 @@ proc transformReturnStmt(ctx: var Ctx, n: PNode): PNode =
     result.add(ctx.newJumpAlongFinallyChain(finallyChain, n.info))
   else:
     # There are no (split) finallies on the path, so we can return right away
+    result.add(ctx.restoreExternExc)
     result.add(n)
 
 proc transformBreaksAndReturns(ctx: var Ctx, n: PNode): PNode =
@@ -960,7 +962,7 @@ proc transformBreaksAndReturns(ctx: var Ctx, n: PNode): PNode =
   # of nkContinueStmt: # By this point all relevant continues should be
   # lowered to breaks in transf.nim.
   of nkReturnStmt:
-    if ctx.curFinallyLevel > 0 and nfNoRewrite notin n.flags:
+    if nfNoRewrite notin n.flags:
       result = ctx.transformReturnStmt(n)
   else:
     for i in 0..<n.len:
@@ -994,8 +996,7 @@ 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)
+    result = newTree(nkStmtList, ctx.restoreExternExc, result)
 
   of nkElse, nkElseExpr:
     result[0] = addGotoOut(result[0], gotoOut)
@@ -1107,7 +1108,6 @@ proc transformClosureIteratorBody(ctx: var Ctx, n: PNode, gotoOut: PNode): PNode
       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)
@@ -1116,7 +1116,6 @@ proc transformClosureIteratorBody(ctx: var Ctx, n: PNode, gotoOut: PNode): PNode
         exceptBody = ctx.addElseToExcept(exceptBody, normalOut)
         # echo "EXCEPT: ", renderTree(exceptBody)
         exceptBody = ctx.transformClosureIteratorBody(exceptBody, tryOut)
-        inc ctx.curExceptLevel
 
       ctx.curExcLandingState = oldExcLandingState
 
@@ -1469,13 +1468,17 @@ proc transformClosureIterator*(g: ModuleGraph; idgen: IdGenerator; fn: PSym, n:
 
   ctx.curExcLandingState = ctx.newStateLabel()
   ctx.stateLoopLabel = newSym(skLabel, getIdent(ctx.g.cache, ":stateLoop"), idgen, fn, fn.info)
+
+
+  ctx.nullifyCurExc = newTree(nkStmtList)
+  ctx.restoreExternExc = newTree(nkStmtList)
+
   var n = n.toStmtList
   # echo "transformed into ", n
 
   discard ctx.newState(n, false, nil)
 
-  let finalState = ctx.newStateLabel()
-  let gotoOut = newTree(nkGotoState, finalState)
+  let gotoOut = newTree(nkGotoState, g.newIntLit(n.info, -1))
 
   var ns = false
   n = ctx.lowerStmtListExprs(n, ns)
@@ -1487,11 +1490,9 @@ 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)
+    ctx.nullifyCurExc.add(ctx.newNullifyCurExc(fn.info))
+    ctx.restoreExternExc.add(ctx.newRestoreExternException())
 
   # Assign state label indexes
   for i in 0 .. ctx.states.high:
@@ -1510,7 +1511,9 @@ proc transformClosureIterator*(g: ModuleGraph; idgen: IdGenerator; fn: PSym, n:
     let body = ctx.transformStateAssignments(s.body)
     caseDispatcher.add newTreeI(nkOfBranch, body.info, s.label, body)
 
-  caseDispatcher.add newTreeI(nkElse, n.info, newTreeI(nkReturnStmt, n.info, g.emptyNode))
+  caseDispatcher.add newTreeI(nkElse, n.info,
+                              newTree(nkStmtList, ctx.restoreExternExc,
+                                      newTreeI(nkReturnStmt, n.info, g.emptyNode)))
 
   result = wrapIntoStateLoop(ctx, caseDispatcher)
   result = liftLocals(ctx, result)

tests/iter/tyieldintry.nim

@@ -798,3 +798,52 @@ block: #25202
     doAssert(checkpoints1 == checkpoints2)
 
   p()
+
+block: #25261
+  iterator y(): int {.closure.} =
+    try:
+      try:
+        raise newException(CatchableError, "Error")
+      except CatchableError:
+        return 123
+      yield 0
+    finally:
+      discard
+
+  let w = y
+  doAssert(w() == 123)
+  doAssert(getCurrentExceptionMsg() == "")
+
+  try:
+    raise newException(ValueError, "Outer error")
+  except:
+    doAssert(getCurrentExceptionMsg() == "Outer error")
+    let w = y
+    doAssert(w() == 123)
+    doAssert(getCurrentExceptionMsg() == "Outer error")
+  doAssert(getCurrentExceptionMsg() == "")
+
+block:
+  # Looks almost like above, but last finally changed to except
+  iterator y(): int {.closure.} =
+    try:
+      try:
+        raise newException(CatchableError, "Error")
+      except CatchableError:
+        return 123
+      yield 0
+    except:
+      discard
+
+  let w = y
+  doAssert(w() == 123)
+  doAssert(getCurrentExceptionMsg() == "")
+
+  try:
+    raise newException(ValueError, "Outer error")
+  except:
+    doAssert(getCurrentExceptionMsg() == "Outer error")
+    let w = y
+    doAssert(w() == 123)
+    doAssert(getCurrentExceptionMsg() == "Outer error")
+  doAssert(getCurrentExceptionMsg() == "")