[mypyc] Add GetElement op and other minor ir tweaks (#20721)

The `GetElement` op gets a struct field from a struct value, unlike
`GetElementPtr` which takes a pointer to a struct.

These changes were extracted from mypyc vec type support branch, and
aren't very useful by themselves. I'm splitting the branch into multiple
PRs to simplify reviewing.

Author: JukkaL

mypyc/analysis/dataflow.py

@@ -24,6 +24,7 @@
     FloatNeg,
     FloatOp,
     GetAttr,
+    GetElement,
     GetElementPtr,
     Goto,
     IncRef,
@@ -271,6 +272,9 @@ def visit_float_comparison_op(self, op: FloatComparisonOp) -> GenAndKill[T]:
     def visit_load_mem(self, op: LoadMem) -> GenAndKill[T]:
         return self.visit_register_op(op)
 
+    def visit_get_element(self, op: GetElement) -> GenAndKill[T]:
+        return self.visit_register_op(op)
+
     def visit_get_element_ptr(self, op: GetElementPtr) -> GenAndKill[T]:
         return self.visit_register_op(op)
 

mypyc/analysis/ircheck.py

@@ -22,6 +22,7 @@
     FloatNeg,
     FloatOp,
     GetAttr,
+    GetElement,
     GetElementPtr,
     Goto,
     IncRef,
@@ -449,6 +450,9 @@ def visit_load_mem(self, op: LoadMem) -> None:
     def visit_set_mem(self, op: SetMem) -> None:
         pass
 
+    def visit_get_element(self, op: GetElement) -> None:
+        pass
+
     def visit_get_element_ptr(self, op: GetElementPtr) -> None:
         pass
 

mypyc/analysis/selfleaks.py

@@ -16,6 +16,7 @@
     FloatNeg,
     FloatOp,
     GetAttr,
+    GetElement,
     GetElementPtr,
     Goto,
     InitStatic,
@@ -179,6 +180,9 @@ def visit_float_comparison_op(self, op: FloatComparisonOp) -> GenAndKill:
     def visit_load_mem(self, op: LoadMem) -> GenAndKill:
         return CLEAN
 
+    def visit_get_element(self, op: GetElement) -> GenAndKill:
+        return CLEAN
+
     def visit_get_element_ptr(self, op: GetElementPtr) -> GenAndKill:
         return CLEAN
 

mypyc/codegen/emitfunc.py

@@ -42,6 +42,7 @@
     FloatNeg,
     FloatOp,
     GetAttr,
+    GetElement,
     GetElementPtr,
     Goto,
     IncRef,
@@ -795,6 +796,12 @@ def visit_set_mem(self, op: SetMem) -> None:
         if dest != src:
             self.emit_line(f"*({dest_type} *){dest} = {src};")
 
+    def visit_get_element(self, op: GetElement) -> None:
+        dest = self.reg(op)
+        src = self.reg(op.src)
+        dest_type = self.ctype(op.type)
+        self.emit_line(f"{dest} = ({dest_type}){src}.{op.field};")
+
     def visit_get_element_ptr(self, op: GetElementPtr) -> None:
         dest = self.reg(op)
         src = self.reg(op.src)

mypyc/ir/ops.py

@@ -337,6 +337,8 @@ def set_sources(self, new: list[Value]) -> None:
         (self.src,) = new
 
     def stolen(self) -> list[Value]:
+        if not self.dest.type.is_refcounted:
+            return []
         return [self.src]
 
     def accept(self, visitor: OpVisitor[T]) -> T:
@@ -1679,9 +1681,36 @@ def accept(self, visitor: OpVisitor[T]) -> T:
         return visitor.visit_set_mem(self)
 
 
+@final
+class GetElement(RegisterOp):
+    """Get the value of a struct element from a struct value."""
+
+    error_kind = ERR_NEVER
+    is_borrowed = True
+
+    def __init__(self, src: Value, field: str, line: int = -1) -> None:
+        super().__init__(line)
+        assert isinstance(src.type, RStruct)
+        self.type = src.type.field_type(field)
+        self.src = src
+        self.src_type = src.type
+        self.field = field
+
+    def sources(self) -> list[Value]:
+        return [self.src]
+
+    def set_sources(self, new: list[Value]) -> None:
+        (self.src,) = new
+
+    def accept(self, visitor: OpVisitor[T]) -> T:
+        return visitor.visit_get_element(self)
+
+
 @final
 class GetElementPtr(RegisterOp):
-    """Get the address of a struct element.
+    """Get the address of a struct element from a pointer to a struct.
+
+    If you have a struct value, use GetElement instead.
 
     Note that you may need to use KeepAlive to avoid the struct
     being freed, if it's reference counted, such as PyObject *.
@@ -1691,6 +1720,7 @@ class GetElementPtr(RegisterOp):
 
     def __init__(self, src: Value, src_type: RType, field: str, line: int = -1) -> None:
         super().__init__(line)
+        assert not isinstance(src.type, RStruct)
         self.type = pointer_rprimitive
         self.src = src
         self.src_type = src_type
@@ -2008,6 +2038,10 @@ def visit_load_mem(self, op: LoadMem) -> T:
     def visit_set_mem(self, op: SetMem) -> T:
         raise NotImplementedError
 
+    @abstractmethod
+    def visit_get_element(self, op: GetElement) -> T:
+        raise NotImplementedError
+
     @abstractmethod
     def visit_get_element_ptr(self, op: GetElementPtr) -> T:
         raise NotImplementedError

mypyc/ir/pprint.py

@@ -29,6 +29,7 @@
     FloatNeg,
     FloatOp,
     GetAttr,
+    GetElement,
     GetElementPtr,
     Goto,
     IncRef,
@@ -280,6 +281,9 @@ def visit_load_mem(self, op: LoadMem) -> str:
     def visit_set_mem(self, op: SetMem) -> str:
         return self.format("set_mem %r, %r :: %t*", op.dest, op.src, op.dest_type)
 
+    def visit_get_element(self, op: GetElement) -> str:
+        return self.format("%r = %r.%s", op, op.src, op.field)
+
     def visit_get_element_ptr(self, op: GetElementPtr) -> str:
         return self.format("%r = get_element_ptr %r %s :: %t", op, op.src, op.field, op.src_type)
 

mypyc/ir/rtypes.py

@@ -865,13 +865,21 @@ def __init__(self, name: str, names: list[str], types: list[RType]) -> None:
         self.name = name
         self.names = names
         self.types = types
+        self.is_refcounted = any(t.is_refcounted for t in self.types)
+
         # generate dummy names
         if len(self.names) < len(self.types):
             for i in range(len(self.types) - len(self.names)):
                 self.names.append("_item" + str(i))
         self.offsets, self.size = compute_aligned_offsets_and_size(types)
         self._ctype = name
 
+    def field_type(self, name: str) -> RType:
+        for n, t in zip(self.names, self.types):
+            if n == name:
+                return t
+        assert False, f"{self.name} has no field '{name}'"
+
     def accept(self, visitor: RTypeVisitor[T]) -> T:
         return visitor.visit_rstruct(self)
 

mypyc/transform/ir_transform.py

@@ -20,6 +20,7 @@
     FloatNeg,
     FloatOp,
     GetAttr,
+    GetElement,
     GetElementPtr,
     Goto,
     IncRef,
@@ -212,6 +213,9 @@ def visit_load_mem(self, op: LoadMem) -> Value | None:
     def visit_set_mem(self, op: SetMem) -> Value | None:
         return self.add(op)
 
+    def visit_get_element(self, op: GetElement) -> Value | None:
+        return self.add(op)
+
     def visit_get_element_ptr(self, op: GetElementPtr) -> Value | None:
         return self.add(op)
 
@@ -355,6 +359,9 @@ def visit_set_mem(self, op: SetMem) -> None:
         op.dest = self.fix_op(op.dest)
         op.src = self.fix_op(op.src)
 
+    def visit_get_element(self, op: GetElement) -> None:
+        op.src = self.fix_op(op.src)
+
     def visit_get_element_ptr(self, op: GetElementPtr) -> None:
         op.src = self.fix_op(op.src)