"""Generate IR for generator functions. A generator function is represented by a class that implements the generator protocol and keeps track of the generator state, including local variables. The top-level logic for dealing with generator functions is in mypyc.irbuild.function. """ from typing import List from mypy.nodes import Var, ARG_OPT from mypyc.common import SELF_NAME, NEXT_LABEL_ATTR_NAME, ENV_ATTR_NAME from mypyc.ir.ops import ( BasicBlock, Call, Return, Goto, Integer, SetAttr, Unreachable, RaiseStandardError, Value, Register, MethodCall, TupleSet, Branch, NO_TRACEBACK_LINE_NO ) from mypyc.ir.rtypes import RInstance, int_rprimitive, object_rprimitive from mypyc.ir.func_ir import FuncIR, FuncDecl, FuncSignature, RuntimeArg from mypyc.ir.class_ir import ClassIR from mypyc.irbuild.nonlocalcontrol import ExceptNonlocalControl from mypyc.primitives.exc_ops import ( raise_exception_with_tb_op, error_catch_op, exc_matches_op, reraise_exception_op, restore_exc_info_op ) from mypyc.irbuild.env_class import ( add_args_to_env, load_outer_env, load_env_registers, finalize_env_class ) from mypyc.irbuild.builder import IRBuilder, gen_arg_defaults from mypyc.irbuild.context import FuncInfo, GeneratorClass def gen_generator_func(builder: IRBuilder) -> None: setup_generator_class(builder) load_env_registers(builder) gen_arg_defaults(builder) finalize_env_class(builder) builder.add(Return(instantiate_generator_class(builder))) def instantiate_generator_class(builder: IRBuilder) -> Value: fitem = builder.fn_info.fitem generator_reg = builder.add(Call(builder.fn_info.generator_class.ir.ctor, [], fitem.line)) # Get the current environment register. If the current function is nested, then the # generator class gets instantiated from the callable class' '__call__' method, and hence # we use the callable class' environment register. Otherwise, we use the original # function's environment register. if builder.fn_info.is_nested: curr_env_reg = builder.fn_info.callable_class.curr_env_reg else: curr_env_reg = builder.fn_info.curr_env_reg # Set the generator class' environment attribute to point at the environment class # defined in the current scope. builder.add(SetAttr(generator_reg, ENV_ATTR_NAME, curr_env_reg, fitem.line)) # Set the generator class' environment class' NEXT_LABEL_ATTR_NAME attribute to 0. zero = Integer(0) builder.add(SetAttr(curr_env_reg, NEXT_LABEL_ATTR_NAME, zero, fitem.line)) return generator_reg def setup_generator_class(builder: IRBuilder) -> ClassIR: name = f'{builder.fn_info.namespaced_name()}_gen' generator_class_ir = ClassIR(name, builder.module_name, is_generated=True) generator_class_ir.attributes[ENV_ATTR_NAME] = RInstance(builder.fn_info.env_class) generator_class_ir.mro = [generator_class_ir] builder.classes.append(generator_class_ir) builder.fn_info.generator_class = GeneratorClass(generator_class_ir) return generator_class_ir def create_switch_for_generator_class(builder: IRBuilder) -> None: builder.add(Goto(builder.fn_info.generator_class.switch_block)) block = BasicBlock() builder.fn_info.generator_class.continuation_blocks.append(block) builder.activate_block(block) def populate_switch_for_generator_class(builder: IRBuilder) -> None: cls = builder.fn_info.generator_class line = builder.fn_info.fitem.line builder.activate_block(cls.switch_block) for label, true_block in enumerate(cls.continuation_blocks): false_block = BasicBlock() comparison = builder.binary_op( cls.next_label_reg, Integer(label), '==', line ) builder.add_bool_branch(comparison, true_block, false_block) builder.activate_block(false_block) builder.add(RaiseStandardError(RaiseStandardError.STOP_ITERATION, None, line)) builder.add(Unreachable()) def add_raise_exception_blocks_to_generator_class(builder: IRBuilder, line: int) -> None: """Add error handling blocks to a generator class. Generates blocks to check if error flags are set while calling the helper method for generator functions, and raises an exception if those flags are set. """ cls = builder.fn_info.generator_class assert cls.exc_regs is not None exc_type, exc_val, exc_tb = cls.exc_regs # Check to see if an exception was raised. error_block = BasicBlock() ok_block = BasicBlock() comparison = builder.translate_is_op(exc_type, builder.none_object(), 'is not', line) builder.add_bool_branch(comparison, error_block, ok_block) builder.activate_block(error_block) builder.call_c(raise_exception_with_tb_op, [exc_type, exc_val, exc_tb], line) builder.add(Unreachable()) builder.goto_and_activate(ok_block) def add_methods_to_generator_class(builder: IRBuilder, fn_info: FuncInfo, sig: FuncSignature, arg_regs: List[Register], blocks: List[BasicBlock], is_coroutine: bool) -> None: helper_fn_decl = add_helper_to_generator_class(builder, arg_regs, blocks, sig, fn_info) add_next_to_generator_class(builder, fn_info, helper_fn_decl, sig) add_send_to_generator_class(builder, fn_info, helper_fn_decl, sig) add_iter_to_generator_class(builder, fn_info) add_throw_to_generator_class(builder, fn_info, helper_fn_decl, sig) add_close_to_generator_class(builder, fn_info) if is_coroutine: add_await_to_generator_class(builder, fn_info) def add_helper_to_generator_class(builder: IRBuilder, arg_regs: List[Register], blocks: List[BasicBlock], sig: FuncSignature, fn_info: FuncInfo) -> FuncDecl: """Generates a helper method for a generator class, called by '__next__' and 'throw'.""" sig = FuncSignature((RuntimeArg(SELF_NAME, object_rprimitive), RuntimeArg('type', object_rprimitive), RuntimeArg('value', object_rprimitive), RuntimeArg('traceback', object_rprimitive), RuntimeArg('arg', object_rprimitive) ), sig.ret_type) helper_fn_decl = FuncDecl('__mypyc_generator_helper__', fn_info.generator_class.ir.name, builder.module_name, sig) helper_fn_ir = FuncIR(helper_fn_decl, arg_regs, blocks, fn_info.fitem.line, traceback_name=fn_info.fitem.name) fn_info.generator_class.ir.methods['__mypyc_generator_helper__'] = helper_fn_ir builder.functions.append(helper_fn_ir) return helper_fn_decl def add_iter_to_generator_class(builder: IRBuilder, fn_info: FuncInfo) -> None: """Generates the '__iter__' method for a generator class.""" with builder.enter_method(fn_info.generator_class.ir, '__iter__', object_rprimitive, fn_info): builder.add(Return(builder.self())) def add_next_to_generator_class(builder: IRBuilder, fn_info: FuncInfo, fn_decl: FuncDecl, sig: FuncSignature) -> None: """Generates the '__next__' method for a generator class.""" with builder.enter_method(fn_info.generator_class.ir, '__next__', object_rprimitive, fn_info): none_reg = builder.none_object() # Call the helper function with error flags set to Py_None, and return that result. result = builder.add(Call(fn_decl, [builder.self(), none_reg, none_reg, none_reg, none_reg], fn_info.fitem.line)) builder.add(Return(result)) def add_send_to_generator_class(builder: IRBuilder, fn_info: FuncInfo, fn_decl: FuncDecl, sig: FuncSignature) -> None: """Generates the 'send' method for a generator class.""" with builder.enter_method(fn_info.generator_class.ir, 'send', object_rprimitive, fn_info): arg = builder.add_argument('arg', object_rprimitive) none_reg = builder.none_object() # Call the helper function with error flags set to Py_None, and return that result. result = builder.add(Call( fn_decl, [builder.self(), none_reg, none_reg, none_reg, builder.read(arg)], fn_info.fitem.line)) builder.add(Return(result)) def add_throw_to_generator_class(builder: IRBuilder, fn_info: FuncInfo, fn_decl: FuncDecl, sig: FuncSignature) -> None: """Generates the 'throw' method for a generator class.""" with builder.enter_method(fn_info.generator_class.ir, 'throw', object_rprimitive, fn_info): typ = builder.add_argument('type', object_rprimitive) val = builder.add_argument('value', object_rprimitive, ARG_OPT) tb = builder.add_argument('traceback', object_rprimitive, ARG_OPT) # Because the value and traceback arguments are optional and hence # can be NULL if not passed in, we have to assign them Py_None if # they are not passed in. none_reg = builder.none_object() builder.assign_if_null(val, lambda: none_reg, builder.fn_info.fitem.line) builder.assign_if_null(tb, lambda: none_reg, builder.fn_info.fitem.line) # Call the helper function using the arguments passed in, and return that result. result = builder.add(Call( fn_decl, [builder.self(), builder.read(typ), builder.read(val), builder.read(tb), none_reg], fn_info.fitem.line)) builder.add(Return(result)) def add_close_to_generator_class(builder: IRBuilder, fn_info: FuncInfo) -> None: """Generates the '__close__' method for a generator class.""" with builder.enter_method(fn_info.generator_class.ir, 'close', object_rprimitive, fn_info): except_block, else_block = BasicBlock(), BasicBlock() builder.builder.push_error_handler(except_block) builder.goto_and_activate(BasicBlock()) generator_exit = builder.load_module_attr_by_fullname('builtins.GeneratorExit', fn_info.fitem.line) builder.add(MethodCall( builder.self(), 'throw', [generator_exit, builder.none_object(), builder.none_object()])) builder.goto(else_block) builder.builder.pop_error_handler() builder.activate_block(except_block) old_exc = builder.call_c(error_catch_op, [], fn_info.fitem.line) builder.nonlocal_control.append( ExceptNonlocalControl(builder.nonlocal_control[-1], old_exc)) stop_iteration = builder.load_module_attr_by_fullname('builtins.StopIteration', fn_info.fitem.line) exceptions = builder.add( TupleSet([generator_exit, stop_iteration], fn_info.fitem.line)) matches = builder.call_c( exc_matches_op, [exceptions], fn_info.fitem.line) match_block, non_match_block = BasicBlock(), BasicBlock() builder.add(Branch(matches, match_block, non_match_block, Branch.BOOL)) builder.activate_block(match_block) builder.call_c(restore_exc_info_op, [builder.read(old_exc)], fn_info.fitem.line) builder.add(Return(builder.none_object())) builder.activate_block(non_match_block) builder.call_c(reraise_exception_op, [], NO_TRACEBACK_LINE_NO) builder.add(Unreachable()) builder.nonlocal_control.pop() builder.activate_block(else_block) builder.add(RaiseStandardError(RaiseStandardError.RUNTIME_ERROR, 'generator ignored GeneratorExit', fn_info.fitem.line)) builder.add(Unreachable()) def add_await_to_generator_class(builder: IRBuilder, fn_info: FuncInfo) -> None: """Generates the '__await__' method for a generator class.""" with builder.enter_method(fn_info.generator_class.ir, '__await__', object_rprimitive, fn_info): builder.add(Return(builder.self())) 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)