Repository: numba/numba

Issue:
Cannot cache functions with callable arguments
<!--

Thanks for opening an issue! To help the Numba team handle your information
efficiently, please first ensure that there is no other issue present that
already describes the issue you have
(search at https://github.com/numba/numba/issues?&q=is%3Aissue).

-->

## Reporting a bug

<!--

Before submitting a bug report please ensure that you can check off these boxes:

-->

- [x] I have tried using ~~the latest released version of Numba (most recent is
 visible in the change log (https://github.com/numba/numba/blob/master/CHANGE_LOG)~~ d58a63a870587bf8732038ca5b2229a4de218378.
- [x] I have included below a minimal working reproducer (if you are unsure how
 to write one see http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports).

<!--

Please include details of the bug here, including, if applicable, what you
expected to happen!

-->
```python
import numba

@numba.njit(cache=True)
def mul_self(x):
	return x * x

@numba.njit(cache=True)
def compose_self_apply(f, x):
	return f(f(x))

compose_self_apply(mul_self, 2)
```
gives
<details>

<summary>Traceback</summary>

```python
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
c:\users\wiese\repos\ga\numba\numba\core\caching.py in IndexDataCacheFile.save(self, key, data)
    482 try:
    483     # If key already exists, we will overwrite the file
--> 484     data_name = overloads[key]
    485 except KeyError:
    486     # Find an available name for the data file

KeyError: ((type(CPUDispatcher(<function mul_self at 0x0000025918224C10>)), int64), ('x86_64-pc-windows-msvc', 'broadwell', '+64bit,+adx,-aes,+avx,+avx2,-avx512bf16,-avx512bitalg,-avx512bw,-avx512cd,-avx512dq,-avx512er,-avx512f,-avx512ifma,-avx512pf,-avx512vbmi,-avx512vbmi2,-avx512vl,-avx512vnni,-avx512vpopcntdq,+bmi,+bmi2,-cldemote,-clflushopt,-clwb,-clzero,+cmov,+cx16,+cx8,-enqcmd,+f16c,+fma,-fma4,+fsgsbase,+fxsr,-gfni,+invpcid,-lwp,+lzcnt,+mmx,+movbe,-movdir64b,-movdiri,-mpx,-mwaitx,+pclmul,-pconfig,-pku,+popcnt,-prefetchwt1,+prfchw,-ptwrite,-rdpid,+rdrnd,+rdseed,-rtm,+sahf,-sgx,-sha,-shstk,+sse,+sse2,+sse3,+sse4.1,+sse4.2,-sse4a,+ssse3,-tbm,-vaes,-vpclmulqdq,-waitpkg,-wbnoinvd,-xop,+xsave,-xsavec,+xsaveopt,-xsaves'), ('95d9e4857cb9755b29e175e557e615d8467821eb42c66a521b4270d86f6b8c06', 'e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'))

During handling of the above exception, another exception occurred:

TypeError                                 Traceback (most recent call last)
<ipython-input-7-74c10e900da8> in <module>
      7 @numba.njit(cache=True)
      8 def compose_self_apply(f, x):
      9         return f(f(x))
---> 11 compose_self_apply(mul_self, 2)

c:\users\wiese\repos\ga\numba\numba\core\dispatcher.py in _DispatcherBase._compile_for_args(self, *args, **kws)
    431         e.patch_message('\n'.join((str(e).rstrip(), help_msg)))
    432 # ignore the FULL_TRACEBACKS config, this needs reporting!
--> 433 raise e

c:\users\wiese\repos\ga\numba\numba\core\dispatcher.py in _DispatcherBase._compile_for_args(self, *args, **kws)
    364         argtypes.append(self.typeof_pyval(a))
    365 try:
--> 366     return self.compile(tuple(argtypes))
    367 except errors.ForceLiteralArg as e:
    368     # Received request for compiler re-entry with the list of arguments
    369     # indicated by e.requested_args.
    370     # First, check if any of these args are already Literal-ized
    371     already_lit_pos = [i for i in e.requested_args
    372                        if isinstance(args[i], types.Literal)]

c:\users\wiese\repos\ga\numba\numba\core\compiler_lock.py in _CompilerLock.__call__.<locals>._acquire_compile_lock(*args, **kwargs)
     29 @functools.wraps(func)
     30 def _acquire_compile_lock(*args, **kwargs):
     31     with self:
---> 32         return func(*args, **kwargs)

c:\users\wiese\repos\ga\numba\numba\core\dispatcher.py in Dispatcher.compile(self, sig)
    822     raise e.bind_fold_arguments(folded)
    823 self.add_overload(cres)
--> 824 self._cache.save_overload(sig, cres)
    825 return cres.entry_point

c:\users\wiese\repos\ga\numba\numba\core\caching.py in Cache.save_overload(self, sig, data)
    664 """
    665 Save the data for the given signature in the cache.
    666 """
    667 with self._guard_against_spurious_io_errors():
--> 668     self._save_overload(sig, data)

c:\users\wiese\repos\ga\numba\numba\core\caching.py in Cache._save_overload(self, sig, data)
    676 key = self._index_key(sig, _get_codegen(data))
    677 data = self._impl.reduce(data)
--> 678 self._cache_file.save(key, data)

c:\users\wiese\repos\ga\numba\numba\core\caching.py in IndexDataCacheFile.save(self, key, data)
    491             break
    492     overloads[key] = data_name
--> 493     self._save_index(overloads)
    494 self._save_data(data_name, data)

c:\users\wiese\repos\ga\numba\numba\core\caching.py in IndexDataCacheFile._save_index(self, overloads)
    537 def _save_index(self, overloads):
    538     data = self._source_stamp, overloads
--> 539     data = self._dump(data)
    540     with self._open_for_write(self._index_path) as f:
    541         pickle.dump(self._version, f, protocol=-1)

c:\users\wiese\repos\ga\numba\numba\core\caching.py in IndexDataCacheFile._dump(self, obj)
    566 def _dump(self, obj):
--> 567     return pickle.dumps(obj, protocol=-1)

TypeError: cannot pickle 'weakref' object
```

</details>


Discussion:

Patch:
diff --git a/numba/core/types/functions.py b/numba/core/types/functions.py
index 0c3926a0e..eafe45a3b 100644
--- a/numba/core/types/functions.py
+++ b/numba/core/types/functions.py
@@ -457,6 +457,20 @@ class MakeFunctionLiteral(Literal, Opaque):
     pass
 
 
+class _PickleableWeakRef(weakref.ref):
+    """
+    Allow a weakref to be pickled.
+
+    Note that if the object referred to is not kept alive elsewhere in the
+    pickle, the weakref will immediately expire after being constructed.
+    """
+    def __getnewargs__(self):
+        obj = self()
+        if obj is None:
+            raise ReferenceError("underlying object has vanished")
+        return (obj,)
+
+
 class WeakType(Type):
     """
     Base class for types parametered by a mortal object, to which only
@@ -464,7 +478,7 @@ class WeakType(Type):
     """
 
     def _store_object(self, obj):
-        self._wr = weakref.ref(obj)
+        self._wr = _PickleableWeakRef(obj)
 
     def _get_object(self):
         obj = self._wr()


Test patch:
diff --git a/numba/tests/cache_usecases.py b/numba/tests/cache_usecases.py
index 27517f3d4..68602852c 100644
--- a/numba/tests/cache_usecases.py
+++ b/numba/tests/cache_usecases.py
@@ -163,6 +163,21 @@ class _TestModule(TestCase):
         raise NotImplementedError
 
 
+@jit(cache=True)
+def first_class_function_mul(x):
+    return x * x
+
+
+@jit(cache=True)
+def first_class_function_add(x):
+    return x + x
+
+
+@jit(cache=True)
+def first_class_function_usecase(f, x):
+    return f(x)
+
+
 def self_test():
     mod = sys.modules[__name__]
     _TestModule().check_module(mod)
diff --git a/numba/tests/test_dispatcher.py b/numba/tests/test_dispatcher.py
index b70c6804a..65f4c07ed 100644
--- a/numba/tests/test_dispatcher.py
+++ b/numba/tests/test_dispatcher.py
@@ -1266,6 +1266,17 @@ class TestCache(BaseCacheUsecasesTest):
             self.assertPreciseEqual(f(3), 12) # 3 + 9 = 12
             self.check_pycache(5) # 1 nbi, 4 nbc
 
+    def test_first_class_function(self):
+        mod = self.import_module()
+        f = mod.first_class_function_usecase
+        self.assertEqual(f(mod.first_class_function_mul, 1), 1)
+        self.assertEqual(f(mod.first_class_function_mul, 10), 100)
+        self.assertEqual(f(mod.first_class_function_add, 1), 2)
+        self.assertEqual(f(mod.first_class_function_add, 10), 20)
+        # 1 + 1 + 1 nbi, 1 + 1 + 2 nbc - a separate cache for each call to `f`
+        # with a different callback.
+        self.check_pycache(7)
+
     def test_cache_reuse(self):
         mod = self.import_module()
         mod.add_usecase(2, 3)