-mm=own: a Shared<T> parameter the body assigns is a counted local of its own

tslang/lib/TypeScript/MLIRGenFunctions.cpp

@@ -1566,8 +1566,13 @@ namespace mlirgen
     // string - starts a local of its own the same way, which own makes from a copy of the caller's string (spec
     // 22.3) and which then owns each string assigned to it. Left as the parameter's slot, what was assigned to it
     // had no owner: it leaked, and `return s` gave back a copy of it - read up to its first NUL, so the default
-    // library's `"".clone().resize(n)` returned a string of one byte that its caller wrote n into. Any other
-    // parameter own refuses to assign (OwnershipInferencePass), and a collector needs none of this.
+    // library's `"".clone().resize(n)` returned a string of one byte that its caller wrote n into.
+    //
+    // -mm=own: a Shared<T> parameter the body assigns is copied too. Own steps aside for Shared handles - their
+    // counts are rc's - so it accepts the assignment, and the slot held a handle that only borrowed the new
+    // value: assigned in a loop, the handle was released at the end of the loop body and a later `p.value` read
+    // freed memory. Any other parameter own refuses to assign (OwnershipInferencePass), and a collector needs
+    // none of this.
     mlir::LogicalResult MLIRGenImpl::mlirGenFunctionOwnAssignedParams(mlir::Location location,
                                                                       FunctionLikeDeclarationBase functionLikeDeclarationBaseAST,
                                                                       FunctionPrototypeDOM::TypePtr funcProto,
@@ -1585,7 +1590,7 @@ namespace mlirgen
         for (const auto &param : funcProto->getParams())
         {
             if (param->getBindingPattern() || !mth.ownsHeapMemory(location, param->getType()) ||
-                (own && !mlir::isa<mlir_ts::StringType>(param->getType())))
+                (own && !mlir::isa<mlir_ts::StringType, mlir_ts::SharedType>(param->getType())))
             {
                 continue;
             }

tslang/test/tester/CMakeLists.txt

@@ -3861,7 +3861,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_argument_assigned
-                        own_shared_global_argument own_shared_global_argument_throw)
+                        own_shared_global_argument own_shared_global_argument_throw own_shared_param_assigned)
         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()
@@ -3876,7 +3876,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_argument_assigned
-                        own_shared_global_argument own_shared_global_argument_throw)
+                        own_shared_global_argument own_shared_global_argument_throw own_shared_param_assigned)
         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_param_assigned.ts

@@ -0,0 +1,79 @@
+// A Shared<T> parameter the body gives a new value - in a loop, on one branch, in a try, then
+// captured by a closure - and reads afterwards. Under own the parameter's slot only borrowed the new
+// handle, which was released at the end of the block that made it, so the read after it saw freed
+// memory. The parameter is a counted local of its own, as under rc. Each read follows allocations
+// of the same size, so a freed block is reused; the caller's handle keeps its count.
+class V {
+    tag = "";
+    constructor(public x: number) {
+        this.tag = "v" + x;
+    }
+}
+
+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;
+}
+
+function check(p: Shared<V>, x: number, what: string) {
+    assert(p.value.x == x && p.value.tag == "v" + x, what);
+}
+
+function inLoop(p: Shared<V>) {
+    for (let i = 0; i < 3; i++) {
+        p = new Shared(new V(10 + i));
+    }
+    churn();
+    check(p, 12, "assigned in a loop");
+    return p.value.x;
+}
+
+function onBranch(p: Shared<V>, b: boolean) {
+    if (b) {
+        p = new Shared(new V(20));
+    }
+    churn();
+    check(p, b ? 20 : 1, "assigned on one branch");
+    return p.value.x;
+}
+
+function inTry(p: Shared<V>) {
+    try {
+        p = new Shared(new V(30));
+    } finally {
+        churn();
+    }
+    check(p, 30, "assigned in a try");
+    return p.value.x;
+}
+
+function captured(p: Shared<V>) {
+    for (let i = 0; i < 2; i++) {
+        p = new Shared(new V(40 + i));
+    }
+    const f = () => p.value.x;
+    churn();
+    check(p, 41, "assigned, then captured");
+    return f();
+}
+
+function main() {
+    const s = new Shared(new V(1));
+    assert(inLoop(s) == 12, "loop");
+    assert(onBranch(s, true) == 20, "branch taken");
+    assert(onBranch(s, false) == 1, "branch not taken");
+    assert(inTry(s) == 30, "try");
+    assert(captured(s) == 41, "closure");
+    churn();
+    check(s, 1, "the caller's handle");
+    const count = Shared.count(s);
+    assert(count == -1 || count == 1, "the caller's handle keeps its one count");
+    print("done.");
+}