Log some implicit increfs

Author: JukkaL

mypyc/transform/log_trace.py

@@ -9,6 +9,8 @@
 
 from __future__ import annotations
 
+from typing import Final
+
 from mypyc.ir.func_ir import FuncIR
 from mypyc.ir.ops import (
     Box,
@@ -54,6 +56,18 @@ def get_load_global_name(op: CallC) -> str | None:
     return None
 
 
+# These primitives perform an implicit IncRef for the return value. Only some of the most common ones
+# are included.
+primitives_that_inc_ref: Final = {
+    "list_get_item_unsafe",
+    "CPyList_GetItemShort",
+    "CPyDict_GetWithNone",
+    "CPyList_GetItem",
+    "CPyDict_GetItem",
+    "CPyList_PopLast",
+}
+
+
 class LogTraceEventTransform(IRTransform):
     def __init__(self, builder: LowLevelIRBuilder, fullname: str) -> None:
         super().__init__(builder)
@@ -64,7 +78,10 @@ def visit_call(self, op: Call) -> Value:
         return self.log(op, "call", op.fn.fullname)
 
     def visit_primitive_op(self, op: PrimitiveOp) -> Value:
-        return self.log(op, "primitive_op", op.desc.name)
+        value = self.log(op, "primitive_op", op.desc.name)
+        if op.desc.name in primitives_that_inc_ref:
+            self.log_inc_ref(value)
+        return value
 
     def visit_call_c(self, op: CallC) -> Value:
         if global_name := get_load_global_name(op):
@@ -79,34 +96,47 @@ def visit_call_c(self, op: CallC) -> Value:
         elif func_name == "PyObject_VectorcallMethod" and isinstance(op.args[0], LoadLiteral):
             return self.log(op, "python_call_method", str(op.args[0].value))
 
-        return self.log(op, "call_c", func_name)
+        value = self.log(op, "call_c", func_name)
+        if func_name in primitives_that_inc_ref:
+            self.log_inc_ref(value)
+        return value
 
-    def visit_get_attr(self, op: GetAttr) -> Value | None:
-        return self.log(op, "get_attr", f"{op.class_type.name}.{op.attr}")
+    def visit_get_attr(self, op: GetAttr) -> Value:
+        value = self.log(op, "get_attr", f"{op.class_type.name}.{op.attr}")
+        if not op.is_borrowed and op.type.is_refcounted:
+            self.log_inc_ref(op)
+        return value
 
-    def visit_set_attr(self, op: SetAttr) -> Value | None:
+    def visit_set_attr(self, op: SetAttr) -> Value:
         name = "set_attr" if not op.is_init else "set_attr_init"
         return self.log(op, name, f"{op.class_type.name}.{op.attr}")
 
-    def visit_box(self, op: Box) -> Value | None:
+    def visit_box(self, op: Box) -> Value:
         if op.src.type is none_rprimitive:
             # Boxing 'None' is a very quick operation, so we don't log it.
             return self.add(op)
         else:
             return self.log(op, "box", str(op.src.type))
 
-    def visit_unbox(self, op: Unbox) -> Value | None:
+    def visit_unbox(self, op: Unbox) -> Value:
         return self.log(op, "unbox", str(op.type))
 
-    def visit_inc_ref(self, op: IncRef) -> Value | None:
+    def visit_inc_ref(self, op: IncRef) -> Value:
         return self.log(op, "inc_ref", str(op.src.type))
 
-    def visit_dec_ref(self, op: DecRef) -> Value | None:
+    def visit_dec_ref(self, op: DecRef) -> Value:
         return self.log(op, "dec_ref", str(op.src.type))
 
+    def log_inc_ref(self, value: Value) -> None:
+        self.log_event("inc_ref", str(value.type), value.line)
+
     def log(self, op: Op, name: str, details: str) -> Value:
-        if op.line >= 0:
-            line_str = str(op.line)
+        self.log_event(name, details, op.line)
+        return self.add(op)
+
+    def log_event(self, name: str, details: str, line: int) -> None:
+        if line >= 0:
+            line_str = str(line)
         else:
             line_str = ""
         self.builder.primitive_op(
@@ -117,6 +147,5 @@ def log(self, op: Op, name: str, details: str) -> Value:
                 CString(name.encode("utf-8")),
                 CString(details.encode("utf-8")),
             ],
-            op.line,
+            line,
         )
-        return self.add(op)