rc: a ternary branch yielding a fresh object widened to a union hands its reference out (#465)

tslang/lib/TypeScript/OwnedReturnConsumptionPass.cpp

@@ -389,9 +389,43 @@ class OwnedReturnConsumptionPass
 
             // Only worth doing when some branch actually carries an unclaimed +1. Where every
             // branch borrows, the receiver's retain is already the right and only answer.
+            // What a branch yields, seen through the casts that only widen it: `new C()` yielded as
+            // `C | null` or `C | undefined` is still that block, so its reference goes out with it
+            // (#465). A cast that makes a new value (a box) carries its own reference instead.
+            auto producerOf = [](mlir::Value value) {
+                while (auto *definingOp = value.getDefiningOp())
+                {
+                    if (definingOp->hasAttr(OWNED_RESULT_ATTR_NAME))
+                    {
+                        break;
+                    }
+
+                    if (auto castOp = mlir::dyn_cast<mlir_ts::CastOp>(definingOp);
+                        castOp && mlir::isa<mlir_ts::UnionType, mlir_ts::OptionalType>(castOp.getType()))
+                    {
+                        value = castOp.getIn();
+                    }
+                    else if (auto optionalValueOp = mlir::dyn_cast<mlir_ts::OptionalValueOp>(definingOp))
+                    {
+                        value = optionalValueOp.getIn();
+                    }
+                    else if (auto createUnionInstanceOp = mlir::dyn_cast<mlir_ts::CreateUnionInstanceOp>(definingOp))
+                    {
+                        value = createUnionInstanceOp.getIn();
+                    }
+                    else
+                    {
+                        break;
+                    }
+                }
+
+                return value;
+            };
+
             auto carriesOwned = [&](mlir_ts::ResultOp resultOp) {
-                auto *definingOp = resultOp->getOperand(0).getDefiningOp();
-                return ::typescript::mayTakeOverReference(resultOp->getOperand(0)) &&
+                auto produced = producerOf(resultOp->getOperand(0));
+                auto *definingOp = produced.getDefiningOp();
+                return ::typescript::mayTakeOverReference(produced) &&
                        definingOp->getParentRegion() == resultOp->getParentRegion();
             };
 
@@ -406,8 +440,9 @@ class OwnedReturnConsumptionPass
                 {
                     // The `ts.If` result now stands for this reference, so the producer's +1 has
                     // an owner and must not also be released as a discarded temporary.
-                    resultOp->getOperand(0).getDefiningOp()->setAttr(OWNED_RESULT_CONSUMED_ATTR_NAME,
-                                                                    mlir::UnitAttr::get(&getContext()));
+                    producerOf(resultOp->getOperand(0))
+                        .getDefiningOp()
+                        ->setAttr(OWNED_RESULT_CONSUMED_ATTR_NAME, mlir::UnitAttr::get(&getContext()));
                     continue;
                 }
 

tslang/test/tester/CMakeLists.txt

@@ -470,6 +470,7 @@ tslang_add_test(NAME test-compile-00-class-structural-extends COMMAND test-runne
 tslang_add_test(NAME test-compile-00-instanceof COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00instanceof.ts")
 tslang_add_test(NAME test-compile-00-instanceof-union COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00instanceof_union.ts")
 tslang_add_test(NAME test-compile-00-nullable-to-any COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00nullable_to_any.ts")
+tslang_add_test(NAME test-compile-00-ternary-new-or-null COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00ternary_new_or_null.ts")
 tslang_add_test(NAME test-compile-00-class COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00class.ts")
 tslang_add_test(NAME test-compile-00-class-new COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00class_new.ts")
 tslang_add_test(NAME test-compile-00-capture-in-new-arguments COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00capture_in_new_arguments.ts")
@@ -954,6 +955,7 @@ tslang_add_test(NAME test-jit-00-class-structural-extends COMMAND test-runner -j
 tslang_add_test(NAME test-jit-00-instanceof COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00instanceof.ts")
 tslang_add_test(NAME test-jit-00-instanceof-union COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00instanceof_union.ts")
 tslang_add_test(NAME test-jit-00-nullable-to-any COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00nullable_to_any.ts")
+tslang_add_test(NAME test-jit-00-ternary-new-or-null COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00ternary_new_or_null.ts")
 tslang_add_test(NAME test-jit-00-class COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00class.ts")
 tslang_add_test(NAME test-jit-00-class-new COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00class_new.ts")
 tslang_add_test(NAME test-jit-00-capture-in-new-arguments COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00capture_in_new_arguments.ts")
@@ -1770,6 +1772,7 @@ set(TSLANG_CORPUS
     00instanceof_namespace_narrowing.ts
     00instanceof_union.ts
     00nullable_to_any.ts
+    00ternary_new_or_null.ts
     00interface_captures.ts
     00interface_conjunction.ts
     00interface_function_typed_field.ts

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

@@ -0,0 +1,53 @@
+// `cond ? new C() : null` hands the new object to whoever takes the result (#465). Under rc the
+// branch released it as a discarded temporary before the merge, so the result pointed at freed
+// memory. Each result is read after other allocations have reused whatever was freed.
+class Node {
+    v = 5;
+    name = "";
+}
+
+function churn() {
+    let keep: Node[] = [];
+    for (let i = 0; i < 1000; i++) {
+        const n = new Node();
+        n.v = -i;
+        n.name = "junk" + i;
+        keep.push(n);
+    }
+
+    return keep.length;
+}
+
+function nodeOrNull(set: boolean): Node | null {
+    return set ? new Node() : null;
+}
+
+function named(set: boolean, name: string): Node | undefined {
+    const n = new Node();
+    n.name = name;
+    return set ? n : undefined;
+}
+
+function main() {
+    const set = true;
+    const n: Node | null = set ? new Node() : null;
+    assert(churn() == 1000);
+    assert(n !== null && n.v == 5, "a local from a ternary keeps its node");
+
+    const m = nodeOrNull(true);
+    assert(churn() == 1000);
+    assert(m !== null && m.v == 5, "a returned ternary keeps its node");
+
+    assert(nodeOrNull(false) === null, "the null side is null");
+
+    const k = named(true, "kept" + 1);
+    assert(churn() == 1000);
+    assert(k !== undefined && k.name == "kept1", "an optional ternary keeps its node");
+
+    for (let i = 0; i < 1000; i++) {
+        const t = i % 2 == 0 ? new Node() : null;
+        if (t) t.v = i;
+    }
+
+    print("done.");
+}