rc: a local read for an argument is held when a later argument assigns it

tslang/lib/TypeScript/OwnedReturnConsumptionPass.cpp

@@ -309,6 +309,12 @@ class OwnedReturnConsumptionPass
     // count (`Shared.count`) as the number of handles - which a parameter, a borrow, is not one of.
     // So a handle is held, under both, only across a call that may give the global it was read
     // from a new value, the one case where the hold is what keeps the block alive.
+    //
+    // A read of a local is held too, under rc, where the local is given a new value after the read
+    // and before the call: by a later argument (`f(n, (n = mk(), 0))`, in a conditional one too),
+    // or, for a captured local, by a closure called there. The assignment released the block the
+    // argument read. Under own only a handle is held so, being counted there too; own rejects the
+    // others (OwnershipInferencePass).
     void holdGlobalArguments(MLIRTypeHelper &mth, mlir::ModuleOp module)
     {
         if (!compileOptions.tracksOwnership())
@@ -617,8 +623,18 @@ class OwnedReturnConsumptionPass
                 }
 
                 auto global = globalOf(argument, false);
+                if (global.empty())
+                {
+                    if (localAssignedBefore(op, argument))
+                    {
+                        held.push_back({op, argument});
+                    }
+
+                    continue;
+                }
+
                 auto isHandle = MLIRTypeHelper::isSharedHandleType(argument.getType());
-                if (global.empty() || (isHandle ? !mayAssign(op, global) : !compileOptions.isRefCounted()))
+                if (isHandle ? !mayAssign(op, global) : !compileOptions.isRefCounted())
                 {
                     continue;
                 }
@@ -640,10 +656,13 @@ class OwnedReturnConsumptionPass
 
             // 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.
+            // before a retain at the call could keep it. A conditional argument between them puts
+            // the call in a block of its own, run once for each read unless a loop leads back.
             mlir::OpBuilder builder(op);
             auto *readOp = argument.getDefiningOp();
-            if (readOp && readOp->getBlock() == op->getBlock())
+            if (readOp && (readOp->getBlock() == op->getBlock() ||
+                           (readOp->getParentRegion() == op->getParentRegion() &&
+                            !blockReaches(op->getBlock(), readOp->getBlock()))))
             {
                 builder.setInsertionPointAfter(readOp);
             }
@@ -700,6 +719,91 @@ class OwnedReturnConsumptionPass
         }
     }
 
+    // Whether a read of a local, `argument`, is followed by a new value given to that local (or its
+    // release) before `call` takes it (holdGlobalArguments). Under rc, and for a handle under own.
+    bool localAssignedBefore(mlir::Operation *call, mlir::Value argument)
+    {
+        if (!compileOptions.isRefCounted() && !MLIRTypeHelper::isSharedHandleType(argument.getType()))
+        {
+            return false;
+        }
+
+        auto *read = argument.getDefiningOp();
+        if (!read)
+        {
+            return false;
+        }
+
+        if (auto local = mlir::dyn_cast_or_null<mlir_ts::VariableOp>(own_facts::slotOf(argument)))
+        {
+            auto captured = local.getCaptured().value_or(false);
+            return own_facts::findBetween(read, call, [&](mlir::Operation *op) {
+                       return (captured && own_facts::isCall(op)) ||
+                              own_facts::sameSlot(own_facts::slotWritten(op), local);
+                   }) != nullptr;
+        }
+
+        // a field or an element of an object no local's slot holds (`const h = new H()`, folded):
+        // a store into the same field, or any element, of that object (`f(h.v, (h.v = mk(), 0))`)
+        auto value = argument;
+        while (auto castOp = value.getDefiningOp<mlir_ts::CastOp>())
+        {
+            value = castOp.getIn();
+        }
+
+        auto loadOp = value.getDefiningOp<mlir_ts::LoadOp>();
+        if (!loadOp || !own_facts::isPlace(loadOp.getReference()))
+        {
+            return false;
+        }
+
+        auto baseOf = [](mlir::Operation *place) {
+            auto base = place->getOperand(0);
+            while (auto castOp = base.getDefiningOp<mlir_ts::CastOp>())
+            {
+                base = castOp.getIn();
+            }
+
+            return base;
+        };
+
+        auto *place = loadOp.getReference().getDefiningOp();
+        auto base = baseOf(place);
+        auto propertyRef = mlir::dyn_cast<mlir_ts::PropertyRefOp>(place);
+        return own_facts::findBetween(read, call, [&](mlir::Operation *op) {
+                   auto storeOp = mlir::dyn_cast<mlir_ts::StoreOp>(op);
+                   auto *written = storeOp ? storeOp.getReference().getDefiningOp() : nullptr;
+                   if (!written || written->getName() != place->getName() || baseOf(written) != base)
+                   {
+                       return false;
+                   }
+
+                   auto writtenField = mlir::dyn_cast<mlir_ts::PropertyRefOp>(written);
+                   return !propertyRef || writtenField.getPosition() == propertyRef.getPosition();
+               }) != nullptr;
+    }
+
+    static bool blockReaches(mlir::Block *from, mlir::Block *to)
+    {
+        llvm::SmallPtrSet<mlir::Block *, 16> seen;
+        llvm::SmallVector<mlir::Block *> work{from};
+        while (!work.empty())
+        {
+            auto *block = work.pop_back_val();
+            if (block == to)
+            {
+                return true;
+            }
+
+            if (seen.insert(block).second)
+            {
+                work.append(block->getSuccessors().begin(), block->getSuccessors().end());
+            }
+        }
+
+        return false;
+    }
+
     static mlir::Value stripCasts(mlir::Value value)
     {
         while (auto castOp = value.getDefiningOp<mlir_ts::CastOp>())

tslang/lib/TypeScript/OwnershipFacts.h

@@ -425,6 +425,178 @@ inline bool isPlace(mlir::Value ref)
            isInterfaceField(ref);
 }
 
+// The storage a reference addresses - a local's `ts.Variable` or a global's `ts.AddressOf` - itself
+// or through the places in what it holds (a field, an element, a handle's value): `n`, `n.v`,
+// `G.items[0]`. Null for anything else.
+inline mlir::Operation *slotOfRef(mlir::Value ref)
+{
+    for (auto depth = 0; depth < 16 && ref; depth++)
+    {
+        auto *def = ref.getDefiningOp();
+        if (mlir::isa_and_nonnull<mlir_ts::VariableOp, mlir_ts::AddressOfOp>(def))
+        {
+            return def;
+        }
+
+        if (!isPlace(ref))
+        {
+            return nullptr;
+        }
+
+        // what the place is in: a read of a slot, or the slot itself (a field of a local tuple)
+        auto in = def->getOperand(0);
+        while (auto castOp = in.getDefiningOp<mlir_ts::CastOp>())
+        {
+            in = castOp.getIn();
+        }
+
+        auto loadOp = in.getDefiningOp<mlir_ts::LoadOp>();
+        ref = loadOp ? loadOp.getReference() : in;
+    }
+
+    return nullptr;
+}
+
+// The storage a value is read out of (slotOfRef), through casts.
+inline mlir::Operation *slotOf(mlir::Value value)
+{
+    while (auto castOp = value.getDefiningOp<mlir_ts::CastOp>())
+    {
+        value = castOp.getIn();
+    }
+
+    auto loadOp = value.getDefiningOp<mlir_ts::LoadOp>();
+    return loadOp ? slotOfRef(loadOp.getReference()) : nullptr;
+}
+
+// One storage: a local's slot, or a global - each access of which is a `ts.AddressOf` of its own.
+inline bool sameSlot(mlir::Operation *a, mlir::Operation *b)
+{
+    if (a == b)
+    {
+        return true;
+    }
+
+    auto globalA = mlir::dyn_cast_or_null<mlir_ts::AddressOfOp>(a);
+    auto globalB = mlir::dyn_cast_or_null<mlir_ts::AddressOfOp>(b);
+    return globalA && globalB && globalA.getGlobalName() == globalB.getGlobalName();
+}
+
+// The storage an op gives a new value to or releases: a store, `ts.ReleaseSlot`.
+inline mlir::Operation *slotWritten(mlir::Operation *op)
+{
+    if (auto storeOp = mlir::dyn_cast<mlir_ts::StoreOp>(op))
+    {
+        return slotOfRef(storeOp.getReference());
+    }
+
+    if (auto releaseSlotOp = mlir::dyn_cast<mlir_ts::ReleaseSlotOp>(op))
+    {
+        return slotOfRef(releaseSlotOp.getSlot());
+    }
+
+    return nullptr;
+}
+
+// An op for which `found` holds that may run after `read` and before `call`: the ops between them
+// in a block, those of every block on a path from one to the other, and what each of them nests.
+// The arguments of a call are read in order, so one read early may be freed by a later one before
+// the call takes it - `f(n, (n = mk(), 0))` - and a conditional argument puts that in a block of its
+// own. Null when there is none.
+inline mlir::Operation *findBetween(mlir::Operation *read, mlir::Operation *call,
+                                    llvm::function_ref<bool(mlir::Operation *)> found)
+{
+    // the two in the innermost region holding both
+    mlir::Operation *readTop = nullptr;
+    mlir::Operation *callTop = nullptr;
+    for (auto *region = call->getParentRegion(); region && !readTop; region = region->getParentRegion())
+    {
+        readTop = region->findAncestorOpInRegion(*read);
+        callTop = readTop ? region->findAncestorOpInRegion(*call) : nullptr;
+    }
+
+    if (!readTop || !callTop || readTop == callTop)
+    {
+        return nullptr;
+    }
+
+    mlir::Operation *at = nullptr;
+    auto scan = [&](mlir::Operation *op) {
+        op->walk([&](mlir::Operation *inner) {
+            if (!at && inner != call && found(inner))
+            {
+                at = inner;
+            }
+        });
+    };
+
+    // the statement the call is in, when the call is not one itself
+    if (callTop != call)
+    {
+        scan(callTop);
+    }
+
+    auto *fromBlock = readTop->getBlock();
+    auto *toBlock = callTop->getBlock();
+    if (fromBlock == toBlock && readTop->isBeforeInBlock(callTop))
+    {
+        for (auto *op = readTop->getNextNode(); op != callTop; op = op->getNextNode())
+        {
+            scan(op);
+        }
+
+        return at;
+    }
+
+    for (auto *op = readTop->getNextNode(); op; op = op->getNextNode())
+    {
+        scan(op);
+    }
+
+    for (auto *op = &toBlock->front(); op != callTop; op = op->getNextNode())
+    {
+        scan(op);
+    }
+
+    // the blocks on a path: reached from the read's, and reaching the call's
+    llvm::SmallPtrSet<mlir::Block *, 16> after;
+    llvm::SmallVector<mlir::Block *> work(fromBlock->getSuccessors().begin(), fromBlock->getSuccessors().end());
+    while (!work.empty())
+    {
+        auto *block = work.pop_back_val();
+        if (after.insert(block).second)
+        {
+            work.append(block->getSuccessors().begin(), block->getSuccessors().end());
+        }
+    }
+
+    llvm::SmallPtrSet<mlir::Block *, 16> before;
+    work.assign(toBlock->getPredecessors().begin(), toBlock->getPredecessors().end());
+    while (!work.empty())
+    {
+        auto *block = work.pop_back_val();
+        if (before.insert(block).second)
+        {
+            work.append(block->getPredecessors().begin(), block->getPredecessors().end());
+        }
+    }
+
+    for (auto *block : after)
+    {
+        if (!before.contains(block))
+        {
+            continue;
+        }
+
+        for (auto &op : *block)
+        {
+            scan(&op);
+        }
+    }
+
+    return at;
+}
+
 // A handle (spec 23.2): counted in every model that tracks ownership, and left alone by own.
 inline bool isHandle(mlir::Value value)
 {

tslang/test/tester/CMakeLists.txt

@@ -410,6 +410,7 @@ tslang_add_test(NAME test-compile-00-argument-place-overwritten COMMAND test-run
 tslang_add_test(NAME test-compile-00-temporary-read-in-finally COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00temporary_read_in_finally.ts")
 tslang_add_test(NAME test-compile-00-unowned-local-store COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00unowned_local_store.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-argument-slot-assigned COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00argument_slot_assigned.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")
@@ -994,6 +995,7 @@ tslang_add_test(NAME test-jit-00-argument-place-overwritten COMMAND test-runner
 tslang_add_test(NAME test-jit-00-temporary-read-in-finally COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00temporary_read_in_finally.ts")
 tslang_add_test(NAME test-jit-00-unowned-local-store COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00unowned_local_store.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-argument-slot-assigned COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00argument_slot_assigned.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")
@@ -1968,6 +1970,7 @@ set(TSLANG_CORPUS
     00temporary_read_in_finally.ts
     00unowned_local_store.ts
     00array_argument_held.ts
+    00argument_slot_assigned.ts
     00array_data_string_alias.ts
     00array7.ts
     00array8_tuple_spread.ts
@@ -3751,7 +3754,7 @@ set(shared_models gc none rc own)
 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_place_two_uses own_shared_global_field own_shared_global_copy own_shared_argument_assigned
                         own_shared_global_argument own_shared_global_argument_throw)
         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")
@@ -3766,7 +3769,7 @@ endforeach()
 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_place_two_uses own_shared_global_field own_shared_global_copy own_shared_argument_assigned
                         own_shared_global_argument own_shared_global_argument_throw)
         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

tslang/test/tester/own/own_shared_argument_assigned.ts

@@ -0,0 +1,36 @@
+// A Shared<T> handle read for an argument, and given a new value by a later argument before the call
+// takes it: the assignment gave back the count the old handle held, the last one, and the callee
+// read the freed block. The caller holds a count across the call, in every model that counts. 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 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 g(s: Shared<V>, b: number) {
+    churn();
+    assert(s.value.tag == "v" + s.value.x, "the handle's block");
+    return s.value.x * 10 + b;
+}
+
+function main() {
+    let s = new Shared(new V(1));
+    assert(g(s, (s = new Shared(new V(2)), 0)) == 10, "a handle given a new value in the next argument");
+    churn();
+    assert(s.value.x == 2 && s.value.tag == "v2", "the handle's new value");
+    print("done.");
+}