An exception through a block's using releases the owned locals of the scopes around it

tslang/include/TypeScript/MLIRLogic/MLIRGenContext.h

@@ -83,6 +83,7 @@ struct GenContext
         capturedVars = nullptr;
         usingVars = nullptr;
         ownedVars = nullptr;
+        isUnwindCleanupScope = false;
 
         currentOperation = nullptr;
         allocateVarsOutsideOfOperation = false;
@@ -234,6 +235,13 @@ struct GenContext
     // is, for every scope nested in it. An unlabelled `break` walking outwards stops here; a
     // `continue` passes through it to the loop.
     bool isSwitchScope = false;
+    // Set on a scope that is the body of a TryOp - a `try` statement's, or the catch-less one
+    // wrapped around a block or a switch for its `using` - and cleared, as isSwitchScope is, for
+    // every scope nested in it. An exception leaving such a scope lands on that TryOp's own
+    // landing pad, whose cleanup region gives back what the scope owns; the unwind leg of a
+    // scope nested further in releases the scopes between, and stops here
+    // (mlirGenReleaseOwnedOnUnwind).
+    bool isUnwindCleanupScope = false;
     // 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/MLIRGenFunctions.cpp

@@ -1751,6 +1751,8 @@ namespace mlirgen
         // context has a list of its own, since the walk used to stop at the first context
         // without one.
         funcGenContext.parentBlockContext = nullptr;
+        // nor is it the body of an enclosing function's TryOp
+        funcGenContext.isUnwindCleanupScope = false;
 
         // if we need params only we do not need to process body
         auto discoverParamsOnly = funcGenContext.allowPartialResolve && funcGenContext.discoverParamsOnly;

tslang/lib/TypeScript/MLIRGenImpl.h

@@ -807,23 +807,7 @@ class MLIRGenImpl
 
         if (genContext->ownedVars != nullptr)
         {
-            // reverse declaration order, the order a scope is unwound in: a later local may
-            // hold the only other reference to what an earlier one points at
-            for (auto storage : llvm::reverse(*genContext->ownedVars))
-            {
-                // Two different debts, and a slot can carry either, both or neither. A cell is
-                // released as a cell whatever its contents are owned by - the block has to go
-                // back regardless. A slot that is neither is here only because it might have
-                // turned into a cell and did not; see trackPossibleCell.
-                if (isCapturedVariableCell(storage))
-                {
-                    builder.create<mlir_ts::ReleaseCellOp>(location, storage);
-                }
-                else if (isOwnedLocalSlot(storage))
-                {
-                    builder.create<mlir_ts::ReleaseSlotOp>(location, storage);
-                }
-            }
+            releaseScopeOwnedVars(location, *genContext->ownedVars);
 
             // Process-once, as for usingVars: CurrentScopeKeepAfterUse is what the try body
             // passes so that the cleanup region, generated after it, still sees the list.
@@ -852,6 +836,57 @@ class MLIRGenImpl
         return mlir::success();
     }
 
+    void releaseScopeOwnedVars(mlir::Location location, const SmallVector<mlir::Value> &ownedVars)
+    {
+        // reverse declaration order, the order a scope is unwound in: a later local may
+        // hold the only other reference to what an earlier one points at
+        for (auto storage : llvm::reverse(ownedVars))
+        {
+            // Two different debts, and a slot can carry either, both or neither. A cell is
+            // released as a cell whatever its contents are owned by - the block has to go
+            // back regardless. A slot that is neither is here only because it might have
+            // turned into a cell and did not; see trackPossibleCell.
+            if (isCapturedVariableCell(storage))
+            {
+                builder.create<mlir_ts::ReleaseCellOp>(location, storage);
+            }
+            else if (isOwnedLocalSlot(storage))
+            {
+                builder.create<mlir_ts::ReleaseSlotOp>(location, storage);
+            }
+        }
+    }
+
+    // The rest of the unwind leg of a catch-less TryOp's cleanup region - the one wrapped around
+    // a block or a switch for its `using`. Once the region has given back its own scope, the
+    // exception keeps unwinding, and the scopes around it that have no cleanup region of their
+    // own get no other chance: a function's body holding `let s = a + b` and, further in,
+    // `{ using r = new R(); ... }` leaked `s` on every exception through the inner block.
+    //
+    // Walks outwards from the scope around the wrapped one, releasing as it goes, and stops at
+    // the first scope that is itself a TryOp's body (isUnwindCleanupScope) without touching its
+    // list: the exception lands on that TryOp's landing pad, whose own cleanup releases that
+    // scope - and, when it has a catch, everything further out is still live after it. Otherwise
+    // the walk ends at the function (whose context ends the chain), parameters included, as a
+    // `return` does.
+    //
+    // Release only, and nothing is marked processed. No disposal: a scope with a `using` it can
+    // dispose on unwind is wrapped, and so is where the walk stops. The scopes' lists stay as
+    // they are for their own exits, which come later in the IR and on other paths. Only the
+    // locals declared ahead of the wrapped scope are in the lists yet, so every slot released
+    // here has had its store - or, hoisted in front of a switch whose case did not run, is
+    // null, which a release does nothing with (VariableOpLowering).
+    void mlirGenReleaseOwnedOnUnwind(mlir::Location location, const GenContext *genContext)
+    {
+        for (; genContext != nullptr && !genContext->isUnwindCleanupScope; genContext = genContext->parentBlockContext)
+        {
+            if (genContext->ownedVars != nullptr)
+            {
+                releaseScopeOwnedVars(location, *genContext->ownedVars);
+            }
+        }
+    }
+
     // Does this reference address a captured variable's cell - a heap block the frame shares
     // with the closures that captured it - rather than ordinary frame storage?
     //

tslang/lib/TypeScript/MLIRGenStatements.cpp

@@ -189,6 +189,7 @@ namespace mlirgen
         genContextUsing.labelledScopeLabel.clear();
         // nor the switch's, when this block is a case's
         genContextUsing.isSwitchScope = false;
+        genContextUsing.isUnwindCleanupScope = false;
         genContextUsing.isLoop = false;
 
         DITableScopeT debugBlockScope(debugScope);
@@ -264,6 +265,7 @@ namespace mlirgen
             tryBodyGenContext.loopBodyLabel = tryGenContext.loopLabel;
             tryBodyGenContext.labelledScopeLabel.clear();
             tryBodyGenContext.isSwitchScope = false;
+            tryBodyGenContext.isUnwindCleanupScope = true;
             tryBodyGenContext.isLoop = false;
 
             auto usingVars = std::make_unique<SmallVector<ts::VariableDeclarationDOM::TypePtr>>();
@@ -283,6 +285,8 @@ namespace mlirgen
             // well as in the body.
             builder.setInsertionPointToStart(&tryOp.getCleanup().front());
             EXIT_IF_FAILED(mlirGenScopeExit(location, DisposeDepth::CurrentScope, {}, &tryBodyGenContext));
+            // and, with no catch here, what the scopes around it owe
+            mlirGenReleaseOwnedOnUnwind(location, tryBodyGenContext.parentBlockContext);
 
             builder.create<mlir_ts::ResultOp>(location);
         }
@@ -690,6 +694,7 @@ namespace mlirgen
         branchContext.isLoop = false;
         branchContext.labelledScopeLabel.clear();
         branchContext.isSwitchScope = false;
+        branchContext.isUnwindCleanupScope = false;
 
         auto usingVars = std::make_unique<SmallVector<ts::VariableDeclarationDOM::TypePtr>>();
         branchContext.usingVars = usingVars.get();
@@ -1076,6 +1081,7 @@ namespace mlirgen
         labelledContext.labelledScopeLabel = label;
         labelledContext.isLoopBodyScope = false;
         labelledContext.isSwitchScope = false;
+        labelledContext.isUnwindCleanupScope = false;
 
         auto usingVars = std::make_unique<SmallVector<ts::VariableDeclarationDOM::TypePtr>>();
         labelledContext.usingVars = usingVars.get();
@@ -1197,6 +1203,7 @@ namespace mlirgen
         // had, which crashed the compiler.
         switchGenContext.parentBlockContext = &genContext;
         switchGenContext.isSwitchScope = true;
+        switchGenContext.isUnwindCleanupScope = static_cast<bool>(tryOp);
         switchGenContext.isLoopBodyScope = false;
         switchGenContext.labelledScopeLabel.clear();
         switchGenContext.isLoop = false;
@@ -1281,6 +1288,7 @@ namespace mlirgen
 
             builder.setInsertionPointToStart(&tryOp.getCleanup().front());
             EXIT_IF_FAILED(mlirGenScopeExit(location, DisposeDepth::CurrentScope, {}, &switchGenContext));
+            mlirGenReleaseOwnedOnUnwind(location, switchGenContext.parentBlockContext);
             builder.create<mlir_ts::ResultOp>(location);
 
             builder.setInsertionPointToStart(&tryOp.getCatches().front());
@@ -1464,6 +1472,8 @@ namespace mlirgen
             tryBodyGenContext.loopBodyLabel = tryGenContext.loopLabel;
             tryBodyGenContext.labelledScopeLabel.clear();
             tryBodyGenContext.isSwitchScope = false;
+            // where the unwind leg of a wrapped scope inside stops (mlirGenReleaseOwnedOnUnwind)
+            tryBodyGenContext.isUnwindCleanupScope = true;
             tryBodyGenContext.isLoop = false;
 
             auto usingVars = std::make_unique<SmallVector<ts::VariableDeclarationDOM::TypePtr>>();

tslang/test/tester/CMakeLists.txt

@@ -336,6 +336,7 @@ tslang_add_test(NAME test-compile-00-label-chain-loop COMMAND test-runner "${PRO
 tslang_add_test(NAME test-compile-00-labelled-scope-exit COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00labelled_scope_exit.ts")
 tslang_add_test(NAME test-compile-00-switch-case-scope-exit COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00switch_case_scope_exit.ts")
 tslang_add_test(NAME test-compile-00-using-dispose-order COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00using_dispose_order.ts")
+tslang_add_test(NAME test-compile-00-using-unwind-outer-locals COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00using_unwind_outer_locals.ts")
 tslang_add_test(NAME test-compile-00-vars COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00vars.ts")
 tslang_add_test(NAME test-compile-00-var-bindings COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00var_bindings.ts")
 tslang_add_test(NAME test-compile-00-globals COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00globals.ts")
@@ -878,6 +879,7 @@ tslang_add_test(NAME test-jit-00-label-chain-loop COMMAND test-runner -jit "${PR
 tslang_add_test(NAME test-jit-00-labelled-scope-exit COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00labelled_scope_exit.ts")
 tslang_add_test(NAME test-jit-00-switch-case-scope-exit COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00switch_case_scope_exit.ts")
 tslang_add_test(NAME test-jit-00-using-dispose-order COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00using_dispose_order.ts")
+tslang_add_test(NAME test-jit-00-using-unwind-outer-locals COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00using_unwind_outer_locals.ts")
 tslang_add_test(NAME test-jit-00-vars COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00vars.ts")
 tslang_add_test(NAME test-jit-00-var-bindings COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00var_bindings.ts")
 tslang_add_test(NAME test-jit-00-globals COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00globals.ts")
@@ -1831,6 +1833,7 @@ set(TSLANG_CORPUS
     00labelled_scope_exit.ts
     00switch_case_scope_exit.ts
     00using_dispose_order.ts
+    00using_unwind_outer_locals.ts
     00capture_in_new_arguments.ts
     00cast_to_f32.ts
     00catch_untyped_value.ts

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

@@ -0,0 +1,167 @@
+// An exception through a block's `using` gives back what the scopes around the block took, too.
+//
+// A block with a `using` is wrapped in a catch-less TryOp whose cleanup region disposes and
+// releases the block's own scope, then lets the exception keep unwinding. A scope around it with
+// no cleanup region of its own - here, mostly, the function body - got no other chance, so
+// `let s = a + b; { using r = new Res(); throw 1; }` leaked `s` under -mm=rc. The ownership
+// verifier (--verify-ownership, swept over this corpus in ctest) reported every such shape; this
+// file is among them, and it is quiet now.
+//
+// The walk outwards stops at the first scope that is itself a TryOp's body: that TryOp's own
+// cleanup gives its scope back, and when it has a catch, everything further out is live after it.
+// The `caughtFurtherOut` cases check that side - a release too many there frees what is still read.
+
+let trail = "";
+
+class Res {
+    constructor(public name: string) {}
+
+    [Symbol.dispose]() {
+        trail += this.name + ";";
+    }
+}
+
+// allocates over whatever was freed last, so that a string released one time too many no longer
+// reads as itself
+function churn() {
+    let junk = "";
+    for (let i = 0; i < 64; i++) {
+        junk = junk + "#";
+    }
+
+    return junk.length;
+}
+
+function returnFromBlock(pre: string) {
+    let s = pre + "1";
+    {
+        using a = new Res("a");
+        return s;
+    }
+}
+
+function throwFromBlock(pre: string) {
+    let s = pre + "2";
+    {
+        using a = new Res("a");
+        trail += s + ";";
+        throw 1;
+    }
+}
+
+function throwThroughPlainScope(pre: string) {
+    let s = pre + "3";
+    {
+        let t = s + "t";
+        {
+            using a = new Res("a");
+            trail += t + ";";
+            throw 1;
+        }
+    }
+}
+
+function throwFromSwitchCase(pre: string, n: number) {
+    let s = pre + "4";
+    switch (n) {
+        case 1:
+            using a = new Res("a");
+            trail += s + ";";
+            throw 1;
+    }
+
+    return s;
+}
+
+function throwFromLoopBody(pre: string) {
+    let s = pre + "5";
+    for (let i = 0; i < 3; i++) {
+        let t = s + i;
+        {
+            using a = new Res("a" + i);
+            if (i == 1) {
+                throw 1;
+            }
+        }
+    }
+}
+
+function caughtFurtherOut(pre: string) {
+    let s = pre + "6";
+    try {
+        let t = s + "t";
+        {
+            using a = new Res("a");
+            throw 1;
+        }
+    }
+    catch (e: TypeOf<1>) {
+        trail += "caught;";
+    }
+
+    churn();
+    return s;
+}
+
+function caughtFurtherOutSwitch(pre: string, n: number) {
+    let s = pre + "7";
+    try {
+        switch (n) {
+            case 1:
+                using a = new Res("a");
+                throw 1;
+        }
+    }
+    catch (e: TypeOf<1>) {
+        trail += "caught;";
+    }
+
+    churn();
+    return s;
+}
+
+function expectThrow(f: () => void) {
+    try {
+        f();
+    }
+    catch (e: TypeOf<1>) {
+        return true;
+    }
+
+    return false;
+}
+
+function check(f: () => void, expected: string, what: string) {
+    trail = "";
+    f();
+    churn();
+    print(what + ": " + trail);
+    assert(trail == expected, what);
+}
+
+function main() {
+    let pre = "p";
+
+    trail = "";
+    let r = returnFromBlock(pre);
+    churn();
+    print("return out of a block: " + r + " " + trail);
+    assert(r == "p1" && trail == "a;", "return out of a block");
+
+    check(() => assert(expectThrow(() => throwFromBlock(pre)), "must throw"), "p2;a;", "throw out of a block");
+    check(() => assert(expectThrow(() => throwThroughPlainScope(pre)), "must throw"), "p3t;a;", "throw through a plain scope");
+    check(() => assert(expectThrow(() => { throwFromSwitchCase(pre, 1); }), "must throw"), "p4;a;", "throw out of a switch case");
+    check(() => assert(expectThrow(() => throwFromLoopBody(pre)), "must throw"), "a0;a1;", "throw out of a loop body");
+
+    trail = "";
+    let c = caughtFurtherOut(pre);
+    print("caught further out: " + c + " " + trail);
+    assert(c == "p6" && trail == "a;caught;", "caught further out");
+
+    trail = "";
+    let cs = caughtFurtherOutSwitch(pre, 1);
+    print("caught further out, switch: " + cs + " " + trail);
+    assert(cs == "p7" && trail == "a;caught;", "caught further out, switch");
+
+    print("done.");
+}