Repository: numba/numba

Issue:
Problems with numpy array iteration
I am getting bad behavior dealing with numpy-array iteration. The following shows an error when invoking `ndenumerate`. I get the same behavior when calling `.flat` on an array.  

## ndenumerate
```python
import numpy as np
import numba

def f():
    x = np.arange(6).reshape(2, 3)
    for index, item in np.ndenumerate(x):
        return index, item

jitted_f = numba.jit(f)
```
```python
>>> print(numba.__version__)
0.29.0

>>> print(f())
((0, 0), 0)

>>> print(jitted_f())
((0, 0), -2387225703656530210)
```

## flat
```python
def f():
    x = np.arange(6).reshape(2, 3)
    for item in x.flat:
        return item

jitted_f = numba.jit(f)
```
```python
>>> print(numba.__version__)
0.29.0

>>> print(f())
0

>>> print(jitted_f())
-2387225703656530210
```

This occurs in Python 3.5 for `numba 0.29.0`, and in Python 2.7 in `numba 0.29.0+0.gead0dd3.dirty`


Discussion:

Resolving patch:
diff --git a/numba/datamodel/models.py b/numba/datamodel/models.py
--- a/numba/datamodel/models.py
+++ b/numba/datamodel/models.py
@@ -984,7 +984,7 @@ def __init__(self, dmm, fe_type, need_indices):
         array_type = fe_type.array_type
         dtype = array_type.dtype
         ndim = array_type.ndim
-        members = [('array', types.EphemeralPointer(array_type)),
+        members = [('array', array_type),
                    ('stride', types.intp),
                    ('index', types.EphemeralPointer(types.intp)),
                    ]
@@ -999,7 +999,7 @@ def __init__(self, dmm, fe_type):
         array_type = fe_type.array_type
         dtype = array_type.dtype
         ndim = array_type.ndim
-        members = [('array', types.EphemeralPointer(array_type)),
+        members = [('array', array_type),
                    ('pointers', types.EphemeralArray(types.CPointer(dtype), ndim)),
                    ('indices', types.EphemeralArray(types.intp, ndim)),
                    ('exhausted', types.EphemeralPointer(types.boolean)),
diff --git a/numba/targets/arrayobj.py b/numba/targets/arrayobj.py
--- a/numba/targets/arrayobj.py
+++ b/numba/targets/arrayobj.py
@@ -2790,11 +2790,10 @@ def make_array_flatiter(context, builder, arrty, arr):
     flatitercls = make_array_flat_cls(types.NumpyFlatType(arrty))
     flatiter = flatitercls(context, builder)
 
-    arrayptr = cgutils.alloca_once_value(builder, arr)
-    flatiter.array = arrayptr
+    flatiter.array = arr
 
     arrcls = context.make_array(arrty)
-    arr = arrcls(context, builder, ref=arrayptr)
+    arr = arrcls(context, builder, ref=flatiter._get_ptr_by_name('array'))
 
     flatiter.init_specific(context, builder, arrty, arr)
 
@@ -2813,7 +2812,7 @@ def iternext_numpy_flatiter(context, builder, sig, args, result):
 
     arrty = flatiterty.array_type
     arrcls = context.make_array(arrty)
-    arr = arrcls(context, builder, value=builder.load(flatiter.array))
+    arr = arrcls(context, builder, value=flatiter.array)
 
     flatiter.iternext_specific(context, builder, arrty, arr, result)
 
@@ -2828,7 +2827,7 @@ def iternext_numpy_getitem(context, builder, sig, args):
 
     arrty = flatiterty.array_type
     arrcls = context.make_array(arrty)
-    arr = arrcls(context, builder, value=builder.load(flatiter.array))
+    arr = arrcls(context, builder, value=flatiter.array)
 
     res = flatiter.getitem(context, builder, arrty, arr, index)
     return impl_ret_borrowed(context, builder, sig.return_type, res)
@@ -2845,7 +2844,7 @@ def iternext_numpy_getitem(context, builder, sig, args):
 
     arrty = flatiterty.array_type
     arrcls = context.make_array(arrty)
-    arr = arrcls(context, builder, value=builder.load(flatiter.array))
+    arr = arrcls(context, builder, value=flatiter.array)
 
     res = flatiter.setitem(context, builder, arrty, arr, index, value)
     return context.get_dummy_value()
@@ -2858,7 +2857,7 @@ def iternext_numpy_getitem(context, builder, sig, args):
     flatiter = flatitercls(context, builder, value=args[0])
 
     arrcls = context.make_array(flatiterty.array_type)
-    arr = arrcls(context, builder, value=builder.load(flatiter.array))
+    arr = arrcls(context, builder, value=flatiter.array)
     return arr.nitems
 
 
@@ -2869,11 +2868,10 @@ def make_array_ndenumerate(context, builder, sig, args):
     nditercls = make_array_ndenumerate_cls(types.NumpyNdEnumerateType(arrty))
     nditer = nditercls(context, builder)
 
-    arrayptr = cgutils.alloca_once_value(builder, arr)
-    nditer.array = arrayptr
+    nditer.array = arr
 
     arrcls = context.make_array(arrty)
-    arr = arrcls(context, builder, ref=arrayptr)
+    arr = arrcls(context, builder, ref=nditer._get_ptr_by_name('array'))
 
     nditer.init_specific(context, builder, arrty, arr)
 
@@ -2892,7 +2890,7 @@ def iternext_numpy_nditer(context, builder, sig, args, result):
 
     arrty = nditerty.array_type
     arrcls = context.make_array(arrty)
-    arr = arrcls(context, builder, value=builder.load(nditer.array))
+    arr = arrcls(context, builder, value=nditer.array)
 
     nditer.iternext_specific(context, builder, arrty, arr, result)
 


Test patch:
diff --git a/numba/tests/support.py b/numba/tests/support.py
--- a/numba/tests/support.py
+++ b/numba/tests/support.py
@@ -596,6 +596,7 @@ def setUp(self):
 
     def tearDown(self):
         super(MemoryLeakMixin, self).tearDown()
+        gc.collect()
         self.memory_leak_teardown()
 
 
diff --git a/numba/tests/test_array_iterators.py b/numba/tests/test_array_iterators.py
--- a/numba/tests/test_array_iterators.py
+++ b/numba/tests/test_array_iterators.py
@@ -92,6 +92,30 @@ def iter_next(arr):
     return next(it), next(it), next(it2)
 
 
+#
+# Test premature free (see issue #2112).
+# The following test allocates an array ``x`` inside the body.
+# The compiler will put a ``del x`` right after the last use of ``x``,
+# which is right after the creation of the array iterator and
+# before the loop is entered.  If the iterator does not incref the array,
+# the iterator will be reading garbage data of free'ed memory.
+#
+
+def array_flat_premature_free(size):
+    x = np.arange(size)
+    res = np.zeros_like(x, dtype=np.intp)
+    for i, v in enumerate(x.flat):
+        res[i] = v
+    return res
+
+def array_ndenumerate_premature_free(size):
+    x = np.arange(size)
+    res = np.zeros_like(x, dtype=np.intp)
+    for i, v in np.ndenumerate(x):
+        res[i] = v
+    return res
+
+
 class TestArrayIterators(MemoryLeakMixin, TestCase):
     """
     Test array.flat, np.ndenumerate(), etc.
@@ -285,6 +309,14 @@ def check(arr):
         arr = np.array([42]).reshape(())
         check(arr)
 
+    def test_array_flat_premature_free(self):
+        cres = compile_isolated(array_flat_premature_free, [types.intp])
+        cfunc = cres.entry_point
+        expect = array_flat_premature_free(6)
+        got = cfunc(6)
+        self.assertTrue(got.sum())
+        self.assertPreciseEqual(expect, got)
+
     @tag('important')
     def test_array_ndenumerate_2d(self):
         arr = np.arange(12).reshape(4, 3)
@@ -329,6 +361,14 @@ def test_array_ndenumerate_empty(self):
         arrty = types.Array(types.int32, 2, layout='A')
         self.check_array_flat_sum(arr, arrty)
 
+    def test_array_ndenumerate_premature_free(self):
+        cres = compile_isolated(array_ndenumerate_premature_free, [types.intp])
+        cfunc = cres.entry_point
+        expect = array_ndenumerate_premature_free(6)
+        got = cfunc(6)
+        self.assertTrue(got.sum())
+        self.assertPreciseEqual(expect, got)
+
     def test_np_ndindex(self):
         func = np_ndindex
         cres = compile_isolated(func, [types.int32, types.int32])