Repository: numba/numba

Issue:
Regression with `literal_unroll` in 0.57.0
The following code used to work on `0.56.4`, but fails on the new version `0.57.0` (installed via `pip` on macOS if that matters):

```python
from numba import literal_unroll, njit
import numpy as np

test = np.array([(1, 2), (2, 3)], dtype=[("a1", "f8"), ("a2", "f8")])
fields = tuple(test.dtype.fields.keys())

@njit
def f(arr):
    res = 0
    for k in literal_unroll(fields):
        res = res + np.abs(arr[k]).sum()
    return res


print(f(test))
```

The script is supposed to print `8.0`, but it raised the following error in the newest version:

```
numba.core.errors.LoweringError: Failed in nopython mode pipeline (step: native lowering)
'vLiteral[str](a1)__28binary_subscr_7' is not a valid parameter name

File "literal_unroll.py", line 11:
def f(arr):
    <source elided>
    for k in literal_unroll(fields):
        res = res + np.abs(arr[k]).sum()
        ^

During: lowering "vLiteral[str](a1)_$30call_method.8 = arrayexpr(expr=(<ufunc 'absolute'>, [Var(vLiteral[str](a1)_$28binary_subscr.7, literal_unroll.py:11)]), ty=array(float64, 1d, C))" at [...]/literal_unroll.py (11)
```



Discussion:
Thanks for the report, this reproduces as above in 0.57.0 but not in 0.56.x. I think this was introduced by c6a25069.
I had a long phone conversation with @stuartarchibald about this one and we have concocted a path forward. I will summarize our findings here.

The original variant of the code was:

```python
                try:
                    scope.get_exact(name)
                except errors.NotDefinedError:
                    # is this correct? In case the scope doesn't have the
                    # variable, we need to define it prior creating new
                    # copies of it!
                    scope.define(name, var.loc)
                new_var_dict[name] = scope.redefine(name, var.loc).name
```

The question here was: why do we need a define and then a redefine? The answer is that this is a result of scoping in Numba IR. Specifically the scope of the function and the scope of the unrolled loop are different, so the redefine takes care of defining the variable inside the correct scope.

This was then changed to be:

```python
                new_var_dict[name] = scope.define(
                    f"v{branch_ty}_{name}", var.loc).name
```

This variant is designed to include the `branch_ty` in the variable name to aid in debugging. Note that this change has no test to accompany it.

Then we received this bug report.

The problem, in case of `branch_ty` now being a literal, is that the `branch_ty` is then:  

```
Literal[str](a1)
```

I.e. the variable name now contains both a `[` and a `]` which then are rejected by the call to `isidentifier` in the `inspect.py` module when attempting to determine the signature of the function during lowering. This is evidenced by the stack trace here:

```
Traceback (most recent call last):
  File "/Users/esc/git/numba/numba/core/errors.py", line 823, in new_error_context
    yield
  File "/Users/esc/git/numba/numba/core/lowering.py", line 285, in lower_block
    self.lower_inst(inst)
  File "/Users/esc/git/numba/numba/core/lowering.py", line 463, in lower_inst
    val = self.lower_assign(ty, inst)
  File "/Users/esc/git/numba/numba/core/lowering.py", line 675, in lower_assign
    return self.lower_expr(ty, value)
  File "/Users/esc/git/numba/numba/core/lowering.py", line 1417, in lower_expr
    res = self.context.special_ops[expr.op](self, expr)
  File "/Users/esc/git/numba/numba/np/ufunc/array_exprs.py", line 406, in _lower_array_expr
    cres = context.compile_subroutine(builder, impl, inner_sig, flags=flags,
  File "/Users/esc/git/numba/numba/core/base.py", line 863, in compile_subroutine
    cres = self._compile_subroutine_no_cache(builder, impl, sig,
  File "/Users/esc/git/numba/numba/core/base.py", line 834, in _compile_subroutine_no_cache
    cres = compiler.compile_internal(self.typing_context, self,
  File "/Users/esc/git/numba/numba/core/compiler.py", line 836, in compile_internal
    return pipeline.compile_extra(func)
  File "/Users/esc/git/numba/numba/core/compiler.py", line 455, in compile_extra
    self.state.func_id = bytecode.FunctionIdentity.from_function(func)
  File "/Users/esc/git/numba/numba/core/bytecode.py", line 386, in from_function
    pysig = utils.pysignature(func)
  File "/Users/esc/miniconda3-arm64/envs/numba_3.10/lib/python3.10/inspect.py", line 3247, in signature
    return Signature.from_callable(obj, follow_wrapped=follow_wrapped,
  File "/Users/esc/miniconda3-arm64/envs/numba_3.10/lib/python3.10/inspect.py", line 2995, in from_callable
    return _signature_from_callable(obj, sigcls=cls,
  File "/Users/esc/miniconda3-arm64/envs/numba_3.10/lib/python3.10/inspect.py", line 2456, in _signature_from_callable
    return _signature_from_function(sigcls, obj,
  File "/Users/esc/miniconda3-arm64/envs/numba_3.10/lib/python3.10/inspect.py", line 2321, in _signature_from_function
    parameters.append(Parameter(name, annotation=annotation,
  File "/Users/esc/miniconda3-arm64/envs/numba_3.10/lib/python3.10/inspect.py", line 2666, in __init__
    raise ValueError('{!r} is not a valid parameter name'.format(name))
ValueError: 'vLiteral[str](a1)__28binary_subscr_7' is not a valid parameter name

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/Users/esc/git/numba/issue_8939.py", line 17, in <module>
    print(f(test))
  File "/Users/esc/git/numba/numba/core/dispatcher.py", line 487, in _compile_for_args
    raise e
  File "/Users/esc/git/numba/numba/core/dispatcher.py", line 420, in _compile_for_args
    return_val = self.compile(tuple(argtypes))
  File "/Users/esc/git/numba/numba/core/dispatcher.py", line 965, in compile
    cres = self._compiler.compile(args, return_type)
  File "/Users/esc/git/numba/numba/core/dispatcher.py", line 125, in compile
    status, retval = self._compile_cached(args, return_type)
  File "/Users/esc/git/numba/numba/core/dispatcher.py", line 139, in _compile_cached
    retval = self._compile_core(args, return_type)
  File "/Users/esc/git/numba/numba/core/dispatcher.py", line 152, in _compile_core
    cres = compiler.compile_extra(self.targetdescr.typing_context,
  File "/Users/esc/git/numba/numba/core/compiler.py", line 762, in compile_extra
    return pipeline.compile_extra(func)
  File "/Users/esc/git/numba/numba/core/compiler.py", line 460, in compile_extra
    return self._compile_bytecode()
  File "/Users/esc/git/numba/numba/core/compiler.py", line 528, in _compile_bytecode
    return self._compile_core()
  File "/Users/esc/git/numba/numba/core/compiler.py", line 507, in _compile_core
    raise e
  File "/Users/esc/git/numba/numba/core/compiler.py", line 494, in _compile_core
    pm.run(self.state)
  File "/Users/esc/git/numba/numba/core/compiler_machinery.py", line 368, in run
    raise patched_exception
  File "/Users/esc/git/numba/numba/core/compiler_machinery.py", line 356, in run
    self._runPass(idx, pass_inst, state)
  File "/Users/esc/git/numba/numba/core/compiler_lock.py", line 35, in _acquire_compile_lock
    return func(*args, **kwargs)
  File "/Users/esc/git/numba/numba/core/compiler_machinery.py", line 311, in _runPass
    mutated |= check(pss.run_pass, internal_state)
  File "/Users/esc/git/numba/numba/core/compiler_machinery.py", line 273, in check
    mangled = func(compiler_state)
  File "/Users/esc/git/numba/numba/core/typed_passes.py", line 466, in run_pass
    lower.lower()
  File "/Users/esc/git/numba/numba/core/lowering.py", line 187, in lower
    self.lower_normal_function(self.fndesc)
  File "/Users/esc/git/numba/numba/core/lowering.py", line 241, in lower_normal_function
    entry_block_tail = self.lower_function_body()
  File "/Users/esc/git/numba/numba/core/lowering.py", line 271, in lower_function_body
    self.lower_block(block)
  File "/Users/esc/git/numba/numba/core/lowering.py", line 283, in lower_block
    with new_error_context('lowering "{inst}" at {loc}', inst=inst,
  File "/Users/esc/miniconda3-arm64/envs/numba_3.10/lib/python3.10/contextlib.py", line 153, in __exit__
    self.gen.throw(typ, value, traceback)
  File "/Users/esc/git/numba/numba/core/errors.py", line 837, in new_error_context
    raise newerr.with_traceback(tb)
numba.core.errors.LoweringError: Failed in nopython mode pipeline (step: native lowering)
'vLiteral[str](a1)__28binary_subscr_7' is not a valid parameter name

File "issue_8939.py", line 13:
def f(arr):
    <source elided>
    for k in literal_unroll(fields):
        res = res + np.abs(arr[k]).sum()
        ^

During: lowering "vLiteral[str](a1)_$30call_method.8 = arrayexpr(expr=(<ufunc 'absolute'>, [Var(vLiteral[str](a1)_$28binary_subscr.7, issue_8939.py:13)]), ty=array(float64, 1d, C))" at /Users/esc/git/numba/issue_8939.py (13)
```

So, the next thing is the fix from #8948, which reads:

```
                if isinstance(branch_ty, types.Literal):
                    new_var_dict[name] = scope.define(
                        f"v{branch_ty.literal_value}_{name}", var.loc).name
                else:
                    new_var_dict[name] = scope.define(
                        f"v{branch_ty}_{name}", var.loc).name
```

This works (more or less, but see below) because now `branch_ty.literal_value` is simply `a1` -- i.e. w/o the `Literal[str](` prefix. This then eliminates the square brackets and so the variable name is a legal identifier once again.

Now, moving along, if we adapt the reproducer from this issue to read, (force inclusion of square brackets in variable name):

```
test = np.array([(1, 2), (2, 3)], dtype=[("a1[0]", "f8"), ("a2", "f8")])
```

We get:

```
numba.core.errors.LoweringError: Failed in nopython mode pipeline (step: native lowering)
'va1[0]__28binary_subscr_7' is not a valid parameter name

File "issue_8939.py", line 13:
def f(arr):
    <source elided>
    for k in literal_unroll(fields):
        res = res + np.abs(arr[k]).sum()
        ^

During: lowering "va1[0]_$30call_method.8 = arrayexpr(expr=(<ufunc 'absolute'>, [Var(va1[0]_$28binary_subscr.7, issue_8939.py:13)]), ty=array(float64, 1d, C))" at /Users/vhaenel/git/numba/issue_8939.py (13)
```

which indicates that the fix from #8948 isn't sufficient. This test case had worked before c6a25069.

@stuartarchibald and I agreed that reverting  https://github.com/numba/numba/commit/c6a25069deabe982bbc589bc98c4290808f8a4e4 as proposed in https://github.com/numba/numba/pull/8948#issuecomment-1571559865 is going to be way forward. I will also add two tests, the original from this issue and the variant using `a1[0]` from above and add a description as to why we need a define followed by a redefine. This will fix this issue and allow us to get 0.57.1 out the door. Then, we can look at adding back in `branch_ty` while also legalizing the original variable name, such that `inspect.py` will not complain further down the line.

Resolving patch:
diff --git a/numba/core/untyped_passes.py b/numba/core/untyped_passes.py
--- a/numba/core/untyped_passes.py
+++ b/numba/core/untyped_passes.py
@@ -906,8 +906,16 @@ def inject_loop_body(self, switch_ir, loop_ir, caller_max_label,
             new_var_dict = {}
             for name, var in var_table.items():
                 scope = switch_ir.blocks[lbl].scope
-                new_var_dict[name] = scope.define(
-                    f"v{branch_ty}_{name}", var.loc).name
+                try:
+                    scope.get_exact(name)
+                except errors.NotDefinedError:
+                    # In case the scope doesn't have the variable, we need to
+                    # define it prior creating new copies of it! This is
+                    # because the scope of the function and the scope of the
+                    # loop are different and the variable needs to be redefined
+                    # within the scope of the loop.
+                    scope.define(name, var.loc)
+                new_var_dict[name] = scope.redefine(name, var.loc).name
             replace_var_names(loop_blocks, new_var_dict)
 
             # clobber the sentinel body and then stuff in the rest


Test patch:
diff --git a/numba/tests/test_mixed_tuple_unroller.py b/numba/tests/test_mixed_tuple_unroller.py
--- a/numba/tests/test_mixed_tuple_unroller.py
+++ b/numba/tests/test_mixed_tuple_unroller.py
@@ -1830,6 +1830,35 @@ def foo(tup):
             while x == 1:
                 x = 0
 
+    def test_literal_unroll_legalize_var_names01(self):
+        # See issue #8939
+        test = np.array([(1, 2), (2, 3)], dtype=[("a1", "f8"), ("a2", "f8")])
+        fields = tuple(test.dtype.fields.keys())
+
+        @njit
+        def foo(arr):
+            res = 0
+            for k in literal_unroll(fields):
+                res = res + np.abs(arr[k]).sum()
+            return res
+
+        self.assertEqual(foo(test), 8.0)
+
+    def test_literal_unroll_legalize_var_names02(self):
+        # See issue #8939
+        test = np.array([(1, 2), (2, 3)],
+                        dtype=[("a1[0]", "f8"), ("a2[1]", "f8")])
+        fields = tuple(test.dtype.fields.keys())
+
+        @njit
+        def foo(arr):
+            res = 0
+            for k in literal_unroll(fields):
+                res = res + np.abs(arr[k]).sum()
+            return res
+
+        self.assertEqual(foo(test), 8.0)
+
 
 def capture(real_pass):
     """ Returns a compiler pass that captures the mutation state reported