[mypyc] Support undefined locals with native int types (#13616)

JukkaL · web-flow · commit 80dfb36b292f · 2022-09-07T16:31:47.000+01:00
There's no reserved value we can use to track if a local variable with a native int type hasn't been initialized yet. Instead, use bitmaps to track whether a local is defined, but only if the local has a native int type and can be undefined. Work on mypyc/mypyc#837.
diff --git a/mypyc/test-data/exceptions.test b/mypyc/test-data/exceptions.test
@@ -598,3 +598,44 @@ L0:
     r0 = c.x
     r1 = c.y
     return 1
+
+[case testConditionallyUndefinedI64]
+from mypy_extensions import i64
+
+def f(x: i64) -> i64:
+    if x:
+        y: i64 = 2
+    return y
+[out]
+def f(x):
+    x, r0, y :: int64
+    __locals_bitmap0 :: uint32
+    r1 :: bit
+    r2, r3 :: uint32
+    r4 :: bit
+    r5 :: bool
+    r6 :: int64
+L0:
+    r0 = <error> :: int64
+    y = r0
+    __locals_bitmap0 = 0
+    r1 = x != 0
+    if r1 goto L1 else goto L2 :: bool
+L1:
+    y = 2
+    r2 = __locals_bitmap0 | 1
+    __locals_bitmap0 = r2
+L2:
+    r3 = __locals_bitmap0 & 1
+    r4 = r3 == 0
+    if r4 goto L3 else goto L5 :: bool
+L3:
+    r5 = raise UnboundLocalError('local variable "y" referenced before assignment')
+    if not r5 goto L6 (error at f:-1) else goto L4 :: bool
+L4:
+    unreachable
+L5:
+    return y
+L6:
+    r6 = <error> :: int64
+    return r6
diff --git a/mypyc/test-data/run-i64.test b/mypyc/test-data/run-i64.test
@@ -920,3 +920,119 @@ def test_magic_default() -> None:
     assert a() == MAGIC
     assert a(1) == 1
     assert a(MAGIC) == MAGIC
+
+[case testI64UndefinedLocal]
+from typing_extensions import Final
+
+MYPY = False
+if MYPY:
+    from mypy_extensions import i64, i32
+
+from testutil import assertRaises
+
+MAGIC: Final = -113
+
+
+def test_conditionally_defined_local() -> None:
+    x = not int()
+    if x:
+        y: i64 = 5
+        z: i32 = 6
+    assert y == 5
+    assert z == 6
+
+def test_conditionally_undefined_local() -> None:
+    x = int()
+    if x:
+        y: i64 = 5
+        z: i32 = 6
+    else:
+        ok: i64 = 7
+    assert ok == 7
+    try:
+        print(y)
+    except NameError as e:
+        assert str(e) == 'local variable "y" referenced before assignment'
+    else:
+        assert False
+    try:
+        print(z)
+    except NameError as e:
+        assert str(e) == 'local variable "z" referenced before assignment'
+    else:
+        assert False
+
+def test_assign_error_value_conditionally() -> None:
+    x = int()
+    if not x:
+        y: i64 = MAGIC
+        z: i32 = MAGIC
+    assert y == MAGIC
+    assert z == MAGIC
+
+def test_many_locals() -> None:
+    x = int()
+    if x:
+        a0: i64 = 0
+        a1: i64 = 1
+        a2: i64 = 2
+        a3: i64 = 3
+        a4: i64 = 4
+        a5: i64 = 5
+        a6: i64 = 6
+        a7: i64 = 7
+        a8: i64 = 8
+        a9: i64 = 9
+        a10: i64 = 10
+        a11: i64 = 11
+        a12: i64 = 12
+        a13: i64 = 13
+        a14: i64 = 14
+        a15: i64 = 15
+        a16: i64 = 16
+        a17: i64 = 17
+        a18: i64 = 18
+        a19: i64 = 19
+        a20: i64 = 20
+        a21: i64 = 21
+        a22: i64 = 22
+        a23: i64 = 23
+        a24: i64 = 24
+        a25: i64 = 25
+        a26: i64 = 26
+        a27: i64 = 27
+        a28: i64 = 28
+        a29: i64 = 29
+        a30: i64 = 30
+        a31: i64 = 31
+        a32: i64 = 32
+        a33: i64 = 33
+    with assertRaises(NameError):
+        print(a0)
+    with assertRaises(NameError):
+        print(a31)
+    with assertRaises(NameError):
+        print(a32)
+    with assertRaises(NameError):
+        print(a33)
+    a0 = 5
+    assert a0 == 5
+    with assertRaises(NameError):
+        print(a31)
+    with assertRaises(NameError):
+        print(a32)
+    with assertRaises(NameError):
+        print(a33)
+    a32 = 55
+    assert a0 == 5
+    assert a32 == 55
+    with assertRaises(NameError):
+        print(a31)
+    with assertRaises(NameError):
+        print(a33)
+    a31 = 10
+    a33 = 20
+    assert a0 == 5
+    assert a31 == 10
+    assert a32 == 55
+    assert a33 == 20
diff --git a/mypyc/transform/uninit.py b/mypyc/transform/uninit.py
@@ -3,11 +3,15 @@
 from __future__ import annotations
 
 from mypyc.analysis.dataflow import AnalysisDict, analyze_must_defined_regs, cleanup_cfg, get_cfg
+from mypyc.common import BITMAP_BITS
 from mypyc.ir.func_ir import FuncIR, all_values
 from mypyc.ir.ops import (
     Assign,
     BasicBlock,
     Branch,
+    ComparisonOp,
+    Integer,
+    IntOp,
     LoadAddress,
     LoadErrorValue,
     Op,
@@ -16,6 +20,7 @@
     Unreachable,
     Value,
 )
+from mypyc.ir.rtypes import bitmap_rprimitive, is_fixed_width_rtype
 
 
 def insert_uninit_checks(ir: FuncIR) -> None:
@@ -38,6 +43,8 @@ def split_blocks_at_uninits(
 
     init_registers = []
     init_registers_set = set()
+    bitmap_registers: list[Register] = []  # Init status bitmaps
+    bitmap_backed: list[Register] = []  # These use bitmaps to track init status
 
     # First split blocks on ops that may raise.
     for block in blocks:
@@ -70,15 +77,28 @@ def split_blocks_at_uninits(
                         init_registers.append(src)
                         init_registers_set.add(src)
 
-                    cur_block.ops.append(
-                        Branch(
+                    if not is_fixed_width_rtype(src.type):
+                        cur_block.ops.append(
+                            Branch(
+                                src,
+                                true_label=error_block,
+                                false_label=new_block,
+                                op=Branch.IS_ERROR,
+                                line=op.line,
+                            )
+                        )
+                    else:
+                        # We need to use bitmap for this one.
+                        check_for_uninit_using_bitmap(
+                            cur_block.ops,
                             src,
-                            true_label=error_block,
-                            false_label=new_block,
-                            op=Branch.IS_ERROR,
-                            line=op.line,
+                            bitmap_registers,
+                            bitmap_backed,
+                            error_block,
+                            new_block,
+                            op.line,
                         )
-                    )
+
                     raise_std = RaiseStandardError(
                         RaiseStandardError.UNBOUND_LOCAL_ERROR,
                         f'local variable "{src.name}" referenced before assignment',
@@ -89,12 +109,82 @@ def split_blocks_at_uninits(
                     cur_block = new_block
             cur_block.ops.append(op)
 
+    if bitmap_backed:
+        update_register_assignments_to_set_bitmap(new_blocks, bitmap_registers, bitmap_backed)
+
     if init_registers:
         new_ops: list[Op] = []
         for reg in init_registers:
             err = LoadErrorValue(reg.type, undefines=True)
             new_ops.append(err)
             new_ops.append(Assign(reg, err))
+        for reg in bitmap_registers:
+            new_ops.append(Assign(reg, Integer(0, bitmap_rprimitive)))
         new_blocks[0].ops[0:0] = new_ops
 
     return new_blocks
+
+
+def check_for_uninit_using_bitmap(
+    ops: list[Op],
+    src: Register,
+    bitmap_registers: list[Register],
+    bitmap_backed: list[Register],
+    error_block: BasicBlock,
+    ok_block: BasicBlock,
+    line: int,
+) -> None:
+    """Check if src is defined using a bitmap.
+
+    Modifies ops, bitmap_registers and bitmap_backed.
+    """
+    if src not in bitmap_backed:
+        # Set up a new bitmap backed register.
+        bitmap_backed.append(src)
+        n = (len(bitmap_backed) - 1) // BITMAP_BITS
+        if len(bitmap_registers) <= n:
+            bitmap_registers.append(Register(bitmap_rprimitive, f"__locals_bitmap{n}"))
+
+    index = bitmap_backed.index(src)
+    masked = IntOp(
+        bitmap_rprimitive,
+        bitmap_registers[index // BITMAP_BITS],
+        Integer(1 << (index & (BITMAP_BITS - 1)), bitmap_rprimitive),
+        IntOp.AND,
+        line,
+    )
+    ops.append(masked)
+    chk = ComparisonOp(masked, Integer(0, bitmap_rprimitive), ComparisonOp.EQ)
+    ops.append(chk)
+    ops.append(Branch(chk, error_block, ok_block, Branch.BOOL))
+
+
+def update_register_assignments_to_set_bitmap(
+    blocks: list[BasicBlock], bitmap_registers: list[Register], bitmap_backed: list[Register]
+) -> None:
+    """Update some assignments to registers to also set a bit in a bitmap.
+
+    The bitmaps are used to track if a local variable has been assigned to.
+
+    Modifies blocks.
+    """
+    for block in blocks:
+        if any(isinstance(op, Assign) and op.dest in bitmap_backed for op in block.ops):
+            new_ops: list[Op] = []
+            for op in block.ops:
+                if isinstance(op, Assign) and op.dest in bitmap_backed:
+                    index = bitmap_backed.index(op.dest)
+                    new_ops.append(op)
+                    reg = bitmap_registers[index // BITMAP_BITS]
+                    new = IntOp(
+                        bitmap_rprimitive,
+                        reg,
+                        Integer(1 << (index & (BITMAP_BITS - 1)), bitmap_rprimitive),
+                        IntOp.OR,
+                        op.line,
+                    )
+                    new_ops.append(new)
+                    new_ops.append(Assign(reg, new))
+                else:
+                    new_ops.append(op)
+            block.ops = new_ops
