1"""
2This file provides internal compiler utilities that support certain special
3operations with numpy.
4"""
5from numba.core import types, typing
6from numba.core.cgutils import unpack_tuple
7from numba.core.extending import intrinsic
8from numba.core.imputils import impl_ret_new_ref
9from numba.core.errors import RequireLiteralValue, TypingError
10
11from numba.cpython.unsafe.tuple import tuple_setitem
12
13
14@intrinsic
15def empty_inferred(typingctx, shape):
16    """A version of numpy.empty whose dtype is inferred by the type system.
17
18    Expects `shape` to be a int-tuple.
19
20    There is special logic in the type-inferencer to handle the "refine"-ing
21    of undefined dtype.
22    """
23    from numba.np.arrayobj import _empty_nd_impl
24
25    def codegen(context, builder, signature, args):
26        # check that the return type is now defined
27        arrty = signature.return_type
28        assert arrty.is_precise()
29        shapes = unpack_tuple(builder, args[0])
30        # redirect implementation to np.empty
31        res = _empty_nd_impl(context, builder, arrty, shapes)
32        return impl_ret_new_ref(context, builder, arrty, res._getvalue())
33
34    # make function signature
35    nd = len(shape)
36    array_ty = types.Array(ndim=nd, layout='C', dtype=types.undefined)
37    sig = array_ty(shape)
38    return sig, codegen
39
40
41@intrinsic
42def to_fixed_tuple(typingctx, array, length):
43    """Convert *array* into a tuple of *length*
44
45    Returns ``UniTuple(array.dtype, length)``
46
47    ** Warning **
48    - No boundchecking.
49      If *length* is longer than *array.size*, the behavior is undefined.
50    """
51    if not isinstance(length, types.IntegerLiteral):
52        raise RequireLiteralValue('*length* argument must be a constant')
53
54    if array.ndim != 1:
55        raise TypingError("Not supported on array.ndim={}".format(array.ndim))
56
57    # Determine types
58    tuple_size = int(length.literal_value)
59    tuple_type = types.UniTuple(dtype=array.dtype, count=tuple_size)
60    sig = tuple_type(array, length)
61
62    def codegen(context, builder, signature, args):
63        def impl(array, length, empty_tuple):
64            out = empty_tuple
65            for i in range(length):
66                out = tuple_setitem(out, i, array[i])
67            return out
68
69        inner_argtypes = [signature.args[0], types.intp, tuple_type]
70        inner_sig = typing.signature(tuple_type, *inner_argtypes)
71        ll_idx_type = context.get_value_type(types.intp)
72        # Allocate an empty tuple
73        empty_tuple = context.get_constant_undef(tuple_type)
74        inner_args = [args[0], ll_idx_type(tuple_size), empty_tuple]
75
76        res = context.compile_internal(builder, impl, inner_sig, inner_args)
77        return res
78
79    return sig, codegen
80
81