Repository: numba/numba

Issue:
Segfault when using generator in nopython mode
This code segfaults if I create an instance of the generator from Python before using it in a nopython mode function:

``` python
from numba import jit
import numpy as np

@jit
def complex_grid(rmin, rmax, nr, imin, imax, ni):
    for r in np.linspace(rmin, rmax, nr):
        for i in np.linspace(imin, imax, ni):
            yield complex(r, i)

#### COMMENTING OUT THE NEXT TWO LINES FIXES THE SEGFAULT
for z in complex_grid(-2, 2, 100, -2, 2, 100):
    pass

@jit(nopython=True)
def mag2(z):
    return z.real * z.real + z.imag * z.imag

MAX_ITERS=1000

@jit(nopython=True)
def mandel(c):
    z = 0j
    for i in range(MAX_ITERS):
        z = z*z + c
        if mag2(z) >= 4:
            return 255 * i // MAX_ITERS
    return 255

@jit(nogil=False, nopython=True)
def mandel_patch(args):
    rmin, rmax, nr, imin, imax, ni = args
    points = np.empty(nr*ni, dtype=np.complex128)
    values = np.empty(nr*ni, dtype=np.uint8)

    for i, c in enumerate(complex_grid(rmin, rmax, nr, imin, imax, ni)):
        points[i] = c
        values[i] = mandel(c)

    return points, values

n = 1000
patches = [(-2, -1, n, -1, 0, n), (-1, 0, n, -1, 0, n), (0, 1, n, -1, 0, n),
           (-2, -1, n,  0, 1, n), (-1, 0, n,  0, 1, n), (0, 1, n,  0, 1, n)]

print(mandel_patch(patches[0]))
```



Discussion:

Resolving patch:
diff --git a/numba/generators.py b/numba/generators.py
--- a/numba/generators.py
+++ b/numba/generators.py
@@ -96,6 +96,7 @@ def lower_init_func(self, lower):
         resume_index = self.context.get_constant(types.int32, 0)
         # Structure index #1: the function arguments
         argsty = retty.elements[1]
+        statesty = retty.elements[2]
 
         # Incref all NRT objects before storing into generator states
         if self.context.enable_nrt:
@@ -104,9 +105,14 @@ def lower_init_func(self, lower):
 
         argsval = cgutils.make_anonymous_struct(builder, lower.fnargs,
                                                 argsty)
+
+        # Zero initialize states
+        statesval = Constant.null(statesty)
         gen_struct = cgutils.make_anonymous_struct(builder,
-                                                   [resume_index, argsval],
+                                                   [resume_index, argsval,
+                                                    statesval],
                                                    retty)
+
         retval = self.box_generator_struct(lower, gen_struct)
         self.call_conv.return_value(builder, retval)
 
@@ -206,24 +212,11 @@ def lower_finalize_func_body(self, builder, genptr):
         state variables.
         """
         if self.context.enable_nrt:
-            resume_index_ptr = self.get_resume_index_ptr(builder, genptr)
-            resume_index = builder.load(resume_index_ptr)
-
-            # If resume_index is 0, next() was never called
-            # Note: The init func has to acquire the reference
-            #       of all arguments; otherwise, they will be destroyed at the
-            #       end of the init func.  The proper release of these
-            #       references relies on the actual NPM function, which is
-            #       never called if the generator exit before any entry into
-            #       the NPM function.
-            need_args_cleanup = builder.icmp_signed(
-                '==', resume_index, Constant.int(resume_index.type, 0))
-
-            with builder.if_then(need_args_cleanup):
-                for elem_idx, argty in enumerate(self.fndesc.argtypes):
-                    argptr = self.get_arg_ptr(builder, genptr, elem_idx)
-                    argval = builder.load(argptr)
-                    self.context.nrt_decref(builder, argty, argval)
+            # Always dereference all arguments
+            for elem_idx, argty in enumerate(self.fndesc.argtypes):
+                argptr = self.get_arg_ptr(builder, genptr, elem_idx)
+                argval = builder.load(argptr)
+                self.context.nrt_decref(builder, argty, argval)
 
             # Always run the finalizer to clear the block
             gen_state_ptr = self.get_state_ptr(builder, genptr)
@@ -326,6 +319,13 @@ def lower_yield_suspend(self):
                                               0, state_index)
             ty = self.gentype.state_types[state_index]
             val = self.lower.loadvar(name)
+            # IncRef newly stored value
+            if self.context.enable_nrt:
+                self.lower.incref(ty, val)
+            # Load and DecRef previously stored value
+            oldval = self.context.unpack_value(self.builder, ty, state_slot)
+            if self.context.enable_nrt:
+                self.lower.decref(ty, oldval)
             self.context.pack_value(self.builder, ty, val, state_slot)
         # Save resume index
         indexval = Constant.int(self.resume_index_ptr.type.pointee,


Test patch:
diff --git a/numba/tests/test_generators.py b/numba/tests/test_generators.py
--- a/numba/tests/test_generators.py
+++ b/numba/tests/test_generators.py
@@ -387,6 +387,68 @@ def test_nrt_gen0_no_iter(self):
         self.assertEqual(sys.getrefcount(py_ary),
                          sys.getrefcount(c_ary))
 
+    def test_nrt_nested_gen(self):
+
+        def gen0(arr):
+            for i in range(arr.size):
+                yield arr
+
+        def factory(gen0):
+            def gen1(arr):
+                out = np.zeros_like(arr)
+                for x in gen0(arr):
+                    out = out + x
+                return out, arr
+
+            return gen1
+
+        py_arr = np.arange(10)
+        c_arr = py_arr.copy()
+        py_res, py_old = factory(gen0)(py_arr)
+        c_gen = jit(nopython=True)(factory(jit(nopython=True)(gen0)))
+        c_res, c_old = c_gen(c_arr)
+
+        self.assertIsNot(py_arr, c_arr)
+        self.assertIs(py_old, py_arr)
+        self.assertIs(c_old, c_arr)
+
+        np.testing.assert_equal(py_res, c_res)
+
+        self.assertEqual(sys.getrefcount(py_res),
+                         sys.getrefcount(c_res))
+
+        # The below test will fail due to generator finalizer not invoked.
+        # This kept a reference of the c_old.
+        #
+        # self.assertEqual(sys.getrefcount(py_old),
+        #                  sys.getrefcount(c_old))
+
+    @unittest.expectedFailure
+    def test_nrt_nested_gen_refct(self):
+        def gen0(arr):
+            yield arr
+
+        def factory(gen0):
+            def gen1(arr):
+                for out in gen0(arr):
+                    return out
+
+            return gen1
+
+        py_arr = np.arange(10)
+        c_arr = py_arr.copy()
+        py_old = factory(gen0)(py_arr)
+        c_gen = jit(nopython=True)(factory(jit(nopython=True)(gen0)))
+        c_old = c_gen(c_arr)
+
+        self.assertIsNot(py_arr, c_arr)
+        self.assertIs(py_old, py_arr)
+        self.assertIs(c_old, c_arr)
+
+        self.assertEqual(sys.getrefcount(py_old),
+                         sys.getrefcount(c_old))
+
+
 
 class TestGeneratorWithNRT(TestCase):
     def test_issue_1254(self):
@@ -409,6 +471,48 @@ def random_directions(n):
         for got in outputs:
             np.testing.assert_equal(expect, got)
 
+    def test_issue_1265(self):
+        """
+        Double-free for locally allocated, non escaping NRT objects
+        """
+        def py_gen(rmin, rmax, nr):
+            a = np.linspace(rmin, rmax, nr)
+            yield a[0]
+            yield a[1]
+
+        c_gen = jit(nopython=True)(py_gen)
+
+        py_res = list(py_gen(-2, 2, 100))
+        c_res = list(c_gen(-2, 2, 100))
+
+        self.assertEqual(py_res, c_res)
+
+        def py_driver(args):
+            rmin, rmax, nr = args
+            points = np.empty(nr, dtype=np.complex128)
+            for i, c in enumerate(py_gen(rmin, rmax, nr)):
+                points[i] = c
+
+            return points
+
+        @jit(nopython=True)
+        def c_driver(args):
+            rmin, rmax, nr = args
+            points = np.empty(nr, dtype=np.complex128)
+            for i, c in enumerate(c_gen(rmin, rmax, nr)):
+                points[i] = c
+
+            return points
+
+        n = 2
+        patches = (-2, -1, n)
+
+        py_res = py_driver(patches)
+        # The error will cause a segfault here
+        c_res = c_driver(patches)
+
+        np.testing.assert_equal(py_res, c_res)
+
 
 if __name__ == '__main__':
     unittest.main()