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)