Repository: numba/numba

Issue:
Reference counting bug with NRT+Generators
the cleanup block in the python wrapper for generators removed all NRT reference before the execution of the actual function.



Discussion:
**Details**

Bug observed as:
- Entry wrapper of generator decreases refct to zero.
- The wrapped function (the NPM function) acquire a object with refct==0 and deference accordingly, causing a double-call to destructor and deallocation.


Resolving patch:
diff --git a/numba/generators.py b/numba/generators.py
--- a/numba/generators.py
+++ b/numba/generators.py
@@ -85,6 +85,12 @@ 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]
+
+        # Incref all NRT objects before storing into generator states
+        if self.context.enable_nrt:
+            for argty, argval in zip(self.fndesc.argtypes, lower.fnargs):
+                self.context.nrt_incref(builder, argty, argval)
+
         argsval = cgutils.make_anonymous_struct(builder, lower.fnargs,
                                                 argsty)
         gen_struct = cgutils.make_anonymous_struct(builder,
@@ -174,7 +180,6 @@ def create_resumption_block(self, lower, index):
         self.resume_blocks[index] = block
 
 
-
 class GeneratorLower(BaseGeneratorLower):
     """
     Support class for lowering nopython generators.
@@ -186,6 +191,46 @@ def get_generator_type(self):
     def box_generator_struct(self, lower, gen_struct):
         return gen_struct
 
+    def lower_finalize_func_body(self, builder, genptr):
+        """
+        Lower the body of the generator's finalizer: decref all live
+        state variables.
+        """
+        if self.context.enable_nrt:
+            resume_index_ptr = cgutils.gep(builder, genptr, 0, 0,
+                                           name='gen.resume_index')
+            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 cgutils.ifthen(builder, need_args_cleanup):
+                gen_args_ptr = cgutils.gep(builder, genptr, 0, 1, name="gen_args")
+                assert len(self.fndesc.argtypes) == len(gen_args_ptr.type.pointee)
+                for elem_idx, argty in enumerate(self.fndesc.argtypes):
+                    argptr = cgutils.gep(builder, gen_args_ptr, 0, elem_idx)
+                    argval = builder.load(argptr)
+                    self.context.nrt_decref(builder, argty, argval)
+
+            # Always run the finalizer to clear the block
+            gen_state_ptr = cgutils.gep(builder, genptr, 0, 2, name='gen.state')
+
+            for state_index in range(len(self.gentype.state_types)):
+                state_slot = cgutils.gep(builder, gen_state_ptr,
+                                         0, state_index)
+                ty = self.gentype.state_types[state_index]
+                val = self.context.unpack_value(builder, ty, state_slot)
+                if self.context.enable_nrt:
+                    self.context.nrt_decref(builder, ty, val)
+
+        builder.ret_void()
 
 class PyGeneratorLower(BaseGeneratorLower):
     """
@@ -293,3 +338,6 @@ def lower_yield_resume(self):
             ty = self.gentype.state_types[state_index]
             val = self.context.unpack_value(self.builder, ty, state_slot)
             self.lower.storevar(val, name)
+            # Previous storevar is making an extra incref
+            if self.context.enable_nrt:
+                self.lower.decref(ty, val)
diff --git a/numba/typeinfer.py b/numba/typeinfer.py
--- a/numba/typeinfer.py
+++ b/numba/typeinfer.py
@@ -438,9 +438,8 @@ def get_generator_type(self, typdict, retty):
         if not yield_types:
             raise TypingError("Cannot type generator: it does not yield any value")
         yield_type = self.context.unify_types(*yield_types)
-        # No finalizer in nopython mode
         return types.Generator(self.py_func, yield_type, arg_types, state_types,
-                               has_finalizer=False)
+                               has_finalizer=True)
 
     def get_function_types(self, typemap):
         calltypes = utils.UniqueDict()


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
@@ -1,5 +1,6 @@
 from __future__ import print_function
 
+import sys
 import numpy as np
 
 from numba.compiler import compile_isolated, Flags
@@ -214,5 +215,121 @@ def test_return_generator_expr(self):
         self.assertEqual(sum(cfunc([1, 2, 3])), sum(pyfunc([1, 2, 3])))
 
 
+
+def nrt_gen0(ary):
+    for elem in ary:
+        yield elem
+
+def nrt_gen1(ary1, ary2):
+    for e1, e2 in zip(ary1, ary2):
+        yield e1
+        yield e2
+
+
+class TestNrtArrayGen(TestCase):
+    def test_nrt_gen0(self):
+        pygen = nrt_gen0
+        cgen = jit(nopython=True)(pygen)
+
+        py_ary = np.arange(10)
+        c_ary = py_ary.copy()
+
+        py_res = list(pygen(py_ary))
+        c_res = list(cgen(c_ary))
+
+        np.testing.assert_equal(py_ary, c_ary)
+        self.assertEqual(py_res, c_res)
+        # Check reference count
+        self.assertEqual(sys.getrefcount(py_ary),
+                         sys.getrefcount(c_ary))
+
+
+    def test_nrt_gen1(self):
+        pygen = nrt_gen1
+        cgen = jit(nopython=True)(pygen)
+
+        py_ary1 = np.arange(10)
+        py_ary2 = py_ary1 + 100
+
+        c_ary1 = py_ary1.copy()
+        c_ary2 = py_ary2.copy()
+
+        py_res = list(pygen(py_ary1, py_ary2))
+        c_res = list(cgen(c_ary1, c_ary2))
+
+        np.testing.assert_equal(py_ary1, c_ary1)
+        np.testing.assert_equal(py_ary2, c_ary2)
+        self.assertEqual(py_res, c_res)
+        # Check reference count
+        self.assertEqual(sys.getrefcount(py_ary1),
+                         sys.getrefcount(c_ary1))
+        self.assertEqual(sys.getrefcount(py_ary2),
+                          sys.getrefcount(c_ary2))
+
+    def test_combine_gen0_gen1(self):
+        """
+        Issue #1163 is observed when two generator with NRT object arguments
+        is ran in sequence.  The first one does a invalid free and corrupts
+        the NRT memory subsystem.  The second generator is likely to segfault
+        due to corrupted NRT data structure (an invalid MemInfo).
+        """
+        self.test_nrt_gen0()
+        self.test_nrt_gen1()
+
+    def test_nrt_gen0_stop_iteration(self):
+        """
+        Test cleanup on StopIteration
+        """
+        pygen = nrt_gen0
+        cgen = jit(nopython=True)(pygen)
+
+        py_ary = np.arange(1)
+        c_ary = py_ary.copy()
+
+        py_iter = pygen(py_ary)
+        c_iter = cgen(c_ary)
+
+        py_res = next(py_iter)
+        c_res = next(c_iter)
+
+        with self.assertRaises(StopIteration):
+            py_res = next(py_iter)
+
+        with self.assertRaises(StopIteration):
+            c_res = next(c_iter)
+
+        del py_iter
+        del c_iter
+
+        np.testing.assert_equal(py_ary, c_ary)
+        self.assertEqual(py_res, c_res)
+        # Check reference count
+        self.assertEqual(sys.getrefcount(py_ary),
+                         sys.getrefcount(c_ary))
+
+    def test_nrt_gen0_no_iter(self):
+        """
+        Test cleanup for a initialized but never iterated (never call next())
+        generator.
+        """
+        pygen = nrt_gen0
+        cgen = jit(nopython=True)(pygen)
+
+        py_ary = np.arange(1)
+        c_ary = py_ary.copy()
+
+        py_iter = pygen(py_ary)
+        c_iter = cgen(c_ary)
+
+        del py_iter
+        del c_iter
+
+        np.testing.assert_equal(py_ary, c_ary)
+
+        # Check reference count
+        self.assertEqual(sys.getrefcount(py_ary),
+                         sys.getrefcount(c_ary))
+
+
 if __name__ == '__main__':
     unittest.main()