python · JukkaL · Dec 28, 2020 · Nov 24, 2020 · Nov 27, 2020 · Nov 27, 2020
diff --git a/mypyc/analysis/dataflow.py b/mypyc/analysis/dataflow.py
@@ -11,6 +11,7 @@
     LoadStatic, InitStatic, MethodCall, RaiseStandardError, CallC, LoadGlobal,
     Truncate, BinaryIntOp, LoadMem, GetElementPtr, LoadAddress, ComparisonOp, SetMem
 )
+from mypyc.ir.func_ir import all_values
 
 
 class CFG:
@@ -361,8 +362,10 @@ def analyze_undefined_regs(blocks: List[BasicBlock],
 
     A register is undefined if there is some path from initial block
     where it has an undefined value.
+
+    Function arguments are assumed to be always defined.
     """
-    initial_undefined = set(env.regs()) - initial_defined
+    initial_undefined = set(all_values([], blocks)) - initial_defined
     return run_analysis(blocks=blocks,
                         cfg=cfg,
                         gen_and_kill=UndefinedVisitor(),

diff --git a/mypyc/codegen/emitfunc.py b/mypyc/codegen/emitfunc.py
@@ -18,7 +18,7 @@
     RType, RTuple, is_tagged, is_int32_rprimitive, is_int64_rprimitive, RStruct,
     is_pointer_rprimitive
 )
-from mypyc.ir.func_ir import FuncIR, FuncDecl, FUNC_STATICMETHOD, FUNC_CLASSMETHOD
+from mypyc.ir.func_ir import FuncIR, FuncDecl, FUNC_STATICMETHOD, FUNC_CLASSMETHOD, all_values
 from mypyc.ir.class_ir import ClassIR
 from mypyc.ir.const_int import find_constant_integer_registers
 from mypyc.ir.pprint import generate_names_for_env
@@ -55,18 +55,20 @@ def generate_native_function(fn: FuncIR,
     else:
         const_int_regs = {}
     declarations = Emitter(emitter.context, fn.env)
-    names = generate_names_for_env(fn.env)
+    names = generate_names_for_env(fn.arg_regs, fn.blocks)
     body = Emitter(emitter.context, fn.env, names)
     visitor = FunctionEmitterVisitor(body, declarations, source_path, module_name, const_int_regs)
 
     declarations.emit_line('{} {{'.format(native_function_header(fn.decl, emitter)))
     body.indent()
 
-    for r, i in fn.env.indexes.items():
+    for r in all_values(fn.arg_regs, fn.blocks):
         if isinstance(r.type, RTuple):
             emitter.declare_tuple_struct(r.type)
-        if i < len(fn.args):
-            continue  # skip the arguments
+
+        if r in fn.arg_regs:
+            continue  # Skip the arguments
+
         ctype = emitter.ctype_spaced(r.type)
         init = ''
         if r in fn.env.vars_needing_init:

diff --git a/mypyc/ir/func_ir.py b/mypyc/ir/func_ir.py
@@ -6,7 +6,9 @@
 from mypy.nodes import FuncDef, Block, ARG_POS, ARG_OPT, ARG_NAMED_OPT
 
 from mypyc.common import JsonDict
-from mypyc.ir.ops import DeserMaps, BasicBlock, Environment
+from mypyc.ir.ops import (
+    DeserMaps, BasicBlock, Environment, Value, Register, Assign, ControlOp, LoadAddress
+)
 from mypyc.ir.rtypes import RType, deserialize_type
 from mypyc.namegen import NameGenerator
 
@@ -151,11 +153,13 @@ class FuncIR:
 
     def __init__(self,
                  decl: FuncDecl,
+                 arg_regs: List[Register],
                  blocks: List[BasicBlock],
                  env: Environment,
                  line: int = -1,
                  traceback_name: Optional[str] = None) -> None:
         self.decl = decl
+        self.arg_regs = arg_regs
         self.blocks = blocks
         self.env = env
         self.line = line
@@ -210,10 +214,66 @@ def deserialize(cls, data: JsonDict, ctx: DeserMaps) -> 'FuncIR':
         return FuncIR(
             FuncDecl.deserialize(data['decl'], ctx),
             [],
+            [],
             Environment(),
             data['line'],
             data['traceback_name'],
         )
 
 
 INVALID_FUNC_DEF = FuncDef('<INVALID_FUNC_DEF>', [], Block([]))  # type: Final
+
+
+def all_values(args: List[Register], blocks: List[BasicBlock]) -> List[Value]:
+    """Return the set of all values that may be initialized in the blocks.
+
+    This omits registers that are only read.
+    """
+    values = list(args)  # type: List[Value]
+    seen_registers = set(args)
+
+    for block in blocks:
+        for op in block.ops:
+            if not isinstance(op, ControlOp):
+                if isinstance(op, Assign):
+                    if op.dest not in seen_registers:
+                        values.append(op.dest)
+                        seen_registers.add(op.dest)
+                elif op.is_void:
+                    continue
+                else:
+                    # If we take the address of a register, it might get initialized.
+                    if (isinstance(op, LoadAddress)
+                            and isinstance(op.src, Register)
+                            and op.src not in seen_registers):
+                        values.append(op.src)
+                        seen_registers.add(op.src)
+                    values.append(op)
+
+    return values
+
+
+def all_values_full(args: List[Register], blocks: List[BasicBlock]) -> List[Value]:
+    """Return set of all values that are initialized or accessed."""
+    values = list(args)  # type: List[Value]
+    seen_registers = set(args)
+
+    for block in blocks:
+        for op in block.ops:
+            for source in op.sources():
+                # Look for unitialized registers that are accessed. Ignore
+                # non-registers since we don't allow ops outside basic blocks.
+                if isinstance(source, Register) and source not in seen_registers:
+                    values.append(source)
+                    seen_registers.add(source)
+            if not isinstance(op, ControlOp):
+                if isinstance(op, Assign):
+                    if op.dest not in seen_registers:
+                        values.append(op.dest)
+                        seen_registers.add(op.dest)
+                elif op.is_void:
+                    continue
+                else:
+                    values.append(op)
+
+    return values
diff --git a/mypyc/ir/ops.py b/mypyc/ir/ops.py
@@ -12,15 +12,12 @@
 
 from abc import abstractmethod
 from typing import (
-    List, Sequence, Dict, Generic, TypeVar, Optional, NamedTuple, Tuple, Union, Iterable, Set
+    List, Sequence, Dict, Generic, TypeVar, Optional, NamedTuple, Tuple, Union, Set
 )
-from mypy.ordered_dict import OrderedDict
 
 from typing_extensions import Final, Type, TYPE_CHECKING
 from mypy_extensions import trait
 
-from mypy.nodes import SymbolNode
-
 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,
@@ -114,59 +111,11 @@ def __init__(self, items: List[AssignmentTarget],
 
 
 class Environment:
-    """Maintain the register symbol table and manage temp generation"""
+    # TODO: Remove this class
 
     def __init__(self) -> None:
-        self.indexes = OrderedDict()  # type: Dict[Value, int]
-        self.symtable = OrderedDict()  # type: OrderedDict[SymbolNode, AssignmentTarget]
         self.vars_needing_init = set()  # type: Set[Value]
 
-    def regs(self) -> Iterable['Value']:
-        return self.indexes.keys()
-
-    def add(self, value: 'Value') -> None:
-        self.indexes[value] = len(self.indexes)
-
-    def add_local(self, symbol: SymbolNode, typ: RType, is_arg: bool = False) -> 'Register':
-        """Add register that represents a symbol to the symbol table.
-
-        Args:
-            is_arg: is this a function argument
-        """
-        assert isinstance(symbol, SymbolNode)
-        reg = Register(typ, symbol.line, is_arg=is_arg, name=symbol.name)
-        self.symtable[symbol] = AssignmentTargetRegister(reg)
-        self.add(reg)
-        return reg
-
-    def add_local_reg(self, symbol: SymbolNode,
-                      typ: RType, is_arg: bool = False) -> AssignmentTargetRegister:
-        """Like add_local, but return an assignment target instead of value."""
-        self.add_local(symbol, typ, is_arg)
-        target = self.symtable[symbol]
-        assert isinstance(target, AssignmentTargetRegister)
-        return target
-
-    def add_target(self, symbol: SymbolNode, target: AssignmentTarget) -> AssignmentTarget:
-        self.symtable[symbol] = target
-        return target
-
-    def lookup(self, symbol: SymbolNode) -> AssignmentTarget:
-        return self.symtable[symbol]
-
-    def add_temp(self, typ: RType) -> 'Register':
-        """Add register that contains a temporary value with the given type."""
-        assert isinstance(typ, RType)
-        reg = Register(typ)
-        self.add(reg)
-        return reg
-
-    def add_op(self, reg: 'RegisterOp') -> None:
-        """Record the value of an operation."""
-        if reg.is_void:
-            return
-        self.add(reg)
-
 
 class BasicBlock:
     """Basic IR block.
@@ -244,9 +193,9 @@ class Register(Value):
     (but not all) temporary values.
     """
 
-    def __init__(self, type: RType, line: int = -1, is_arg: bool = False, name: str = '') -> None:
-        self.name = name
+    def __init__(self, type: RType, name: str = '', is_arg: bool = False, line: int = -1) -> None:
         self.type = type
+        self.name = name
         self.is_arg = is_arg
         self.is_borrowed = is_arg
         self.line = line
@@ -255,6 +204,9 @@ def __init__(self, type: RType, line: int = -1, is_arg: bool = False, name: str
     def is_void(self) -> bool:
         return False
 
+    def __repr__(self) -> str:
+        return '<Register %r at %s>' % (self.name, hex(id(self)))
+
 
 class Op(Value):
     """Abstract base class for all operations (as opposed to values)."""

diff --git a/mypyc/ir/pprint.py b/mypyc/ir/pprint.py
@@ -10,9 +10,9 @@
     Goto, Branch, Return, Unreachable, Assign, LoadInt, LoadErrorValue, GetAttr, SetAttr,
     LoadStatic, InitStatic, TupleGet, TupleSet, IncRef, DecRef, Call, MethodCall, Cast, Box, Unbox,
     RaiseStandardError, CallC, Truncate, LoadGlobal, BinaryIntOp, ComparisonOp, LoadMem, SetMem,
-    GetElementPtr, LoadAddress, Register, Value, OpVisitor, BasicBlock, Environment
+    GetElementPtr, LoadAddress, Register, Value, OpVisitor, BasicBlock, Environment, ControlOp
 )
-from mypyc.ir.func_ir import FuncIR
+from mypyc.ir.func_ir import FuncIR, all_values_full
 from mypyc.ir.module_ir import ModuleIRs
 from mypyc.ir.rtypes import is_bool_rprimitive, is_int_rprimitive, RType
 from mypyc.ir.const_int import find_constant_integer_registers
@@ -132,13 +132,13 @@ def visit_unbox(self, op: Unbox) -> str:
     def visit_raise_standard_error(self, op: RaiseStandardError) -> str:
         if op.value is not None:
             if isinstance(op.value, str):
-                return 'raise %s(%r)' % (op.class_name, op.value)
+                return self.format('%r = raise %s(%s)', op, op.class_name, repr(op.value))
             elif isinstance(op.value, Value):
-                return self.format('raise %s(%r)', op.class_name, op.value)
+                return self.format('%r = raise %s(%r)', op, op.class_name, op.value)
             else:
                 assert False, 'value type must be either str or Value'
         else:
-            return 'raise %s' % op.class_name
+            return self.format('%r = raise %s', op, op.class_name)
 
     def visit_call_c(self, op: CallC) -> str:
         args_str = ', '.join(self.format('%r', arg) for arg in op.args)
@@ -244,12 +244,12 @@ def format(self, fmt: str, *args: Any) -> str:
         return ''.join(result)
 
 
-def env_to_lines(env: Environment,
-                 names: Dict[Value, str],
-                 const_regs: Optional[Dict[LoadInt, int]] = None) -> List[str]:
+def format_registers(func_ir: FuncIR,
+                     names: Dict[Value, str],
+                     const_regs: Optional[Dict[LoadInt, int]] = None) -> List[str]:
     result = []
     i = 0
-    regs = list(env.regs())
+    regs = all_values_full(func_ir.arg_regs, func_ir.blocks)
     if const_regs is None:
         const_regs = {}
     regs = [reg for reg in regs if reg not in const_regs]
@@ -323,8 +323,8 @@ def format_func(fn: FuncIR) -> List[str]:
                                         ', '.join(arg.name for arg in fn.args)))
     # compute constants
     const_regs = find_constant_integer_registers(fn.blocks)
-    names = generate_names_for_env(fn.env)
-    for line in env_to_lines(fn.env, names, const_regs):
+    names = generate_names_for_env(fn.arg_regs, fn.blocks)
+    for line in format_registers(fn, names, const_regs):
         lines.append('    ' + line)
     code = format_blocks(fn.blocks, fn.env, names, const_regs)
     lines.extend(code)
@@ -340,39 +340,60 @@ def format_modules(modules: ModuleIRs) -> List[str]:
     return ops
 
 
-def generate_names_for_env(env: Environment) -> Dict[Value, str]:
+def generate_names_for_env(args: List[Register], blocks: List[BasicBlock]) -> Dict[Value, str]:
     """Generate unique names for values in an environment.
 
     Give names such as 'r5' or 'i0' to temp values in IR which are useful
     when pretty-printing or generating C. Ensure generated names are unique.
     """
-    names = {}
+    names = {}  # type: Dict[Value, str]
     used_names = set()
 
     temp_index = 0
     int_index = 0
 
-    for value in env.indexes:
-        if isinstance(value, Register) and value.name:
-            name = value.name
-        elif isinstance(value, LoadInt):
-            name = 'i%d' % int_index
-            int_index += 1
-        else:
-            name = 'r%d' % temp_index
-            temp_index += 1
-
-        # Append _2, _3, ... if needed to make the name unique.
-        if name in used_names:
-            n = 2
-            while True:
-                candidate = '%s_%d' % (name, n)
-                if candidate not in used_names:
-                    name = candidate
-                    break
-                n += 1
-
-        names[value] = name
-        used_names.add(name)
+    for arg in args:
+        names[arg] = arg.name
+        used_names.add(arg.name)
+
+    for block in blocks:
+        for op in block.ops:
+            values = []
+
+            for source in op.sources():
+                if source not in names:
+                    values.append(source)
+
+            if isinstance(op, Assign):
+                values.append(op.dest)
+            elif isinstance(op, ControlOp) or op.is_void:
+                continue
+            elif op not in names:
+                values.append(op)
+
+            for value in values:
+                if value in names:
+                    continue
+                if isinstance(value, Register) and value.name:
+                    name = value.name
+                elif isinstance(value, LoadInt):
+                    name = 'i%d' % int_index
+                    int_index += 1
+                else:
+                    name = 'r%d' % temp_index
+                    temp_index += 1
+
+                # Append _2, _3, ... if needed to make the name unique.
+                if name in used_names:
+                    n = 2
+                    while True:
+                        candidate = '%s_%d' % (name, n)
+                        if candidate not in used_names:
+                            name = candidate
+                            break
+                        n += 1
+
+                names[value] = name
+                used_names.add(name)
 
     return names