Tragic loss of one-shot information

Loss of one-shot info

The problem

Consider these definitions, gotten from Note [The one-shot state monad trick] in GHC.Utils.Monad:

   newtype M a = MkM' { unM :: State -> (State, a) }

   -- The pattern has the old constructor name.
   pattern MkM f <- MkM' f
      where
        MkM f = MkM' (oneShot f)

    instance Applicative M where
          pure a = MkM (\s -> (s, a))
          (<*>)  = ap

    instance Monad UM where
      m >>= k  = MkM (\s -> case unM m s of (s',v) -> unM (k v) s')

From these we get Core like this (omitting cases to avoid clutter):

$c>>=  =  \a b s{one-shot). blah
$c>>   =  \a b. $c>>= a (\_. b)

The one-shot info in $c>>= comes from that oneShot call. Now, if we inlined $c>>= immediately, all would be well. But sadly we don't. Instead we eta-expand $c>> thus:

$c>>   =  \a b s. $c>>= a (\_. b) s

Note: no one-shot info on that \s!! Now if we later inline $c>>= we'll get simply

$c>>   =  \a b s. blah

Result: we get one-shot info on $c>>= but not on $c>>. Sigh.

This actually happened in GHC.Core.Unify.bindTv, which ended up with a smaller arity than it should have had.

Possible solutions

• One solution is to add an INLINE pragma on (>>=) so that it'll inline when given only two arguments:

    instance Monad UM where
      {-# INLINE (>>=) #-}
      m >>= k  = MkM (\s -> case unM m s of (s',v) -> unM (k v) s')

  Now, we'll inline before eta-expanding.

• Another solution is to write out (>>) explicitly, using that MkM data constructor directly, insteaed of defining (>>) via (>>=):

    instance Monad UM where
      m >>= k  = MkM (\s -> case unM m  s of (s',v) -> unM (k v) s')
      m1 >> m2 = MkM (\s -> case unM m1 s of (s',_) -> unM m2    s')

• Both these solutions are delicate. A more robust solution would be to store the ArityType of the Id, rather than the Arity, so that we can drive the eta expansion from that