Let the specialiser work on dicts under lambdas

Following the discussion under #16473, this change allows the
specializer to work on any dicts in a lambda, not just those that occur
at the beginning.

For example, if you use data types which contain dictionaries and
higher-rank functions then once these are erased by the optimiser you
end up with functions such as:

```
  go_s4K9
  Int#
  -> forall (m :: * -> *).
     Monad m =>
     (forall x. Union '[State (Sum Int)] x -> m x) -> m ()
```

The dictionary argument is after the Int# value argument, this patch
allows `go` to be specialised.

4 files changed, 248 insertions, 1 deletions

diff --git a/testsuite/tests/perf/compiler/Makefile b/testsuite/tests/perf/compiler/Makefile
index 7d8e96fd44..b27c842e91 100644
--- a/testsuite/tests/perf/compiler/Makefile
+++ b/testsuite/tests/perf/compiler/Makefile
@@ -2,8 +2,12 @@ TOP=../../..
 include $(TOP)/mk/boilerplate.mk
 include $(TOP)/mk/test.mk
 
-.PHONY: T4007
+.PHONY: T4007 T16473
 T4007:
 	$(RM) -f T4007.hi T4007.o
 	'$(TEST_HC)' $(TEST_HC_OPTS) -c -O -ddump-rule-firings T4007.hs
 
+T16473:
+	$(RM) -f T16473.hi T16473.o
+	'$(TEST_HC)' $(TEST_HC_OPTS) -c -O -ddump-rule-firings T16473.hs
+
diff --git a/testsuite/tests/perf/compiler/T16473.hs b/testsuite/tests/perf/compiler/T16473.hs
new file mode 100644
index 0000000000..8a9751e306
--- /dev/null
+++ b/testsuite/tests/perf/compiler/T16473.hs
@@ -0,0 +1,102 @@
+{-# LANGUAGE BangPatterns        #-}
+{-# LANGUAGE DataKinds           #-}
+{-# LANGUAGE DeriveFunctor       #-}
+{-# LANGUAGE GADTs               #-}
+{-# LANGUAGE KindSignatures      #-}
+{-# LANGUAGE LambdaCase          #-}
+{-# LANGUAGE RankNTypes          #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeOperators       #-}
+
+{-# OPTIONS_GHC -flate-specialise -O2 #-}
+
+module Main (main) where
+
+import qualified Control.Monad.State.Strict as S
+import           Data.Foldable
+import           Data.Functor.Identity
+import           Data.Kind
+import           Data.Monoid
+import           Data.Tuple
+
+main :: IO ()
+main = print $ badCore 100
+
+badCore :: Int -> Int
+badCore n  = getSum $ fst $ run  $ runState mempty $ for_ [0..n] $ \i ->   modify (<> Sum i)
+
+data Union (r :: [Type -> Type]) a where
+  Union :: e a -> Union '[e] a
+
+decomp :: Union (e ': r) a -> e a
+decomp (Union a) = a
+{-# INLINE decomp #-}
+
+absurdU :: Union '[] a -> b
+absurdU = absurdU
+
+newtype Semantic r a = Semantic
+  { runSemantic
+        :: forall m
+         . Monad m
+        => (forall x. Union r x -> m x)
+        -> m a
+  }
+
+instance Functor (Semantic f) where
+  fmap f (Semantic m) = Semantic $ \k -> fmap f $ m k
+  {-# INLINE fmap #-}
+
+instance Applicative (Semantic f) where
+  pure a = Semantic $ const $ pure a
+  {-# INLINE pure #-}
+  Semantic f <*> Semantic a = Semantic $ \k -> f k <*> a k
+  {-# INLINE (<*>) #-}
+
+instance Monad (Semantic f) where
+  return = pure
+  {-# INLINE return #-}
+  Semantic ma >>= f = Semantic $ \k -> do
+    z <- ma k
+    runSemantic (f z) k
+  {-# INLINE (>>=) #-}
+
+data State s a
+  = Get (s -> a)
+  | Put s a
+  deriving Functor
+
+get :: Semantic '[State s] s
+get = Semantic $ \k -> k $ Union $ Get id
+{-# INLINE get #-}
+
+put :: s -> Semantic '[State s] ()
+put !s = Semantic $ \k -> k $ Union $! Put s ()
+{-# INLINE put #-}
+
+modify :: (s -> s) -> Semantic '[State s] ()
+modify f = do
+  !s <- get
+  put $! f s
+{-# INLINE modify #-}
+
+runState :: s -> Semantic (State s ': r) a -> Semantic r (s, a)
+runState = interpretInStateT $ \case
+  Get k   -> fmap k S.get
+  Put s k -> S.put s >> pure k
+{-# INLINE[3] runState #-}
+
+run :: Semantic '[] a -> a
+run (Semantic m) = runIdentity $ m absurdU
+{-# INLINE run #-}
+
+interpretInStateT
+    :: (forall x. e x -> S.StateT s (Semantic r) x)
+    -> s
+    -> Semantic (e ': r) a
+    -> Semantic r (s, a)
+interpretInStateT f s (Semantic m) = Semantic $ \k ->
+  fmap swap $ flip S.runStateT s $ m $ \u ->
+    S.mapStateT (\z -> runSemantic z k) $ f $ decomp u
+{-# INLINE interpretInStateT #-}
+
diff --git a/testsuite/tests/perf/compiler/T16473.stdout b/testsuite/tests/perf/compiler/T16473.stdout
new file mode 100644
index 0000000000..3a1f5a571b
--- /dev/null
+++ b/testsuite/tests/perf/compiler/T16473.stdout
@@ -0,0 +1,139 @@
+Rule fired: Class op fmap (BUILTIN)
+Rule fired: Class op liftA2 (BUILTIN)
+Rule fired: Class op $p1Applicative (BUILTIN)
+Rule fired: Class op <*> (BUILTIN)
+Rule fired: Class op <$ (BUILTIN)
+Rule fired: Class op $p1Applicative (BUILTIN)
+Rule fired: Class op <*> (BUILTIN)
+Rule fired: Class op fmap (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op fmap (BUILTIN)
+Rule fired: Class op get (BUILTIN)
+Rule fired: Class op return (BUILTIN)
+Rule fired: Class op fmap (BUILTIN)
+Rule fired: Class op >> (BUILTIN)
+Rule fired: Class op put (BUILTIN)
+Rule fired: Class op return (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: Class op return (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op fmap (BUILTIN)
+Rule fired: Class op get (BUILTIN)
+Rule fired: Class op return (BUILTIN)
+Rule fired: Class op fmap (BUILTIN)
+Rule fired: Class op >> (BUILTIN)
+Rule fired: Class op put (BUILTIN)
+Rule fired: Class op return (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: Class op return (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op show (BUILTIN)
+Rule fired: Class op mempty (BUILTIN)
+Rule fired: Class op fromInteger (BUILTIN)
+Rule fired: integerToInt (BUILTIN)
+Rule fired: Class op <> (BUILTIN)
+Rule fired: Class op + (BUILTIN)
+Rule fired: Class op enumFromTo (BUILTIN)
+Rule fired: Class op *> (BUILTIN)
+Rule fired: Class op *> (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: Class op *> (BUILTIN)
+Rule fired: Class op *> (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: Class op *> (BUILTIN)
+Rule fired: Class op *> (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: fold/build (GHC.Base)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op $p1Applicative (BUILTIN)
+Rule fired: Class op fmap (BUILTIN)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op $p1Applicative (BUILTIN)
+Rule fired: Class op fmap (BUILTIN)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: ># (BUILTIN)
+Rule fired: ==# (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op $p1Applicative (BUILTIN)
+Rule fired: SPEC/Main $fApplicativeStateT @ Identity _ (Main)
+Rule fired: SPEC/Main $fMonadStateT_$c>>= @ Identity _ (Main)
+Rule fired: SPEC/Main $fMonadStateT_$c>> @ Identity _ (Main)
+Rule fired: Class op return (BUILTIN)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op $p1Applicative (BUILTIN)
+Rule fired: SPEC/Main $fApplicativeStateT @ Identity _ (Main)
+Rule fired: SPEC/Main $fMonadStateT_$c>>= @ Identity _ (Main)
+Rule fired: SPEC/Main $fMonadStateT_$c>> @ Identity _ (Main)
+Rule fired: Class op return (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op $p1Applicative (BUILTIN)
+Rule fired: SPEC/Main $fApplicativeStateT @ Identity _ (Main)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op $p1Applicative (BUILTIN)
+Rule fired: SPEC/Main $fApplicativeStateT @ Identity _ (Main)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op $p1Applicative (BUILTIN)
+Rule fired: Class op fmap (BUILTIN)
+Rule fired: Class op fmap (BUILTIN)
+Rule fired: Class op return (BUILTIN)
+Rule fired: Class op return (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op return (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op return (BUILTIN)
+Rule fired: Class op fmap (BUILTIN)
+Rule fired: Class op fmap (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op fmap (BUILTIN)
+Rule fired: SPEC/Main $fFunctorStateT @ Identity _ (Main)
+Rule fired:
+    SPEC/Main $fApplicativeStateT_$cpure @ Identity _ (Main)
+Rule fired: SPEC/Main $fApplicativeStateT_$c<*> @ Identity _ (Main)
+Rule fired: Class op fmap (BUILTIN)
+Rule fired: SPEC/Main $fApplicativeStateT_$c*> @ Identity _ (Main)
+Rule fired: Class op fmap (BUILTIN)
+Rule fired: SPEC/Main $fFunctorStateT @ Identity _ (Main)
+Rule fired:
+    SPEC/Main $fApplicativeStateT_$cpure @ Identity _ (Main)
+Rule fired: SPEC/Main $fApplicativeStateT_$c<*> @ Identity _ (Main)
+Rule fired: SPEC/Main $fApplicativeStateT_$c*> @ Identity _ (Main)
+Rule fired: SPEC/Main $fMonadStateT @ Identity _ (Main)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op <*> (BUILTIN)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op $p1Applicative (BUILTIN)
+Rule fired: Class op fmap (BUILTIN)
+Rule fired: Class op fmap (BUILTIN)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op <*> (BUILTIN)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op $p1Applicative (BUILTIN)
+Rule fired: Class op fmap (BUILTIN)
+Rule fired: Class op >>= (BUILTIN)
+Rule fired: Class op fmap (BUILTIN)
+Rule fired: SPEC go @ (StateT (Sum Int) Identity) (Main)
+Rule fired: Class op $p1Monad (BUILTIN)
+Rule fired: Class op pure (BUILTIN)
+Rule fired: SPEC/Main $fMonadStateT @ Identity _ (Main)
+Rule fired: SPEC go @ (StateT (Sum Int) Identity) (Main)
+Rule fired: Class op fmap (BUILTIN)
diff --git a/testsuite/tests/perf/compiler/all.T b/testsuite/tests/perf/compiler/all.T
index 44216f4075..0db9bcf0ee 100644
--- a/testsuite/tests/perf/compiler/all.T
+++ b/testsuite/tests/perf/compiler/all.T
@@ -404,3 +404,5 @@ test('T16190',
       collect_stats(),
       multimod_compile,
       ['T16190.hs', '-v0'])
+
+test('T16473', normal, makefile_test, ['T16473'])