A return, break or continue through a finally disposes and releases the scopes the finally reads after it runs

tslang/include/TypeScript/Defines.h

@@ -70,13 +70,21 @@
 // rest.
 #define UNWIND_ONLY_ATTR_NAME "__unwind_only"
 
-// The `return` counterpart: a release MLIRGen put at the end of a hand-written `finally` for what
-// the scopes around the `try` owe once a `return` from its body or catch clause has run the
-// finally. The `return` itself stops short of those scopes, since the finally may still read them
-// (GenContext::finallyGuardScope). TryOpLowering keeps these in the copy a `return` runs and
-// erases them from the rest.
+// The `return` counterpart: the `ts.If` MLIRGen put at the end of a hand-written `finally` around
+// what the scopes around the `try` owe - their `using`s to dispose and their locals to release -
+// once a `return` from its body or catch clause has run the finally. The `return` itself stops
+// short of those scopes, since the finally may still read them (GenContext::finallyGuardScope).
+// TryOpLowering keeps it in the copy a `return` runs and erases it from the rest.
 #define RETURN_ONLY_ATTR_NAME "__return_only"
 
+// The same for a `break` or `continue` out of the `try` body or catch clause, one `ts.If` per
+// place they go. Its value names the jump: JUMP_ONLY_BREAK_KEY or JUMP_ONLY_CONTINUE_KEY followed
+// by the label (empty when there is none), as the jump's BreakOp or ContinueOp carries it.
+// TryOpLowering keeps it in the copy that jump runs and erases it from the rest.
+#define JUMP_ONLY_ATTR_NAME "__jump_only"
+#define JUMP_ONLY_BREAK_KEY "break:"
+#define JUMP_ONLY_CONTINUE_KEY "continue:"
+
 // Marks a `ts.CreateBoundFunction` whose `this` is a capture box built for it a moment earlier,
 // rather than a receiver that belongs to somebody else. Only such a closure owns its `this`, and
 // only it gets the type tag at CLOSURE_TYPE_INDEX - a bound method must not take ownership of

tslang/include/TypeScript/MLIRLogic/MLIRGenContext.h

@@ -244,12 +244,17 @@ struct GenContext
     bool isUnwindCleanupScope = false;
     // Points at itself on the context that stands, in the walk outwards, for the scope around a
     // `try` with a `finally` - the parent of its body's scope, and its catch clause's context. A
-    // `return` from inside the body or the catch releases the scopes up to here and no further:
-    // the finally still runs before the function returns, and it can read the scopes from here
-    // outwards, so the copy of the finally that a `return` runs releases them once it is done
-    // (mlirGenReleaseOwnedOnReturn). A pointer rather than a flag, so that it needs no clearing:
-    // every context copied from this one points at this one, not at itself.
+    // `return`, `break` or `continue` from inside the body or the catch that gets this far
+    // disposes and releases the scopes up to here and no further: the finally still runs first,
+    // and it can read the scopes from here outwards, so the copy of the finally that the jump
+    // runs disposes and releases them once it is done (mlirGenScopeExitAfterFinally). A pointer
+    // rather than a flag, so that it needs no clearing: every context copied from this one points
+    // at this one, not at itself.
     const GenContext *finallyGuardScope = nullptr;
+    // On a finallyGuardScope: the `break`s (true) and `continue`s (false), by label, that stopped
+    // there - one entry for each place they go, and the finally gets the rest of the scope exit
+    // for each. Shared by the body's guard and the catch clause's, which is a copy of it.
+    mlir::SmallVector<std::pair<bool, std::string>> *finallyEscapingJumps = nullptr;
     // out-of-band cancellation signal; mutable so stop() stays callable through the const& threading
     mutable bool stopProcess = false;
     mlir::SmallVector<std::unique_ptr<mlir::Diagnostic>> *postponedMessages = nullptr;

tslang/lib/TypeScript/LowerToAffineLoops.cpp

@@ -1629,10 +1629,13 @@ struct TryOpLowering : public TsPattern<mlir_ts::TryOp>
             // the last copy is made: the rewriter defers an erasure, so an op erased from the
             // region still goes into every copy cloned from it afterwards.
             //
-            // The releases a `return` left for after the finally (RETURN_ONLY_ATTR_NAME) the other
-            // way round: kept in the copy a `return` runs, and gone from every other.
+            // The scope exit a `return` left for after the finally (RETURN_ONLY_ATTR_NAME) the other
+            // way round: kept in the copy a `return` runs, and gone from every other. And the one
+            // each `break` or `continue` left (JUMP_ONLY_ATTR_NAME), kept in the copy for that
+            // jump alone. Each is a `ts.If` holding the lot, so erasing it erases them all.
             SmallVector<Operation *> unwindOnlyOps;
             SmallVector<Operation *> returnOnlyOps;
+            SmallVector<Operation *> jumpOnlyOps;
             tryOp.getFinally().walk<mlir::WalkOrder::PreOrder>([&](Operation *op) {
                 if (isa<mlir_ts::TryOp>(op))
                 {
@@ -1647,6 +1650,13 @@ struct TryOpLowering : public TsPattern<mlir_ts::TryOp>
                 if (op->hasAttr(RETURN_ONLY_ATTR_NAME))
                 {
                     returnOnlyOps.push_back(op);
+                    return WalkResult::skip();
+                }
+
+                if (op->hasAttr(JUMP_ONLY_ATTR_NAME))
+                {
+                    jumpOnlyOps.push_back(op);
+                    return WalkResult::skip();
                 }
 
                 return WalkResult::advance();
@@ -1659,7 +1669,44 @@ struct TryOpLowering : public TsPattern<mlir_ts::TryOp>
                 }
             };
 
+            // the key MLIRGen tagged a jump's scope exit with: the kind of jump and its label
+            auto jumpKey = [](Operation *op) {
+                auto label = op->getAttrOfType<StringAttr>("label");
+                return std::string(isa<mlir_ts::BreakOp>(op) ? JUMP_ONLY_BREAK_KEY : JUMP_ONLY_CONTINUE_KEY) +
+                       (label ? label.getValue().str() : std::string());
+            };
+
+            // In the copy it belongs to, the `ts.If` goes and what it holds stays in its place:
+            // lowered as it is, the constant condition became a conditional branch, and to the
+            // ownership verifier its other edge was a path that never released.
+            auto unwrapOps = [&](ArrayRef<Operation *> ops, const mlir::IRMapping &mapping) {
+                for (auto *op : ops)
+                {
+                    auto ifOp = cast<mlir_ts::IfOp>(mapping.lookup(op));
+                    auto &thenBlock = ifOp.getThenRegion().front();
+                    auto *terminator = thenBlock.getTerminator();
+                    rewriter.inlineBlockBefore(&thenBlock, ifOp);
+                    rewriter.eraseOp(terminator);
+                    rewriter.eraseOp(ifOp);
+                }
+            };
+
+            auto keepJumpOnlyOps = [&](StringRef keepKey, const mlir::IRMapping &mapping) {
+                for (auto *op : jumpOnlyOps)
+                {
+                    if (op->getAttrOfType<StringAttr>(JUMP_ONLY_ATTR_NAME).getValue() == keepKey)
+                    {
+                        unwrapOps(op, mapping);
+                    }
+                    else
+                    {
+                        rewriter.eraseOp(mapping.lookup(op));
+                    }
+                }
+            };
+
             eraseOps(returnOnlyOps, finallyMapping);
+            eraseOps(jumpOnlyOps, finallyMapping);
 
             // add clone for 'return'
             if (returns.size() > 0)
@@ -1669,6 +1716,8 @@ struct TryOpLowering : public TsPattern<mlir_ts::TryOp>
                                             returnFinallyMapping);
                 propagateTsContextEntries(returnFinallyMapping);
                 eraseOps(unwindOnlyOps, returnFinallyMapping);
+                eraseOps(jumpOnlyOps, returnFinallyMapping);
+                unwrapOps(returnOnlyOps, returnFinallyMapping);
                 auto returnFinallyBlockLast = continuation->getPrevNode();
                 rewriter.setInsertionPoint(returnFinallyBlockLast->getTerminator());
                 auto resultOpOfReturnFinallyBlock = cast<mlir_ts::ResultOp>(returnFinallyBlockLast->getTerminator());
@@ -1677,7 +1726,7 @@ struct TryOpLowering : public TsPattern<mlir_ts::TryOp>
                 // A `try` around this one lowered first and sent these returns to its own finally
                 // copy. The return this copy ends with goes there in their stead, so the outer
                 // finally still runs - pointing the returns at this copy alone skipped it - and
-                // releases the scopes it guards (mlirGenReleaseOwnedOnReturn).
+                // releases the scopes it guards (mlirGenScopeExitAfterFinally).
                 mlir::Block *outerCleanup = nullptr;
                 for (auto retOp : returns)
                 {
@@ -1705,12 +1754,18 @@ struct TryOpLowering : public TsPattern<mlir_ts::TryOp>
             // 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.
+            //
+            // A copy ends with the rest of its jump's scope exit (JUMP_ONLY_ATTR_NAME), and that is
+            // MLIRGen's to say by the jump's kind and label, so a copy is per target and key: an
+            // unlabelled `break` and a `break outer` for the same loop go to one place, and each
+            // finds its own scope exit in its copy.
             if (escapingJumpsBody.size() > 0 || escapingJumpsCatches.size() > 0)
             {
-                mlir::DenseMap<mlir::Block *, mlir::Block *> finallyCopyPerTarget;
+                std::map<std::pair<mlir::Block *, std::string>, mlir::Block *> finallyCopyPerTarget;
                 auto routeThroughFinally = [&](Operation *op) {
                     auto target = tsContext->jumps.lookup(op);
-                    auto &finallyCopy = finallyCopyPerTarget[target];
+                    auto key = jumpKey(op);
+                    auto &finallyCopy = finallyCopyPerTarget[{target, key}];
                     if (!finallyCopy)
                     {
                         mlir::IRMapping jumpFinallyMapping;
@@ -1719,6 +1774,7 @@ struct TryOpLowering : public TsPattern<mlir_ts::TryOp>
                         propagateTsContextEntries(jumpFinallyMapping);
                         eraseOps(unwindOnlyOps, jumpFinallyMapping);
                         eraseOps(returnOnlyOps, jumpFinallyMapping);
+                        keepJumpOnlyOps(key, jumpFinallyMapping);
                         auto jumpFinallyBlockLast = continuation->getPrevNode();
                         rewriter.setInsertionPoint(jumpFinallyBlockLast->getTerminator());
                         auto resultOpOfJumpFinallyBlock = cast<mlir_ts::ResultOp>(jumpFinallyBlockLast->getTerminator());
@@ -1743,7 +1799,7 @@ struct TryOpLowering : public TsPattern<mlir_ts::TryOp>
                 }
             }
 
-            for (auto *op : llvm::concat<Operation *const>(unwindOnlyOps, returnOnlyOps))
+            for (auto *op : llvm::concat<Operation *const>(unwindOnlyOps, returnOnlyOps, jumpOnlyOps))
             {
                 rewriter.eraseOp(op);
             }

tslang/lib/TypeScript/MLIRGenImpl.h

@@ -817,6 +817,41 @@ class MLIRGenImpl
         return !isTarget;
     }
 
+    // Whether a scope exit walking outwards stops here, at a `try` whose `finally` still has to
+    // run first (GenContext::finallyGuardScope): a `return`, or a `break` or `continue` that
+    // leaves the `try`. The finally can read this scope and the ones outside it, so the copy of
+    // the finally that the jump runs disposes and releases them once it is done
+    // (mlirGenScopeExitAfterFinally). A `break` or `continue` is noted on the guard, for the
+    // finally to know where it goes; the two walks of one scope exit both come by, hence the
+    // check for one already there.
+    static bool scopeExitStopsAtFinallyGuard(DisposeDepth disposeDepth, const std::string &loopLabel,
+                                             const GenContext *genContext)
+    {
+        if (genContext->finallyGuardScope != genContext)
+        {
+            return false;
+        }
+
+        if (disposeDepth == DisposeDepth::FullStack)
+        {
+            return true;
+        }
+
+        if (disposeDepth != DisposeDepth::LoopScope && disposeDepth != DisposeDepth::BreakScope)
+        {
+            return false;
+        }
+
+        std::pair<bool, std::string> jump(disposeDepth == DisposeDepth::BreakScope, loopLabel);
+        auto &escapingJumps = *genContext->finallyEscapingJumps;
+        if (!llvm::is_contained(escapingJumps, jump))
+        {
+            escapingJumps.push_back(jump);
+        }
+
+        return true;
+    }
+
     // Drops the reference each local of this scope took when it was declared. Shaped after
     // mlirGenDisposable, and walks outwards on the same terms, so that a `return` from a
     // nested block releases every scope it leaves and a `break` releases up to the loop.
@@ -828,9 +863,7 @@ class MLIRGenImpl
             return mlir::success();
         }
 
-        // a `return` stops at a `try` whose `finally` still has to run: that finally's copy for
-        // the `return` releases the scopes from here outwards (mlirGenReleaseOwnedOnReturn)
-        if (disposeDepth == DisposeDepth::FullStack && genContext->finallyGuardScope == genContext)
+        if (scopeExitStopsAtFinallyGuard(disposeDepth, loopLabel, genContext))
         {
             return mlir::success();
         }
@@ -931,25 +964,42 @@ class MLIRGenImpl
         }
     }
 
-    // The end of the copy of a hand-written `finally` that a `return` from its `try` body or catch
-    // clause runs (RETURN_ONLY_ATTR_NAME): the rest of that `return`'s scope exit. The `return`
-    // released the scopes inside the `try` and stopped at its finallyGuardScope, because the
+    // The end of the copy of a hand-written `finally` that a `return`, `break` or `continue` from
+    // its `try` body or catch clause runs: the rest of that jump's scope exit. The jump disposed
+    // and released the scopes inside the `try` and stopped at its finallyGuardScope, because the
     // finally can still read every scope from there outwards - `let s = ...; try { return x; }
     // finally { log(s); }` read `s` after the `return` had given it back.
     //
-    // Walks outwards from the scope around the `try`, as a `return` does, and stops at the
-    // finallyGuardScope of the next `try` out with a `finally`: the `return` goes on through that
-    // finally's own copy (TryOpLowering chains them), which releases from there.
-    void mlirGenReleaseOwnedOnReturn(mlir::Location location, const GenContext *genContext)
+    // Walks outwards from the scope around the `try` on the jump's own terms, and so stops where
+    // it goes, or at the finallyGuardScope of the next `try` out with a `finally`: the jump goes
+    // on through that finally's own copy (TryOpLowering chains them), which goes on from there.
+    //
+    // Inside a `ts.If` on a constant true, tagged `onlyOnAttr` (RETURN_ONLY_ATTR_NAME or
+    // JUMP_ONLY_ATTR_NAME) with `onlyOnValue`, so that TryOpLowering can keep the lot in the one
+    // copy of the finally it belongs to and erase it from every other. Nothing at all when the
+    // scopes owe nothing.
+    mlir::LogicalResult mlirGenScopeExitAfterFinally(mlir::Location location, DisposeDepth disposeDepth,
+                                                     std::string loopLabel, const GenContext *genContext,
+                                                     StringRef onlyOnAttr, mlir::Attribute onlyOnValue)
     {
-        for (; genContext != nullptr && genContext->finallyGuardScope != genContext;
-             genContext = genContext->parentBlockContext)
+        auto alwaysTrue = builder.create<mlir_ts::ConstantOp>(location, getBooleanType(), builder.getBoolAttr(true));
+        auto ifOp = builder.create<mlir_ts::IfOp>(location, alwaysTrue, /*withElseRegion=*/false);
+        ifOp->setAttr(onlyOnAttr, onlyOnValue);
+
         {
-            if (genContext->ownedVars != nullptr)
-            {
-                releaseScopeOwnedVars(location, *genContext->ownedVars, RETURN_ONLY_ATTR_NAME);
-            }
+            mlir::OpBuilder::InsertionGuard insertGuard(builder);
+            builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
+            EXIT_IF_FAILED(mlirGenScopeExit(location, disposeDepth, loopLabel, genContext));
         }
+
+        auto &thenBlock = ifOp.getThenRegion().front();
+        if (&thenBlock.front() == thenBlock.getTerminator())
+        {
+            ifOp->erase();
+            alwaysTrue->erase();
+        }
+
+        return mlir::success();
     }
 
     // Whether an exception out of a scope opened in this context lands on a TryOp's landing pad
@@ -1345,6 +1395,13 @@ class MLIRGenImpl
             return mlir::success();
         }
 
+        // and at a `try` whose `finally` reads this scope before it may be disposed - `{ using r =
+        // ...; try { return; } finally { r.read(); } }` disposed `r` and then ran the finally
+        if (scopeExitStopsAtFinallyGuard(disposeDepth, loopLabel, genContext))
+        {
+            return mlir::success();
+        }
+
         if (genContext->usingVars != nullptr)
         {
             // reverse declaration order, as JS disposes a scope (and as mlirGenReleaseOwned releases

tslang/lib/TypeScript/MLIRGenStatements.cpp

@@ -1468,14 +1468,17 @@ namespace mlirgen
         tryGenContext.currentOperation = tryOp;
 
         // The scope around the `try` as its body and catch clause see it on the way out, when it
-        // has a `finally`: a `return` from either stops here, and the finally's copy for the
-        // `return` releases the rest (GenContext::finallyGuardScope). A copy of tryGenContext,
-        // because the finally itself is generated in tryGenContext, and a `return` in the
-        // finally is not one this finally runs after.
+        // has a `finally`: a `return`, `break` or `continue` from either stops here, and the
+        // finally's copy for that jump disposes and releases the rest
+        // (GenContext::finallyGuardScope). A copy of tryGenContext, because the finally itself is
+        // generated in tryGenContext, and a jump in the finally is not one this finally runs
+        // after.
+        SmallVector<std::pair<bool, std::string>> finallyEscapingJumps;
         GenContext finallyGuardGenContext(tryGenContext);
         if (tryStatementAST->finallyBlock)
         {
             finallyGuardGenContext.finallyGuardScope = &finallyGuardGenContext;
+            finallyGuardGenContext.finallyEscapingJumps = &finallyEscapingJumps;
         }
 
         SmallVector<mlir::Type, 0> types;
@@ -1591,8 +1594,19 @@ namespace mlirgen
 
             // And the copy a `return` from the body or the catch clause runs: the rest of that
             // `return`'s scope exit, which it left for after the finally (finallyGuardScope).
-            // Marked return-only, so TryOpLowering drops them from the other copies.
-            mlirGenReleaseOwnedOnReturn(location, &tryGenContext);
+            // Marked return-only, so TryOpLowering drops it from the other copies.
+            EXIT_IF_FAILED(mlirGenScopeExitAfterFinally(location, DisposeDepth::FullStack, {}, &tryGenContext,
+                                                        RETURN_ONLY_ATTR_NAME, builder.getUnitAttr()));
+
+            // The same for each place a `break` or `continue` from them goes, marked with that
+            // place, so TryOpLowering keeps it in the copy that jump runs and no other
+            for (auto &[isBreak, label] : finallyEscapingJumps)
+            {
+                auto key = std::string(isBreak ? JUMP_ONLY_BREAK_KEY : JUMP_ONLY_CONTINUE_KEY) + label;
+                EXIT_IF_FAILED(mlirGenScopeExitAfterFinally(
+                    location, isBreak ? DisposeDepth::BreakScope : DisposeDepth::LoopScope, label, &tryGenContext,
+                    JUMP_ONLY_ATTR_NAME, builder.getStringAttr(key)));
+            }
         }
 
         // terminator