Tipragot
628be439b8
Cela permet de ne pas avoir de problèmes de compatibilité car python est dans le git.
347 lines
13 KiB
Python
347 lines
13 KiB
Python
"""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 __future__ import annotations
|
|
|
|
from mypy.nodes import ARG_OPT, Var
|
|
from mypyc.common import ENV_ATTR_NAME, NEXT_LABEL_ATTR_NAME, SELF_NAME
|
|
from mypyc.ir.class_ir import ClassIR
|
|
from mypyc.ir.func_ir import FuncDecl, FuncIR, FuncSignature, RuntimeArg
|
|
from mypyc.ir.ops import (
|
|
NO_TRACEBACK_LINE_NO,
|
|
BasicBlock,
|
|
Branch,
|
|
Call,
|
|
Goto,
|
|
Integer,
|
|
MethodCall,
|
|
RaiseStandardError,
|
|
Register,
|
|
Return,
|
|
SetAttr,
|
|
TupleSet,
|
|
Unreachable,
|
|
Value,
|
|
)
|
|
from mypyc.ir.rtypes import RInstance, int_rprimitive, object_rprimitive
|
|
from mypyc.irbuild.builder import IRBuilder, gen_arg_defaults
|
|
from mypyc.irbuild.context import FuncInfo, GeneratorClass
|
|
from mypyc.irbuild.env_class import (
|
|
add_args_to_env,
|
|
finalize_env_class,
|
|
load_env_registers,
|
|
load_outer_env,
|
|
)
|
|
from mypyc.irbuild.nonlocalcontrol import ExceptNonlocalControl
|
|
from mypyc.primitives.exc_ops import (
|
|
error_catch_op,
|
|
exc_matches_op,
|
|
raise_exception_with_tb_op,
|
|
reraise_exception_op,
|
|
restore_exc_info_op,
|
|
)
|
|
|
|
|
|
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)
|