208 lines
8.5 KiB
Python
208 lines
8.5 KiB
Python
|
"""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 typing import Dict, Optional, Union
|
||
|
|
|
||
|
|
from mypy.nodes import FuncDef, SymbolNode
|
||
|
|
|
||
|
|
from mypyc.common import SELF_NAME, ENV_ATTR_NAME
|
||
|
|
from mypyc.ir.ops import Call, GetAttr, SetAttr, Value
|
||
|
|
from mypyc.ir.rtypes import RInstance, object_rprimitive
|
||
|
|
from mypyc.ir.class_ir import ClassIR
|
||
|
|
from mypyc.irbuild.builder import IRBuilder, SymbolTarget
|
||
|
|
from mypyc.irbuild.targets import AssignmentTargetAttr
|
||
|
|
from mypyc.irbuild.context import FuncInfo, ImplicitClass, GeneratorClass
|
||
|
|
|
||
|
|
|
||
|
|
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):
|
||
|
|
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 add_args_to_env(builder: IRBuilder,
|
||
|
|
local: bool = True,
|
||
|
|
base: Optional[Union[FuncInfo, ImplicitClass]] = None,
|
||
|
|
reassign: bool = True) -> None:
|
||
|
|
fn_info = builder.fn_info
|
||
|
|
if local:
|
||
|
|
for arg in fn_info.fitem.arguments:
|
||
|
|
rtype = builder.type_to_rtype(arg.variable.type)
|
||
|
|
builder.add_local_reg(arg.variable, rtype, is_arg=True)
|
||
|
|
else:
|
||
|
|
for arg in fn_info.fitem.arguments:
|
||
|
|
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]
|
||
|
|
)
|