Repository: numba/numba

Issue:
Typer does not propagate return type to all return variables
- [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).
- [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).

When there are multiple return statements in a function the typer does not propagate the unified return type to all return variables. This often works because return values are implicitly casted to the return type [here](https://github.com/numba/numba/blob/626b40eddc9a29a1e89bfb6b8706e4f390cfcc4a/numba/core/lowering.py#L357). However, it does make a difference for memory managed types (the implicit cast does not neccessarily get assigned hence a borrowed reference can escape).

In the example below the first return variable is of type `int64`, the second is `float64`. IMHO both should have been explicitly casted to `float64`.

```python
from numba import jit, types
from numba.core import ir


@jit(nopython=True)
def confuse_typer(x):
    # --- LINE 16 ---
    # label 0
    #   x = arg(0, name=x)  :: float64
    #   $6compare_op.2 = x == x  :: bool
    #   bool8 = global(bool: <class 'bool'>)  :: Function(<class 'bool'>)
    #   $8pred = call bool8($6compare_op.2, func=bool8, args=(Var($6compare_op.2, cast_quantity_to_adf_crash.py:5),), kws=(), vararg=None)  :: (bool,) -> bool
    #   del bool8
    #   del $6compare_op.2
    #   branch $8pred, 10, 18
    if x == x:
        # --- LINE 27 ---
        # label 10
        #   del $8pred
        #   $10load_global.0 = global(int: <class 'int'>)  :: Function(<class 'int'>)
        #   $14call_function.2 = call $10load_global.0(x, func=$10load_global.0, args=[Var(x, cast_quantity_to_adf_crash.py:5)], kws=(), vararg=None)  :: (float64,) -> int64
        #   del x
        #   del $10load_global.0
        #   $16return_value.3 = cast(value=$14call_function.2)  :: int64    <- cast to float64 should happen here
        #   del $14call_function.2
        #   return $16return_value.3    <- returns int64
        return int(x)
    # --- LINE 29 ---
    else:
        # --- LINE 36 ---
        # label 18
        #   del $8pred
        #   $20return_value.1 = cast(value=x)  :: float64
        #   del x
        #   return $20return_value.1    <- returns float64
        return x

confuse_typer.compile((types.float64,))
cres = confuse_typer.overloads[(types.float64,)]
typemap = cres.type_annotation.typemap
return_vars = {}


for block in cres.type_annotation.blocks.values():
    for inst in block.body:
        if isinstance(inst, ir.Return):
            varname = inst.value.name
            return_vars[varname] = typemap[varname]

# fails b/c return_vars = {'$16return_value.3': types.int64, '$20return_value.1': types.float64}
assert all(vt == types.float64 for vt in return_vars.values())
```


Discussion:

Resolving patch:
diff --git a/numba/core/lowering.py b/numba/core/lowering.py
--- a/numba/core/lowering.py
+++ b/numba/core/lowering.py
@@ -353,8 +353,8 @@ def lower_inst(self, inst):
                 # If returning an optional type
                 self.call_conv.return_optional_value(self.builder, ty, oty, val)
                 return
-            if ty != oty:
-                val = self.context.cast(self.builder, val, oty, ty)
+            assert ty == oty, (
+                "type '{}' does not match return type '{}'".format(oty, ty))
             retval = self.context.get_return_value(self.builder, ty, val)
             self.call_conv.return_value(self.builder, retval)
 
diff --git a/numba/core/typeinfer.py b/numba/core/typeinfer.py
--- a/numba/core/typeinfer.py
+++ b/numba/core/typeinfer.py
@@ -1214,6 +1214,9 @@ def check_var(name):
                 raise e
             else:
                 retty = None
+        else:
+            typdict = utils.UniqueDict(
+                typdict, **{v.name: retty for v in self._get_return_vars()})
 
         try:
             fntys = self.get_function_types(typdict)


Test patch:
diff --git a/numba/tests/test_typeinfer.py b/numba/tests/test_typeinfer.py
--- a/numba/tests/test_typeinfer.py
+++ b/numba/tests/test_typeinfer.py
@@ -5,7 +5,7 @@
 
 from numba.core.compiler import compile_isolated
 from numba import jit
-from numba.core import types, typing, errors, typeinfer, utils
+from numba.core import errors, ir, types, typing, typeinfer, utils
 from numba.core.typeconv import Conversion
 
 from numba.tests.support import TestCase, tag
@@ -678,6 +678,31 @@ def test_issue_1394(self):
             res = cfunc(v)
             self.assertEqual(res, pyfunc(v))
 
+    def test_issue_6293(self):
+        """https://github.com/numba/numba/issues/6293
+
+        Typer does not propagate return type to all return variables
+        """
+        @jit(nopython=True)
+        def confuse_typer(x):
+            if x == x:
+                return int(x)
+            else:
+                return x
+
+        confuse_typer.compile((types.float64,))
+        cres = confuse_typer.overloads[(types.float64,)]
+        typemap = cres.type_annotation.typemap
+        return_vars = {}
+
+        for block in cres.type_annotation.blocks.values():
+            for inst in block.body:
+                if isinstance(inst, ir.Return):
+                    varname = inst.value.name
+                    return_vars[varname] = typemap[varname]
+
+        self.assertTrue(all(vt == types.float64 for vt in return_vars.values()))
+
 
 class TestFoldArguments(unittest.TestCase):
     def check_fold_arguments_list_inputs(self, func, args, kws):