Static argument transformation should also run after specialisation

Consider the following program where we eventually want thepayload to simplify to the same code as direct.

{-# LANGUAGE ExistentialQuantification, RankNTypes, DeriveFunctor #-}
module Foo where

newtype Q a b t = Q { getQ :: forall f . Applicative f => (a -> f b) -> f t }
  deriving Functor

instance Applicative (Q a b) where
  pure a = Q (\_ -> pure a)
  (Q ab) <*> (Q a) = (Q (\v -> ab v <*> a v))

singleQ :: a -> Q a b b
singleQ a = Q (\f -> f a)

data L a = Nil | L a (L a) deriving Show

traverseList :: Applicative f => (a -> f b) -> L a -> f (L b)
traverseList f Nil = pure Nil
traverseList f (L a la) = L <$> f a <*> traverseList f la

newtype Identity a = Identity { runIdentity :: a } deriving (Functor, Show)

instance Applicative Identity where
  pure = Identity
  (Identity f) <*> (Identity x) = Identity (f x)


thepayload :: L String -> L String
thepayload l = runIdentity $ (getQ $ (traverseList singleQ l)) Identity

direct :: L String -> L String
direct Nil = Nil
direct (L a b) = L a (direct b)

With ghc-8.0.2 and -fstatic-argument-transformation, the specialiser will specialise the call to traverseList and leave us with a definition like,

rec
(20)
$ssat_worker= λ sg sc l →
   case l of
      Nil→ pure sc Nil
      L a la → <*> sc (fmap ($p1Applicative sc) (L a)) ($ssat_worker sg sc la)

(4) thepayload = λl→ $ssat_worker $fApplicativeIdentity l

$ssat_worked is recursive in l but not in the other two arguments so we can also apply SAT here.

Notice that $ssat_worker is called with a statically known dictionary in thepayload and so if we can inline $ssat_worker we would get the same code as the naive direct, as desired.

I verified that inserting another SAT pass later in the compilation pipeline does have the desired effect but am not sure where exactly the right place would be or whether it is in general desirable.