def setup_env_for_generator_class(builder: IRBuilder) -> None:
"""Populates the environment for a generator class."""
fitem = builder.fn_info.fitem
cls = builder.fn_info.generator_class
self_target = builder.add_self_to_env(cls.ir)
# Add the type, value, and traceback variables to the environment.
exc_type = builder.add_local(Var('type'), object_rprimitive, is_arg=True)
exc_val = builder.add_local(Var('value'), object_rprimitive, is_arg=True)
exc_tb = builder.add_local(Var('traceback'),
object_rprimitive,
is_arg=True)
# TODO: Use the right type here instead of object?
exc_arg = builder.add_local(Var('arg'), object_rprimitive, is_arg=True)
cls.exc_regs = (exc_type, exc_val, exc_tb)
cls.send_arg_reg = exc_arg
cls.self_reg = builder.read(self_target, fitem.line)
cls.curr_env_reg = load_outer_env(builder, cls.self_reg,
builder.symtables[-1])
# Define a variable representing the label to go to the next time
# the '__next__' function of the generator is called, and add it
# as an attribute to the environment class.
cls.next_label_target = builder.add_var_to_env_class(
Var(NEXT_LABEL_ATTR_NAME), int_rprimitive, cls, reassign=False)
# Add arguments from the original generator function to the
# environment of the generator class.
add_args_to_env(builder, local=False, base=cls, reassign=False)
# Set the next label register for the generator class.
cls.next_label_reg = builder.read(cls.next_label_target, fitem.line)
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.environment.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 gen_func_item(builder: IRBuilder,
fitem: FuncItem,
name: str,
sig: FuncSignature,
cdef: Optional[ClassDef] = None,
) -> Tuple[FuncIR, Optional[Value]]:
"""Generate and return the FuncIR for a given FuncDef.
If the given FuncItem is a nested function, then we generate a
callable class representing the function and use that instead of
the actual function. if the given FuncItem contains a nested
function, then we generate an environment class so that inner
nested functions can access the environment of the given FuncDef.
Consider the following nested function:
def a() -> None:
def b() -> None:
def c() -> None:
return None
return None
return None
The classes generated would look something like the following.
has pointer to +-------+
+--------------------------> | a_env |
| +-------+
| ^
| | has pointer to
+-------+ associated with +-------+
| b_obj | -------------------> | b_env |
+-------+ +-------+
^
|
+-------+ has pointer to |
| c_obj | --------------------------+
+-------+
"""
# TODO: do something about abstract methods.
func_reg = None # type: Optional[Value]
# We treat lambdas as always being nested because we always generate
# a class for lambdas, no matter where they are. (It would probably also
# work to special case toplevel lambdas and generate a non-class function.)
is_nested = fitem in builder.nested_fitems or isinstance(fitem, LambdaExpr)
contains_nested = fitem in builder.encapsulating_funcs.keys()
is_decorated = fitem in builder.fdefs_to_decorators
in_non_ext = False
class_name = None
if cdef:
ir = builder.mapper.type_to_ir[cdef.info]
in_non_ext = not ir.is_ext_class
class_name = cdef.name
builder.enter(FuncInfo(fitem, name, class_name, gen_func_ns(builder),
is_nested, contains_nested, is_decorated, in_non_ext))
# Functions that contain nested functions need an environment class to store variables that
# are free in their nested functions. Generator functions need an environment class to
# store a variable denoting the next instruction to be executed when the __next__ function
# is called, along with all the variables inside the function itself.
if builder.fn_info.contains_nested or builder.fn_info.is_generator:
setup_env_class(builder)
if builder.fn_info.is_nested or builder.fn_info.in_non_ext:
setup_callable_class(builder)
if builder.fn_info.is_generator:
# Do a first-pass and generate a function that just returns a generator object.
gen_generator_func(builder)
blocks, env, ret_type, fn_info = builder.leave()
func_ir, func_reg = gen_func_ir(builder, blocks, sig, env, fn_info, cdef)
# Re-enter the FuncItem and visit the body of the function this time.
builder.enter(fn_info)
setup_env_for_generator_class(builder)
load_outer_envs(builder, builder.fn_info.generator_class)
if builder.fn_info.is_nested and isinstance(fitem, FuncDef):
setup_func_for_recursive_call(builder, fitem, builder.fn_info.generator_class)
create_switch_for_generator_class(builder)
add_raise_exception_blocks_to_generator_class(builder, fitem.line)
else:
load_env_registers(builder)
gen_arg_defaults(builder)
if builder.fn_info.contains_nested and not builder.fn_info.is_generator:
finalize_env_class(builder)
builder.ret_types[-1] = sig.ret_type
# Add all variables and functions that are declared/defined within this
# function and are referenced in functions nested within this one to this
# function's environment class so the nested functions can reference
# them even if they are declared after the nested function's definition.
# Note that this is done before visiting the body of this function.
env_for_func = builder.fn_info # type: Union[FuncInfo, ImplicitClass]
if builder.fn_info.is_generator:
env_for_func = builder.fn_info.generator_class
elif builder.fn_info.is_nested or builder.fn_info.in_non_ext:
env_for_func = builder.fn_info.callable_class
if builder.fn_info.fitem in builder.free_variables:
# Sort the variables to keep things deterministic
for var in sorted(builder.free_variables[builder.fn_info.fitem],
key=lambda x: x.name):
if isinstance(var, Var):
rtype = builder.type_to_rtype(var.type)
builder.add_var_to_env_class(var, rtype, env_for_func, reassign=False)
if builder.fn_info.fitem in builder.encapsulating_funcs:
for nested_fn in builder.encapsulating_funcs[builder.fn_info.fitem]:
if isinstance(nested_fn, FuncDef):
# The return type is 'object' instead of an RInstance of the
# callable class because differently defined functions with
# the same name and signature across conditional blocks
# will generate different callable classes, so the callable
# class that gets instantiated must be generic.
builder.add_var_to_env_class(
nested_fn, object_rprimitive, env_for_func, reassign=False
)
builder.accept(fitem.body)
builder.maybe_add_implicit_return()
if builder.fn_info.is_generator:
populate_switch_for_generator_class(builder)
blocks, env, ret_type, fn_info = builder.leave()
if fn_info.is_generator:
add_methods_to_generator_class(builder, fn_info, sig, env, blocks, fitem.is_coroutine)
else:
func_ir, func_reg = gen_func_ir(builder, blocks, sig, env, fn_info, cdef)
calculate_arg_defaults(builder, fn_info, env, func_reg)
return (func_ir, func_reg)
