Wanted: better arity analysis
@tomjaguarpaw1, in this email showed an example where arity analysis fails badly.
His program
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
module Main where
import Control.Monad.Trans.Class (MonadTrans)
import Control.Monad.Trans.State.Strict (StateT)
import qualified Control.Monad.Trans.State.Strict as State
import Data.Coerce (coerce)
import Data.Functor.Identity (Identity (runIdentity))
import Data.Kind (Type)
import Data.Foldable (for_)
-- These "myModify" functions all compile to the same code --
-- *exactly* the same code, i.e. they share the implementation (which
-- is just a constant)! Good.
myModifyMTL :: Int
myModifyMTL =
runIdentity $
flip State.evalStateT 0 $ do
s <- State.get
State.put $! s + 1
State.get
myModifyNewtype :: Int
myModifyNewtype =
runIdentity $
unE $
flip State.evalStateT 0 $
unL $
unB $ do
s <- Bm $ Lm State.get
Bm $ Lm $ State.put $! s + 1
Bm $ Lm State.get
myModifyEff :: Int
myModifyEff =
runIdentity $
unMkEff $
flip State.evalStateT 0 $
unMkEff $
unMkEff @(Branch (Leaf (StateT Int)) Empty) $ do
s <- MkEff $ MkEff State.get
MkEff $ MkEff $ State.put $! s + 1
MkEff $ MkEff State.get
-- These "mySum" functions don't all compile to the same code. Bad!
-- "mySumMTL" and "mySumNewtype" compile to the same code but it's not
-- shared. That's fine. I don't care. I really care that "mySumEff"
-- compiles to the same code, but it doesn't. It compiles to an
-- inefficient loop with allocation.
mySumMTL :: Int
mySumMTL =
runIdentity $
flip State.evalStateT 0 $ do
for_ [1 .. 10] $ \i -> do
s <- State.get
State.put $! s + i
State.get
mySumNewtype :: Int
mySumNewtype =
runIdentity $
unE $
flip State.evalStateT 0 $
unL $
unB $ do
for_ [1.. 10] $ \i -> do
s <- Bm $ Lm State.get
Bm $ Lm $ State.put $! s + i
Bm $ Lm State.get
mySumEff :: Int
mySumEff =
runIdentity $
unMkEff $
flip State.evalStateT 0 $
unMkEff $
unMkEff @(Branch (Leaf (StateT Int)) Empty) $
do
for_ [1..10] $ \i -> do
s <- MkEff $ MkEff State.get
MkEff $ MkEff $ State.put $! s + i
MkEff $ MkEff State.get
-- Definitions for the "Newtype" versions
newtype Lt t m a = Lm {unL :: t m a}
deriving newtype (Functor, Applicative, Monad)
newtype Et m a = Em {unE :: m a}
deriving newtype (Functor, Applicative, Monad)
newtype Bt t1 t2 m a = Bm {unB :: t1 (t2 m) a}
deriving newtype (Functor, Applicative, Monad)
-- Definitions for the "Eff" versions, that use Eff, an "effect
-- system"
-- A type that encodes the structure of a composed series of monad
-- transformers
data Effects
= Branch Effects Effects
| Leaf ((Type -> Type) -> Type -> Type)
| Empty
-- A singleton type for Effects
data SEffects i where
SBranch :: (SingI i1, SingI i2) => SEffects (Branch i1 i2)
SLeaf :: (MonadTrans t) => SEffects (Leaf t)
SEmpty :: SEffects Empty
-- For passing the singleton implicitly, like in the singletons
-- library
class SingI i where
sing :: SEffects i
instance (MonadTrans t) => SingI (Leaf t) where
sing = SLeaf
instance (SingI t1, SingI t2) => SingI (Branch t1 t2) where
sing = SBranch
instance SingI Empty where
sing = SEmpty
newtype Eff es m a = MkEff {unMkEff :: EffF es m a}
type family EffF (es :: Effects) m where
EffF Empty m = m
EffF (Leaf t) m = t m
EffF (Branch s1 s2) m = Eff s1 (Eff s2 m)
coerceFmapEmpty ::
forall a b m. (Functor m) => (a -> b) -> Eff Empty m a -> Eff Empty m b
coerceFmapEmpty = coerce (fmap @m @a @b)
coerceFmapLeaf ::
forall a b t m.
(MonadTrans t, Monad m) =>
(a -> b) ->
Eff (Leaf t) m a ->
Eff (Leaf t) m b
coerceFmapLeaf = coerce (fmap @(t m) @a @b)
coerceFmapBranch ::
forall a b s1 s2 m.
(SingI s1, SingI s2, Monad m) =>
(a -> b) ->
Eff (Branch s1 s2) m a ->
Eff (Branch s1 s2) m b
coerceFmapBranch = coerce (fmap @(Eff s1 (Eff s2 m)) @a @b)
coercePureEmpty :: forall a m. (Applicative m) => a -> Eff Empty m a
coercePureEmpty = coerce (pure :: a -> m a)
coercePureLeaf ::
forall a t m.
(MonadTrans t, Monad m) =>
a ->
Eff (Leaf t) m a
coercePureLeaf = coerce (pure :: a -> t m a)
coercePureBranch ::
forall a s1 s2 m.
(Monad m, SingI s1, SingI s2) =>
a ->
Eff (Branch s1 s2) m a
coercePureBranch = coerce (pure :: a -> Eff s1 (Eff s2 m) a)
coerceAndThenEmpty ::
forall m a b. (Applicative m) => Eff Empty m a -> Eff Empty m b -> Eff Empty m b
coerceAndThenEmpty = coerce ((*>) :: m a -> m b -> m b)
coerceAndThenLeaf ::
forall t m a b.
(MonadTrans t, Monad m) =>
Eff (Leaf t) m a ->
Eff (Leaf t) m b ->
Eff (Leaf t) m b
coerceAndThenLeaf = coerce ((*>) :: t m a -> t m b -> t m b)
coerceAndThenBranch ::
forall s1 s2 m a b.
(SingI s1, SingI s2, Monad m) =>
Eff (Branch s1 s2) m a ->
Eff (Branch s1 s2) m b ->
Eff (Branch s1 s2) m b
coerceAndThenBranch =
coerce
( (*>) ::
Eff s1 (Eff s2 m) a ->
Eff s1 (Eff s2 m) b ->
Eff s1 (Eff s2 m) b
)
coerceAppEmpty ::
forall m a b. (Applicative m) => Eff Empty m (a -> b) -> Eff Empty m a -> Eff Empty m b
coerceAppEmpty = coerce ((<*>) :: m (a -> b) -> m a -> m b)
coerceAppLeaf ::
forall t m a b.
(MonadTrans t, Monad m) =>
Eff (Leaf t) m (a -> b) ->
Eff (Leaf t) m a ->
Eff (Leaf t) m b
coerceAppLeaf = coerce ((<*>) :: t m (a -> b) -> t m a -> t m b)
coerceAppBranch ::
forall s1 s2 m a b.
(SingI s1, SingI s2, Monad m) =>
Eff (Branch s1 s2) m (a -> b) ->
Eff (Branch s1 s2) m a ->
Eff (Branch s1 s2) m b
coerceAppBranch =
coerce
( (<*>) ::
Eff s1 (Eff s2 m) (a -> b) ->
Eff s1 (Eff s2 m) a ->
Eff s1 (Eff s2 m) b
)
coerceBindEmpty ::
forall m a b. (Monad m) => Eff Empty m a -> (a -> Eff Empty m b) -> Eff Empty m b
coerceBindEmpty = coerce ((>>=) @m @a @b)
coerceBindLeaf ::
forall t m a b.
(MonadTrans t, Monad m) =>
Eff (Leaf t) m a ->
(a -> Eff (Leaf t) m b) ->
Eff (Leaf t) m b
coerceBindLeaf = coerce ((>>=) @(t m) @a @b)
coerceBindBranch ::
forall s1 s2 m a b.
(SingI s1, SingI s2, Monad m) =>
Eff (Branch s1 s2) m a ->
(a -> Eff (Branch s1 s2) m b) ->
Eff (Branch s1 s2) m b
coerceBindBranch =
coerce ((>>=) @(Eff s1 (Eff s2 m)) @a @b)
instance (SingI es, Monad m) => Functor (Eff es m) where
fmap = case sing @es of
SEmpty -> coerceFmapEmpty
SBranch -> coerceFmapBranch
SLeaf -> coerceFmapLeaf
{-# INLINE fmap #-}
instance (SingI es, Monad m) => Applicative (Eff es m) where
pure = case sing @es of
SEmpty -> coercePureEmpty
SLeaf -> coercePureLeaf
SBranch -> coercePureBranch
{-# INLINE pure #-}
(<*>) = case sing @es of
SEmpty -> coerceAppEmpty
SLeaf -> coerceAppLeaf
SBranch -> coerceAppBranch
{-# INLINE (<*>) #-}
(*>) = case sing @es of
SEmpty -> coerceAndThenEmpty
SLeaf -> coerceAndThenLeaf
SBranch -> coerceAndThenBranch
{-# INLINE (*>) #-}
instance (SingI es, Monad m) => Monad (Eff es m) where
(>>=) = case sing @es of
SEmpty -> coerceBindEmpty
SLeaf -> coerceBindLeaf
SBranch -> coerceBindBranch
{-# INLINE (>>=) #-}
main :: IO ()
main = pure ()Focus on the code for mySumEff and you get this monstrosity
Main.mySumEff_go3 :: GHC.Prim.Int# -> Eff (Branch (Leaf (StateT Int)) Empty) Identity ()
Main.mySumEff_go3
= \ (x_a3In :: GHC.Prim.Int#) ->
let {
k_a2Fj :: Eff (Branch (Leaf (StateT Int)) Empty) Identity ()
k_a2Fj
= case x_a3In of wild_X1H {
__DEFAULT -> Main.mySumEff_go3 (GHC.Prim.+# wild_X1H 1#);
10# ->
n_r3O8
`cast` (co1 :: (Int -> ((), Int))
~R# Eff (Branch (Leaf (StateT Int)) Empty) Identity ())
} } in
(\ (eta2_a2GN :: Int) ->
case eta2_a2GN of { GHC.Types.I# x1_a2Yy ->
(k_a2Fj
`cast` (co2 :: Eff (Branch (Leaf (StateT Int)) Empty) Identity ()
~R# (Int -> Eff Empty Identity ((), Int))))
(GHC.Types.I# (GHC.Prim.+# x1_a2Yy x_a3In))
})
`cast` (co3 :: (Int -> Eff Empty Identity ((), Int))
~R# Eff (Branch (Leaf (StateT Int)) Empty) Identity ())```Not efficient at all
Diagnosis
GHC is reluctant to inline k_a2Fj inside the \eta in case that would duplicate work: inlining
k would push the call mySumEff_go3 (wild_X1H + 1) inside the lambda. What we want is to recognise that k has arity 1 -- it could have a lambda at the top. Absent the casts, GHC will work that out -- look at GHC.Core.Opt.Arity.findRhsArity.
But the casts get in the way. What we'd like is
k_a2Fj = (\eta. case x of wild
DEFAULT -> ((go3 (wild + 1)) `cast` sym co1) eta
10 -> n eta)
`cast` co1
where
co1 :: (Int -> ((), Int))
~R# Eff (Branch (Leaf (StateT Int)) Empty) Identity ())Now, GHC's arity analyser is clever enough to do this when co1 is just newtypes.
See Note [The EtaInfo mechanism] in GHC.Core.Opt.Arity.
But here it isn't. We have to execute the recursive EffF type-level function
Eff (Branch (Leaf (StateT Int)) Empty) Identity ()
--> EffF (Branch (Leaf (StateT Int)) Empty) Identity ()
--> Eff (Leaf (StateT Int)) (Eff Empty Identy) ()
--> EffF (Leaf (StateT Int)) (EffF Empty Identy) ()
--> StateT Int Identity ()Running type functions is (currently) only done in the type checker, not the Simplifier.
Tantalisingly, though, the casts we need are lying around in the tree, so it ought be be possible to do this. But it would take some careful work, for example