Skip to content
Cannot retrieve contributors at this time
Support for lowering generators.
from llvmlite.llvmpy.core import Constant, Type, Builder
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__ = ()
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
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.call_conv = lower.call_conv
self.func_ir = lower.func_ir
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 = {}
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,
def get_state_ptr(self, builder, genptr):
return cgutils.gep_inbounds(builder, genptr, 0, 2,
def lower_init_func(self, lower):
Lower the generator's initialization function (which will fill up
the passed-by-reference generator structure).
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,
# Init argument values
# 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.null(statesty)
gen_struct = cgutils.make_anonymous_struct(builder,
[resume_index, argsval,
retval = self.box_generator_struct(lower, gen_struct)
lower.debug_print("# low_init_func before return")
self.call_conv.return_value(builder, retval)
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.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.get_args_ptr(builder, genptr),
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")
# Add prologue switch to resume blocks
# 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),
for index, block in self.resume_blocks.items():
switch.add_case(index, block)
# Close tail of entry block
def lower_finalize_func(self, lower):
Lower the generator's finalizer.
fnty = Type.function(Type.void(),
function = cgutils.get_or_insert_function(
lower.module, fnty, self.gendesc.llvm_finalizer_name)
entry_block = function.append_basic_block('entry')
builder = Builder(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 =, -1), self.resume_index_ptr)
def create_resumption_block(self, lower, index):
block_name = "generator_resume%d" % (index,)
block = lower.function.append_basic_block(block_name)
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")
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(
arg_types=(types.pyobject,) * self.func_ir.arg_count,
state_types=(types.pyobject,) * len(self.geninfo.state_vars),
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.
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,, 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)
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 =,
self.inst.index), 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")