{- |
  Naive Free monads suffer from a quadratic complexity,
  as explained in

  * Janis Voigtlander, /Asymptotic Improvement of Computations over Free Monads, MPC'08/

  The solution is to redefine the Free datatype in CPS,
  similar to what is done in difference lists to solve the problem on quadratic append
  for lists.
-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE FlexibleContexts, FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}

module Control.Monad.Free.Improve (
   C(..), rep, improve
  ) where

import Control.Applicative
import Control.Monad
import Control.Monad.Free
import Control.Monad.Trans.Class

newtype C mu a = C (forall b. (a -> mu b) -> mu b)

rep :: Monad mu => mu a -> C mu a
rep :: forall (mu :: * -> *) a. Monad mu => mu a -> C mu a
rep mu a
m = (forall b. (a -> mu b) -> mu b) -> C mu a
forall (mu :: * -> *) a. (forall b. (a -> mu b) -> mu b) -> C mu a
C (mu a
m mu a -> (a -> mu b) -> mu b
forall a b. mu a -> (a -> mu b) -> mu b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>=)

improve :: Monad mu => C mu a -> mu a
improve :: forall (mu :: * -> *) a. Monad mu => C mu a -> mu a
improve (C forall b. (a -> mu b) -> mu b
p) = (a -> mu a) -> mu a
forall b. (a -> mu b) -> mu b
p a -> mu a
forall a. a -> mu a
forall (m :: * -> *) a. Monad m => a -> m a
return

instance Functor (C mu) where
  fmap :: forall a b. (a -> b) -> C mu a -> C mu b
fmap a -> b
f (C forall b. (a -> mu b) -> mu b
m) = (forall b. (b -> mu b) -> mu b) -> C mu b
forall (mu :: * -> *) a. (forall b. (a -> mu b) -> mu b) -> C mu a
C (\b -> mu b
h -> (a -> mu b) -> mu b
forall b. (a -> mu b) -> mu b
m (b -> mu b
h(b -> mu b) -> (a -> b) -> a -> mu b
forall b c a. (b -> c) -> (a -> b) -> a -> c
.a -> b
f))
--  fmap f (C m) = C (m . (.f))

instance Monad (C mu) where
  return :: forall a. a -> C mu a
return a
a = (forall b. (a -> mu b) -> mu b) -> C mu a
forall (mu :: * -> *) a. (forall b. (a -> mu b) -> mu b) -> C mu a
C (\a -> mu b
h -> a -> mu b
h a
a)
  C forall b. (a -> mu b) -> mu b
p >>= :: forall a b. C mu a -> (a -> C mu b) -> C mu b
>>= a -> C mu b
k = (forall b. (b -> mu b) -> mu b) -> C mu b
forall (mu :: * -> *) a. (forall b. (a -> mu b) -> mu b) -> C mu a
C (\b -> mu b
h -> (a -> mu b) -> mu b
forall b. (a -> mu b) -> mu b
p (\a
a -> case a -> C mu b
k a
a of C forall b. (b -> mu b) -> mu b
q -> (b -> mu b) -> mu b
forall b. (b -> mu b) -> mu b
q b -> mu b
h))

instance Applicative (C mu) where
  pure :: forall a. a -> C mu a
pure = a -> C mu a
forall a. a -> C mu a
forall (m :: * -> *) a. Monad m => a -> m a
return
  <*> :: forall a b. C mu (a -> b) -> C mu a -> C mu b
(<*>) = C mu (a -> b) -> C mu a -> C mu b
forall (m :: * -> *) a b. Monad m => m (a -> b) -> m a -> m b
ap

instance Functor f => MonadFree f (C (Free f)) where
  wrap :: forall a. f (C (Free f) a) -> C (Free f) a
wrap f (C (Free f) a)
t = (forall b. (a -> Free f b) -> Free f b) -> C (Free f) a
forall (mu :: * -> *) a. (forall b. (a -> mu b) -> mu b) -> C mu a
C (\a -> Free f b
h -> f (Free f b) -> Free f b
forall a. f (Free f a) -> Free f a
forall (f :: * -> *) (m :: * -> *) a.
MonadFree f m =>
f (m a) -> m a
wrap ((C (Free f) a -> Free f b) -> f (C (Free f) a) -> f (Free f b)
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(C forall b. (a -> Free f b) -> Free f b
p) -> (a -> Free f b) -> Free f b
forall b. (a -> Free f b) -> Free f b
p a -> Free f b
h) f (C (Free f) a)
t))
  free :: forall a. C (Free f) a -> C (Free f) (Either a (f (C (Free f) a)))
free   = Free f (Either a (f (C (Free f) a)))
-> C (Free f) (Either a (f (C (Free f) a)))
forall (mu :: * -> *) a. Monad mu => mu a -> C mu a
rep (Free f (Either a (f (C (Free f) a)))
 -> C (Free f) (Either a (f (C (Free f) a))))
-> (C (Free f) a -> Free f (Either a (f (C (Free f) a))))
-> C (Free f) a
-> C (Free f) (Either a (f (C (Free f) a)))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Either a (f (Free f a)) -> Either a (f (C (Free f) a)))
-> Free f (Either a (f (Free f a)))
-> Free f (Either a (f (C (Free f) a)))
forall a b. (a -> b) -> Free f a -> Free f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap((Either a (f (Free f a)) -> Either a (f (C (Free f) a)))
 -> Free f (Either a (f (Free f a)))
 -> Free f (Either a (f (C (Free f) a))))
-> ((Free f a -> C (Free f) a)
    -> Either a (f (Free f a)) -> Either a (f (C (Free f) a)))
-> (Free f a -> C (Free f) a)
-> Free f (Either a (f (Free f a)))
-> Free f (Either a (f (C (Free f) a)))
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(f (Free f a) -> f (C (Free f) a))
-> Either a (f (Free f a)) -> Either a (f (C (Free f) a))
forall a b. (a -> b) -> Either a a -> Either a b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap((f (Free f a) -> f (C (Free f) a))
 -> Either a (f (Free f a)) -> Either a (f (C (Free f) a)))
-> ((Free f a -> C (Free f) a) -> f (Free f a) -> f (C (Free f) a))
-> (Free f a -> C (Free f) a)
-> Either a (f (Free f a))
-> Either a (f (C (Free f) a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(Free f a -> C (Free f) a) -> f (Free f a) -> f (C (Free f) a)
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap) Free f a -> C (Free f) a
forall (mu :: * -> *) a. Monad mu => mu a -> C mu a
rep (Free f (Either a (f (Free f a)))
 -> Free f (Either a (f (C (Free f) a))))
-> (C (Free f) a -> Free f (Either a (f (Free f a))))
-> C (Free f) a
-> Free f (Either a (f (C (Free f) a)))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Free f a -> Free f (Either a (f (Free f a)))
forall a. Free f a -> Free f (Either a (f (Free f a)))
forall (f :: * -> *) (m :: * -> *) a.
MonadFree f m =>
m a -> m (Either a (f (m a)))
free (Free f a -> Free f (Either a (f (Free f a))))
-> (C (Free f) a -> Free f a)
-> C (Free f) a
-> Free f (Either a (f (Free f a)))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. C (Free f) a -> Free f a
forall (mu :: * -> *) a. Monad mu => C mu a -> mu a
improve

instance (Monad m, Functor f) => MonadFree f (C (FreeT f m)) where
  wrap :: forall a. f (C (FreeT f m) a) -> C (FreeT f m) a
wrap f (C (FreeT f m) a)
t = (forall b. (a -> FreeT f m b) -> FreeT f m b) -> C (FreeT f m) a
forall (mu :: * -> *) a. (forall b. (a -> mu b) -> mu b) -> C mu a
C (\a -> FreeT f m b
h -> f (FreeT f m b) -> FreeT f m b
forall a. f (FreeT f m a) -> FreeT f m a
forall (f :: * -> *) (m :: * -> *) a.
MonadFree f m =>
f (m a) -> m a
wrap ((C (FreeT f m) a -> FreeT f m b)
-> f (C (FreeT f m) a) -> f (FreeT f m b)
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(C forall b. (a -> FreeT f m b) -> FreeT f m b
p) -> (a -> FreeT f m b) -> FreeT f m b
forall b. (a -> FreeT f m b) -> FreeT f m b
p a -> FreeT f m b
h) f (C (FreeT f m) a)
t))
  free :: forall a.
C (FreeT f m) a -> C (FreeT f m) (Either a (f (C (FreeT f m) a)))
free   = FreeT f m (Either a (f (C (FreeT f m) a)))
-> C (FreeT f m) (Either a (f (C (FreeT f m) a)))
forall (mu :: * -> *) a. Monad mu => mu a -> C mu a
rep (FreeT f m (Either a (f (C (FreeT f m) a)))
 -> C (FreeT f m) (Either a (f (C (FreeT f m) a))))
-> (C (FreeT f m) a -> FreeT f m (Either a (f (C (FreeT f m) a))))
-> C (FreeT f m) a
-> C (FreeT f m) (Either a (f (C (FreeT f m) a)))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Either a (f (FreeT f m a)) -> Either a (f (C (FreeT f m) a)))
-> FreeT f m (Either a (f (FreeT f m a)))
-> FreeT f m (Either a (f (C (FreeT f m) a)))
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM((Either a (f (FreeT f m a)) -> Either a (f (C (FreeT f m) a)))
 -> FreeT f m (Either a (f (FreeT f m a)))
 -> FreeT f m (Either a (f (C (FreeT f m) a))))
-> ((FreeT f m a -> C (FreeT f m) a)
    -> Either a (f (FreeT f m a)) -> Either a (f (C (FreeT f m) a)))
-> (FreeT f m a -> C (FreeT f m) a)
-> FreeT f m (Either a (f (FreeT f m a)))
-> FreeT f m (Either a (f (C (FreeT f m) a)))
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(f (FreeT f m a) -> f (C (FreeT f m) a))
-> Either a (f (FreeT f m a)) -> Either a (f (C (FreeT f m) a))
forall a b. (a -> b) -> Either a a -> Either a b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap((f (FreeT f m a) -> f (C (FreeT f m) a))
 -> Either a (f (FreeT f m a)) -> Either a (f (C (FreeT f m) a)))
-> ((FreeT f m a -> C (FreeT f m) a)
    -> f (FreeT f m a) -> f (C (FreeT f m) a))
-> (FreeT f m a -> C (FreeT f m) a)
-> Either a (f (FreeT f m a))
-> Either a (f (C (FreeT f m) a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(FreeT f m a -> C (FreeT f m) a)
-> f (FreeT f m a) -> f (C (FreeT f m) a)
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap) FreeT f m a -> C (FreeT f m) a
forall (mu :: * -> *) a. Monad mu => mu a -> C mu a
rep (FreeT f m (Either a (f (FreeT f m a)))
 -> FreeT f m (Either a (f (C (FreeT f m) a))))
-> (C (FreeT f m) a -> FreeT f m (Either a (f (FreeT f m a))))
-> C (FreeT f m) a
-> FreeT f m (Either a (f (C (FreeT f m) a)))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FreeT f m a -> FreeT f m (Either a (f (FreeT f m a)))
forall a. FreeT f m a -> FreeT f m (Either a (f (FreeT f m a)))
forall (f :: * -> *) (m :: * -> *) a.
MonadFree f m =>
m a -> m (Either a (f (m a)))
free (FreeT f m a -> FreeT f m (Either a (f (FreeT f m a))))
-> (C (FreeT f m) a -> FreeT f m a)
-> C (FreeT f m) a
-> FreeT f m (Either a (f (FreeT f m a)))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. C (FreeT f m) a -> FreeT f m a
forall (mu :: * -> *) a. Monad mu => C mu a -> mu a
improve

instance MonadPlus mu => MonadPlus (C mu) where
  mzero :: forall a. C mu a
mzero       = mu a -> C mu a
forall (mu :: * -> *) a. Monad mu => mu a -> C mu a
rep mu a
forall a. mu a
forall (m :: * -> *) a. MonadPlus m => m a
mzero
  mplus :: forall a. C mu a -> C mu a -> C mu a
mplus C mu a
p1 C mu a
p2 = mu a -> C mu a
forall (mu :: * -> *) a. Monad mu => mu a -> C mu a
rep (mu a -> mu a -> mu a
forall a. mu a -> mu a -> mu a
forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
mplus (C mu a -> mu a
forall (mu :: * -> *) a. Monad mu => C mu a -> mu a
improve C mu a
p1) (C mu a -> mu a
forall (mu :: * -> *) a. Monad mu => C mu a -> mu a
improve C mu a
p2))

instance MonadPlus mu => Alternative (C mu) where
  empty :: forall a. C mu a
empty = C mu a
forall a. C mu a
forall (m :: * -> *) a. MonadPlus m => m a
mzero
  <|> :: forall a. C mu a -> C mu a -> C mu a
(<|>) = C mu a -> C mu a -> C mu a
forall a. C mu a -> C mu a -> C mu a
forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
mplus

instance MonadTrans C where lift :: forall (mu :: * -> *) a. Monad mu => mu a -> C mu a
lift m a
m = (forall b. (a -> m b) -> m b) -> C m a
forall (mu :: * -> *) a. (forall b. (a -> mu b) -> mu b) -> C mu a
C (m a
m m a -> (a -> m b) -> m b
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>=)