-mm=own: calls in a global's initializer get their callees' facts

tslang/lib/TypeScript/OwnershipSignaturePass.cpp

@@ -251,6 +251,7 @@ class OwnershipSignaturePass : public mlir::PassWrapper<OwnershipSignaturePass,
         }
 
         exportNoDrops();
+        ownGlobalInitializerCalls(module);
 
         module.walk([&](mlir_ts::CreateBoundFunctionOp boundOp) {
             if (!boundOp->hasAttr(OWNS_CAPTURE_ATTR_NAME))
@@ -2353,6 +2354,107 @@ class OwnershipSignaturePass : public mlir::PassWrapper<OwnershipSignaturePass,
                llvm::any_of(callees.funcs, [&](mlir_ts::FuncOp callee) { return !noDrops.contains(callee); });
     }
 
+    // ---- Calls in a global's initializer ----
+
+    // A global's region is no function: its calls get no facts pinned, and the inference never sees
+    // it, so its calls were read with the defaults - an owned result, borrowed arguments. A result
+    // that borrows an argument (`const c = child(n)`, `id(new V(1))`) then made the global a second
+    // owner of a block its argument's owner keeps, or kept one the region freed; and a fresh
+    // argument a callee keeps (`new N(new V(4))`, the constructor keeping `c`) was released at the
+    // end of the region, under the object that kept it. A kept argument made in the region and used
+    // by nothing after the call moves into the callee: its release goes. Anything else is an error.
+    void ownGlobalInitializerCalls(mlir::ModuleOp module)
+    {
+        module.walk([&](mlir_ts::GlobalOp globalOp) {
+            // collected first: a move erases a release the walk has not reached yet
+            llvm::SmallVector<mlir::Operation *> calls;
+            globalOp.walk([&](mlir::Operation *op) {
+                if (isCall(op))
+                {
+                    calls.push_back(op);
+                }
+            });
+
+            for (auto *op : calls)
+            {
+                auto callees = resolve(op);
+                auto facts = callFacts(callees);
+                auto calleeName = callees.funcs.empty() ? llvm::StringRef("the callee") : callees.funcs.front().getSymName();
+                if (facts.borrows >= 0)
+                {
+                    op->emitError("the initializer of global '")
+                        << globalOp.getSymName() << "' holds the result of '" << calleeName
+                        << "', which borrows its argument; a global cannot hold a borrow under -mm=own";
+                    signalPassFailure();
+                    continue;
+                }
+
+                auto args = callArgs(op);
+                llvm::SmallPtrSet<mlir::Value, 4> given;
+                for (auto index : facts.owned)
+                {
+                    if (index < 0 || index >= static_cast<int32_t>(args.size()) ||
+                        !given.insert(rootOf(args[index])).second || !moveIntoCallee(op, args[index]))
+                    {
+                        op->emitError("the initializer of global '")
+                            << globalOp.getSymName() << "' gives '" << calleeName
+                            << "' a value it keeps, but -mm=own cannot move it there";
+                        signalPassFailure();
+                        break;
+                    }
+                }
+            }
+        });
+    }
+
+    // `value` was made in the call's block and nothing but its release, if it has one, comes after
+    // the call: the release goes, and the callee owns the block. Constant data (a string literal)
+    // is nobody's, and is given as it is.
+    static bool moveIntoCallee(mlir::Operation *call, mlir::Value value)
+    {
+        auto root = rootOf(value);
+        if (isConstantData(root))
+        {
+            return true;
+        }
+
+        if (!isFresh(root) || root.getDefiningOp()->getBlock() != call->getBlock())
+        {
+            return false;
+        }
+
+        mlir::Operation *release = nullptr;
+        auto usedAfter = false;
+        forEachUse(root, [&](mlir::Operation *user, mlir::Value) {
+            if (user == call)
+            {
+                return;
+            }
+
+            if (user->getBlock() != call->getBlock())
+            {
+                usedAfter = true;
+            }
+            else if (call->isBeforeInBlock(user))
+            {
+                usedAfter = usedAfter || release || !mlir::isa<mlir_ts::ReleaseOp>(user);
+                release = mlir::isa<mlir_ts::ReleaseOp>(user) ? user : release;
+            }
+        });
+
+        if (usedAfter)
+        {
+            return false;
+        }
+
+        if (release)
+        {
+            release->erase();
+        }
+
+        return true;
+    }
+
     // ---- Export ----
 
     // A library says, beside its memory-model marker, which of its exported functions destroy

tslang/test/tester/CMakeLists.txt

@@ -3503,6 +3503,13 @@ foreach(global_mm rc none gc)
     tslang_add_test(NAME test-compile-${global_mm}-own-global-script COMMAND test-runner -mm=${global_mm} "${PROJECT_SOURCE_DIR}/test/tester/own/own_global_script.ts")
 endforeach()
 
+# A global made by a call that keeps an argument made for it owns that argument through what it was
+# given; the same program under every model.
+foreach(kept_mm own rc none gc)
+    tslang_add_test(NAME test-jit-${kept_mm}-own-global-kept-argument COMMAND test-runner -jit -mm=${kept_mm} "${PROJECT_SOURCE_DIR}/test/tester/own/own_global_kept_argument.ts")
+    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 `let` given a new value after its value moved owns the new one (#528); the same program under
 # rc and gc (not none: nothing is freed there, and the loops would hold about 800 MB).
 foreach(reassign_mm rc gc)
@@ -3700,6 +3707,9 @@ set(own_error_cases
     "own_err_param_assigned|is a parameter and cannot be assigned under -mm=own yet"
     "own_err_local_owns_nothing|owns nothing under -mm=own .* so it cannot be given a new value here"
     "own_err_merged_result|is merged from several branches"
+    "own_err_global_borrowed_result|the initializer of global 'c' holds the result of 'child', which borrows its argument"
+    "own_err_global_borrowed_temporary|the initializer of global 'c' holds the result of 'id', which borrows its argument"
+    "own_err_global_kept_not_fresh|the initializer of global 'n' gives 'N.constructor' a value it keeps, but -mm=own cannot move it there"
     "own_err_global_returned|borrows the global 'V.def' and cannot be stored, returned or captured"
     "own_err_global_let_call|borrows the global 'names', but -mm=own cannot prove the global keeps its value"
     "own_err_captured_returned|borrows a captured variable and cannot be stored, returned or captured"

tslang/test/tester/own/own_err_global_borrowed_result.ts

@@ -0,0 +1,22 @@
+// -mm=own rejects: `child` returns a borrow of its argument's field, so `c` would be a second owner
+// of the block `n.c` owns, and `n.c = new V(5)` frees what `c` still reads. A global cannot borrow.
+class V {
+    constructor(public x: number) {}
+}
+
+class N {
+    constructor(public c: V) {}
+}
+
+const n = new N(new V(4));
+
+function child(m: N) {
+    return m.c;
+}
+
+const c = child(n);
+
+function main() {
+    n.c = new V(5);
+    print(c.x, n.c.x);
+}

tslang/test/tester/own/own_err_global_borrowed_temporary.ts

@@ -0,0 +1,15 @@
+// -mm=own rejects: `id` returns a borrow of its argument, a temporary the initializer frees at its
+// end, so `c` would read a freed block. A global cannot borrow.
+class V {
+    constructor(public x: number) {}
+}
+
+function id(v: V) {
+    return v;
+}
+
+const c = id(new V(1));
+
+function main() {
+    print(c.x);
+}

tslang/test/tester/own/own_err_global_kept_not_fresh.ts

@@ -0,0 +1,16 @@
+// -mm=own rejects: the constructor keeps `c`, but what it is given is the global `v`, which keeps
+// its value, so `n` and `v` would both own one block. Only a value made for the call moves.
+class V {
+    constructor(public x: number) {}
+}
+
+class N {
+    constructor(public c: V) {}
+}
+
+const v = new V(1);
+const n = new N(v);
+
+function main() {
+    print(n.c.x, v.x);
+}

tslang/test/tester/own/own_global_kept_argument.ts

@@ -0,0 +1,45 @@
+// A global made by a call that keeps an argument made for it (`new N(new V(4))`, the constructor
+// keeping `c`; a function that stores its parameter in what it returns) owns that argument through
+// what it was given. Under own the argument used to be released at the end of the global's
+// initializer, under the object that kept it. Each read follows allocations of the same size, so a
+// freed block is reused.
+class V {
+    tag = "";
+    constructor(public x: number) {
+        this.tag = "v" + x;
+    }
+}
+
+class N {
+    constructor(public c: V) {}
+}
+
+function wrap(v: V) {
+    const n = new N(v);
+    return n;
+}
+
+class Junk {
+    a = -1;
+    b = -2;
+    tag = "junk";
+}
+
+function churn() {
+    let keep: Junk[] = [];
+    for (let i = 0; i < 1000; i++) keep.push(new Junk());
+    return keep.length;
+}
+
+const n = new N(new V(4));
+let w = wrap(new V(5));
+
+function main() {
+    churn();
+    assert(n.c.x == 4 && n.c.tag == "v4", "a constructor's kept argument");
+    assert(w.c.x == 5 && w.c.tag == "v5", "a function's kept argument");
+    w = wrap(new V(6));
+    churn();
+    assert(w.c.x == 6 && n.c.x == 4, "the global given a new value");
+    print("done.");
+}