A global held across a call is given back when the call throws

tslang/lib/TypeScript/OwnedReturnConsumptionPass.cpp

@@ -298,9 +298,9 @@ class OwnedReturnConsumptionPass
     // other argument, and the callee may give the global a new value while it still uses the
     // parameter (`function run(h: C) { G = mk(); print(h.x); }`): the global's release freed the
     // block the parameter read. Under rc the caller holds a count across the call - retained right
-    // before it, released right after it. This runs before the affine lowering, so a call in a try
-    // body is still a plain call here (it becomes a `ts.Invoke` later, with the release first in its
-    // normal successor); a throw out of the call skips the release, and the count leaks.
+    // before it, released right after it - and, when a try around the call takes what it throws,
+    // released where that try's handler starts (unwindRegionOf). This runs before the affine
+    // lowering, so a call in a try body is still a plain call here (a `ts.Invoke` later).
     //
     // A Shared<T> handle (spec 23.1) is counted under rc and own alike, and a program reads its
     // count (`Shared.count`) as the number of handles - which a parameter, a borrow, is not one of.
@@ -494,11 +494,108 @@ class OwnedReturnConsumptionPass
 
         for (auto [op, argument] : held)
         {
+            auto location = op->getLoc();
             mlir::OpBuilder builder(op);
-            builder.create<mlir_ts::RetainOp>(op->getLoc(), argument);
+            builder.create<mlir_ts::RetainOp>(location, argument);
+
+            // A call a try around it takes the exceptions of gives the hold back on that path too,
+            // through a slot that holds the value only while the call runs: null before and after
+            // it, so the handler's release does nothing when something else threw.
+            auto *unwindRegion = unwindRegionOf(op);
+            mlir::Value slot;
+            if (unwindRegion)
+            {
+                auto type = argument.getType();
+                auto &entryBlock = op->getParentOfType<mlir_ts::FuncOp>().getBody().front();
+                mlir::OpBuilder entryBuilder(&entryBlock, entryBlock.begin());
+                if (auto entryOp = dyn_cast<mlir_ts::EntryOp>(entryBlock.front()))
+                {
+                    entryBuilder.setInsertionPointAfter(entryOp);
+                }
+
+                auto none = entryBuilder.create<mlir_ts::DefaultOp>(location, type);
+                slot = entryBuilder.create<mlir_ts::VariableOp>(location, mlir_ts::RefType::get(type), none,
+                                                                entryBuilder.getBoolAttr(false),
+                                                                entryBuilder.getIndexAttr(0));
+                builder.create<mlir_ts::StoreOp>(location, argument, slot);
+            }
+
             builder.setInsertionPointAfter(op);
-            builder.create<mlir_ts::ReleaseOp>(op->getLoc(), argument);
+            builder.create<mlir_ts::ReleaseOp>(location, argument);
+            if (!slot)
+            {
+                continue;
+            }
+
+            builder.create<mlir_ts::StoreOp>(location, builder.create<mlir_ts::DefaultOp>(location, argument.getType()),
+                                             slot);
+
+            // in a catch clause, after the catch variable is set
+            auto &handlerBlock = unwindRegion->front();
+            mlir::OpBuilder handlerBuilder(&handlerBlock, handlerBlock.begin());
+            for (auto &handlerOp : handlerBlock)
+            {
+                if (isa<mlir_ts::CatchOp>(handlerOp))
+                {
+                    handlerBuilder.setInsertionPointAfter(&handlerOp);
+                    break;
+                }
+            }
+
+            handlerBuilder.create<mlir_ts::ReleaseSlotOp>(location, slot);
+            handlerBuilder.create<mlir_ts::StoreOp>(
+                location, handlerBuilder.create<mlir_ts::DefaultOp>(location, argument.getType()), slot);
+        }
+    }
+
+    // Where an exception out of `op` is taken first within its function: the region of the try
+    // around it that runs then, in the order its landing pad takes them (TryOpLowering) - the
+    // cleanup, else the catch clause, else the finally, for a call in its body; the finally for a
+    // call in its catch clause. Null when the exception leaves the function.
+    static mlir::Region *unwindRegionOf(mlir::Operation *op)
+    {
+        auto hasOps = [](mlir::Region &region) {
+            return llvm::any_of(region, [](mlir::Block &block) { return &block.front() != block.getTerminator(); });
+        };
+
+        for (auto *region = op->getParentRegion(); region; region = region->getParentOp()->getParentRegion())
+        {
+            auto *parentOp = region->getParentOp();
+            if (isa<mlir::FunctionOpInterface>(parentOp))
+            {
+                return nullptr;
+            }
+
+            auto tryOp = dyn_cast<mlir_ts::TryOp>(parentOp);
+            if (!tryOp)
+            {
+                continue;
+            }
+
+            if (region == &tryOp.getBody())
+            {
+                if (hasOps(tryOp.getCleanup()))
+                {
+                    return &tryOp.getCleanup();
+                }
+
+                if (tryOp->hasAttr(HAS_CATCH_ATTR_NAME) || hasOps(tryOp.getCatches()))
+                {
+                    return &tryOp.getCatches();
+                }
+
+                if (hasOps(tryOp.getFinally()))
+                {
+                    return &tryOp.getFinally();
+                }
+            }
+            else if (region == &tryOp.getCatches() && hasOps(tryOp.getFinally()))
+            {
+                return &tryOp.getFinally();
+            }
         }
+
+        return nullptr;
     }
 
     // A global is born holding the reference its initializer carries.

tslang/test/tester/CMakeLists.txt

@@ -3598,7 +3598,7 @@ foreach(shared_mm ${shared_models})
     foreach(shared_test own_shared_basic own_shared_user_named own_shared_graph own_shared_containers own_shared_cycle own_shared_union
                         own_shared_global own_shared_unlink own_shared_global_init
                         own_shared_place_two_uses own_shared_global_field own_shared_global_copy
-                        own_shared_global_argument)
+                        own_shared_global_argument own_shared_global_argument_throw)
         tslang_add_test(NAME test-jit-${shared_mm}-${shared_test} COMMAND test-runner -jit -mm=${shared_mm} "${PROJECT_SOURCE_DIR}/test/tester/own/${shared_test}.ts")
         tslang_add_test(NAME test-compile-${shared_mm}-${shared_test} COMMAND test-runner -mm=${shared_mm} "${PROJECT_SOURCE_DIR}/test/tester/own/${shared_test}.ts")
     endforeach()
@@ -3613,7 +3613,7 @@ foreach(shared_verify_mm own rc)
     foreach(shared_test own_shared_basic own_shared_graph own_shared_count own_shared_containers own_shared_cycle own_shared_union
                         own_shared_global own_shared_unlink own_shared_global_init
                         own_shared_place_two_uses own_shared_global_field own_shared_global_copy
-                        own_shared_global_argument)
+                        own_shared_global_argument own_shared_global_argument_throw)
         add_test(NAME test-${shared_verify_mm}-verify-ownership-${shared_test}
                  COMMAND $<TARGET_FILE:tslang> --emit=mlir-affine --no-default-lib -mm=${shared_verify_mm} --verify-ownership
                          "${PROJECT_SOURCE_DIR}/test/tester/own/${shared_test}.ts")

tslang/test/tester/own/own_shared_global_argument_throw.ts

@@ -0,0 +1,123 @@
+// A Shared<T> global given to a call that may assign it is held by the caller across the call
+// (own_shared_global_argument.ts). When the call throws instead of returning, the hold is given
+// back on the way to the handler: before, only the normal path released it, and every throw left
+// the global's block one count it never lost. Counts are checked where there is one (rc, own);
+// gc and none report -1.
+class Node {
+    v = 0;
+    tag = "";
+}
+
+function countIs(s: Shared<Node>, expected: number) {
+    const c = Shared.count(s);
+    return c < 0 || c == expected;
+}
+
+function mk(tag: string, v: number) {
+    const n = new Node();
+    n.tag = tag;
+    n.v = v;
+    return n;
+}
+
+let g = new Shared(mk("g", 1));
+
+// may assign the global, so a caller holds what it gives it; throws before it does
+function maybeAssign(h: Shared<Node>, assign: boolean) {
+    if (assign) {
+        g = new Shared(mk("replaced", 2));
+    }
+
+    if (h.value.v > 0) {
+        throw "thrown";
+    }
+
+    return h.value.v;
+}
+
+function inTryCatch() {
+    try {
+        maybeAssign(g, false);
+    } catch (e) {
+        assert(countIs(g, 1), "the hold given back before the catch runs");
+    }
+
+    assert(countIs(g, 1), "try/catch");
+}
+
+function inTryFinally() {
+    try {
+        maybeAssign(g, false);
+    } finally {
+        assert(countIs(g, 1), "the hold given back before the finally runs");
+    }
+}
+
+function throughFinally() {
+    try {
+        inTryFinally();
+    } catch (e) {
+    }
+
+    assert(countIs(g, 1), "try/finally");
+}
+
+function inBlockWithLocal() {
+    try {
+        {
+            const local = mk("local", 3);
+            maybeAssign(g, false);
+            print(local.v);
+        }
+    } catch (e) {
+    }
+
+    assert(countIs(g, 1), "a block with a local in a try body");
+}
+
+function inCatchWithFinally() {
+    let caught = 0;
+    try {
+        try {
+            throw "first";
+        } catch (e) {
+            maybeAssign(g, false);
+        } finally {
+            caught++;
+        }
+    } catch (e) {
+        caught++;
+    }
+
+    assert(caught == 2 && countIs(g, 1), "a call in a catch clause with a finally");
+}
+
+function inLoop() {
+    for (let i = 0; i < 3; i++) {
+        try {
+            maybeAssign(g, false);
+        } catch (e) {
+        }
+    }
+
+    assert(countIs(g, 1), "a try in a loop");
+}
+
+function main() {
+    inTryCatch();
+    throughFinally();
+    inBlockWithLocal();
+    inCatchWithFinally();
+    inLoop();
+
+    // and the normal path still holds across an assignment
+    let ok = false;
+    try {
+        g = new Shared(mk("neg", -1));
+        ok = maybeAssign(g, true) == -1;
+    } catch (e) {
+    }
+
+    assert(ok && countIs(g, 1) && g.value.tag == "replaced", "the normal path");
+    print("done.");
+}