rc: hold a global argument from its read, and give a const read of a global a count

tslang/lib/TypeScript/MLIRGenImpl.h

@@ -1737,6 +1737,18 @@ class MLIRGenImpl
             return loadOp && loadOp.getReference().getDefiningOp<mlir_ts::PropertyRefOp>();
         }
 
+        // A value read out of a global variable, through any casts: what `G` is in `const n = G`.
+        static bool readsGlobal(mlir::Value value)
+        {
+            while (auto castOp = value ? value.getDefiningOp<mlir_ts::CastOp>() : mlir_ts::CastOp())
+            {
+                value = castOp.getIn();
+            }
+
+            auto loadOp = value ? value.getDefiningOp<mlir_ts::LoadOp>() : mlir_ts::LoadOp();
+            return loadOp && loadOp.getReference().getDefiningOp<mlir_ts::AddressOfOp>();
+        }
+
         // A value read out of an element, through any casts: what `a[0]` is in `const c = a[0]`, and
         // what `b` is in `const [b] = a`, read there only when `a` is long enough (the `then` of an if).
         static bool readsElement(mlir::Value value)
@@ -1886,11 +1898,12 @@ class MLIRGenImpl
             // interface or a bound method made over such a read (`const i: I = a`, `const m = a.get`):
             // it holds the block as `b` would.
             //
-            // Under rc, so does a const read out of a field (`const c = h.c`) or an element (`const c =
-            // a[0]`, `const [c] = a`): folded, `h.c = d` or `a[0] = d` freed what `c` still read. So
-            // does an interface or a bound method made over such a read (`const i: I = h.v`,
-            // `const m = a[0].get`), which holds the block as `c` would. Under own the folded read is a
-            // borrow, and a use after the overwrite is an error there.
+            // Under rc, so does a const read out of a field (`const c = h.c`), an element (`const c =
+            // a[0]`, `const [c] = a`) or a global (`const c = G`): folded, `h.c = d`, `a[0] = d` or
+            // `G = d` - here or in a call - freed what `c` still read. So does an interface or a bound
+            // method made over such a read (`const i: I = h.v`, `const m = a[0].get`), which holds the
+            // block as `c` would. Under own the folded read is a borrow, and a use after the overwrite
+            // is an error there.
             //
             // And a const record whose fields own blocks (`const o = { s: p.child }`): the literal
             // hands its fields over unretained - its receiver is the one that retains them, as a
@@ -1910,8 +1923,9 @@ class MLIRGenImpl
                 || ((isa<mlir_ts::ArrayType>(type) || MLIRTypeHelper::isSharedHandleType(type)
                      || ((readsLocalSlot(initial) || viewsLocalSlot(initial)) && mth.ownsHeapMemory(location, type))
                      || (compileOptions.isRefCounted()
-                         && (readsField(initial) || readsElement(initial) || readsField(viewedValue(initial))
-                             || readsElement(viewedValue(initial)))
+                         && (readsField(initial) || readsElement(initial) || readsGlobal(initial)
+                             || readsField(viewedValue(initial)) || readsElement(viewedValue(initial))
+                             || readsGlobal(viewedValue(initial)))
                          && mth.ownsHeapMemory(location, type))
                      || (readsLiteralSlot(mth, location, initial) && isa<mlir_ts::TupleType>(type)
                          && mth.ownsHeapMemory(location, type))

tslang/lib/TypeScript/OwnedReturnConsumptionPass.cpp

@@ -297,10 +297,10 @@ class OwnedReturnConsumptionPass
     // A call's argument read out of a global (`run(G)`, `run(G.c)`, `G.m()`) is borrowed like any
     // 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 - 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).
+    // block the parameter read. Under rc the caller holds a count across the call - retained where
+    // the argument is read, released right after the call - 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.
@@ -495,12 +495,22 @@ class OwnedReturnConsumptionPass
         for (auto [op, argument] : held)
         {
             auto location = op->getLoc();
+
+            // From the read on, not from the call: an argument after it may give the global a new
+            // value first (`f(G, (G = mk(), x))`), and its release freed the block this one read
+            // before a retain at the call could keep it.
             mlir::OpBuilder builder(op);
+            auto *readOp = argument.getDefiningOp();
+            if (readOp && readOp->getBlock() == op->getBlock())
+            {
+                builder.setInsertionPointAfter(readOp);
+            }
+
             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.
+            // through a slot that holds the value only while the hold lasts: null before the read
+            // and after the call, so the handler's release does nothing when something else threw.
             auto *unwindRegion = unwindRegionOf(op);
             mlir::Value slot;
             if (unwindRegion)

tslang/test/tester/CMakeLists.txt

@@ -402,6 +402,8 @@ tslang_add_test(NAME test-compile-00-global-copy-counted COMMAND test-runner "${
 tslang_add_test(NAME test-compile-00-c-string-copy COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00c_string_copy.ts")
 tslang_add_test(NAME test-compile-00-const-field-read COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00const_field_read.ts")
 tslang_add_test(NAME test-compile-00-const-element-read COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00const_element_read.ts")
+tslang_add_test(NAME test-compile-00-const-global-read COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00const_global_read.ts")
+tslang_add_test(NAME test-compile-00-global-argument-reassigned COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00global_argument_reassigned.ts")
 tslang_add_test(NAME test-compile-00-array-data-string-alias COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00array_data_string_alias.ts")
 tslang_add_test(NAME test-compile-00-arrays7 COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00array7.ts")
 tslang_add_test(NAME test-compile-00-arrays8-tuple-spread COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00array8_tuple_spread.ts")
@@ -978,6 +980,8 @@ tslang_add_test(NAME test-jit-00-global-copy-counted COMMAND test-runner -jit "$
 tslang_add_test(NAME test-jit-00-c-string-copy COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00c_string_copy.ts")
 tslang_add_test(NAME test-jit-00-const-field-read COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00const_field_read.ts")
 tslang_add_test(NAME test-jit-00-const-element-read COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00const_element_read.ts")
+tslang_add_test(NAME test-jit-00-const-global-read COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00const_global_read.ts")
+tslang_add_test(NAME test-jit-00-global-argument-reassigned COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00global_argument_reassigned.ts")
 tslang_add_test(NAME test-jit-00-array-data-string-alias COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00array_data_string_alias.ts")
 tslang_add_test(NAME test-jit-00-arrays7 COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00array7.ts")
 tslang_add_test(NAME test-jit-00-arrays8-tuple-spread COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00array8_tuple_spread.ts")
@@ -1944,6 +1948,8 @@ set(TSLANG_CORPUS
     00c_string_copy.ts
     00const_field_read.ts
     00const_element_read.ts
+    00const_global_read.ts
+    00global_argument_reassigned.ts
     00array_data_string_alias.ts
     00array7.ts
     00array8_tuple_spread.ts

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

@@ -0,0 +1,103 @@
+// A `const` read out of a global keeps what it read alive when the global is given a new value,
+// here or in a call. Folded, the const was the read itself and held no reference of its own: under
+// rc the overwrite freed the block it still read, and a value read through it and lent to a call
+// that assigns the global (`run(n.v)`) was retained and released after the block was gone.
+// Allocations of the same size reuse anything freed before each read.
+class V {
+    tag = "";
+    constructor(public x: number) {
+        this.tag = "v" + x;
+    }
+
+    get() {
+        return this.x;
+    }
+}
+
+class N {
+    v: V;
+    constructor(x: number) {
+        this.v = new V(x);
+    }
+}
+
+interface IV {
+    x: number;
+    get(): number;
+}
+
+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;
+}
+
+let G = new N(1);
+let GV = new V(1);
+
+function assignG(x: number) {
+    G = new N(x);
+    churn();
+}
+
+// assigns the global its argument was read out of
+function run(v: V, x: number) {
+    assignG(100 + x);
+    return v.x == x && v.tag == "v" + x;
+}
+
+function reassignedHere(i: number) {
+    G = new N(i);
+    const n = G;
+    G = new N(9);
+    churn();
+    assert(n.v.x == i && n.v.tag == "v" + i, "a const read of a global, reassigned here");
+}
+
+function reassignedInCall(i: number) {
+    G = new N(i);
+    const n = G;
+    assignG(9);
+    assert(n.v.x == i && n.v.get() == i, "a const read of a global, reassigned in a call");
+}
+
+function lentFromConst(i: number) {
+    G = new N(i);
+    const n = G;
+    for (let j = 0; j < 3; j++) {
+        assert(run(n.v, i), "a field of a const read of a global, lent to a call that assigns it");
+        try {
+            churn();
+        } finally {
+            assert(n.v.tag == "v" + i, "read in a finally");
+        }
+    }
+}
+
+function views(i: number) {
+    GV = new V(i);
+    const iv: IV = GV;
+    const m = GV.get;
+    GV = new V(9);
+    churn();
+    assert(iv.x == i && iv.get() == i && m() == i, "an interface and a bound method over a global");
+}
+
+function main() {
+    for (let i = 2; i < 5; i++) {
+        reassignedHere(i);
+        reassignedInCall(i);
+        lentFromConst(i);
+        views(i);
+    }
+
+    churn();
+    assert(G.v.x > 100 && GV.x == 9, "the globals' last values");
+    print("done.");
+}

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

@@ -0,0 +1,81 @@
+// An argument read out of a global, or out of a global's field, stays alive when a later argument
+// of the same call gives that global a new value: `check(G, (G = mk(2), G), 1)`. Under rc the
+// caller holds a count across a call that may assign the global (00global_argument_held.ts), and
+// took it at the call - after the later argument had already freed the block, so the hold
+// retained and released a freed block. Each read follows allocations of the same size, so a freed
+// block is reused.
+// @strict-null false
+class C {
+    x = 5;
+    tag = "";
+    v: number[] = [];
+}
+
+class H {
+    c: C = mk(3);
+}
+
+function mk(x: number): C {
+    const r = new C();
+    r.x = x;
+    r.tag = "c" + x;
+    r.v.push(x);
+    return r;
+}
+
+function churn() {
+    const junk: C[] = [];
+    for (let i = 0; i < 64; i++) {
+        const c = new C();
+        c.x = 999;
+        c.tag = "junk";
+        c.v.push(99);
+        junk.push(c);
+    }
+}
+
+let G: C = mk(1);
+const GH = new H();
+
+function thrower(): C {
+    throw "thrown";
+}
+
+// assigns both globals, so a caller holds what it reads out of them
+function check(c: C, other: C, x: number) {
+    G = mk(100 + x);
+    GH.c = mk(200 + x);
+    churn();
+    return c.x == x && c.tag == "c" + x && c.v[0] == x && other.x > 0;
+}
+
+function inLoop() {
+    for (let i = 0; i < 3; i++) {
+        const x = G.x;
+        assert(check(G, (G = mk(10 + i), G), x), "a global, in a loop");
+    }
+}
+
+function inTryCatch() {
+    G = mk(30);
+    let caught = false;
+    try {
+        check(G, (G = mk(31), thrower()), 30);
+    } catch (e) {
+        caught = true;
+    }
+
+    churn();
+    assert(caught && G.x == 31 && G.tag == "c31", "a later argument throws after the assignment");
+}
+
+function main() {
+    assert(check(G, (G = mk(2), G), 1), "a global given before an argument that assigns it");
+    const fx = GH.c.x;
+    assert(check(GH.c, (GH.c = mk(4), GH.c), fx), "a global's field given before an argument that assigns it");
+    inLoop();
+    inTryCatch();
+    churn();
+    assert(G.x == 31 && GH.c.x == 701, "the globals' last values");
+    print("done.");
+}