[mypyc] Introduce GetElementPtr (python#9260)

TH3CHARLie · web-flow · commit ffd9d1cdff4a · 2020-08-04T16:48:03.000+01:00
Related to mypyc/mypyc#741. This PR introduces GetElementPtr, which computes the address of an element in an aggregate type (currently RStruct is supported). Part of efforts to support the len and other macro-related primitives.
diff --git a/mypyc/analysis/dataflow.py b/mypyc/analysis/dataflow.py
@@ -9,7 +9,7 @@
     BasicBlock, OpVisitor, Assign, LoadInt, LoadErrorValue, RegisterOp, Goto, Branch, Return, Call,
     Environment, Box, Unbox, Cast, Op, Unreachable, TupleGet, TupleSet, GetAttr, SetAttr,
     LoadStatic, InitStatic, PrimitiveOp, MethodCall, RaiseStandardError, CallC, LoadGlobal,
-    Truncate, BinaryIntOp, LoadMem
+    Truncate, BinaryIntOp, LoadMem, GetElementPtr
 )
 
 
@@ -211,6 +211,9 @@ def visit_binary_int_op(self, op: BinaryIntOp) -> GenAndKill:
     def visit_load_mem(self, op: LoadMem) -> GenAndKill:
         return self.visit_register_op(op)
 
+    def visit_get_element_ptr(self, op: GetElementPtr) -> GenAndKill:
+        return self.visit_register_op(op)
+
 
 class DefinedVisitor(BaseAnalysisVisitor):
     """Visitor for finding defined registers.
diff --git a/mypyc/codegen/emitfunc.py b/mypyc/codegen/emitfunc.py
@@ -12,10 +12,10 @@
     LoadStatic, InitStatic, TupleGet, TupleSet, Call, IncRef, DecRef, Box, Cast, Unbox,
     BasicBlock, Value, MethodCall, PrimitiveOp, EmitterInterface, Unreachable, NAMESPACE_STATIC,
     NAMESPACE_TYPE, NAMESPACE_MODULE, RaiseStandardError, CallC, LoadGlobal, Truncate,
-    BinaryIntOp, LoadMem
+    BinaryIntOp, LoadMem, GetElementPtr
 )
 from mypyc.ir.rtypes import (
-    RType, RTuple, is_tagged, is_int32_rprimitive, is_int64_rprimitive
+    RType, RTuple, is_tagged, is_int32_rprimitive, is_int64_rprimitive, RStruct
 )
 from mypyc.ir.func_ir import FuncIR, FuncDecl, FUNC_STATICMETHOD, FUNC_CLASSMETHOD
 from mypyc.ir.class_ir import ClassIR
@@ -471,6 +471,14 @@ def visit_load_mem(self, op: LoadMem) -> None:
         type = self.ctype(op.type)
         self.emit_line('%s = *(%s *)%s;' % (dest, type, src))
 
+    def visit_get_element_ptr(self, op: GetElementPtr) -> None:
+        dest = self.reg(op)
+        src = self.reg(op.src)
+        # TODO: support tuple type
+        assert isinstance(op.src_type, RStruct)
+        assert op.field in op.src_type.names, "Invalid field name."
+        self.emit_line('%s = &%s.%s;' % (dest, src, op.field))
+
     # Helpers
 
     def label(self, label: BasicBlock) -> str:
diff --git a/mypyc/ir/ops.py b/mypyc/ir/ops.py
@@ -25,7 +25,8 @@
 from mypyc.ir.rtypes import (
     RType, RInstance, RTuple, RVoid, is_bool_rprimitive, is_int_rprimitive,
     is_short_int_rprimitive, is_none_rprimitive, object_rprimitive, bool_rprimitive,
-    short_int_rprimitive, int_rprimitive, void_rtype, is_c_py_ssize_t_rprimitive
+    short_int_rprimitive, int_rprimitive, void_rtype, is_c_py_ssize_t_rprimitive,
+    c_pyssize_t_rprimitive
 )
 from mypyc.common import short_name
 
@@ -1372,6 +1373,28 @@ def accept(self, visitor: 'OpVisitor[T]') -> T:
         return visitor.visit_load_mem(self)
 
 
+class GetElementPtr(RegisterOp):
+    """Get the address of a struct element"""
+    error_kind = ERR_NEVER
+
+    def __init__(self, src: Value, src_type: RType, field: str, line: int = -1) -> None:
+        super().__init__(line)
+        self.type = c_pyssize_t_rprimitive
+        self.src = src
+        self.src_type = src_type
+        self.field = field
+
+    def sources(self) -> List[Value]:
+        return [self.src]
+
+    def to_str(self, env: Environment) -> str:
+        return env.format("%r = get_element_ptr %r %r :: %r", self, self.src,
+                          self.field, self.src_type)
+
+    def accept(self, visitor: 'OpVisitor[T]') -> T:
+        return visitor.visit_get_element_ptr(self)
+
+
 @trait
 class OpVisitor(Generic[T]):
     """Generic visitor over ops (uses the visitor design pattern)."""
@@ -1482,6 +1505,10 @@ def visit_binary_int_op(self, op: BinaryIntOp) -> T:
     def visit_load_mem(self, op: LoadMem) -> T:
         raise NotImplementedError
 
+    @abstractmethod
+    def visit_get_element_ptr(self, op: GetElementPtr) -> T:
+        raise NotImplementedError
+
 
 # TODO: Should this live somewhere else?
 LiteralsMap = Dict[Tuple[Type[object], Union[int, float, str, bytes, complex]], str]
diff --git a/mypyc/test/test_emitfunc.py b/mypyc/test/test_emitfunc.py
@@ -10,12 +10,12 @@
 from mypyc.ir.ops import (
     Environment, BasicBlock, Goto, Return, LoadInt, Assign, IncRef, DecRef, Branch,
     Call, Unbox, Box, TupleGet, GetAttr, PrimitiveOp, RegisterOp,
-    SetAttr, Op, Value, CallC, BinaryIntOp, LoadMem
+    SetAttr, Op, Value, CallC, BinaryIntOp, LoadMem, GetElementPtr
 )
 from mypyc.ir.rtypes import (
     RTuple, RInstance, int_rprimitive, bool_rprimitive, list_rprimitive,
     dict_rprimitive, object_rprimitive, c_int_rprimitive, short_int_rprimitive, int32_rprimitive,
-    int64_rprimitive
+    int64_rprimitive, StructInfo, RStruct
 )
 from mypyc.ir.func_ir import FuncIR, FuncDecl, RuntimeArg, FuncSignature
 from mypyc.ir.class_ir import ClassIR
@@ -282,6 +282,17 @@ def test_load_mem(self) -> None:
             self.assert_emit(LoadMem(bool_rprimitive, self.i64),
                              """cpy_r_r0 = *(char *)cpy_r_i64;""")
 
+    def test_get_element_ptr(self) -> None:
+        info = StructInfo("Foo", ["b", "i32", "i64"], [bool_rprimitive,
+                                                       int32_rprimitive, int64_rprimitive])
+        r = RStruct(info)
+        self.assert_emit(GetElementPtr(self.o, r, "b"),
+                        """cpy_r_r0 = &cpy_r_o.b;""")
+        self.assert_emit(GetElementPtr(self.o, r, "i32"),
+                        """cpy_r_r00 = &cpy_r_o.i32;""")
+        self.assert_emit(GetElementPtr(self.o, r, "i64"),
+                        """cpy_r_r01 = &cpy_r_o.i64;""")
+
     def assert_emit(self, op: Op, expected: str) -> None:
         self.emitter.fragments = []
         self.declarations.fragments = []
