Skip to content

Should GHC be able to figure out that this eta-expansion would be beneficial?

In the process of investigating #23172 (closed), I've run into a deficiency in optimizing the following program:

{-# LANGUAGE NoImplicitPrelude, CPP #-}
module Ap where

class Monad m where
    pure :: a -> m a

    (<*>) :: m (a -> b) -> m a -> m b
    infixl 4 <*>

    (>>=)  :: m a -> (a -> m b) -> m b
    infixl 1 >>=

{-# INLINE ap #-}
ap :: (Monad m) => m (a -> b) -> m a -> m b
ap ff fx =
    ff >>= \ f ->
    fx >>= \ x ->
    pure (f x)

data T a = MkT

instance Monad T where
    pure _ = MkT
    _ >>= _ = MkT

    (<*>) =
#if ETA_EXPANDED_AP
        \ff fx -> ap ff fx
#else
        ap
#endif

If I compile this with -O2 -fno-specialise, then with ETA_EXPANDED_AP=0 I get the following output:

$ ./_build/stage1/bin/ghc \
    -dsuppress-all -dsuppress-uniques -dno-suppress-type-signatures -dno-typeable-binds \
    -fforce-recomp \
    -O2 -fno-specialise \
    -ddump-ds-preopt -ddump-prep \
    input/ap7.hs -DETA_EXPANDED_AP=0

==================== Desugar (before optimization) ====================

$c<*> :: forall a b. T (a -> b) -> T a -> T b
$c<*>
  = \ (@a) (@b) ->
      let {
        $dMonad :: Monad T
        $dMonad = $fMonadT } in
      ap $dMonad

==================== CorePrep ====================

Rec {
$fMonadT :: Monad T
$fMonadT = C:Monad $fMonadT_$cpure $fMonadT_$c<*> $fMonadT_$c>>=

$fMonadT_$c<*> :: forall a b. T (a -> b) -> T a -> T b
$fMonadT_$c<*>
  = \ (@a) (@b) (eta :: T (a -> b)) (eta :: T a) ->
      ap $fMonadT eta eta
end Rec }

However, if the input is manually eta-expanded (by setting ETA_EXPANDED_AP=1), the change shows up in the desugared output, and eventually leads to much nicer output after the optimisation pipeline:

$ ./_build/stage1/bin/ghc \
    -dsuppress-all -dsuppress-uniques -dno-suppress-type-signatures -dno-typeable-binds \
    -fforce-recomp \
    -O2 -fno-specialise \
    -ddump-ds-preopt -ddump-prep \
    input/ap7.hs -DETA_EXPANDED_AP=

==================== Desugar (before optimization) ====================

$c<*> :: forall a b. T (a -> b) -> T a -> T b
$c<*>
  = \ (@a) (@b) ->
      let {
        $dMonad :: Monad T
        $dMonad = $fMonadT } in
      \ (ff :: T (a -> b)) (fx :: T a) -> ap $dMonad ff fx

==================== CorePrep ====================

$fMonadT_$c<*> :: forall a b. T (a -> b) -> T a -> T b
$fMonadT_$c<*> = \ (@a) (@b) _ _ -> MkT

$fMonadT :: Monad T
$fMonadT = C:Monad $fMonadT_$cpure $fMonadT_$c<*> $fMonadT_$c>>=

Would it be possible/reasonable(/desirable?) to expect GHC to generate this second version even from the eta-contracted input?

Edited by Gergő Érdi
To upload designs, you'll need to enable LFS and have an admin enable hashed storage. More information