141 lines
4.8 KiB
Python
141 lines
4.8 KiB
Python
|
"""
|
||
|
|
Implementation of various iterable and iterator types.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from numba.core import types, cgutils
|
||
|
|
from numba.core.imputils import (
|
||
|
|
lower_builtin, iternext_impl, call_iternext, call_getiter,
|
||
|
|
impl_ret_borrowed, impl_ret_new_ref, RefType)
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
@lower_builtin('getiter', types.IteratorType)
|
||
|
|
def iterator_getiter(context, builder, sig, args):
|
||
|
|
[it] = args
|
||
|
|
return impl_ret_borrowed(context, builder, sig.return_type, it)
|
||
|
|
|
||
|
|
#-------------------------------------------------------------------------------
|
||
|
|
# builtin `enumerate` implementation
|
||
|
|
|
||
|
|
@lower_builtin(enumerate, types.IterableType)
|
||
|
|
@lower_builtin(enumerate, types.IterableType, types.Integer)
|
||
|
|
def make_enumerate_object(context, builder, sig, args):
|
||
|
|
assert len(args) == 1 or len(args) == 2 # enumerate(it) or enumerate(it, start)
|
||
|
|
srcty = sig.args[0]
|
||
|
|
|
||
|
|
if len(args) == 1:
|
||
|
|
src = args[0]
|
||
|
|
start_val = context.get_constant(types.intp, 0)
|
||
|
|
elif len(args) == 2:
|
||
|
|
src = args[0]
|
||
|
|
start_val = context.cast(builder, args[1], sig.args[1], types.intp)
|
||
|
|
|
||
|
|
iterobj = call_getiter(context, builder, srcty, src)
|
||
|
|
|
||
|
|
enum = context.make_helper(builder, sig.return_type)
|
||
|
|
|
||
|
|
countptr = cgutils.alloca_once(builder, start_val.type)
|
||
|
|
builder.store(start_val, countptr)
|
||
|
|
|
||
|
|
enum.count = countptr
|
||
|
|
enum.iter = iterobj
|
||
|
|
|
||
|
|
res = enum._getvalue()
|
||
|
|
return impl_ret_new_ref(context, builder, sig.return_type, res)
|
||
|
|
|
||
|
|
@lower_builtin('iternext', types.EnumerateType)
|
||
|
|
@iternext_impl(RefType.NEW)
|
||
|
|
def iternext_enumerate(context, builder, sig, args, result):
|
||
|
|
[enumty] = sig.args
|
||
|
|
[enum] = args
|
||
|
|
|
||
|
|
enum = context.make_helper(builder, enumty, value=enum)
|
||
|
|
|
||
|
|
count = builder.load(enum.count)
|
||
|
|
ncount = builder.add(count, context.get_constant(types.intp, 1))
|
||
|
|
builder.store(ncount, enum.count)
|
||
|
|
|
||
|
|
srcres = call_iternext(context, builder, enumty.source_type, enum.iter)
|
||
|
|
is_valid = srcres.is_valid()
|
||
|
|
result.set_valid(is_valid)
|
||
|
|
|
||
|
|
with builder.if_then(is_valid):
|
||
|
|
srcval = srcres.yielded_value()
|
||
|
|
result.yield_(context.make_tuple(builder, enumty.yield_type,
|
||
|
|
[count, srcval]))
|
||
|
|
|
||
|
|
|
||
|
|
#-------------------------------------------------------------------------------
|
||
|
|
# builtin `zip` implementation
|
||
|
|
|
||
|
|
@lower_builtin(zip, types.VarArg(types.Any))
|
||
|
|
def make_zip_object(context, builder, sig, args):
|
||
|
|
zip_type = sig.return_type
|
||
|
|
|
||
|
|
assert len(args) == len(zip_type.source_types)
|
||
|
|
|
||
|
|
zipobj = context.make_helper(builder, zip_type)
|
||
|
|
|
||
|
|
for i, (arg, srcty) in enumerate(zip(args, sig.args)):
|
||
|
|
zipobj[i] = call_getiter(context, builder, srcty, arg)
|
||
|
|
|
||
|
|
res = zipobj._getvalue()
|
||
|
|
return impl_ret_new_ref(context, builder, sig.return_type, res)
|
||
|
|
|
||
|
|
@lower_builtin('iternext', types.ZipType)
|
||
|
|
@iternext_impl(RefType.NEW)
|
||
|
|
def iternext_zip(context, builder, sig, args, result):
|
||
|
|
[zip_type] = sig.args
|
||
|
|
[zipobj] = args
|
||
|
|
|
||
|
|
zipobj = context.make_helper(builder, zip_type, value=zipobj)
|
||
|
|
|
||
|
|
if len(zipobj) == 0:
|
||
|
|
# zip() is an empty iterator
|
||
|
|
result.set_exhausted()
|
||
|
|
return
|
||
|
|
|
||
|
|
p_ret_tup = cgutils.alloca_once(builder,
|
||
|
|
context.get_value_type(zip_type.yield_type))
|
||
|
|
p_is_valid = cgutils.alloca_once_value(builder, value=cgutils.true_bit)
|
||
|
|
|
||
|
|
for i, (iterobj, srcty) in enumerate(zip(zipobj, zip_type.source_types)):
|
||
|
|
is_valid = builder.load(p_is_valid)
|
||
|
|
# Avoid calling the remaining iternext if a iterator has been exhausted
|
||
|
|
with builder.if_then(is_valid):
|
||
|
|
srcres = call_iternext(context, builder, srcty, iterobj)
|
||
|
|
is_valid = builder.and_(is_valid, srcres.is_valid())
|
||
|
|
builder.store(is_valid, p_is_valid)
|
||
|
|
val = srcres.yielded_value()
|
||
|
|
ptr = cgutils.gep_inbounds(builder, p_ret_tup, 0, i)
|
||
|
|
builder.store(val, ptr)
|
||
|
|
|
||
|
|
is_valid = builder.load(p_is_valid)
|
||
|
|
result.set_valid(is_valid)
|
||
|
|
|
||
|
|
with builder.if_then(is_valid):
|
||
|
|
result.yield_(builder.load(p_ret_tup))
|
||
|
|
|
||
|
|
|
||
|
|
#-------------------------------------------------------------------------------
|
||
|
|
# generator implementation
|
||
|
|
|
||
|
|
@lower_builtin('iternext', types.Generator)
|
||
|
|
@iternext_impl(RefType.BORROWED)
|
||
|
|
def iternext_zip(context, builder, sig, args, result):
|
||
|
|
genty, = sig.args
|
||
|
|
gen, = args
|
||
|
|
impl = context.get_generator_impl(genty)
|
||
|
|
status, retval = impl(context, builder, sig, args)
|
||
|
|
context.add_linking_libs(getattr(impl, 'libs', ()))
|
||
|
|
|
||
|
|
with cgutils.if_likely(builder, status.is_ok):
|
||
|
|
result.set_valid(True)
|
||
|
|
result.yield_(retval)
|
||
|
|
with cgutils.if_unlikely(builder, status.is_stop_iteration):
|
||
|
|
result.set_exhausted()
|
||
|
|
with cgutils.if_unlikely(builder,
|
||
|
|
builder.and_(status.is_error,
|
||
|
|
builder.not_(status.is_stop_iteration))):
|
||
|
|
context.call_conv.return_status_propagate(builder, status)
|