Repository: numba/numba

Issue:
Tuple getitem cannot compile with IntegerLiterals as index
- [x] I have tried using the latest released version of Numba (most recent is
 visible in the change log (https://github.com/numba/numba/blob/master/CHANGE_LOG).
- [x] I have included below a minimal working reproducer (if you are unsure how
 to write one see http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports).

Hi,

We recently found some code in our code base, that was obviously used to extend operator.getitem for BaseTuple types and IntegerLiterals as index. Without it the following fails to compile:
```
import operator
import numpy as np

import numba
from numba import types
from numba.core import cgutils
from numba import literally
from numba.extending import (
    make_attribute_wrapper,
    overload,
    models,
    register_model,
    type_callable,
    typeof_impl,
    lower_builtin,
)

def make_dummy_type():
    class Dummy(object):
        def __init__(self, value):
            self.value = value

    class DummyType(types.Type):
        def __init__(self, data):
            self.data = data
            super(DummyType, self).__init__(name="Dummy({})".format(data))

    @register_model(DummyType)
    class DummyModel(models.StructModel):
        def __init__(self, dmm, fe_type):
            data_type = types.Tuple(fe_type.data)
            members = [("data", data_type)]
            models.StructModel.__init__(self, dmm, fe_type, members)

    make_attribute_wrapper(DummyType, "data", "_data")

    @type_callable(Dummy)
    def type_dummy(context):
        def typer(value):
            return DummyType(value)

        return typer

    @lower_builtin(Dummy, types.Tuple)
    def impl_dummy(context, builder, sig, args):
        typ = sig.return_type
        [value] = args
        dummy = cgutils.create_struct_proxy(typ)(context, builder)
        dummy.data = value
        return dummy._getvalue()


    @typeof_impl.register(Dummy)
    def typeof_dummy(val, c):
        return DummyType(val)

    # Store attributes
    return Dummy, DummyType

Dummy, DummyType = make_dummy_type()

@overload(operator.getitem)
def dummy_getitem_ovld(self, idx):
    if not isinstance(self, DummyType):
        return None

    if isinstance(idx, types.IntegerLiteral):
        def dummy_getitem_impl(self, idx):
            return self._data[idx]

        return dummy_getitem_impl

@numba.njit
def test_impl_issue(x):
    d = Dummy(x)
    return d[0]


n = 11
a1 = np.ones(n, dtype='int')
a2 = np.arange(n, dtype=np.float64)
a3 = np.random.ranf(n)
print("DEBUG:", test_impl_issue((a1, a2, a3)))
```

And the error is following:
```
No implementation of function Function(<built-in function getitem>) found for signature:

 >>> getitem(Tuple(array(int32, 1d, C), array(float64, 1d, C), array(float64, 1d, C)), Literal[int](0))

There are 18 candidate implementations:
   - Of which 18 did not match due to:
   Overload of function 'getitem': File: <numerous>: Line N/A.
     With argument(s): '(Tuple(array(int32, 1d, C), array(float64, 1d, C), array(float64, 1d, C)), Literal[int](0))':
    No match.
```

As far as I see from the func_ir 'getitem' expr is not transformed to 'static_getitem', since idx is not const and so this continues to type_inference_stage and fails, since probably there are no templates that can handle this:
```
------------IR DUMP: dummy_getitem_ovld.<locals>.dummy_getitem_impl-------------
label 0:
    self = arg(0, name=self)                 ['self']
    idx = arg(1, name=idx)                   ['idx']
    $4load_attr.1 = getattr(value=self, attr=_data) ['$4load_attr.1', 'self']
    $8binary_subscr.3 = getitem(value=$4load_attr.1, index=idx) ['$4load_attr.1', '$8binary_subscr.3', 'idx']
    $10return_value.4 = cast(value=$8binary_subscr.3) ['$10return_value.4', '$8binary_subscr.3']
    return $10return_value.4                 ['$10return_value.4']
```

After that I found below simple script also fails:
```
import operator
import numpy as np
import numba
from numba import types
from numba.extending import overload

def bar(tup, idx):
    pass

@overload(bar)
def bar_ovld(tup, idx):
    if not isinstance(idx, types.IntegerLiteral):
        return None

    def _impl(tup, idx):
        return tup[idx]        # it works with attr[0]
    return _impl

@numba.njit
def test_bar(arrs):
    return bar(arrs, 0)

n = 11
a1 = np.ones(n, dtype=np.int32)
a2 = np.arange(n, dtype=np.float64)
a3 = np.random.ranf(n)
tup_of_arrays = (a1, a2, a3,)
print("DEBUG:", test_bar(tup_of_arrays))
```



Do you think we can add below code to Numba to fix this?
```
@overload(operator.getitem)
def base_tuple_getitem_ovld(tup, idx):
    if not (isinstance(tup, types.BaseTuple)
            and isinstance(idx, types.IntegerLiteral)):
        return None
    idx_val = idx.literal_value
    def base_tuple_getitem_impl(tup, idx):
        return tup[idx_val]
    return base_tuple_getitem_impl
```

I was using Numba: 0.51.0.dev0+147.g7bb3157cf

Best Regards,
Alexey.

Discussion:

Patch:
diff --git a/numba/cpython/tupleobj.py b/numba/cpython/tupleobj.py
index 7a547577f..34739623f 100644
--- a/numba/cpython/tupleobj.py
+++ b/numba/cpython/tupleobj.py
@@ -187,6 +187,24 @@ def iternext_unituple(context, builder, sig, args, result):
         builder.store(nidx, iterval.index)
 
 
+@overload(operator.getitem)
+def getitem_literal_idx(tup, idx):
+    """
+    Overloads BaseTuple getitem to cover cases where constant
+    inference and RewriteConstGetitems cannot replace it
+    with a static_getitem.
+    """
+    if not (isinstance(tup, types.BaseTuple)
+            and isinstance(idx, types.IntegerLiteral)):
+        return None
+
+    idx_val = idx.literal_value
+    def getitem_literal_idx_impl(tup, idx):
+        return tup[idx_val]
+
+    return getitem_literal_idx_impl
+
+
 @lower_builtin('typed_getitem', types.BaseTuple, types.Any)
 def getitem_typed(context, builder, sig, args):
     tupty, _ = sig.args


Test patch:
diff --git a/numba/tests/test_tuples.py b/numba/tests/test_tuples.py
index dfd93ad0e..f0ae01167 100644
--- a/numba/tests/test_tuples.py
+++ b/numba/tests/test_tuples.py
@@ -4,7 +4,7 @@ import itertools
 import numpy as np
 
 from numba.core.compiler import compile_isolated
-from numba import njit, jit, typeof
+from numba import njit, jit, typeof, literally
 from numba.core import types, errors, utils
 from numba.tests.support import TestCase, MemoryLeakMixin, tag
 import unittest
@@ -218,6 +218,17 @@ class TestOperations(TestCase):
                               [types.UniTuple(types.int64, 3)])
         self.assertPreciseEqual(cr.entry_point((4, 5, 6)), 3)
 
+    def test_index_literal(self):
+        # issue #6023, test non-static getitem with IntegerLiteral index
+        def pyfunc(tup, idx):
+            idx = literally(idx)
+            return tup[idx]
+        cfunc = njit(pyfunc)
+
+        tup = (4, 3.1, 'sss')
+        for i in range(len(tup)):
+            self.assertPreciseEqual(cfunc(tup, i), tup[i])
+
     def test_index(self):
         pyfunc = tuple_index
         cr = compile_isolated(pyfunc,