Core global uniqueness problems since 9.2.2

Summary

For a core exploration tool I would greatly prefer that all my binding sites have a globally unique, Unique attached to them. Since normally these Uniques are only locally unique, I created my own plugin to take care of this. This worked great until 9.2.2 after which the simplifier is now unable to find a certain variable when calling refineFromInScope. Although this is labelled a bug I realize that the issue is likely on my end, but I need some advice

Steps to reproduce

Below is the essence of my plugin. I basically traverse the AST with a global set that tracks wether a unique was seen at a binding site before, if so it generates a new unique and replaces all its occurences using a scoped mapping.

-- A local scope map and a global unique set
data UniqEnv = UniqEnv
  { uenv_scope :: Map Int Int
  , uenv_global :: Set Int
  }

emptyEnv = UniqEnv
  { uenv_scope = M.empty
  , uenv_global = S.empty
  }

uenvInsertScope :: Int -> Int -> UniqEnv -> UniqEnv
uenvInsertScope from to env = env { uenv_scope = M.insert from to (uenv_scope env) }

uenvInsertGlobal :: Int -> UniqEnv -> UniqEnv
uenvInsertGlobal uid env = env { uenv_global = S.insert uid (uenv_global env) }


type Uniq a = StateT UniqEnv IO a

runUnique :: Uniq a -> IO a
runUnique uq = evalStateT uq emptyEnv

uniqueModule :: ModGuts -> IO ModGuts
uniqueModule guts@ModGuts { mg_binds = mg_binds } = do
    nbinds <- runUnique (uniqProgram mg_binds)
    pure $ guts { mg_binds = nbinds }

-- Prevent scoped substitutions to escape 
limitScope :: Uniq a -> Uniq a
limitScope g = do
    scope_before <- gets uenv_scope
    r <- g
    env_after <- get
    put $ env_after { uenv_scope = scope_before }
    pure r

uniqVar :: Var -> Uniq Var
uniqVar var = 
    if isTyVar var
    then pure var
    else do
      scope <- gets uenv_scope
      let (tag, uid) = unpkUnique (getUnique var)
      -- Lookup if the unique must change
      case M.lookup uid scope of
        Just i -> pure $ setVarUnique var (mkUnique tag i)
        Nothing -> pure var

uniqBndr :: CoreBndr -> Uniq Var
uniqBndr var = do
    if isTyVar var
       then pure var
       else do
        env <- get
        let (tag, uid) = unpkUnique (getUnique var)
        if S.member uid (uenv_global env)
           then do
              -- We've seen a different binding site with this unique before, generate a new one
              new_id <- (`mod` 100000000) <$> randomIO @Int
              modify $ uenvInsertScope uid new_id
              modify $ uenvInsertGlobal new_id
              uniqVar var
           else do
              modify $ uenvInsertGlobal uid
              pure var

uniqProgram :: CoreProgram -> Uniq CoreProgram
uniqProgram = mapM uniqBind

uniqBind :: CoreBind -> Uniq CoreBind
uniqBind (NonRec b e) = NonRec <$> uniqBndr b <*> uniqExpr e
uniqBind (Rec xs) = do
    let (bs, es) = unzip xs
    bs' <- mapM uniqBndr bs
    es' <- mapM uniqExpr es
    pure $ Rec $ zip bs' es'

uniqExpr :: CoreExpr -> Uniq CoreExpr
uniqExpr (Var var) = Var <$> uniqVar var
uniqExpr (Lit lit) = pure $ Lit lit
uniqExpr (App f a) = App <$> uniqExpr f <*> uniqExpr a
uniqExpr (Lam b e) = limitScope $ Lam <$> uniqBndr b <*> uniqExpr e
uniqExpr (Let b e) = limitScope $ Let <$> uniqBind b <*> uniqExpr e
uniqExpr (Case e b t alts) = limitScope $ Case <$> uniqExpr e <*> uniqBndr b <*> pure t <*> mapM uniqAlt alts
uniqExpr (Cast e c) = Cast <$> uniqExpr e <*> pure c
uniqExpr (Tick t e) = Tick t <$> uniqExpr e
uniqExpr (Type t) = pure $ Type t
uniqExpr (Coercion c) = pure $ Coercion c

uniqAlt :: CoreAlt -> Uniq CoreAlt
#if MIN_VERSION_ghc(9,2,0)
uniqAlt (Alt con bs e) = limitScope $ Alt con <$> mapM uniqBndr bs <*> uniqExpr e
#else
uniqAlt (con, bs, e) = limitScope $ (,,) con <$> mapM uniqBndr bs <*> uniqExpr e
#endif

When run over this program:

module Text where

import Data.Text (Text)
import qualified Data.Text as T

{-# INLINE slice #-}
slice :: Int -> Int -> Text -> Text
slice offset len = T.take len . T.drop offset

it results in this error

[11 of 12] Compiling Text             ( app/Text.hs, /home/hugo/repos/hs-comprehension/test-project/dist-newstyle/build/x86_64-linux/ghc-9.2.2/test-project-0.1.0.0/x/hs-plugin-test/build/hs-plugin-test/hs-plugin-test-tmp/Text.o )
ghc: panic! (the 'impossible' happened)
  (GHC version 9.2.2:
	refineFromInScope
  InScope {wild_00 eta_a5bN ds_a5bR wild_a5bS s_a5bV next0_a5c2
           s0_a5c3 len_a5c4 wild1_a5cp dt_a5cq dt1_a5cr wild_aSxFb i_a2T3Uo
           wild_a3qZuE x1_a3RYGz slice $trModule $trModule_s5cL $trModule_s5cM
           $trModule_s5cN $trModule_s5cO $snext_s5dq $stake_s5dF $sdrop_s5eB
           lvl_s5eD lvl_s5eJ lvl_s5eL lvl_s5eN lvl_s5eT lvl_s5eU lvl_s5eV
           lvl_s5eW lvl_s5eX lvl_s5eY $j_s5fw}
  i_a5fo
  Call stack:
      CallStack (from HasCallStack):
        callStackDoc, called at compiler/GHC/Utils/Panic.hs:181:37 in ghc:GHC.Utils.Panic
        pprPanic, called at compiler/GHC/Core/Opt/Simplify/Env.hs:706:30 in ghc:GHC.Core.Opt.Simplify.Env

Please report this as a GHC bug:  https://www.haskell.org/ghc/reportabug

While it does work on my other, far more complicated modules. Is there some AST invariant that I have violated? Advice much appreaciated.

Expected behavior

That this yields a valid Core program (ignoring the probability of random collisions which has been ruled out as the cause in the crash)

Environment

• GHC version used: 9.2.2

Optional:

• Operating System: Ubuntu
• System Architecture: amd64