Repository: numba/numba

Issue:
iteration fails in tuple comprehension
Hello,

As discussed on [gitter](https://gitter.im/numba/numba?at=613a2952d206b85e4f8c176e), I found a bug in list comprehension when passing a tuple as input:
```
@numba.njit
def f1_fine(array_tuple):
    vals = []
    for a in array_tuple:
        vals.append(a[0])
    return np.array(vals)

@numba.njit
def f2_fine(array_tuple):
    return np.array([array_tuple[i][0] for i in range(len(array_tuple))])


@numba.njit
def f3_crash(array_tuple):
    return np.array([a[0] for a in array_tuple])

at = (np.ones(4), np.ones(3))

f1_fine(at)
f2_fine(at)
f3_crash(at)
```

While these functions are pretty much the same, `f3` does not compile with error
```
TypingError: Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<intrinsic range_iter_len>) found for signature:
 
 >>> <unknown function>(iter(UniTuple(array(float64, 1d, C) x 2)))
 
There are 2 candidate implementations:
      - Of which 2 did not match due to:
      Intrinsic of function '<unknown function>': File: numba/cpython/rangeobj.py: Line 183.
        With argument(s): '(iter(UniTuple(array(float64, 1d, C) x 2)))':
       No match.

During: resolving callee type: Function(<intrinsic range_iter_len>)
During: typing of call at <ipython-input-20-4047ac44d0b7> (15)


File "<ipython-input-20-4047ac44d0b7>", line 15:
def f3_crash(array_tuple):
    return np.array([a[0] for a in array_tuple])
    ^
```
numba version `0.53.1`


Discussion:
Seems to be related to #6550

I also poked around a bit and found what I believe to be the source of the error:

https://github.com/numba/numba/blob/e6c4bd1dd3badaf7b3de414b0e730a9f5291d0ce/numba/core/inline_closurecall.py#L1044

The function `_find_iter_range` raises `GuardException` when it does not find a `range` iterator. This exception is turned into `None` by `guard` (as expected). However, the following `if` statements uses that info for different things and essentially assumes that there was a `range_iter` and only differentiates between range iter with and without step (as far as I can tell). A fix for this specific situation is to insert `require(range_def is not None)`, but I am not sure if that works in general.
> Seems to be related to #6550
> 
> I also poked around a bit and found what I believe to be the source of the error:
> 
> https://github.com/numba/numba/blob/e6c4bd1dd3badaf7b3de414b0e730a9f5291d0ce/numba/core/inline_closurecall.py#L1044
> 
> The function `_find_iter_range` raises `GuardException` when it does not find a `range` iterator. This exception is turned into `None` by `guard` (as expected). However, the following `if` statements uses that info for different things and essentially assumes that there was a `range_iter` and only differentiates between range iter with and without step (as far as I can tell). A fix for this specific situation is to insert `require(range_def is not None)`, but I am not sure if that works in general.

Thanks for debugging this @MegaIng, the fix you describe will prevent this problem from occurring and I think given the nature of what the optimisation subsequently does it is an appropriate way to fix it. A PR with unit test based on the bug added to: https://github.com/numba/numba/blob/e6c4bd1dd3badaf7b3de414b0e730a9f5291d0ce/numba/tests/test_comprehension.py#L241 would be welcomed if you have time. Many Thanks!
While writing the tests right now, I noticed that the most recent master apparently  transforms `f1_fine` in such a way that it now also breaks.
The fix does not work. It breaks a few tests in the same class: https://github.com/numba/numba/blob/e6c4bd1dd3badaf7b3de414b0e730a9f5291d0ce/numba/tests/test_comprehension.py#L401, https://github.com/numba/numba/blob/e6c4bd1dd3badaf7b3de414b0e730a9f5291d0ce/numba/tests/test_comprehension.py#L274 and https://github.com/numba/numba/blob/e6c4bd1dd3badaf7b3de414b0e730a9f5291d0ce/numba/tests/test_comprehension.py#L342 .

Apparently the function `range_iter_len` is just misnamed. It is supposed to handle all possible things one might throw at an array comprehension, not just range iterators. So currently the correct fix is to tell `range_iter_len` how to deal with tuple iterators (which is what is suggested as solution in #6550). However, that means that for every iterator type there might be the function has to be implemented.

Alternatively, is there a way to check if a function is implemented for the given parameters? We could check that there, and if no implementation is found reject the optimization.
> While writing the tests right now, I noticed that the most recent master apparently transforms `f1_fine` in such a way that it now also breaks.

Thanks. On which version of Numba does it work? I can't reproduce it working on 0.53.1:
```
$ cat issue7394.py 
from numba import njit, version_info
import numpy as np

print(version_info)

@njit
def f1_fine(array_tuple):
    vals = []
    for a in array_tuple:
        vals.append(a[0])
    return np.array(vals)

at = (np.ones(4), np.ones(3))

f1_fine(at)

$ python issue7394.py 
version_info(major=0, minor=53, patch=1, short=(0, 53), full=(0, 53, 1), string='0.53.1', tuple=('0', '53', '1'), git_revision=None)
Traceback (most recent call last):
  File "issue7394.py", line 15, in <module>
    f1_fine(at)
  File "<path>/numba/core/dispatcher.py", line 420, in _compile_for_args
    error_rewrite(e, 'typing')
  File "<path>/numba/core/dispatcher.py", line 361, in error_rewrite
    raise e.with_traceback(None)
numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<intrinsic range_iter_len>) found for signature:
 
 >>> <unknown function>(iter(UniTuple(array(float64, 1d, C) x 2)))
 
There are 2 candidate implementations:
   - Of which 2 did not match due to:
   Intrinsic of function '<unknown function>': File: numba/cpython/rangeobj.py: Line 183.
     With argument(s): '(iter(UniTuple(array(float64, 1d, C) x 2)))':
    No match.

During: resolving callee type: Function(<intrinsic range_iter_len>)
During: typing of call at issue7394.py (8)


File "issue7394.py", line 8:
def f1_fine(array_tuple):
    vals = []
```
@stuartarchibald For me it works on 0.54.0  and on 0.53.1

Python version differences are probably the problem.
> @stuartarchibald For me it works on 0.54.0 and on 0.53.1
> 
> Python version differences are probably the problem.

Yes, are you on 3.7? I tested this on 3.9.
@stuartarchibald I have two different environments: One for numba development, and another where I have numba installed from pip. Both are Python version 3.9.4, but on the one with numba master `f1_fine` doesn't work, on the other it does.
Ok, no worries, I am stupid. My one version of `f1_fine` did not have a `numba.jit` decorator. It fails everywhere.
I run `f1_fine` with python `3.7.11`, numpy `1.20.3` and numba `0.53.1+8.gddf817af0.dirty`, it works. So yes probably python version.
RE failing tests, thanks for trying this out. I've taken a bit more of a look at the code and agree that the original suggestion will not work for the reasons you noted here: https://github.com/numba/numba/issues/7394#issuecomment-916838990, i.e. if there's no `range` then `range_iter_len` is used and it needs to work out the length of the iterator.

I think the issue here is that as the array comprehension transform is untyped it's not possible to know prior to doing the transform whether `range_iter_len` actually supports computing a length for the type of the iterator.

Two things probably need doing:
1. Write code to handle the tuple case.
2. Raise a `TypingError` if none of the types are matched in `range_iter_len` with a message spelling out the problem.

Untested patch for the above items:
```diff
diff --git a/numba/cpython/rangeobj.py b/numba/cpython/rangeobj.py
index 0f61a4f..8a26cbe 100644
--- a/numba/cpython/rangeobj.py
+++ b/numba/cpython/rangeobj.py
@@ -7,7 +7,7 @@ import operator
 import llvmlite.llvmpy.core as lc
 
 from numba import prange
-from numba.core import types, cgutils
+from numba.core import types, cgutils, errors
 from numba.cpython.listobj import ListIterInstance
 from numba.np.arrayobj import make_array
 from numba.core.imputils import (lower_builtin, lower_cast,
@@ -213,7 +213,18 @@ def range_iter_len(typingctx, val):
             # array iterates along the outer dimension
             return impl_ret_untracked(context, builder, intp_t, shape[0])
         return signature(types.intp, val), codegen
-
+    elif isinstance(val, types.UniTupleIter):
+        def  codegen(context, builder, sig, args):
+            (iterty,) = sig.args
+            tuplety = iterty.container
+            intp_t = context.get_value_type(types.intp)
+            count_const = intp_t(tuplety.count)
+            return impl_ret_untracked(context, builder, intp_t, count_const)
+        return signature(types.intp, val), codegen
+    else:
+        msg = ('Unsupported iterator found in array comprehension, try '
+               'preallocating the array and filling manually.')
+        raise errors.TypingError(msg)
```
I came to the same conclusion. I also think that `range_iter_len` should be renamed (I choose `comprehension_iter_len`). It probably should also be moved to another module, including the tests, but for the moment I am just going to leave it there.

Resolving patch:
diff --git a/numba/core/inline_closurecall.py b/numba/core/inline_closurecall.py
--- a/numba/core/inline_closurecall.py
+++ b/numba/core/inline_closurecall.py
@@ -1068,10 +1068,12 @@ def is_removed(val, removed):
         # this doesn't work in objmode as it's effectively untyped
         if typed:
             len_func_var = ir.Var(scope, mk_unique_var("len_func"), loc)
-            from numba.cpython.rangeobj import range_iter_len
+            from numba.cpython.rangeobj import length_of_iterator
             stmts.append(_new_definition(func_ir, len_func_var,
-                        ir.Global('range_iter_len', range_iter_len, loc=loc),
-                        loc))
+                                         ir.Global('length_of_iterator',
+                                                   length_of_iterator,
+                                                   loc=loc),
+                                         loc))
             size_val = ir.Expr.call(len_func_var, (iter_var,), (), loc=loc)
         else:
             raise GuardException
diff --git a/numba/cpython/rangeobj.py b/numba/cpython/rangeobj.py
--- a/numba/cpython/rangeobj.py
+++ b/numba/cpython/rangeobj.py
@@ -7,7 +7,7 @@
 import llvmlite.llvmpy.core as lc
 
 from numba import prange
-from numba.core import types, cgutils
+from numba.core import types, cgutils, errors
 from numba.cpython.listobj import ListIterInstance
 from numba.np.arrayobj import make_array
 from numba.core.imputils import (lower_builtin, lower_cast,
@@ -181,9 +181,10 @@ def range_to_range(context, builder, fromty, toty, val):
     return cgutils.make_anonymous_struct(builder, items)
 
 @intrinsic
-def range_iter_len(typingctx, val):
+def length_of_iterator(typingctx, val):
     """
-    An implementation of len(range_iter) for internal use.
+    An implementation of len(iter) for internal use.
+    Primary use is for array comprehensions (see inline_closurecall).
     """
     if isinstance(val, types.RangeIteratorType):
         val_type = val.yield_type
@@ -213,7 +214,27 @@ def  codegen(context, builder, sig, args):
             # array iterates along the outer dimension
             return impl_ret_untracked(context, builder, intp_t, shape[0])
         return signature(types.intp, val), codegen
+    elif isinstance(val, types.UniTupleIter):
+        def codegen(context, builder, sig, args):
+            (iterty,) = sig.args
+            tuplety = iterty.container
+            intp_t = context.get_value_type(types.intp)
+            count_const = intp_t(tuplety.count)
+            return impl_ret_untracked(context, builder, intp_t, count_const)
 
+        return signature(types.intp, val), codegen
+    elif isinstance(val, types.ListTypeIteratorType):
+        def codegen(context, builder, sig, args):
+            (value,) = args
+            intp_t = context.get_value_type(types.intp)
+            from numba.typed.listobject import ListIterInstance
+            iterobj = ListIterInstance(context, builder, sig.args[0], value)
+            return impl_ret_untracked(context, builder, intp_t, iterobj.size)
+        return signature(types.intp, val), codegen
+    else:
+        msg = ('Unsupported iterator found in array comprehension, try '
+               'preallocating the array and filling manually.')
+        raise errors.TypingError(msg)
 
 def make_range_attr(index, attribute):
     @intrinsic


Test patch:
diff --git a/numba/tests/test_comprehension.py b/numba/tests/test_comprehension.py
--- a/numba/tests/test_comprehension.py
+++ b/numba/tests/test_comprehension.py
@@ -7,7 +7,7 @@
 import numpy
 
 from numba.core.compiler import compile_isolated
-from numba import jit
+from numba import jit, typed
 from numba.core import types, utils
 from numba.core.errors import TypingError, LoweringError
 from numba.core.types.functions import _header_lead
@@ -370,6 +370,16 @@ def comp_nest_with_dependency(n):
         self.assertIn(_header_lead, str(raises.exception))
         self.assertIn('array(undefined,', str(raises.exception))
 
+    def test_comp_unsupported_iter(self):
+        def comp_unsupported_iter():
+            val = zip([1, 2, 3], [4, 5, 6])
+            return np.array([a for a, b in val])
+        with self.assertRaises(TypingError) as raises:
+            self.check(comp_unsupported_iter)
+        self.assertIn(_header_lead, str(raises.exception))
+        self.assertIn('Unsupported iterator found in array comprehension',
+                      str(raises.exception))
+
     def test_no_array_comp(self):
         def no_array_comp1(n):
             l = [1,2,3,4]
@@ -407,6 +417,10 @@ def array_comp(a):
         self.check(array_comp, l)
         # with array iterator
         self.check(array_comp, np.array(l))
+        # with tuple iterator (issue #7394)
+        self.check(array_comp, tuple(l))
+        # with typed.List iterator (issue #6550)
+        self.check(array_comp, typed.List(l))
 
     def test_array_comp_with_dtype(self):
         def array_comp(n):
diff --git a/numba/tests/test_range.py b/numba/tests/test_range.py
--- a/numba/tests/test_range.py
+++ b/numba/tests/test_range.py
@@ -9,7 +9,7 @@
 from numba.core import types, utils
 from numba.tests.support import tag
 
-from numba.cpython.rangeobj import range_iter_len
+from numba.cpython.rangeobj import length_of_iterator
 def loop1(n):
     s = 0
     for i in range(n):
@@ -40,10 +40,10 @@ def range_len2(a, b):
 def range_len3(a, b, c):
     return len(range(a, b, c))
 def range_iter_len1(a):
-    return range_iter_len(iter(range(a)))
+    return length_of_iterator(iter(range(a)))
 
 def range_iter_len2(a):
-    return range_iter_len(iter(a))
+    return length_of_iterator(iter(a))
 
 def range_attrs(start, stop, step):
     r1 = range(start)