Run enclosing finally when break/continue jumps out of a try block

docs/LowerToAffineLoops-LowerToLLVM-review.md

@@ -237,14 +237,37 @@ Fixed on branch `fix/lowering-passes-bugs-3`, verified via full ctest suite
    times to catch wrong-copy-used-once-vs-every-iteration bugs), JIT +
    AOT green.
 
-## Remaining known issues (not yet fixed, tracked here for follow-up)
+## Fourth pass (2026-07-14) — item 5b fixed
+
+5b. **`LowerToAffineLoops.cpp`** (`TryOpLowering` / loop `jumps[]` targets) —
+    FIXED (branch `fix/break-continue-in-try-finally`), using the same
+    mechanism `return` in try/finally already used (suggested by the repo
+    owner). `TryOpLowering` now collects `BreakOp`/`ContinueOp` in the try
+    body and catches that already have a `tsContext->jumps` entry at
+    try-lowering time — a clean discriminator for "escapes this try": the
+    target loop/label encloses the try, so its lowering (which populates
+    `jumps`) ran first in the pre-order worklist, whereas jumps targeting a
+    loop *inside* the try aren't resolved yet. For each distinct escape
+    target it clones the `finally` region once (mirroring the return-cleanup
+    clone), terminates the clone with a `cf::BranchOp` to the original
+    target, and re-points `jumps[op]` at the clone — so Break/Continue
+    lowering needs no structural change and just branches into the finally
+    copy. Nested try/finally chains compose naturally: the outer try lowers
+    first and re-points the op at its finally copy; a nested try later
+    re-points the same op at *its* copy, which branches to the outer copy —
+    innermost-first execution. Escaping jumps in catch clauses additionally
+    get `tsContext->unwind[op] = catchesBlock`, and Break/Continue lowering
+    now emits `EndCatchOp` when that's set, mirroring return-from-catch.
+
+    Verified by JIT execution: break in try body, continue in try body,
+    two-level nested finally chain (both finallys run, innermost first),
+    break from a catch clause (catch + finally + EndCatch), one try with two
+    distinct escape targets (break + continue → two separate finally
+    copies), and the negative case (break of a loop fully inside the try is
+    not rerouted — finally runs exactly once). All added to
+    `00try_finally_break_continue.ts`.
 
-5b. **`LowerToAffineLoops.cpp`** (loop `jumps[]` target computation) —
-    `break`/`continue` inside a `try` body does not run an enclosing
-    `finally` block before jumping to the loop's continuation/increment.
-    See "Third pass" above for the repro and root cause. Needs the
-    break/continue lowering (or the `TryOp` lowering) to detect that the
-    jump target crosses a `finally` boundary and route through it first.
+## Remaining known issues (not yet fixed, tracked here for follow-up)
 
 8. **`LowerToLLVM.cpp:4159`** (`LandingPadOpLowering`, Windows path) —
    cleanup-only branch reuses the typed-catch filter value instead of an

tslang/lib/TypeScript/LowerToAffineLoops.cpp

@@ -794,6 +794,11 @@ struct BreakOpLowering : public TsPattern<mlir_ts::BreakOp>
         auto jump = tsContext->jumps[breakOp];
         assert(jump);
 
+        if (auto unwind = tsContext->unwind[breakOp])
+        {
+            rewriter.create<mlir_ts::EndCatchOp>(loc);
+        }
+
         rewriter.replaceOpWithNewOp<mlir::cf::BranchOp>(breakOp, jump);
         clh.CutBlock();
 
@@ -815,6 +820,11 @@ struct ContinueOpLowering : public TsPattern<mlir_ts::ContinueOp>
         auto jump = tsContext->jumps[continueOp];
         assert(jump);
 
+        if (auto unwind = tsContext->unwind[continueOp])
+        {
+            rewriter.create<mlir_ts::EndCatchOp>(loc);
+        }
+
         rewriter.replaceOpWithNewOp<mlir::cf::BranchOp>(continueOp, jump);
         clh.CutBlock();
 
@@ -1287,6 +1297,28 @@ struct TryOpLowering : public TsPattern<mlir_ts::TryOp>
         tryOp.getBody().walk(visitorReturnOps);
         tryOp.getCatches().walk(visitorReturnOps);
 
+        // break/continue ops in the body/catches that jump OUT of this try already have their
+        // tsContext->jumps entry resolved (the target loop/label encloses the try, so its
+        // lowering ran first); ones targeting a construct inside the try don't have one yet.
+        // Escaping jumps must run `finally` on the way out, the same way `return` does --
+        // collected here (before the regions are inlined/emptied), rerouted through their own
+        // finally copy below. Ops directly inside the finally region itself are deliberately
+        // not collected: a break/continue there aborts the rest of the finally and jumps
+        // straight out (and its wiring in the cloned copies is handled by the side-table
+        // propagation below).
+        mlir::SmallVector<Operation *> escapingJumpsBody;
+        mlir::SmallVector<Operation *> escapingJumpsCatches;
+        auto collectEscapingJumps = [&](mlir::Region &region, mlir::SmallVector<Operation *> &to) {
+            region.walk([&](Operation *op) {
+                if (isa<mlir_ts::BreakOp, mlir_ts::ContinueOp>(op) && tsContext->jumps.count(op))
+                {
+                    to.push_back(op);
+                }
+            });
+        };
+        collectEscapingJumps(tryOp.getBody(), escapingJumpsBody);
+        collectEscapingJumps(tryOp.getCatches(), escapingJumpsCatches);
+
         // inline structure
         OpBuilder::InsertionGuard guard(rewriter);
         mlir::Block *currentBlock = rewriter.getInsertionBlock();
@@ -1421,6 +1453,49 @@ struct TryOpLowering : public TsPattern<mlir_ts::TryOp>
                 }
             }
 
+            // add clones for escaping break/continue, the same way as for 'return' above:
+            // one finally copy per distinct jump target, terminated by a branch to that
+            // target, and the escaping op is re-pointed at the copy -- so Break/Continue
+            // lowering just branches into the finally copy, which then continues to the
+            // real target. Nested try/finally chains compose naturally: this try lowers
+            // before any nested one, so a nested try later re-points the same op at its
+            // own finally copy, which branches here.
+            if (escapingJumpsBody.size() > 0 || escapingJumpsCatches.size() > 0)
+            {
+                mlir::DenseMap<mlir::Block *, mlir::Block *> finallyCopyPerTarget;
+                auto routeThroughFinally = [&](Operation *op) {
+                    auto target = tsContext->jumps[op];
+                    auto &finallyCopy = finallyCopyPerTarget[target];
+                    if (!finallyCopy)
+                    {
+                        mlir::IRMapping jumpFinallyMapping;
+                        rewriter.cloneRegionBefore(tryOp.getFinally(), *continuation->getParent(),
+                                                    continuation->getIterator(), jumpFinallyMapping);
+                        propagateTsContextEntries(jumpFinallyMapping);
+                        auto jumpFinallyBlockLast = continuation->getPrevNode();
+                        rewriter.setInsertionPoint(jumpFinallyBlockLast->getTerminator());
+                        auto resultOpOfJumpFinallyBlock = cast<mlir_ts::ResultOp>(jumpFinallyBlockLast->getTerminator());
+                        rewriter.replaceOpWithNewOp<mlir::cf::BranchOp>(resultOpOfJumpFinallyBlock, target);
+                        finallyCopy = jumpFinallyMapping.lookup(&tryOp.getFinally().front());
+                    }
+
+                    tsContext->jumps[op] = finallyCopy;
+                };
+
+                for (auto op : escapingJumpsBody)
+                {
+                    routeThroughFinally(op);
+                }
+
+                for (auto op : escapingJumpsCatches)
+                {
+                    routeThroughFinally(op);
+                    // leaving a catch clause abruptly must end the active catch first,
+                    // mirroring the return-from-catch handling
+                    tsContext->unwind[op] = catchesBlock;
+                }
+            }
+
             rewriter.inlineRegionBefore(tryOp.getFinally(), continuation);
             exitBlockLast = continuation->getPrevNode();
         }

tslang/test/tester/tests/00try_finally_break_continue.ts

@@ -1,11 +1,10 @@
-// regression test for TryOpLowering's finally-region cloning: break/continue and a
-// nested try inside `finally` must keep correct wiring in all 3 copies of the region
-// (the two cloned exception-path copies, and the inlined normal-exit copy).
-//
-// Note: break/continue inside a try *body* (rather than directly in `finally`) is a
-// separate, pre-existing bug tracked in docs/LowerToAffineLoops-LowerToLLVM-review.md --
-// intervening finally blocks are not run before the jump, so those cases are not covered
-// here.
+// regression tests for break/continue interacting with try/finally:
+// - break/continue directly inside `finally` must keep correct wiring in all copies of
+//   the cloned finally region (TryOpLowering clones it for the exception/return paths);
+// - break/continue inside the try *body* or a catch clause that jumps out of the try must
+//   run the enclosing `finally` on the way out, the same way `return` does -- routed
+//   through a dedicated finally copy per escape target;
+// - nested try/finally chains must run every finally on the way out, innermost first.
 
 function test_break_in_finally() {
     let iterations = 0;
@@ -75,10 +74,156 @@ function test_nested_try_in_finally() {
     return caught;
 }
 
+function test_break_in_try_body() {
+    let iterations = 0;
+    let finallyRuns = 0;
+
+    for (let i = 0; i < 5; i++) {
+        iterations++;
+        try {
+            if (i == 2) {
+                break;
+            }
+        } finally {
+            finallyRuns++;
+        }
+    }
+
+    assert(iterations == 3, "break-in-body: wrong iteration count");
+    assert(finallyRuns == 3, "break-in-body: finally must run on the break iteration too");
+    return 1;
+}
+
+function test_continue_in_try_body() {
+    let sum = 0;
+    let finallyRuns = 0;
+
+    for (let i = 0; i < 5; i++) {
+        try {
+            if (i == 2) {
+                continue;
+            }
+            sum += i;
+        } finally {
+            finallyRuns++;
+        }
+    }
+
+    assert(finallyRuns == 5, "continue-in-body: finally must run on the continue iteration too");
+
+    // sum skips i == 2: 0 + 1 + 3 + 4 = 8
+    return sum;
+}
+
+function test_nested_finally_chain() {
+    let f1 = 0;
+    let f2 = 0;
+    let iterations = 0;
+
+    for (let i = 0; i < 5; i++) {
+        iterations++;
+        try {
+            try {
+                if (i == 1) {
+                    break;
+                }
+            } finally {
+                f1++;
+            }
+        } finally {
+            f2++;
+        }
+    }
+
+    assert(iterations == 2, "nested-chain: wrong iteration count");
+    assert(f1 == 2, "nested-chain: inner finally skipped");
+    assert(f2 == 2, "nested-chain: outer finally skipped");
+    return 1;
+}
+
+function test_break_in_catch() {
+    let caught = 0;
+    let finallyRuns = 0;
+    let iterations = 0;
+
+    for (let i = 0; i < 5; i++) {
+        iterations++;
+        try {
+            if (i == 2) {
+                throw "boom";
+            }
+        } catch (e: string) {
+            caught++;
+            break;
+        } finally {
+            finallyRuns++;
+        }
+    }
+
+    assert(iterations == 3, "break-in-catch: wrong iteration count");
+    assert(caught == 1, "break-in-catch: catch did not run");
+    assert(finallyRuns == 3, "break-in-catch: finally must run when catch breaks out");
+    return 1;
+}
+
+function test_continue_and_break_two_targets() {
+    // the same try has both a break and a continue escaping -- two distinct targets,
+    // each must get its own finally routing
+    let sum = 0;
+    let finallyRuns = 0;
+
+    for (let i = 0; i < 10; i++) {
+        try {
+            if (i % 2 == 0) {
+                continue;
+            }
+            if (i == 7) {
+                break;
+            }
+            sum += i;
+        } finally {
+            finallyRuns++;
+        }
+    }
+
+    // odd i below 7 summed: 1 + 3 + 5 = 9; loop left at i == 7
+    assert(sum == 9, "two-targets: wrong sum");
+    assert(finallyRuns == 8, "two-targets: finally runs for i = 0..7");
+    return 1;
+}
+
+function test_inner_loop_break_untouched() {
+    // break targeting a loop fully inside the try must NOT be routed through the
+    // finally -- it doesn't escape the try
+    let finallyRuns = 0;
+    let innerRuns = 0;
+
+    try {
+        for (let i = 0; i < 10; i++) {
+            if (i == 3) {
+                break;
+            }
+            innerRuns++;
+        }
+    } finally {
+        finallyRuns++;
+    }
+
+    assert(innerRuns == 3, "inner-loop: wrong inner count");
+    assert(finallyRuns == 1, "inner-loop: finally must run exactly once");
+    return 1;
+}
+
 function main() {
     assert(test_break_in_finally() == 33, "failed. 1");
     assert(test_continue_in_finally() == 8, "failed. 2");
     assert(test_nested_try_in_finally() == 3, "failed. 3");
+    assert(test_break_in_try_body() == 1, "failed. 4");
+    assert(test_continue_in_try_body() == 8, "failed. 5");
+    assert(test_nested_finally_chain() == 1, "failed. 6");
+    assert(test_break_in_catch() == 1, "failed. 7");
+    assert(test_continue_and_break_two_targets() == 1, "failed. 8");
+    assert(test_inner_loop_break_untouched() == 1, "failed. 9");
 
     print("done.");
 }