Infer and propagate constant information earlier in the pipeline.

[stuartarchibald]

Reporting a bug

• I have tried using the latest released version of Numba (most recent is
  visible in the change log (https://github.com/numba/numba/blob/main/CHANGE_LOG).
• I have included a self contained code sample to reproduce the problem.
  i.e. it's possible to run as 'python bug.py'.

I'm wondering if code like this could potentially be handled directly by Numba. The following fails at present as it can't unify agg in the agg.shape call. If the following sequence of passes were run:

• PartialTypeInference (aggs_and_cols[0] inferred as 3d)
• RewriteSemanticConstants (aggs_and_cols[0].ndim baked into IR as const(3)? maybe needs some work on the pass)
• DeadBranchPrune (Remove branching as 3 == 3 is always True)

it might allow elimination of the control flow and permit compilation.

from numba import njit
import numpy as np

@njit
def foo(*aggs_and_cols):
    i = 0
    agg = aggs_and_cols[0] # This agg is 3D
    if aggs_and_cols[0].ndim == 3: # This type could be inferred, semantic rewritten and branch pruned.
        cat_index = aggs_and_cols[1][i]
        agg = agg[:, :, cat_index] # This agg is 2D

    xmax = agg.shape[1] - 1 # What's this agg? 2D or 3D?

args = (np.zeros((3, 4, 5)), (0, 0))

foo(*args)