ghc-9.0.1 RankNTypes inference regression

ghc-8.10.5 accepts this file, but I get the error below with ghc-9.0.1 and 9.2.0.20210422.

{-# LANGUAGE RankNTypes #-}
module Bug where

type L a = forall f. Functor f => a -> f a

f1 :: Int -> Int -> L a
f1 _ _ = undefined

f :: Eq b => Int -> Int -> L b
f i0 = sortArgs f1 i0

sortArgs :: Ord a => (a -> a -> r) -> a -> a -> r
sortArgs g x y
        | x <= y = g x y
        | otherwise = g y x

BugA.hs:10:8: error:
    • Couldn't match type ‘r0’ with ‘L b’
      Expected: Int -> L b
        Actual: Int -> r0
      Cannot instantiate unification variable ‘r0’
      with a type involving polytypes: L b
    • In the expression: sortArgs f1 i0
      In an equation for ‘f’: f i0 = sortArgs f1 i0
    • Relevant bindings include
        f :: Int -> Int -> L b (bound at BugA.hs:10:1)
   |
10 | f i0 = sortArgs f1 i0
   |        ^^^^^^^^^^^^^^

I started with this:

{-# LANGUAGE RankNTypes, TypeFamilies #-}
module Bug where
import qualified Data.IntMap as IM


-- import Control.Lens
-- or just these definitions

type family IxValue (a :: *) :: *
type family Index (a :: *) :: *

type instance IxValue (IM.IntMap a) = a
type instance Index (IM.IntMap a) = Int

at :: Index m -> Lens' m (Maybe (IxValue m))
at _ = undefined

non :: Eq a => a -> Lens' (Maybe a) a
non _ = undefined

type Lens' s a = forall f. Functor f => (a -> f a) -> s -> f s

at2s :: Eq b => Int -> Int -> Lens' (IM.IntMap (IM.IntMap (b,b))) (Maybe (b,b))
at2s i0 = sortArgs (\i j -> at i . non mempty . at j) i0

sortArgs :: Ord a => (a -> a -> r) -> a -> a -> r
sortArgs g x y
        | x <= y = g x y
        | otherwise = g y x

Bug.hs:24:11: error:
    • Couldn't match type: (Maybe (IxValue (IxValue m0))
                            -> f0 (Maybe (IxValue (IxValue m0))))
                           -> m0 -> f0 m0
                     with: forall (f :: * -> *).
                           Functor f =>
                           (Maybe (b, b) -> f (Maybe (b, b)))
                           -> IM.IntMap (IM.IntMap (b, b)) -> f (IM.IntMap (IM.IntMap (b, b)))
      Expected: Int
                -> Lens' (IM.IntMap (IM.IntMap (b, b))) (Maybe (b, b))
        Actual: Int
                -> (Maybe (IxValue (IxValue m0))
                    -> f0 (Maybe (IxValue (IxValue m0))))
                -> m0
                -> f0 m0
    • In the expression:
        sortArgs (\ i j -> at i . non mempty . at j) i0
      In an equation for ‘at2s’:
          at2s i0 = sortArgs (\ i j -> at i . non mempty . at j) i0
    • Relevant bindings include
        at2s :: Int
                -> Int -> Lens' (IM.IntMap (IM.IntMap (b, b))) (Maybe (b, b))
          (bound at Bug.hs:24:1)
   |
24 | at2s i0 = sortArgs (\i j -> at i . non mempty . at j) i0
   |          

ghc-9.0.1 accepts Bug.hs if I comment the at2s type signature, but removing the type signature for f in BugA.hs makes no difference.