rc: a new array given straight to a call is held across it

tslang/lib/TypeScript/OwnedReturnConsumptionPass.cpp

@@ -464,7 +464,9 @@ class OwnedReturnConsumptionPass
             auto value = argument;
             for (;;)
             {
-                if (auto castOp = value.getDefiningOp<mlir_ts::CastOp>())
+                // a constant array cast to an array is copied to the heap there: the copy is the read
+                auto castOp = value.getDefiningOp<mlir_ts::CastOp>();
+                if (castOp && !mlir::isa<mlir_ts::ConstArrayType>(castOp.getIn().getType()))
                 {
                     value = castOp.getIn();
                 }
@@ -523,6 +525,44 @@ class OwnedReturnConsumptionPass
             reads.push_back(value);
         };
 
+        // A new array given straight to a call (`g([mk(1)])`) is born at count 0 (bornUncounted), and
+        // nothing else owns it: a count the callee took and gave back - `const a = p` in a loop, or a
+        // parameter it assigns, copied into an owned local - freed the array the callee still read.
+        // Under rc the caller holds it across the call too, and so gives it back after the call,
+        // which nothing did before.
+        auto arrayMadeFor = [](mlir::Value value, mlir::Operation *call) {
+            auto *definingOp = value.getDefiningOp();
+            auto castOp = mlir::dyn_cast_or_null<mlir_ts::CastOp>(definingOp);
+            auto isNewArray =
+                mlir::isa_and_nonnull<mlir_ts::CreateArrayOp, mlir_ts::NewArrayOp, mlir_ts::NewEmptyArrayOp>(definingOp) ||
+                (castOp && mlir::isa<mlir_ts::ConstArrayType>(castOp.getIn().getType()) &&
+                 mlir::isa<mlir_ts::ArrayType>(value.getType()));
+            if (!isNewArray)
+            {
+                return false;
+            }
+
+            // its one use is the call, through casts
+            while (value.hasOneUse())
+            {
+                auto *user = *value.getUsers().begin();
+                if (user == call)
+                {
+                    return true;
+                }
+
+                auto userCastOp = mlir::dyn_cast<mlir_ts::CastOp>(user);
+                if (!userCastOp)
+                {
+                    return false;
+                }
+
+                value = userCastOp.getRes();
+            }
+
+            return false;
+        };
+
         llvm::SmallVector<std::pair<mlir::Operation *, mlir::Value>> held;
         module.walk([&](mlir::Operation *op) {
             if (!own_facts::isCall(op) || op->hasTrait<mlir::OpTrait::IsTerminator>() ||
@@ -555,6 +595,12 @@ class OwnedReturnConsumptionPass
                     continue;
                 }
 
+                if (compileOptions.isRefCounted() && arrayMadeFor(argument, op))
+                {
+                    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()))

tslang/test/tester/CMakeLists.txt

@@ -405,6 +405,7 @@ tslang_add_test(NAME test-compile-00-const-element-read COMMAND test-runner "${P
 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-global-argument-view COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00global_argument_view.ts")
+tslang_add_test(NAME test-compile-00-array-argument-held COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00array_argument_held.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")
@@ -984,6 +985,7 @@ tslang_add_test(NAME test-jit-00-const-element-read COMMAND test-runner -jit "${
 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-global-argument-view COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00global_argument_view.ts")
+tslang_add_test(NAME test-jit-00-array-argument-held COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00array_argument_held.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")
@@ -1953,6 +1955,7 @@ set(TSLANG_CORPUS
     00const_global_read.ts
     00global_argument_reassigned.ts
     00global_argument_view.ts
+    00array_argument_held.ts
     00array_data_string_alias.ts
     00array7.ts
     00array8_tuple_spread.ts

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

@@ -0,0 +1,83 @@
+// A new array given straight to a call (`g([mk(1)])`) stays alive while the callee takes a count
+// of it and gives it back. Under rc a new array is born at count 0 and nothing else owned it: an
+// alias in a loop or a try (`const a = p`), or the parameter assigned (copied into a local that
+// owns it), freed the array the callee still read. 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 {
+    v: V;
+    constructor(x: number) {
+        this.v = new V(x);
+    }
+}
+
+function mk(x: number) {
+    return new N(x);
+}
+
+function churn() {
+    const junk: N[] = [];
+    for (let i = 0; i < 64; i++) {
+        junk.push(new N(999));
+    }
+}
+
+function aliasInLoop(p: N[]) {
+    for (let i = 0; i < 3; i++) {
+        const a: N[] = p;
+        churn();
+    }
+
+    return p[0].v.x == 1 && p[1].v.tag == "v2";
+}
+
+function aliasInTry(p: N[]) {
+    try {
+        let a: N[] = p;
+    } finally {
+        p.push(mk(3));
+    }
+
+    churn();
+    return p.length == 3 && p[2].v.x == 3 && p[0].v.x == 1;
+}
+
+function assigned(p: N[]) {
+    const first = p[0].v.x;
+    p = [mk(7)];
+    churn();
+    return first == 1 && p[0].v.x == 7;
+}
+
+function passedOn(p: N[]) {
+    const ok = assigned(p);
+    churn();
+    return ok && p[0].v.x == 1 && p[0].v.tag == "v1";
+}
+
+function numbers(p: number[]) {
+    for (let i = 0; i < 3; i++) {
+        const a = p;
+        churn();
+    }
+
+    return p[0] + p[1] + p[2];
+}
+
+function main() {
+    for (let i = 0; i < 3; i++) {
+        assert(aliasInLoop([mk(1), mk(2)]), "an alias of a new array, in a loop");
+        assert(aliasInTry([mk(1), mk(2)]), "an alias of a new array, in a try");
+        assert(assigned([mk(1)]), "a new array whose parameter is assigned");
+        assert(passedOn([mk(1)]), "a new array given on to a call that assigns its parameter");
+        assert(numbers([1, 2, 3]) == 6, "a constant array, copied for the call");
+    }
+
+    print("done.");
+}