"""
Support for lowering generators.
"""
import llvmlite.ir
from llvmlite.ir import Constant, IRBuilder
from numba.core import types, config, cgutils
from numba.core.funcdesc import FunctionDescriptor
class GeneratorDescriptor(FunctionDescriptor):
"""
The descriptor for a generator's next function.
"""
__slots__ = ()
@classmethod
def from_generator_fndesc(cls, func_ir, fndesc, gentype, mangler):
"""
Build a GeneratorDescriptor for the generator returned by the
function described by *fndesc*, with type *gentype*.
The generator inherits the env_name from the *fndesc*.
All emitted functions for the generator shares the same Env.
"""
assert isinstance(gentype, types.Generator)
restype = gentype.yield_type
args = ['gen']
argtypes = (gentype,)
qualname = fndesc.qualname + '.next'
unique_name = fndesc.unique_name + '.next'
self = cls(fndesc.native, fndesc.modname, qualname, unique_name,
fndesc.doc, fndesc.typemap, restype, fndesc.calltypes,
args, fndesc.kws, argtypes=argtypes, mangler=mangler,
inline=False, env_name=fndesc.env_name)
return self
@property
def llvm_finalizer_name(self):
"""
The LLVM name of the generator's finalizer function
(if <generator type>.has_finalizer is true).
"""
return 'finalize_' + self.mangled_name
class BaseGeneratorLower(object):
"""
Base support class for lowering generators.
"""
def __init__(self, lower):
self.context = lower.context
self.fndesc = lower.fndesc
self.library = lower.library
self.func_ir = lower.func_ir
self.lower = lower
self.geninfo = lower.generator_info
self.gentype = self.get_generator_type()
self.gendesc = GeneratorDescriptor.from_generator_fndesc(
lower.func_ir, self.fndesc, self.gentype, self.context.mangler)
# Helps packing non-omitted arguments into a structure
self.arg_packer = self.context.get_data_packer(self.fndesc.argtypes)
self.resume_blocks = {}
@property
def call_conv(self):
return self.lower.call_conv
def get_args_ptr(self, builder, genptr):
return cgutils.gep_inbounds(builder, genptr, 0, 1)
def get_resume_index_ptr(self, builder, genptr):
return cgutils.gep_inbounds(builder, genptr, 0, 0,
name='gen.resume_index')
def get_state_ptr(self, builder, genptr):
return cgutils.gep_inbounds(builder, genptr, 0, 2,
name='gen.state')
def lower_init_func(self, lower):
"""
Lower the generator's initialization function (which will fill up
the passed-by-reference generator structure).
"""
lower.setup_function(self.fndesc)
builder = lower.builder
# Insert the generator into the target context in order to allow
# calling from other Numba-compiled functions.
lower.context.insert_generator(self.gentype, self.gendesc,
[self.library])
# Init argument values
lower.extract_function_arguments()
lower.pre_lower()
# Initialize the return structure (i.e. the generator structure).
retty = self.context.get_return_type(self.gentype)
# Structure index #0: the initial resume index (0 == start of generator)
resume_index = self.context.get_constant(types.int32, 0)
# Structure index #1: the function arguments
argsty = retty.elements[1]
statesty = retty.elements[2]
lower.debug_print("# low_init_func incref")
# Incref all NRT arguments before storing into generator states
if self.context.enable_nrt:
for argty, argval in zip(self.fndesc.argtypes, lower.fnargs):
self.context.nrt.incref(builder, argty, argval)
# Filter out omitted arguments
argsval = self.arg_packer.as_data(builder, lower.fnargs)
# Zero initialize states
statesval = Constant(statesty, None)
gen_struct = cgutils.make_anonymous_struct(builder,
[resume_index, argsval,
statesval],
retty)
retval = self.box_generator_struct(lower, gen_struct)
lower.debug_print("# low_init_func before return")
self.call_conv.return_value(builder, retval)
lower.post_lower()
def lower_next_func(self, lower):
"""
Lower the generator's next() function (which takes the
passed-by-reference generator structure and returns the next
yielded value).
"""
lower.setup_function(self.gendesc)
lower.debug_print("# lower_next_func: {0}".format(self.gendesc.unique_name))
assert self.gendesc.argtypes[0] == self.gentype
builder = lower.builder
function = lower.function
# Extract argument values and other information from generator struct
genptr, = self.call_conv.get_arguments(function)
self.arg_packer.load_into(builder,
self.get_args_ptr(builder, genptr),
lower.fnargs)
self.resume_index_ptr = self.get_resume_index_ptr(builder, genptr)
self.gen_state_ptr = self.get_state_ptr(builder, genptr)
prologue = function.append_basic_block("generator_prologue")
# Lower the generator's Python code
entry_block_tail = lower.lower_function_body()
# Add block for StopIteration on entry
stop_block = function.append_basic_block("stop_iteration")
builder.position_at_end(stop_block)
self.call_conv.return_stop_iteration(builder)
# Add prologue switch to resume blocks
builder.position_at_end(prologue)
# First Python block is also the resume point on first next() call
first_block = self.resume_blocks[0] = lower.blkmap[lower.firstblk]
# Create front switch to resume points
switch = builder.switch(builder.load(self.resume_index_ptr),
stop_block)
for index, block in self.resume_blocks.items():
switch.add_case(index, block)
# Close tail of entry block
builder.position_at_end(entry_block_tail)
builder.branch(prologue)
def lower_finalize_func(self, lower):
"""
Lower the generator's finalizer.
"""
fnty = llvmlite.ir.FunctionType(llvmlite.ir.VoidType(),
[self.context.get_value_type(self.gentype)])
function = cgutils.get_or_insert_function(
lower.module, fnty, self.gendesc.llvm_finalizer_name)
entry_block = function.append_basic_block('entry')
builder = IRBuilder(entry_block)
genptrty = self.context.get_value_type(self.gentype)
genptr = builder.bitcast(function.args[0], genptrty)
self.lower_finalize_func_body(builder, genptr)
def return_from_generator(self, lower):
"""
Emit a StopIteration at generator end and mark the generator exhausted.
"""
indexval = Constant(self.resume_index_ptr.type.pointee, -1)
lower.builder.store(indexval, self.resume_index_ptr)
self.call_conv.return_stop_iteration(lower.builder)
def create_resumption_block(self, lower, index):
block_name = "generator_resume%d" % (index,)
block = lower.function.append_basic_block(block_name)
lower.builder.position_at_end(block)
self.resume_blocks[index] = block
def debug_print(self, builder, msg):
if config.DEBUG_JIT:
self.context.debug_print(builder, "DEBUGJIT: {0}".format(msg))
class GeneratorLower(BaseGeneratorLower):
"""
Support class for lowering nopython generators.
"""
def get_generator_type(self):
return self.fndesc.restype
def box_generator_struct(self, lower, gen_struct):
return gen_struct
def lower_finalize_func_body(self, builder, genptr):
"""
Lower the body of the generator's finalizer: decref all live
state variables.
"""
self.debug_print(builder, "# generator: finalize")
if self.context.enable_nrt:
# Always dereference all arguments
# self.debug_print(builder, "# generator: clear args")
args_ptr = self.get_args_ptr(builder, genptr)
for ty, val in self.arg_packer.load(builder, args_ptr):
self.context.nrt.decref(builder, ty, val)
self.debug_print(builder, "# generator: finalize end")
builder.ret_void()
class PyGeneratorLower(BaseGeneratorLower):
"""
Support class for lowering object mode generators.
"""
def get_generator_type(self):
"""
Compute the actual generator type (the generator function's return
type is simply "pyobject").
"""
return types.Generator(
gen_func=self.func_ir.func_id.func,
yield_type=types.pyobject,
arg_types=(types.pyobject,) * self.func_ir.arg_count,
state_types=(types.pyobject,) * len(self.geninfo.state_vars),
has_finalizer=True,
)
def box_generator_struct(self, lower, gen_struct):
"""
Box the raw *gen_struct* as a Python object.
"""
gen_ptr = cgutils.alloca_once_value(lower.builder, gen_struct)
return lower.pyapi.from_native_generator(gen_ptr, self.gentype, lower.envarg)
def init_generator_state(self, lower):
"""
NULL-initialize all generator state variables, to avoid spurious
decref's on cleanup.
"""
lower.builder.store(Constant(self.gen_state_ptr.type.pointee, None),
self.gen_state_ptr)
def lower_finalize_func_body(self, builder, genptr):
"""
Lower the body of the generator's finalizer: decref all live
state variables.
"""
pyapi = self.context.get_python_api(builder)
resume_index_ptr = self.get_resume_index_ptr(builder, genptr)
resume_index = builder.load(resume_index_ptr)
# If resume_index is 0, next() was never called
# If resume_index is -1, generator terminated cleanly
# (note function arguments are saved in state variables,
# so they don't need a separate cleanup step)
need_cleanup = builder.icmp_signed(
'>', resume_index, Constant(resume_index.type, 0))
with cgutils.if_unlikely(builder, need_cleanup):
# Decref all live vars (some may be NULL)
gen_state_ptr = self.get_state_ptr(builder, genptr)
for state_index in range(len(self.gentype.state_types)):
state_slot = cgutils.gep_inbounds(builder, gen_state_ptr,
0, state_index)
ty = self.gentype.state_types[state_index]
val = self.context.unpack_value(builder, ty, state_slot)
pyapi.decref(val)
builder.ret_void()
class LowerYield(object):
"""
Support class for lowering a particular yield point.
"""
def __init__(self, lower, yield_point, live_vars):
self.lower = lower
self.context = lower.context
self.builder = lower.builder
self.genlower = lower.genlower
self.gentype = self.genlower.gentype
self.gen_state_ptr = self.genlower.gen_state_ptr
self.resume_index_ptr = self.genlower.resume_index_ptr
self.yp = yield_point
self.inst = self.yp.inst
self.live_vars = live_vars
self.live_var_indices = [lower.generator_info.state_vars.index(v)
for v in live_vars]
def lower_yield_suspend(self):
self.lower.debug_print("# generator suspend")
# Save live vars in state
for state_index, name in zip(self.live_var_indices, self.live_vars):
state_slot = cgutils.gep_inbounds(self.builder, self.gen_state_ptr,
0, state_index)
ty = self.gentype.state_types[state_index]
# The yield might be in a loop, in which case the state might
# contain a predicate var that branches back to the loop head, in
# this case the var is live but in sequential lowering won't have
# been alloca'd yet, so do this here.
fetype = self.lower.typeof(name)
self.lower._alloca_var(name, fetype)
val = self.lower.loadvar(name)
# IncRef newly stored value
if self.context.enable_nrt:
self.context.nrt.incref(self.builder, ty, val)
self.context.pack_value(self.builder, ty, val, state_slot)
# Save resume index
indexval = Constant(self.resume_index_ptr.type.pointee,
self.inst.index)
self.builder.store(indexval, self.resume_index_ptr)
self.lower.debug_print("# generator suspend end")
def lower_yield_resume(self):
# Emit resumption point
self.genlower.create_resumption_block(self.lower, self.inst.index)
self.lower.debug_print("# generator resume")
# Reload live vars from state
for state_index, name in zip(self.live_var_indices, self.live_vars):
state_slot = cgutils.gep_inbounds(self.builder, self.gen_state_ptr,
0, state_index)
ty = self.gentype.state_types[state_index]
val = self.context.unpack_value(self.builder, ty, state_slot)
self.lower.storevar(val, name)
# Previous storevar is making an extra incref
if self.context.enable_nrt:
self.context.nrt.decref(self.builder, ty, val)
self.lower.debug_print("# generator resume end")