Repository: numba/numba

Issue:
RewriteArrayExprs pass incorrectly rewrites all Array-returning operations
<!--

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).

For more general "how do I do X?" type questions, please speak to us in real
time on https://gitter.im/numba/numba or post to the Numba mailing list
https://groups.google.com/a/continuum.io/forum/#!forum/numba-users.

-->

## Reporting a bug

<!--

Before submitting a bug report please ensure that you can check off these boxes:

-->

- [x] I am using the latest released version of Numba (most recent is visible in
 the change log (https://github.com/numba/numba/blob/master/CHANGE_LOG).
- [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!

-->
Reproducer:
```python
from numba.extending import (
    overload,
    models,
    lower_builtin,
    register_model,
    make_attribute_wrapper,
    type_callable,
    typeof_impl,
)
from numba import types, cgutils
from numba import njit
import operator
import numpy as np


def make_dummy_type():
    class Dummy(object):
        def __init__(self, value):
            self.value = value

    class DummyType(types.Type):
        def __init__(self):
            super(DummyType, self).__init__(name="Dummy")

    dummy_type = DummyType()

    @register_model(DummyType)
    class DummyModel(models.StructModel):
        def __init__(self, dmm, fe_type):
            members = [("value", types.intp)]
            models.StructModel.__init__(self, dmm, fe_type, members)

    make_attribute_wrapper(DummyType, "value", "value")

    @type_callable(Dummy)
    def type_dummy(context):
        def typer(value):
            return dummy_type

        return typer

    @lower_builtin(Dummy, types.intp)
    def impl_dummy(context, builder, sig, args):
        typ = sig.return_type
        [value] = args
        dummy = cgutils.create_struct_proxy(typ)(context, builder)
        dummy.value = value
        return dummy._getvalue()

    @typeof_impl.register(Dummy)
    def typeof_dummy(val, c):
        return DummyType()

    # Store attributes
    return Dummy, DummyType


Dummy, DummyType = make_dummy_type()


@overload(operator.add)
def ov_add_dummy(lhs, rhs):
    if isinstance(lhs, DummyType) and isinstance(rhs, DummyType):
        print("===", lhs, rhs)

        def imp(lhs, rhs):
            return np.array([lhs.value, rhs.value])

        return imp


@njit
def foo(x, y):
    return Dummy(x) + Dummy(y)


r = foo(1, 2)
print(r)
```


The above raises:
```pytb
numba.errors.TypingError: Failed in nopython mode pipeline (step: nopython mode backend)
Failed in nopython mode pipeline (step: nopython frontend)
No conversion from array(int64, 1d, C) to int64 for '$8return_value.3', defined at None

File "array_expr_issue.py", line 68:
def foo(x, y):
    return Dummy(x) + Dummy(y)
    ^

[1] During: typing of assignment at /Users/siu/dev/numba-git/array_expr_issue.py (68)

File "array_expr_issue.py", line 68:
def foo(x, y):
    return Dummy(x) + Dummy(y)
    ^

[1] During: lowering "$14binary_add.6 = arrayexpr(expr=(<built-in function add>, [Var($6call_function.2, array_expr_issue.py:68), Var($12call_function.5, array_expr_issue.py:68)]), ty=array(int64, 1d, C))" at /Users/siu/dev/numba-git/array_expr_issue.py (68)
```

Expected Behavior:
- `RewriteArrayExprs` pass should not make `arrayexpr` node for the `Dummy(x) + Dummy(y)` expression.
- Disabling the pass will allow the code run correctly.



Discussion:
Would it be sufficient here to ban rewrites on types which don't have the `is_internal` property? The effect being that arrayexpr rewrites would only be done on Numba internally supported types.

Resolving patch:
diff --git a/numba/np/ufunc/array_exprs.py b/numba/np/ufunc/array_exprs.py
--- a/numba/np/ufunc/array_exprs.py
+++ b/numba/np/ufunc/array_exprs.py
@@ -77,7 +77,10 @@ def _match_array_expr(self, instr, expr, target_name):
         if ((expr_op in ('unary', 'binop')) and (
                 expr.fn in npydecl.supported_array_operators)):
             # It is an array operator that maps to a ufunc.
-            array_assigns[target_name] = instr
+            # check that all args have internal types
+            if all(self.typemap[var.name].is_internal
+                   for var in expr.list_vars()):
+                array_assigns[target_name] = instr
 
         elif ((expr_op == 'call') and (expr.func.name in self.typemap)):
             # It could be a match for a known ufunc call.


Test patch:
diff --git a/numba/tests/test_array_exprs.py b/numba/tests/test_array_exprs.py
--- a/numba/tests/test_array_exprs.py
+++ b/numba/tests/test_array_exprs.py
@@ -5,9 +5,27 @@
 
 from numba import njit, vectorize
 from numba import typeof
-from numba.core import utils, types, typing, ir, compiler, cpu
+from numba.core import utils, types, typing, ir, compiler, cpu, cgutils
 from numba.core.compiler import Compiler, Flags
-from numba.tests.support import MemoryLeakMixin, TestCase
+from numba.tests.support import MemoryLeakMixin, TestCase, temp_directory
+from numba.extending import (
+    overload,
+    models,
+    lower_builtin,
+    register_model,
+    make_attribute_wrapper,
+    type_callable,
+    typeof_impl
+)
+import operator
+import contextlib
+import importlib
+import shutil
+import uuid
+import os
+import sys
+import textwrap
+
 import unittest
 
 
@@ -627,5 +645,127 @@ def do_call(x, y):
         self.assertTrue(isinstance(oty.type.dtype, types.Float))
 
 
+class TestExternalTypes(MemoryLeakMixin, unittest.TestCase):
+    """ Tests RewriteArrayExprs with external (user defined) types,
+    see #5157"""
+
+    source_lines = textwrap.dedent("""
+        from numba.core import types
+
+        class FooType(types.Type):
+            def __init__(self):
+                super(FooType, self).__init__(name='Foo')
+        """)
+
+    # copied from test_types
+    @contextlib.contextmanager
+    def create_temp_module(self, source_lines=None, **jit_options):
+        # Use try/finally so cleanup happens even when an exception is raised
+        try:
+            if source_lines is None:
+                source_lines = self.source_lines
+            tempdir = temp_directory('test_extension_type')
+            # Generate random module name
+            temp_module_name = 'test_extension_type_{}'.format(
+                str(uuid.uuid4()).replace('-', '_'))
+            temp_module_path = os.path.join(tempdir, temp_module_name + '.py')
+
+            with open(temp_module_path, 'w') as f:
+                lines = source_lines.format(jit_options=jit_options)
+                f.write(lines)
+            # Add test_module to sys.path so it can be imported
+            sys.path.insert(0, tempdir)
+            test_module = importlib.import_module(temp_module_name)
+            yield test_module
+        finally:
+            sys.modules.pop(temp_module_name, None)
+            sys.path.remove(tempdir)
+            shutil.rmtree(tempdir)
+
+    def make_foo_type(self, FooType):
+        class Foo(object):
+            def __init__(self, value):
+                self.value = value
+
+        @register_model(FooType)
+        class FooModel(models.StructModel):
+            def __init__(self, dmm, fe_type):
+                members = [("value", types.intp)]
+                models.StructModel.__init__(self, dmm, fe_type, members)
+
+        make_attribute_wrapper(FooType, "value", "value")
+
+        @type_callable(Foo)
+        def type_foo(context):
+            def typer(value):
+                return FooType()
+
+            return typer
+
+        @lower_builtin(Foo, types.intp)
+        def impl_foo(context, builder, sig, args):
+            typ = sig.return_type
+            [value] = args
+            foo = cgutils.create_struct_proxy(typ)(context, builder)
+            foo.value = value
+            return foo._getvalue()
+
+        @typeof_impl.register(Foo)
+        def typeof_foo(val, c):
+            return FooType()
+
+        return Foo, FooType
+
+    def test_external_type(self):
+        with self.create_temp_module(self.source_lines) as test_module:
+            Foo, FooType = self.make_foo_type(test_module.FooType)
+
+            # sum of foo class instance and array return an array
+            # binary operation with foo class instance as one of args
+            @overload(operator.add)
+            def overload_foo_add(lhs, rhs):
+                if isinstance(lhs, FooType) and isinstance(rhs, types.Array):
+                    def imp(lhs, rhs):
+                        return np.array([lhs.value, rhs[0]])
+
+                    return imp
+
+            # sum of 2 foo class instances return an array
+            # binary operation with 2 foo class instances as args
+            @overload(operator.add)
+            def overload_foo_add(lhs, rhs):
+                if isinstance(lhs, FooType) and isinstance(rhs, FooType):
+                    def imp(lhs, rhs):
+                        return np.array([lhs.value, rhs.value])
+
+                    return imp
+
+            # neg of foo class instance return an array
+            # unary operation with foo class instance arg
+            @overload(operator.neg)
+            def overload_foo_neg(x):
+                if isinstance(x, FooType):
+                    def imp(x):
+                        return np.array([-x.value])
+
+                    return imp
+
+            @njit
+            def arr_expr_sum1(x, y):
+                return Foo(x) + np.array([y])
+
+            @njit
+            def arr_expr_sum2(x, y):
+                return Foo(x) + Foo(y)
+
+            @njit
+            def arr_expr_neg(x):
+                return -Foo(x)
+
+            np.testing.assert_array_equal(arr_expr_sum1(0, 1), np.array([0, 1]))
+            np.testing.assert_array_equal(arr_expr_sum2(2, 3), np.array([2, 3]))
+            np.testing.assert_array_equal(arr_expr_neg(4), np.array([-4]))
+
+
 if __name__ == "__main__":
     unittest.main()