Repository: numba/numba

Issue:
`zip`: on iterable's exhaustion succeeding iterables still get consumed
Using `Python 3.5.2`, `Numba 0.29.0`.
When an iterable is completely consumed by `zip`, one element from each succeeding (longer) iterable gets consumed.
https://docs.python.org/3/library/functions.html#zip states:

```
The iterator stops when the shortest input iterable is exhausted. 
```

and

```
The left-to-right evaluation order of the iterables is guaranteed.
```

. Numba seems to conform to "left-to-right evaluation" but does not stop consumption of succeeding iterables right away when the (leftmost) shortest iterable is exhausted. See:

``` python
from numba import njit


def zip_long_short(iterable, n):
    it = iter(iterable)
    front = list(zip(it, range(n)))
    back = list(it)
    return front, back


def zip_short_long(iterable, n):
    it = iter(iterable)
    front = list(zip(range(n), it))
    back = list(it)
    return front, back


iterable = list(range(10, 16))
len_front = 3

for pyfunc in (zip_long_short, zip_short_long):
    cfunc = njit(pyfunc)
    print(pyfunc.__name__)
    print(" pyfunc:", pyfunc(iterable, len_front))
    print(" cfunc :",  cfunc(iterable, len_front))
```

output:

``` python
zip_long_short
 pyfunc: ([(10, 0), (11, 1), (12, 2)], [14, 15])
 cfunc : ([(10, 0), (11, 1), (12, 2)], [14, 15])
zip_short_long
 pyfunc: ([(0, 10), (1, 11), (2, 12)], [13, 14, 15])
 cfunc : ([(0, 10), (1, 11), (2, 12)], [14, 15])
```

`zip`: on iterable's exhaustion succeeding iterables still get consumed
Using `Python 3.5.2`, `Numba 0.29.0`.
When an iterable is completely consumed by `zip`, one element from each succeeding (longer) iterable gets consumed.
https://docs.python.org/3/library/functions.html#zip states:

```
The iterator stops when the shortest input iterable is exhausted. 
```

and

```
The left-to-right evaluation order of the iterables is guaranteed.
```

. Numba seems to conform to "left-to-right evaluation" but does not stop consumption of succeeding iterables right away when the (leftmost) shortest iterable is exhausted. See:

``` python
from numba import njit


def zip_long_short(iterable, n):
    it = iter(iterable)
    front = list(zip(it, range(n)))
    back = list(it)
    return front, back


def zip_short_long(iterable, n):
    it = iter(iterable)
    front = list(zip(range(n), it))
    back = list(it)
    return front, back


iterable = list(range(10, 16))
len_front = 3

for pyfunc in (zip_long_short, zip_short_long):
    cfunc = njit(pyfunc)
    print(pyfunc.__name__)
    print(" pyfunc:", pyfunc(iterable, len_front))
    print(" cfunc :",  cfunc(iterable, len_front))
```

output:

``` python
zip_long_short
 pyfunc: ([(10, 0), (11, 1), (12, 2)], [14, 15])
 cfunc : ([(10, 0), (11, 1), (12, 2)], [14, 15])
zip_short_long
 pyfunc: ([(0, 10), (1, 11), (2, 12)], [13, 14, 15])
 cfunc : ([(0, 10), (1, 11), (2, 12)], [14, 15])
```



Discussion:
reproducible since numba 0.27.  I didn't check earlier version.

Correct, this is not a new regression but was also observed in 0.28.1 just before updating to 0.29.

reproducible since numba 0.27.  I didn't check earlier version.

Correct, this is not a new regression but was also observed in 0.28.1 just before updating to 0.29.


Resolving patch:
diff --git a/numba/targets/iterators.py b/numba/targets/iterators.py
--- a/numba/targets/iterators.py
+++ b/numba/targets/iterators.py
@@ -95,17 +95,26 @@ def iternext_zip(context, builder, sig, args, result):
         result.set_exhausted()
         return
 
-    is_valid = context.get_constant(types.boolean, True)
-    values = []
-
-    for iterobj, srcty in zip(zipobj, zip_type.source_types):
-        srcres = call_iternext(context, builder, srcty, iterobj)
-        is_valid = builder.and_(is_valid, srcres.is_valid())
-        values.append(srcres.yielded_value())
-
+    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_(context.make_tuple(builder, zip_type.yield_type, values))
+        result.yield_(builder.load(p_ret_tup))
 
 
 #-------------------------------------------------------------------------------


Test patch:
diff --git a/numba/tests/test_builtins.py b/numba/tests/test_builtins.py
--- a/numba/tests/test_builtins.py
+++ b/numba/tests/test_builtins.py
@@ -20,6 +20,9 @@
 
 no_pyobj_flags = Flags()
 
+nrt_no_pyobj_flags = Flags()
+nrt_no_pyobj_flags.set("nrt")
+
 
 def abs_usecase(x):
     return abs(x)
@@ -182,6 +185,17 @@ def zip_3_usecase():
     return result
 
 
+def zip_first_exhausted():
+    iterable = range(7)
+    n = 3
+    it = iter(iterable)
+    # 1st iterator is shorter
+    front = list(zip(range(n), it))
+    # Make sure that we didn't skip one in `it`
+    back = list(it)
+    return front, back
+
+
 def pow_op_usecase(x, y):
     return x ** y
 
@@ -903,6 +917,17 @@ def test_zip_0(self, flags=forceobj_flags):
     def test_zip_0_npm(self):
         self.test_zip_0(flags=no_pyobj_flags)
 
+    def test_zip_first_exhausted(self, flags=forceobj_flags):
+        """
+        Test side effect to the input iterators when a left iterator has been
+        exhausted before the ones on the right.
+        """
+        self.run_nullary_func(zip_first_exhausted, flags)
+
+    @tag('important')
+    def test_zip_first_exhausted_npm(self):
+        self.test_zip_first_exhausted(flags=nrt_no_pyobj_flags)
+
     def test_pow_op_usecase(self):
         args = [
             (2, 3),