1from numba.core import errors, ir
2from numba.core.rewrites import register_rewrite, Rewrite
3
4
5@register_rewrite('before-inference')
6class DetectStaticBinops(Rewrite):
7    """
8    Detect constant arguments to select binops.
9    """
10
11    # Those operators can benefit from a constant-inferred argument
12    rhs_operators = {'**'}
13
14    def match(self, func_ir, block, typemap, calltypes):
15        self.static_lhs = {}
16        self.static_rhs = {}
17        self.block = block
18        # Find binop expressions with a constant lhs or rhs
19        for expr in block.find_exprs(op='binop'):
20            try:
21                if (expr.fn in self.rhs_operators
22                    and expr.static_rhs is ir.UNDEFINED):
23                    self.static_rhs[expr] = func_ir.infer_constant(expr.rhs)
24            except errors.ConstantInferenceError:
25                continue
26
27        return len(self.static_lhs) > 0 or len(self.static_rhs) > 0
28
29    def apply(self):
30        """
31        Store constant arguments that were detected in match().
32        """
33        for expr, rhs in self.static_rhs.items():
34            expr.static_rhs = rhs
35        return self.block
36