-mm=own: a read of a global ends at a call that may assign the global

tslang/lib/TypeScript/OwnershipFacts.h

@@ -63,6 +63,11 @@ namespace mlir_ts = mlir::typescript;
 // result is the caller's to own, but no use of it may come after any of these arguments is given
 // back, and it may not be kept anywhere that outlives them.
 #define OWN_RESULT_BOUNDED_ATTR_NAME "__own_result_bounded"
+// The globals a call may give a new value, by name: a store to one, or its release, in a callee or
+// in a call it makes. A call that may drop and whose callees are not known may assign any of them
+// (OWN_ASSIGNS_ANY_GLOBAL_ATTR_NAME). A read of a global ends at such a call (checkGlobalRead).
+#define OWN_ASSIGNS_GLOBALS_ATTR_NAME "__own_assigns_globals"
+#define OWN_ASSIGNS_ANY_GLOBAL_ATTR_NAME "__own_assigns_any_global"
 // An immortal global (spec 25): a class-typed global of this module, not exported, whose only
 // write is its initializer, which gives it an object made there (not a call's result) or another
 // immortal global's. Like a

tslang/lib/TypeScript/OwnershipInferencePass.cpp

@@ -4305,8 +4305,9 @@ class OwnershipInferencePass
     // Such a read is a borrow of the value the global owns - `const c = K.w` names the read itself,
     // with no retain - and an assignment of the global releases it. So no use of the read, nor of
     // what it produces (walkBorrowed), may run after one in this function. A path back to a use
-    // through the read itself is a new read. A call that assigns the global is #639's at its
-    // arguments. (C1 of the fix-up review: `const c = K.w; K.w = new V(i + 1); print(c.x)` in a loop
+    // through the read itself is a new read. A call that may assign the global, itself or through
+    // its callees, is such an assignment too (follow-up 11: `const c = G; reset(); print(c.x)`); the
+    // read given to that call as an argument is #639's. (C1 of the fix-up review: `const c = K.w; K.w = new V(i + 1); print(c.x)` in a loop
     // read the block the assignment freed.)
     //
     // Where the walk meets a keep, the read is either followed into what holds it or rejected -
@@ -4477,6 +4478,19 @@ class OwnershipInferencePass
             {
                 ends.push_back(op);
             }
+
+            // a call that may give the global a new value (OwnershipSignaturePass, pinAssignedGlobals):
+            // `const c = G; reset(); print(c.x)`
+            if (isCall(op))
+            {
+                auto assigned = op->getAttrOfType<mlir::ArrayAttr>(OWN_ASSIGNS_GLOBALS_ATTR_NAME);
+                if (op->hasAttr(OWN_ASSIGNS_ANY_GLOBAL_ATTR_NAME) ||
+                    (assigned && llvm::any_of(assigned.getAsRange<mlir::StringAttr>(),
+                                              [&](mlir::StringAttr name) { return name.getValue() == global; })))
+                {
+                    ends.push_back(op);
+                }
+            }
         });
 
         for (auto *end : ends)

tslang/lib/TypeScript/OwnershipSignaturePass.cpp

@@ -453,6 +453,7 @@ class OwnershipSignaturePass : public mlir::PassWrapper<OwnershipSignaturePass,
         findOpen(module);
         computeFacts();
         computeDrops();
+        pinAssignedGlobals();
         rejectGlobalArgumentsDropped();
         findBorrowingStates(module);
         findBoundedForwarders();
@@ -5523,6 +5524,116 @@ class OwnershipSignaturePass : public mlir::PassWrapper<OwnershipSignaturePass,
                llvm::any_of(callees.funcs, [&](mlir_ts::FuncOp callee) { return !noDrops.contains(callee); });
     }
 
+    // ---- Globals a call may assign ----
+
+    // The globals each call may give a new value, pinned on the call for the inference
+    // (OWN_ASSIGNS_GLOBALS_ATTR_NAME, OWN_ASSIGNS_ANY_GLOBAL_ATTR_NAME): a store to a global, or its
+    // release, in a callee's body, or in a call it makes. A call that may drop and whose callees are
+    // not known may assign any; one that drops nothing assigns none. A read of a global ends where
+    // the global is given a new value (checkGlobalRead), and a call that does so ends it as surely as
+    // a store in the same function: `const c = G; reset(); print(c.x)`.
+    void pinAssignedGlobals()
+    {
+        llvm::DenseMap<mlir::Operation *, llvm::SmallVector<Call *>> callsIn;
+        for (auto &call : calls)
+        {
+            callsIn[call.op->getParentOfType<mlir_ts::FuncOp>()].push_back(&call);
+        }
+
+        llvm::DenseMap<mlir::Operation *, llvm::SetVector<mlir::StringRef>> assigns;
+        llvm::DenseSet<mlir::Operation *> assignsAny;
+        for (auto &entry : functions)
+        {
+            auto funcOp = entry.second;
+            funcOp.walk([&](mlir::Operation *op) {
+                mlir::Value slot;
+                if (auto storeOp = mlir::dyn_cast<mlir_ts::StoreOp>(op))
+                {
+                    slot = storeOp.getReference();
+                }
+                else if (auto releaseSlotOp = mlir::dyn_cast<mlir_ts::ReleaseSlotOp>(op))
+                {
+                    slot = releaseSlotOp.getSlot();
+                }
+
+                if (auto addressOfOp = slot ? slot.getDefiningOp<mlir_ts::AddressOfOp>() : mlir_ts::AddressOfOp())
+                {
+                    assigns[funcOp].insert(addressOfOp.getGlobalName());
+                }
+            });
+        }
+
+        // what a call may assign, from what its callees do so far
+        auto assignedBy = [&](Call &call, llvm::SetVector<mlir::StringRef> &globals) {
+            if (!callMayDrop(call.callees))
+            {
+                return false;
+            }
+
+            if (!call.callees.known)
+            {
+                return true;
+            }
+
+            auto any = false;
+            for (auto callee : call.callees.funcs)
+            {
+                any |= assignsAny.contains(callee);
+                if (auto found = assigns.find(callee); found != assigns.end())
+                {
+                    globals.insert(found->second.begin(), found->second.end());
+                }
+            }
+
+            return any;
+        };
+
+        for (auto changed = true; changed;)
+        {
+            changed = false;
+            for (auto &entry : functions)
+            {
+                auto funcOp = entry.second;
+                for (auto *call : callsIn[funcOp])
+                {
+                    llvm::SetVector<mlir::StringRef> globals;
+                    if (assignedBy(*call, globals))
+                    {
+                        changed |= assignsAny.insert(funcOp).second;
+                    }
+
+                    for (auto global : globals)
+                    {
+                        changed |= assigns[funcOp].insert(global);
+                    }
+                }
+            }
+        }
+
+        for (auto &call : calls)
+        {
+            llvm::SetVector<mlir::StringRef> globals;
+            if (assignedBy(call, globals))
+            {
+                call.op->setAttr(OWN_ASSIGNS_ANY_GLOBAL_ATTR_NAME, mlir::UnitAttr::get(&getContext()));
+                continue;
+            }
+
+            if (globals.empty())
+            {
+                continue;
+            }
+
+            llvm::SmallVector<mlir::Attribute> names;
+            for (auto global : globals)
+            {
+                names.push_back(mlir::StringAttr::get(&getContext(), global));
+            }
+
+            call.op->setAttr(OWN_ASSIGNS_GLOBALS_ATTR_NAME, mlir::ArrayAttr::get(&getContext(), names));
+        }
+    }
+
     // An argument read out of a global (`run(G)`, `G.m()`) is lent to the callee, and the callee may
     // free it under its own parameter: `function run(h: C) { G = mk(); print(h.x); }`. The caller's
     // borrow ends at the call that takes it, so nothing in the caller sees that, and in the callee

tslang/test/tester/CMakeLists.txt

@@ -3700,6 +3700,12 @@ foreach(kept_mm own rc none gc)
     tslang_add_test(NAME test-compile-${kept_mm}-own-global-kept-argument COMMAND test-runner -mm=${kept_mm} "${PROJECT_SOURCE_DIR}/test/tester/own/own_global_kept_argument.ts")
 endforeach()
 
+# A read of a global used after calls that give only other globals new values; the same program under
+# every model.
+foreach(read_mm own rc none gc)
+    tslang_add_test(NAME test-jit-${read_mm}-own-global-read-callee-other COMMAND test-runner -jit -mm=${read_mm} "${PROJECT_SOURCE_DIR}/test/tester/own/own_global_read_callee_other.ts")
+    tslang_add_test(NAME test-compile-${read_mm}-own-global-read-callee-other COMMAND test-runner -mm=${read_mm} "${PROJECT_SOURCE_DIR}/test/tester/own/own_global_read_callee_other.ts")
+endforeach()
 # An interface-typed global given a fresh class object owns it; the same program under every model.
 foreach(iface_mm own rc none gc)
     tslang_add_test(NAME test-jit-${iface_mm}-own-global-interface-fresh COMMAND test-runner -jit -mm=${iface_mm} "${PROJECT_SOURCE_DIR}/test/tester/own/own_global_interface_fresh.ts")
@@ -3991,6 +3997,8 @@ set(own_error_cases
     "own_err_global_argument_assigned|a value read out of the global 'G' is lent to a call that may release or overwrite it"
     "own_err_global_method_assigns|a value read out of the global 'G' is lent to a call that may release or overwrite it"
     "own_err_global_argument_transitive|a value read out of the global 'G' is lent to a call that may release or overwrite it"
+    "own_err_global_read_callee_assigns|borrows 'G' but is used here after 'G' is released or overwritten"
+    "own_err_global_read_callee_assigns_transitive|borrows 'G' but is used here after 'G' is released or overwritten"
     "own_err_argument_assigned|is lent to 'g', but is given a new value before the call takes it"
     "own_err_argument_assigned_conditional|is lent to 'g', but is given a new value before the call takes it"
     "own_err_argument_assigned_closure|is lent to 'g', but is given a new value before the call takes it"

tslang/test/tester/own/own_err_global_read_callee_assigns.ts

@@ -0,0 +1,17 @@
+// -mm=own rejects: `c` borrows what `G` holds, and `reset` gives `G` a new value, which frees that
+// object while `c` still reads it.
+class V {
+    constructor(public x: number) {}
+}
+
+let G = new V(0);
+
+function reset() {
+    G = new V(1);
+}
+
+function main() {
+    const c = G;
+    reset();
+    print(c.x);
+}

tslang/test/tester/own/own_err_global_read_callee_assigns_transitive.ts

@@ -0,0 +1,21 @@
+// -mm=own rejects: `c` borrows what `G` holds, and `step` calls `reset`, which gives `G` a new value
+// and frees that object while `c` still reads it.
+class V {
+    constructor(public x: number) {}
+}
+
+let G = new V(0);
+
+function reset() {
+    G = new V(1);
+}
+
+function step() {
+    reset();
+}
+
+function main() {
+    const c = G;
+    step();
+    print(c.x);
+}

tslang/test/tester/own/own_global_read_callee_other.ts

@@ -0,0 +1,34 @@
+// A read of a global is used after calls that give other globals new values, or none: -mm=own accepts
+// it, and it reads the global's own value.
+class V {
+    tag = "";
+    constructor(public x: number) {
+        this.tag = "v" + x;
+    }
+}
+
+let G = new V(1);
+let H = new V(2);
+
+function resetH() {
+    H = new V(3);
+}
+
+function sum(a: number, b: number) {
+    return a + b;
+}
+
+function readG() {
+    return G.x;
+}
+
+function main() {
+    const c = G;
+    resetH();
+    const s = sum(c.x, readG());
+    assert(c.x == 1 && c.tag == "v1" && s == 2, "the global's value, read before the calls");
+    assert(H.x == 3, "the other global's new value");
+    G = new V(4);
+    assert(G.x == 4, "the global's new value");
+    print("done.");
+}