generators.py

"""
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")