[mypyc] Implement SetMem (#9364)

This PR implements SetMem op to set value to a memory address.

Author: TH3CHARLie

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, GetElementPtr, LoadAddress, ComparisonOp
+    Truncate, BinaryIntOp, LoadMem, GetElementPtr, LoadAddress, ComparisonOp, SetMem
 )
 
 
@@ -151,6 +151,10 @@ def visit_register_op(self, op: RegisterOp) -> GenAndKill:
     def visit_assign(self, op: Assign) -> GenAndKill:
         raise NotImplementedError
 
+    @abstractmethod
+    def visit_set_mem(self, op: SetMem) -> GenAndKill:
+        raise NotImplementedError
+
     def visit_call(self, op: Call) -> GenAndKill:
         return self.visit_register_op(op)
 
@@ -248,6 +252,9 @@ def visit_assign(self, op: Assign) -> GenAndKill:
         else:
             return {op.dest}, set()
 
+    def visit_set_mem(self, op: SetMem) -> GenAndKill:
+        return set(), set()
+
 
 def analyze_maybe_defined_regs(blocks: List[BasicBlock],
                                cfg: CFG,
@@ -308,6 +315,9 @@ def visit_assign(self, op: Assign) -> GenAndKill:
             return set(), {op.dest}
         return set(), set()
 
+    def visit_set_mem(self, op: SetMem) -> GenAndKill:
+        return set(), set()
+
 
 def analyze_borrowed_arguments(
         blocks: List[BasicBlock],
@@ -342,6 +352,9 @@ def visit_register_op(self, op: RegisterOp) -> GenAndKill:
     def visit_assign(self, op: Assign) -> GenAndKill:
         return set(), {op.dest}
 
+    def visit_set_mem(self, op: SetMem) -> GenAndKill:
+        return set(), set()
+
 
 def analyze_undefined_regs(blocks: List[BasicBlock],
                            cfg: CFG,
@@ -381,6 +394,9 @@ def visit_register_op(self, op: RegisterOp) -> GenAndKill:
     def visit_assign(self, op: Assign) -> GenAndKill:
         return set(op.sources()), {op.dest}
 
+    def visit_set_mem(self, op: SetMem) -> GenAndKill:
+        return set(op.sources()), set()
+
 
 def analyze_live_regs(blocks: List[BasicBlock],
                       cfg: CFG) -> AnalysisResult[Value]:

mypyc/codegen/emitfunc.py

@@ -12,7 +12,7 @@
     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, GetElementPtr, LoadAddress, ComparisonOp
+    BinaryIntOp, LoadMem, GetElementPtr, LoadAddress, ComparisonOp, SetMem
 )
 from mypyc.ir.rtypes import (
     RType, RTuple, is_tagged, is_int32_rprimitive, is_int64_rprimitive, RStruct,
@@ -478,6 +478,15 @@ def visit_load_mem(self, op: LoadMem) -> None:
         type = self.ctype(op.type)
         self.emit_line('%s = *(%s *)%s;' % (dest, type, src))
 
+    def visit_set_mem(self, op: SetMem) -> None:
+        dest = self.reg(op.dest)
+        src = self.reg(op.src)
+        type = self.ctype(op.type)
+        # clang whines about self assignment (which we might generate
+        # for some casts), so don't emit it.
+        if dest != src:
+            self.emit_line('*(%s *)%s = %s;' % (type, dest, src))
+
     def visit_get_element_ptr(self, op: GetElementPtr) -> None:
         dest = self.reg(op)
         src = self.reg(op.src)

mypyc/ir/ops.py

@@ -1414,6 +1414,55 @@ def accept(self, visitor: 'OpVisitor[T]') -> T:
         return visitor.visit_load_mem(self)
 
 
+class SetMem(Op):
+    """Write a memory location.
+
+    *(type *)dest = src
+
+    Attributes:
+      type: Type of the read value
+      dest: Pointer to memory to write
+      src: Source value
+      base: If not None, the object from which we are reading memory.
+            It's used to avoid the target object from being freed via
+            reference counting. If the target is not in reference counted
+            memory, or we know that the target won't be freed, it can be
+            None.
+    """
+    error_kind = ERR_NEVER
+
+    def __init__(self,
+                 type: RType,
+                 dest: Register,
+                 src: Value,
+                 base: Optional[Value],
+                 line: int = -1) -> None:
+        super().__init__(line)
+        self.type = type
+        self.src = src
+        self.dest = dest
+        self.base = base
+
+    def sources(self) -> List[Value]:
+        if self.base:
+            return [self.src, self.base, self.dest]
+        else:
+            return [self.src, self.dest]
+
+    def stolen(self) -> List[Value]:
+        return [self.src]
+
+    def to_str(self, env: Environment) -> str:
+        if self.base:
+            base = env.format(', %r', self.base)
+        else:
+            base = ''
+        return env.format("%r = set_mem %r%s :: %r*", self.dest, self.src, base, self.type)
+
+    def accept(self, visitor: 'OpVisitor[T]') -> T:
+        return visitor.visit_set_mem(self)
+
+
 class GetElementPtr(RegisterOp):
     """Get the address of a struct element"""
     error_kind = ERR_NEVER
@@ -1569,6 +1618,10 @@ def visit_comparison_op(self, op: ComparisonOp) -> T:
     def visit_load_mem(self, op: LoadMem) -> T:
         raise NotImplementedError
 
+    @abstractmethod
+    def visit_set_mem(self, op: SetMem) -> T:
+        raise NotImplementedError
+
     @abstractmethod
     def visit_get_element_ptr(self, op: GetElementPtr) -> T:
         raise NotImplementedError

mypyc/test/test_emitfunc.py

@@ -10,7 +10,8 @@
 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, GetElementPtr, LoadAddress, ComparisonOp
+    SetAttr, Op, Value, CallC, BinaryIntOp, LoadMem, GetElementPtr, LoadAddress, ComparisonOp,
+    SetMem
 )
 from mypyc.ir.rtypes import (
     RTuple, RInstance, int_rprimitive, bool_rprimitive, list_rprimitive,
@@ -300,6 +301,10 @@ def test_load_mem(self) -> None:
         self.assert_emit(LoadMem(bool_rprimitive, self.ptr, self.s1),
                          """cpy_r_r00 = *(char *)cpy_r_ptr;""")
 
+    def test_set_mem(self) -> None:
+        self.assert_emit(SetMem(bool_rprimitive, self.ptr, self.b, None),
+                         """*(char *)cpy_r_ptr = cpy_r_b;""")
+
     def test_get_element_ptr(self) -> None:
         r = RStruct("Foo", ["b", "i32", "i64"], [bool_rprimitive,
                                                  int32_rprimitive, int64_rprimitive])