A Shared global lent to a call that assigns it: the caller holds a count

tslang/lib/TypeScript/OwnedReturnConsumptionPass.cpp

@@ -13,6 +13,7 @@
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Debug.h"
 
@@ -297,36 +298,154 @@ 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. A call in a try body (`ts.Invoke`) ends its block, and is
-    // left as it was; so is a Shared<T> handle, whose count a program can read (`Shared.count`).
+    // 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.
+    //
+    // 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.
+    // So a handle is held, under both, only across a call that may give the global it was read
+    // from a new value, the one case where the hold is what keeps the block alive.
     void holdGlobalArguments(MLIRTypeHelper &mth, mlir::ModuleOp module)
     {
-        if (!compileOptions.isRefCounted())
+        if (!compileOptions.tracksOwnership())
         {
             return;
         }
 
-        // a read out of a global, or out of a field or an element of one, through casts
-        auto readsGlobal = [](mlir::Value value) {
+        // the global a read comes out of, or a place (a field, an element, a handle's value) is in,
+        // through casts
+        auto globalOf = [](mlir::Value value, bool isPlace) -> mlir::StringRef {
             for (auto depth = 0; depth < 16 && value; depth++)
             {
-                while (auto castOp = value.getDefiningOp<mlir_ts::CastOp>())
+                mlir::Operation *place = nullptr;
+                if (isPlace)
                 {
-                    value = castOp.getIn();
+                    place = value.getDefiningOp();
+                    isPlace = false;
+                }
+                else
+                {
+                    while (auto castOp = value.getDefiningOp<mlir_ts::CastOp>())
+                    {
+                        value = castOp.getIn();
+                    }
+
+                    auto loadOp = value.getDefiningOp<mlir_ts::LoadOp>();
+                    place = loadOp ? loadOp.getReference().getDefiningOp() : nullptr;
                 }
 
-                auto loadOp = value.getDefiningOp<mlir_ts::LoadOp>();
-                auto *place = loadOp ? loadOp.getReference().getDefiningOp() : nullptr;
-                if (mlir::isa_and_nonnull<mlir_ts::AddressOfOp>(place))
+                if (auto addressOfOp = mlir::dyn_cast_or_null<mlir_ts::AddressOfOp>(place))
                 {
-                    return true;
+                    return addressOfOp.getGlobalName();
                 }
 
-                value = mlir::isa_and_nonnull<mlir_ts::PropertyRefOp, mlir_ts::ElementRefOp>(place) ? place->getOperand(0)
-                                                                                                     : mlir::Value();
+                value = mlir::isa_and_nonnull<mlir_ts::PropertyRefOp, mlir_ts::ElementRefOp, mlir_ts::SharedValueRefOp>(place)
+                            ? place->getOperand(0)
+                            : mlir::Value();
             }
 
-            return false;
+            return {};
+        };
+
+        auto readsGlobalItself = [](mlir::Value value) {
+            while (auto castOp = value.getDefiningOp<mlir_ts::CastOp>())
+            {
+                value = castOp.getIn();
+            }
+
+            auto loadOp = value.getDefiningOp<mlir_ts::LoadOp>();
+            return loadOp && loadOp.getReference().getDefiningOp<mlir_ts::AddressOfOp>();
+        };
+
+        // The function a call names, if it can be named.
+        auto calleeOf = [](mlir::Operation *op) -> mlir::StringRef {
+            if (auto callee = own_facts::directCallee(op))
+            {
+                return callee.getValue();
+            }
+
+            auto value = own_facts::calleeValue(op);
+            auto symbolRefOp = value ? value.getDefiningOp<mlir_ts::SymbolRefOp>() : mlir_ts::SymbolRefOp();
+            if (symbolRefOp)
+            {
+                return symbolRefOp.getIdentifier();
+            }
+
+            auto thisSymbolRefOp = value ? value.getDefiningOp<mlir_ts::ThisSymbolRefOp>() : mlir_ts::ThisSymbolRefOp();
+            return thisSymbolRefOp ? thisSymbolRefOp.getIdentifier() : mlir::StringRef();
+        };
+
+        // The globals each function may give a new value: by a store to the global or into a place
+        // in it, its own or a callee's. A call through a value that names no function may assign
+        // any of them; a function with no body here assigns none.
+        llvm::StringMap<llvm::StringSet<>> assigns;
+        llvm::StringMap<llvm::SmallVector<mlir::StringRef>> callees;
+        llvm::StringSet<> assignsAny;
+        module.walk([&](mlir_ts::FuncOp funcOp) {
+            auto name = funcOp.getSymName();
+            funcOp.walk([&](mlir::Operation *op) {
+                if (auto storeOp = mlir::dyn_cast<mlir_ts::StoreOp>(op))
+                {
+                    if (auto global = globalOf(storeOp.getReference(), true); !global.empty())
+                    {
+                        assigns[name].insert(global);
+                    }
+                }
+                else if (own_facts::isCall(op))
+                {
+                    if (auto callee = calleeOf(op); !callee.empty())
+                    {
+                        callees[name].push_back(callee);
+                    }
+                    else
+                    {
+                        assignsAny.insert(name);
+                    }
+                }
+            });
+        });
+
+        for (auto changed = true; changed;)
+        {
+            changed = false;
+            for (auto &entry : callees)
+            {
+                auto caller = entry.getKey();
+                for (auto callee : entry.getValue())
+                {
+                    if (assignsAny.contains(callee) && assignsAny.insert(caller).second)
+                    {
+                        changed = true;
+                    }
+
+                    if (auto it = assigns.find(callee); it != assigns.end())
+                    {
+                        // copied first: `assigns[caller]` may move the entry being read
+                        llvm::SmallVector<mlir::StringRef> globals;
+                        for (auto &global : it->getValue())
+                        {
+                            globals.push_back(global.getKey());
+                        }
+
+                        for (auto global : globals)
+                        {
+                            changed |= assigns[caller].insert(global).second;
+                        }
+                    }
+                }
+            }
+        }
+
+        auto mayAssign = [&](mlir::Operation *op, mlir::StringRef global) {
+            auto callee = calleeOf(op);
+            if (callee.empty() || assignsAny.contains(callee))
+            {
+                return true;
+            }
+
+            auto it = assigns.find(callee);
+            return it != assigns.end() && it->getValue().contains(global);
         };
 
         llvm::SmallVector<std::pair<mlir::Operation *, mlir::Value>> held;
@@ -350,12 +469,26 @@ class OwnedReturnConsumptionPass
                     }
                 }
 
-                if (mth.ownsHeapMemory(op->getLoc(), argument.getType()) &&
-                    !MLIRTypeHelper::isSharedHandleType(argument.getType()) && readsGlobal(argument) &&
-                    seen.insert(argument).second)
+                if (!mth.ownsHeapMemory(op->getLoc(), argument.getType()) || !seen.insert(argument).second)
                 {
-                    held.push_back({op, argument});
+                    continue;
                 }
+
+                auto global = globalOf(argument, false);
+                auto isHandle = MLIRTypeHelper::isSharedHandleType(argument.getType());
+                if (global.empty() || (isHandle ? !mayAssign(op, global) : !compileOptions.isRefCounted()))
+                {
+                    continue;
+                }
+
+                // under own, a handle read out of a place in a global is a borrow of that place,
+                // rejected when the call may overwrite it; only the global itself is held
+                if (!compileOptions.isRefCounted() && !readsGlobalItself(argument))
+                {
+                    continue;
+                }
+
+                held.push_back({op, argument});
             }
         });
 

tslang/test/tester/CMakeLists.txt

@@ -388,6 +388,7 @@ tslang_add_test(NAME test-compile-00-string-empty-falsy COMMAND test-runner "${P
 tslang_add_test(NAME test-compile-00-string-concat-bounds COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00string_concat_bounds.ts")
 tslang_add_test(NAME test-compile-00-parse-number-range COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00parse_number_range.ts")
 tslang_add_test(NAME test-compile-00-global-argument-held COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00global_argument_held.ts")
+tslang_add_test(NAME test-compile-00-shared-global-field-argument COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00shared_global_field_argument.ts")
 tslang_add_test(NAME test-compile-00-narrowed-assign-other-member COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00narrowed_assign_other_member.ts")
 tslang_add_test(NAME test-compile-00-bound-method-count COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00bound_method_count.ts")
 tslang_add_test(NAME test-compile-00-const-record-owned-fields COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00const_record_owned_fields.ts")
@@ -960,6 +961,7 @@ tslang_add_test(NAME test-jit-00-string-empty-falsy COMMAND test-runner -jit "${
 tslang_add_test(NAME test-jit-00-string-concat-bounds COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00string_concat_bounds.ts")
 tslang_add_test(NAME test-jit-00-parse-number-range COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00parse_number_range.ts")
 tslang_add_test(NAME test-jit-00-global-argument-held COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00global_argument_held.ts")
+tslang_add_test(NAME test-jit-00-shared-global-field-argument COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00shared_global_field_argument.ts")
 tslang_add_test(NAME test-jit-00-narrowed-assign-other-member COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00narrowed_assign_other_member.ts")
 tslang_add_test(NAME test-jit-00-bound-method-count COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00bound_method_count.ts")
 tslang_add_test(NAME test-jit-00-const-record-owned-fields COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00const_record_owned_fields.ts")
@@ -1922,6 +1924,7 @@ set(TSLANG_CORPUS
     00string_concat_bounds.ts
     00parse_number_range.ts
     00global_argument_held.ts
+    00shared_global_field_argument.ts
     00narrowed_assign_other_member.ts
     00bound_method_count.ts
     00const_record_owned_fields.ts
@@ -3585,7 +3588,8 @@ set(shared_models gc none rc own)
 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_place_two_uses own_shared_global_field own_shared_global_copy
+                        own_shared_global_argument)
         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()
@@ -3599,7 +3603,8 @@ endforeach()
 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_place_two_uses own_shared_global_field own_shared_global_copy
+                        own_shared_global_argument)
         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.ts

@@ -0,0 +1,72 @@
+// A Shared<T> handle read out of a global and given to a call that gives the global a new value:
+// the parameter is a borrow, so the global's release freed the block the callee still read (rc and
+// own alike). The caller now holds a count across such a call - and only such a call, so
+// `Shared.count` elsewhere is still the number of handles. Counts are checked where there is one
+// (rc, own); gc and none report -1. A handle read out of a global's field, which own rejects, is in
+// 00shared_global_field_argument.ts.
+class Node {
+    v = 0;
+    tag = "";
+}
+
+function churn() {
+    let keep: Node[] = [];
+    for (let i = 0; i < 1000; i++) {
+        const n = new Node();
+        n.v = -i;
+        n.tag = "junk" + i;
+        keep.push(n);
+    }
+
+    return keep.length;
+}
+
+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 bare = new Shared(mk("bare", 2));
+
+// assigns the global it was given
+function run(h: Shared<Node>, v: number) {
+    bare = new Shared(mk("replaced", 5));
+    assert(churn() == 1000);
+    assert(countIs(h, 1), "the caller's hold is the only count left");
+    return h.value.v == v;
+}
+
+// assigns it through a callee
+function replaceBare() {
+    bare = new Shared(mk("again", 6));
+}
+
+function runThrough(h: Shared<Node>) {
+    replaceBare();
+    assert(churn() == 1000);
+    return h.value.v == 5 && h.value.tag == "replaced";
+}
+
+function main() {
+    assert(countIs(bare, 1), "a call that does not assign the global takes no count");
+    assert(run(bare, 2), "a handle given to a call that assigns its global");
+    assert(runThrough(bare), "a handle given to a call that assigns its global through a callee");
+
+    let ok = false;
+    try {
+        ok = run(bare, 6);
+    } catch (e) {
+        print("caught");
+    }
+
+    assert(ok, "a handle given in a try body");
+    assert(countIs(bare, 1) && bare.value.tag == "replaced", "the global's own count after the calls");
+    print("done.");
+}

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

@@ -0,0 +1,61 @@
+// A Shared<T> handle read out of a global's field and given to a call that gives the field a new
+// value, and a value read through a global handle given to a call that replaces it: the caller
+// holds a count across the call under rc, so the block the parameter reads stays alive. (Own rejects the call: the argument borrows the field.) Counts are checked where
+// there is one (rc); gc and none report -1.
+class Node {
+    v = 0;
+    tag = "";
+}
+
+class Holder {
+    s = new Shared(mk("field", 7));
+}
+
+function churn() {
+    let keep: Node[] = [];
+    for (let i = 0; i < 1000; i++) {
+        const n = new Node();
+        n.v = -i;
+        n.tag = "junk" + i;
+        keep.push(n);
+    }
+
+    return keep.length;
+}
+
+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;
+}
+
+const holder = new Holder();
+const G = new Shared(mk("first", 1));
+
+function runField(h: Shared<Node>) {
+    holder.s = new Shared(mk("field2", 8));
+    assert(churn() == 1000);
+    assert(countIs(h, 1), "the caller's hold is the only count left");
+    return h.value.v == 7 && h.value.tag == "field";
+}
+
+// a value read through a global handle, given to a call that replaces it
+function runValue(n: Node) {
+    G.value = mk("second", 2);
+    assert(churn() == 1000);
+    return n.v == 1 && n.tag == "first";
+}
+
+function main() {
+    assert(countIs(holder.s, 1), "a call that does not assign the field takes no count");
+    assert(runField(holder.s), "a handle given to a call that assigns the field it was read from");
+    assert(countIs(holder.s, 1) && holder.s.value.v == 8, "the field's own count after the call");
+    assert(runValue(G.value), "a value read through a global handle given to a call that replaces it");
+    print("done.");
+}