From dc4bd45c611ca0c31cefb091fef2c025e6536422 Mon Sep 17 00:00:00 2001
From: Jukka Lehtosalo <jukka.lehtosalo@iki.fi>
Date: Tue, 25 Oct 2022 17:13:44 +0100
Subject: [PATCH] [mypyc] Fix ircheck if register is initialized through
 LoadAddress

It's okay to take an address of a register, pass it to a function that
initializes the register, and then read the register. Ircheck
complained about an invalid op reference to register in this case. It
affected at least async code.
---
 mypyc/analysis/ircheck.py  |  8 ++++++--
 mypyc/test/test_ircheck.py | 31 +++++++++++++++++++++++++++++--
 2 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/mypyc/analysis/ircheck.py b/mypyc/analysis/ircheck.py
index c2cdd073f62e..b141784ef9ff 100644
--- a/mypyc/analysis/ircheck.py
+++ b/mypyc/analysis/ircheck.py
@@ -129,7 +129,11 @@ def check_op_sources_valid(fn: FuncIR) -> list[FnError]:
     for block in fn.blocks:
         valid_ops.update(block.ops)
 
-        valid_registers.update([op.dest for op in block.ops if isinstance(op, BaseAssign)])
+        for op in block.ops:
+            if isinstance(op, BaseAssign):
+                valid_registers.add(op.dest)
+            elif isinstance(op, LoadAddress) and isinstance(op.src, Register):
+                valid_registers.add(op.src)
 
     valid_registers.update(fn.arg_regs)
 
@@ -150,7 +154,7 @@ def check_op_sources_valid(fn: FuncIR) -> list[FnError]:
                     if source not in valid_registers:
                         errors.append(
                             FnError(
-                                source=op, desc=f"Invalid op reference to register {source.name}"
+                                source=op, desc=f"Invalid op reference to register {source.name!r}"
                             )
                         )
 
diff --git a/mypyc/test/test_ircheck.py b/mypyc/test/test_ircheck.py
index 30ddd39fef0d..008963642272 100644
--- a/mypyc/test/test_ircheck.py
+++ b/mypyc/test/test_ircheck.py
@@ -5,7 +5,17 @@
 from mypyc.analysis.ircheck import FnError, can_coerce_to, check_func_ir
 from mypyc.ir.class_ir import ClassIR
 from mypyc.ir.func_ir import FuncDecl, FuncIR, FuncSignature
-from mypyc.ir.ops import Assign, BasicBlock, Goto, Integer, LoadLiteral, Op, Register, Return
+from mypyc.ir.ops import (
+    Assign,
+    BasicBlock,
+    Goto,
+    Integer,
+    LoadAddress,
+    LoadLiteral,
+    Op,
+    Register,
+    Return,
+)
 from mypyc.ir.pprint import format_func
 from mypyc.ir.rtypes import (
     RInstance,
@@ -16,6 +26,7 @@
     int64_rprimitive,
     none_rprimitive,
     object_rprimitive,
+    pointer_rprimitive,
     str_rprimitive,
 )
 
@@ -88,7 +99,7 @@ def test_invalid_register_source(self) -> None:
         ret = Return(value=Register(type=none_rprimitive, name="r1"))
         block = self.basic_block([ret])
         fn = FuncIR(decl=self.func_decl(name="func_1"), arg_regs=[], blocks=[block])
-        assert_has_error(fn, FnError(source=ret, desc="Invalid op reference to register r1"))
+        assert_has_error(fn, FnError(source=ret, desc="Invalid op reference to register 'r1'"))
 
     def test_invalid_op_source(self) -> None:
         ret = Return(value=LoadLiteral(value="foo", rtype=str_rprimitive))
@@ -170,3 +181,19 @@ def test_pprint(self) -> None:
             "    goto L1",
             "  ERR: Invalid control operation target: 1",
         ]
+
+    def test_load_address_declares_register(self) -> None:
+        rx = Register(str_rprimitive, "x")
+        ry = Register(pointer_rprimitive, "y")
+        load_addr = LoadAddress(pointer_rprimitive, rx)
+        assert_no_errors(
+            FuncIR(
+                decl=self.func_decl(name="func_1"),
+                arg_regs=[],
+                blocks=[
+                    self.basic_block(
+                        ops=[load_addr, Assign(ry, load_addr), Return(value=NONE_VALUE)]
+                    )
+                ],
+            )
+        )