[mypyc] Include more operations in the trace log

mypyc/transform/log_trace.py

@@ -9,8 +9,27 @@
 
 from __future__ import annotations
 
+from typing import Final
+
 from mypyc.ir.func_ir import FuncIR
-from mypyc.ir.ops import Call, CallC, CString, LoadLiteral, LoadStatic, Op, PrimitiveOp, Value
+from mypyc.ir.ops import (
+    Box,
+    Call,
+    CallC,
+    Cast,
+    CString,
+    DecRef,
+    GetAttr,
+    IncRef,
+    LoadLiteral,
+    LoadStatic,
+    Op,
+    PrimitiveOp,
+    SetAttr,
+    Unbox,
+    Value,
+)
+from mypyc.ir.rtypes import none_rprimitive
 from mypyc.irbuild.ll_builder import LowLevelIRBuilder
 from mypyc.options import CompilerOptions
 from mypyc.primitives.misc_ops import log_trace_event
@@ -38,6 +57,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, and mostly ops that could be switched to use borrowing in some contexts.
+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)
@@ -48,7 +79,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):
@@ -63,11 +97,53 @@ 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:
+        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:
+        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:
+        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:
+        return self.log(op, "unbox", str(op.type))
+
+    def visit_cast(self, op: Cast) -> Value | None:
+        value = self.log(op, "cast", str(op.type))
+        if not op.is_borrowed:
+            self.log_inc_ref(value)
+        return value
+
+    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:
+        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(
@@ -78,6 +154,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)