Add missing case for BORROW_TRANSIENT in call

This is leading to inappropriate tracking of temporary lifetimes
numba · Jul 26, 2023 · 240f81c · 240f81c
1 parent 498a547
commit 240f81c
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 0 deletions.
diff --git a/dispyatcher/__init__.py b/dispyatcher/__init__.py
@@ -697,6 +697,12 @@ def call(self, handle: "Handle", args: Sequence[Union[TemporaryValue, Tuple[IRVa
                         self.builder.comment(f"Transferring {caller_arg} argument to {index} for {handle}...")
                         caller_arg_lifetimes[caller_arg] = False
                     callee_argument_lifetimes.append(True)
+                elif arg_management == ArgumentManagement.BORROW_TRANSIENT:
+                    callee_argument_lifetimes.append(set())
+                    for caller_arg in caller_args:
+                        lifetimes_for_arg = caller_arg_lifetimes[caller_arg]
+                        if lifetimes_for_arg:
+                            lifetimes_to_check.update(lifetimes_for_arg)
                 else:
                     combined_caller_lifetimes = set()
                     for caller_arg in caller_args:

diff --git a/tests/test_cpython.py b/tests/test_cpython.py
@@ -104,3 +104,18 @@ def test_multi_unwind(self):
         callsite = CallSite(handle, PythonControlFlowType())
         self.assertEqual(callsite({"a": {"b": 1}}, "a", "b"), 1)
 
+    def test_transient_borrow_not_capture(self):
+        pytuple_pack2 = dispyatcher.general.CurrentProcessFunction(
+            dispyatcher.cpython.PyObjectType(tuple),
+            dispyatcher.ReturnManagement.TRANSFER,
+            "PyTuple_Pack",
+            (MachineType(dispyatcher.cpython.SIZE_T_TYPE), dispyatcher.ArgumentManagement.BORROW_TRANSIENT),
+            (PY_OBJECT_TYPE, dispyatcher.ArgumentManagement.BORROW_TRANSIENT),
+            (PY_OBJECT_TYPE, dispyatcher.ArgumentManagement.BORROW_TRANSIENT))
+        pack2 = pytuple_pack2 << (dispyatcher.general.SimpleConstant(MachineType(dispyatcher.cpython.SIZE_T_TYPE), 2),)
+        lookup_handle = dispyatcher.cpython.PY_DICT_GET_ITEM + dispyatcher.cpython.ThrowIfNull + dispyatcher.Clone
+
+        callsite = CallSite(pack2 << lookup_handle, PythonControlFlowType())
+
+        self.assertEqual(callsite({"a": 1}, "a", "b"), (1, "b"))
+