Test for temp assignment removal in Interpreter

numba/tests/test_copy_propagate.py

@@ -3,14 +3,17 @@
 # SPDX-License-Identifier: BSD-2-Clause
 #
 
+from numba import njit
 from numba.core import types, typing, ir, config, compiler, cpu
 from numba.core.registry import cpu_target
 from numba.core.annotations import type_annotations
 from numba.core.ir_utils import (copy_propagate, apply_copy_propagate,
                             get_name_var_table)
 from numba.core.typed_passes import type_inference_stage
+from numba.tests.test_ir_inlining import InlineTestPipeline
 import unittest
 
+
 def test_will_propagate(b, z, w):
     x = 3
     x1 = x
@@ -21,6 +24,7 @@ def test_will_propagate(b, z, w):
     a = 2 * x1
     return a < b
 
+
 def test_wont_propagate(b, z, w):
     x = 3
     if b > 0:
@@ -31,15 +35,18 @@ def test_wont_propagate(b, z, w):
     a = 2 * x
     return a < b
 
+
 def null_func(a,b,c,d):
     False
 
+
 def inListVar(list_var, var):
     for i in list_var:
         if i.name == var:
             return True
     return False
 
+
 def findAssign(func_ir, var):
     for label, block in func_ir.blocks.items():
         for i, inst in enumerate(block.body):
@@ -50,6 +57,7 @@ def findAssign(func_ir, var):
 
     return False
 
+
 class TestCopyPropagate(unittest.TestCase):
     def test1(self):
         typingctx = typing.Context()
@@ -97,6 +105,30 @@ def test2(self):
 
             self.assertTrue(findAssign(test_ir, "x"))
 
+    def test_input_ir_extra_copies(self):
+        """make sure Interpreter._remove_unused_temporaries() has removed extra copies
+        in the IR in simple cases so copy propagation is faster
+        """
+        def test_impl(a):
+            b = a + 3
+            return b
+
+        j_func = njit(pipeline_class=InlineTestPipeline)(test_impl)
+        self.assertEqual(test_impl(5), j_func(5))
+
+        # make sure b is the target of the expression assignment, not a temporary
+        fir = j_func.overloads[j_func.signatures[0]].metadata['preserved_ir']
+        self.assertTrue(len(fir.blocks) == 1)
+        block = next(iter(fir.blocks.values()))
+        b_found = False
+        for stmt in block.body:
+            if isinstance(stmt, ir.Assign) and stmt.target.name == "b":
+                b_found = True
+                self.assertTrue(isinstance(stmt.value, ir.Expr)
+                    and stmt.value.op == "binop" and stmt.value.lhs.name == "a")
+
+        self.assertTrue(b_found)
+
 
 if __name__ == "__main__":
     unittest.main()