mypy-dev 1.9.0.dev0__py3-none-any.whl

mypyc/irbuild/constant_fold.py

@@ -0,0 +1,95 @@
+"""Constant folding of IR values.
+
+For example, 3 + 5 can be constant folded into 8.
+
+This is mostly like mypy.constant_fold, but we can bind some additional
+NameExpr and MemberExpr references here, since we have more knowledge
+about which definitions can be trusted -- we constant fold only references
+to other compiled modules in the same compilation unit.
+"""
+
+from __future__ import annotations
+
+from typing import Final, Union
+
+from mypy.constant_fold import constant_fold_binary_op, constant_fold_unary_op
+from mypy.nodes import (
+    BytesExpr,
+    ComplexExpr,
+    Expression,
+    FloatExpr,
+    IntExpr,
+    MemberExpr,
+    NameExpr,
+    OpExpr,
+    StrExpr,
+    UnaryExpr,
+    Var,
+)
+from mypyc.irbuild.builder import IRBuilder
+from mypyc.irbuild.util import bytes_from_str
+
+# All possible result types of constant folding
+ConstantValue = Union[int, float, complex, str, bytes]
+CONST_TYPES: Final = (int, float, complex, str, bytes)
+
+
+def constant_fold_expr(builder: IRBuilder, expr: Expression) -> ConstantValue | None:
+    """Return the constant value of an expression for supported operations.
+
+    Return None otherwise.
+    """
+    if isinstance(expr, IntExpr):
+        return expr.value
+    if isinstance(expr, FloatExpr):
+        return expr.value
+    if isinstance(expr, StrExpr):
+        return expr.value
+    if isinstance(expr, BytesExpr):
+        return bytes_from_str(expr.value)
+    if isinstance(expr, ComplexExpr):
+        return expr.value
+    elif isinstance(expr, NameExpr):
+        node = expr.node
+        if isinstance(node, Var) and node.is_final:
+            final_value = node.final_value
+            if isinstance(final_value, (CONST_TYPES)):
+                return final_value
+    elif isinstance(expr, MemberExpr):
+        final = builder.get_final_ref(expr)
+        if final is not None:
+            fn, final_var, native = final
+            if final_var.is_final:
+                final_value = final_var.final_value
+                if isinstance(final_value, (CONST_TYPES)):
+                    return final_value
+    elif isinstance(expr, OpExpr):
+        left = constant_fold_expr(builder, expr.left)
+        right = constant_fold_expr(builder, expr.right)
+        if left is not None and right is not None:
+            return constant_fold_binary_op_extended(expr.op, left, right)
+    elif isinstance(expr, UnaryExpr):
+        value = constant_fold_expr(builder, expr.expr)
+        if value is not None and not isinstance(value, bytes):
+            return constant_fold_unary_op(expr.op, value)
+    return None
+
+
+def constant_fold_binary_op_extended(
+    op: str, left: ConstantValue, right: ConstantValue
+) -> ConstantValue | None:
+    """Like mypy's constant_fold_binary_op(), but includes bytes support.
+
+    mypy cannot use constant folded bytes easily so it's simpler to only support them in mypyc.
+    """
+    if not isinstance(left, bytes) and not isinstance(right, bytes):
+        return constant_fold_binary_op(op, left, right)
+
+    if op == "+" and isinstance(left, bytes) and isinstance(right, bytes):
+        return left + right
+    elif op == "*" and isinstance(left, bytes) and isinstance(right, int):
+        return left * right
+    elif op == "*" and isinstance(left, int) and isinstance(right, bytes):
+        return left * right
+
+    return None

mypyc/irbuild/context.py

@@ -0,0 +1,186 @@
+"""Helpers that store information about functions and the related classes."""
+
+from __future__ import annotations
+
+from mypy.nodes import FuncItem
+from mypyc.ir.class_ir import ClassIR
+from mypyc.ir.func_ir import INVALID_FUNC_DEF
+from mypyc.ir.ops import BasicBlock, Value
+from mypyc.irbuild.targets import AssignmentTarget
+
+
+class FuncInfo:
+    """Contains information about functions as they are generated."""
+
+    def __init__(
+        self,
+        fitem: FuncItem = INVALID_FUNC_DEF,
+        name: str = "",
+        class_name: str | None = None,
+        namespace: str = "",
+        is_nested: bool = False,
+        contains_nested: bool = False,
+        is_decorated: bool = False,
+        in_non_ext: bool = False,
+        add_nested_funcs_to_env: bool = False,
+    ) -> None:
+        self.fitem = fitem
+        self.name = name
+        self.class_name = class_name
+        self.ns = namespace
+        # Callable classes implement the '__call__' method, and are used to represent functions
+        # that are nested inside of other functions.
+        self._callable_class: ImplicitClass | None = None
+        # Environment classes are ClassIR instances that contain attributes representing the
+        # variables in the environment of the function they correspond to. Environment classes are
+        # generated for functions that contain nested functions.
+        self._env_class: ClassIR | None = None
+        # Generator classes implement the '__next__' method, and are used to represent generators
+        # returned by generator functions.
+        self._generator_class: GeneratorClass | None = None
+        # Environment class registers are the local registers associated with instances of an
+        # environment class, used for getting and setting attributes. curr_env_reg is the register
+        # associated with the current environment.
+        self._curr_env_reg: Value | None = None
+        # These are flags denoting whether a given function is nested, contains a nested function,
+        # is decorated, or is within a non-extension class.
+        self.is_nested = is_nested
+        self.contains_nested = contains_nested
+        self.is_decorated = is_decorated
+        self.in_non_ext = in_non_ext
+        self.add_nested_funcs_to_env = add_nested_funcs_to_env
+
+        # TODO: add field for ret_type: RType = none_rprimitive
+
+    def namespaced_name(self) -> str:
+        return "_".join(x for x in [self.name, self.class_name, self.ns] if x)
+
+    @property
+    def is_generator(self) -> bool:
+        return self.fitem.is_generator or self.fitem.is_coroutine
+
+    @property
+    def is_coroutine(self) -> bool:
+        return self.fitem.is_coroutine
+
+    @property
+    def callable_class(self) -> ImplicitClass:
+        assert self._callable_class is not None
+        return self._callable_class
+
+    @callable_class.setter
+    def callable_class(self, cls: ImplicitClass) -> None:
+        self._callable_class = cls
+
+    @property
+    def env_class(self) -> ClassIR:
+        assert self._env_class is not None
+        return self._env_class
+
+    @env_class.setter
+    def env_class(self, ir: ClassIR) -> None:
+        self._env_class = ir
+
+    @property
+    def generator_class(self) -> GeneratorClass:
+        assert self._generator_class is not None
+        return self._generator_class
+
+    @generator_class.setter
+    def generator_class(self, cls: GeneratorClass) -> None:
+        self._generator_class = cls
+
+    @property
+    def curr_env_reg(self) -> Value:
+        assert self._curr_env_reg is not None
+        return self._curr_env_reg
+
+
+class ImplicitClass:
+    """Contains information regarding implicitly generated classes.
+
+    Implicit classes are generated for nested functions and generator
+    functions. They are not explicitly defined in the source code.
+
+    NOTE: This is both a concrete class and used as a base class.
+    """
+
+    def __init__(self, ir: ClassIR) -> None:
+        # The ClassIR instance associated with this class.
+        self.ir = ir
+        # The register associated with the 'self' instance for this generator class.
+        self._self_reg: Value | None = None
+        # Environment class registers are the local registers associated with instances of an
+        # environment class, used for getting and setting attributes. curr_env_reg is the register
+        # associated with the current environment. prev_env_reg is the self.__mypyc_env__ field
+        # associated with the previous environment.
+        self._curr_env_reg: Value | None = None
+        self._prev_env_reg: Value | None = None
+
+    @property
+    def self_reg(self) -> Value:
+        assert self._self_reg is not None
+        return self._self_reg
+
+    @self_reg.setter
+    def self_reg(self, reg: Value) -> None:
+        self._self_reg = reg
+
+    @property
+    def curr_env_reg(self) -> Value:
+        assert self._curr_env_reg is not None
+        return self._curr_env_reg
+
+    @curr_env_reg.setter
+    def curr_env_reg(self, reg: Value) -> None:
+        self._curr_env_reg = reg
+
+    @property
+    def prev_env_reg(self) -> Value:
+        assert self._prev_env_reg is not None
+        return self._prev_env_reg
+
+    @prev_env_reg.setter
+    def prev_env_reg(self, reg: Value) -> None:
+        self._prev_env_reg = reg
+
+
+class GeneratorClass(ImplicitClass):
+    """Contains information about implicit generator function classes."""
+
+    def __init__(self, ir: ClassIR) -> None:
+        super().__init__(ir)
+        # This register holds the label number that the '__next__' function should go to the next
+        # time it is called.
+        self._next_label_reg: Value | None = None
+        self._next_label_target: AssignmentTarget | None = None
+
+        # These registers hold the error values for the generator object for the case that the
+        # 'throw' function is called.
+        self.exc_regs: tuple[Value, Value, Value] | None = None
+
+        # Holds the arg passed to send
+        self.send_arg_reg: Value | None = None
+
+        # The switch block is used to decide which instruction to go using the value held in the
+        # next-label register.
+        self.switch_block = BasicBlock()
+        self.continuation_blocks: list[BasicBlock] = []
+
+    @property
+    def next_label_reg(self) -> Value:
+        assert self._next_label_reg is not None
+        return self._next_label_reg
+
+    @next_label_reg.setter
+    def next_label_reg(self, reg: Value) -> None:
+        self._next_label_reg = reg
+
+    @property
+    def next_label_target(self) -> AssignmentTarget:
+        assert self._next_label_target is not None
+        return self._next_label_target
+
+    @next_label_target.setter
+    def next_label_target(self, target: AssignmentTarget) -> None:
+        self._next_label_target = target

mypyc/irbuild/env_class.py

@@ -0,0 +1,223 @@
+"""Generate classes representing function environments (+ related operations).
+
+If we have a nested function that has non-local (free) variables, access to the
+non-locals is via an instance of an environment class. Example:
+
+    def f() -> int:
+        x = 0  # Make 'x' an attribute of an environment class instance
+
+        def g() -> int:
+            # We have access to the environment class instance to
+            # allow accessing 'x'
+            return x + 2
+
+        x = x + 1  # Modify the attribute
+        return g()
+"""
+
+from __future__ import annotations
+
+from mypy.nodes import Argument, FuncDef, SymbolNode, Var
+from mypyc.common import BITMAP_BITS, ENV_ATTR_NAME, SELF_NAME, bitmap_name
+from mypyc.ir.class_ir import ClassIR
+from mypyc.ir.ops import Call, GetAttr, SetAttr, Value
+from mypyc.ir.rtypes import RInstance, bitmap_rprimitive, object_rprimitive
+from mypyc.irbuild.builder import IRBuilder, SymbolTarget
+from mypyc.irbuild.context import FuncInfo, GeneratorClass, ImplicitClass
+from mypyc.irbuild.targets import AssignmentTargetAttr
+
+
+def setup_env_class(builder: IRBuilder) -> ClassIR:
+    """Generate a class representing a function environment.
+
+    Note that the variables in the function environment are not
+    actually populated here. This is because when the environment
+    class is generated, the function environment has not yet been
+    visited. This behavior is allowed so that when the compiler visits
+    nested functions, it can use the returned ClassIR instance to
+    figure out free variables it needs to access.  The remaining
+    attributes of the environment class are populated when the
+    environment registers are loaded.
+
+    Return a ClassIR representing an environment for a function
+    containing a nested function.
+    """
+    env_class = ClassIR(
+        f"{builder.fn_info.namespaced_name()}_env", builder.module_name, is_generated=True
+    )
+    env_class.attributes[SELF_NAME] = RInstance(env_class)
+    if builder.fn_info.is_nested:
+        # If the function is nested, its environment class must contain an environment
+        # attribute pointing to its encapsulating functions' environment class.
+        env_class.attributes[ENV_ATTR_NAME] = RInstance(builder.fn_infos[-2].env_class)
+    env_class.mro = [env_class]
+    builder.fn_info.env_class = env_class
+    builder.classes.append(env_class)
+    return env_class
+
+
+def finalize_env_class(builder: IRBuilder) -> None:
+    """Generate, instantiate, and set up the environment of an environment class."""
+    instantiate_env_class(builder)
+
+    # Iterate through the function arguments and replace local definitions (using registers)
+    # that were previously added to the environment with references to the function's
+    # environment class.
+    if builder.fn_info.is_nested:
+        add_args_to_env(builder, local=False, base=builder.fn_info.callable_class)
+    else:
+        add_args_to_env(builder, local=False, base=builder.fn_info)
+
+
+def instantiate_env_class(builder: IRBuilder) -> Value:
+    """Assign an environment class to a register named after the given function definition."""
+    curr_env_reg = builder.add(
+        Call(builder.fn_info.env_class.ctor, [], builder.fn_info.fitem.line)
+    )
+
+    if builder.fn_info.is_nested:
+        builder.fn_info.callable_class._curr_env_reg = curr_env_reg
+        builder.add(
+            SetAttr(
+                curr_env_reg,
+                ENV_ATTR_NAME,
+                builder.fn_info.callable_class.prev_env_reg,
+                builder.fn_info.fitem.line,
+            )
+        )
+    else:
+        builder.fn_info._curr_env_reg = curr_env_reg
+
+    return curr_env_reg
+
+
+def load_env_registers(builder: IRBuilder) -> None:
+    """Load the registers for the current FuncItem being visited.
+
+    Adds the arguments of the FuncItem to the environment. If the
+    FuncItem is nested inside of another function, then this also
+    loads all of the outer environments of the FuncItem into registers
+    so that they can be used when accessing free variables.
+    """
+    add_args_to_env(builder, local=True)
+
+    fn_info = builder.fn_info
+    fitem = fn_info.fitem
+    if fn_info.is_nested:
+        load_outer_envs(builder, fn_info.callable_class)
+        # If this is a FuncDef, then make sure to load the FuncDef into its own environment
+        # class so that the function can be called recursively.
+        if isinstance(fitem, FuncDef) and fn_info.add_nested_funcs_to_env:
+            setup_func_for_recursive_call(builder, fitem, fn_info.callable_class)
+
+
+def load_outer_env(
+    builder: IRBuilder, base: Value, outer_env: dict[SymbolNode, SymbolTarget]
+) -> Value:
+    """Load the environment class for a given base into a register.
+
+    Additionally, iterates through all of the SymbolNode and
+    AssignmentTarget instances of the environment at the given index's
+    symtable, and adds those instances to the environment of the
+    current environment. This is done so that the current environment
+    can access outer environment variables without having to reload
+    all of the environment registers.
+
+    Returns the register where the environment class was loaded.
+    """
+    env = builder.add(GetAttr(base, ENV_ATTR_NAME, builder.fn_info.fitem.line))
+    assert isinstance(env.type, RInstance), f"{env} must be of type RInstance"
+
+    for symbol, target in outer_env.items():
+        env.type.class_ir.attributes[symbol.name] = target.type
+        symbol_target = AssignmentTargetAttr(env, symbol.name)
+        builder.add_target(symbol, symbol_target)
+
+    return env
+
+
+def load_outer_envs(builder: IRBuilder, base: ImplicitClass) -> None:
+    index = len(builder.builders) - 2
+
+    # Load the first outer environment. This one is special because it gets saved in the
+    # FuncInfo instance's prev_env_reg field.
+    if index > 1:
+        # outer_env = builder.fn_infos[index].environment
+        outer_env = builder.symtables[index]
+        if isinstance(base, GeneratorClass):
+            base.prev_env_reg = load_outer_env(builder, base.curr_env_reg, outer_env)
+        else:
+            base.prev_env_reg = load_outer_env(builder, base.self_reg, outer_env)
+        env_reg = base.prev_env_reg
+        index -= 1
+
+    # Load the remaining outer environments into registers.
+    while index > 1:
+        # outer_env = builder.fn_infos[index].environment
+        outer_env = builder.symtables[index]
+        env_reg = load_outer_env(builder, env_reg, outer_env)
+        index -= 1
+
+
+def num_bitmap_args(builder: IRBuilder, args: list[Argument]) -> int:
+    n = 0
+    for arg in args:
+        t = builder.type_to_rtype(arg.variable.type)
+        if t.error_overlap and arg.kind.is_optional():
+            n += 1
+    return (n + (BITMAP_BITS - 1)) // BITMAP_BITS
+
+
+def add_args_to_env(
+    builder: IRBuilder,
+    local: bool = True,
+    base: FuncInfo | ImplicitClass | None = None,
+    reassign: bool = True,
+) -> None:
+    fn_info = builder.fn_info
+    args = fn_info.fitem.arguments
+    nb = num_bitmap_args(builder, args)
+    if local:
+        for arg in args:
+            rtype = builder.type_to_rtype(arg.variable.type)
+            builder.add_local_reg(arg.variable, rtype, is_arg=True)
+        for i in reversed(range(nb)):
+            builder.add_local_reg(Var(bitmap_name(i)), bitmap_rprimitive, is_arg=True)
+    else:
+        for arg in args:
+            if is_free_variable(builder, arg.variable) or fn_info.is_generator:
+                rtype = builder.type_to_rtype(arg.variable.type)
+                assert base is not None, "base cannot be None for adding nonlocal args"
+                builder.add_var_to_env_class(arg.variable, rtype, base, reassign=reassign)
+
+
+def setup_func_for_recursive_call(builder: IRBuilder, fdef: FuncDef, base: ImplicitClass) -> None:
+    """Enable calling a nested function (with a callable class) recursively.
+
+    Adds the instance of the callable class representing the given
+    FuncDef to a register in the environment so that the function can
+    be called recursively. Note that this needs to be done only for
+    nested functions.
+    """
+    # First, set the attribute of the environment class so that GetAttr can be called on it.
+    prev_env = builder.fn_infos[-2].env_class
+    prev_env.attributes[fdef.name] = builder.type_to_rtype(fdef.type)
+
+    if isinstance(base, GeneratorClass):
+        # If we are dealing with a generator class, then we need to first get the register
+        # holding the current environment class, and load the previous environment class from
+        # there.
+        prev_env_reg = builder.add(GetAttr(base.curr_env_reg, ENV_ATTR_NAME, -1))
+    else:
+        prev_env_reg = base.prev_env_reg
+
+    # Obtain the instance of the callable class representing the FuncDef, and add it to the
+    # current environment.
+    val = builder.add(GetAttr(prev_env_reg, fdef.name, -1))
+    target = builder.add_local_reg(fdef, object_rprimitive)
+    builder.assign(target, val, -1)
+
+
+def is_free_variable(builder: IRBuilder, symbol: SymbolNode) -> bool:
+    fitem = builder.fn_info.fitem
+    return fitem in builder.free_variables and symbol in builder.free_variables[fitem]