{-# LANGUAGE GADTs              #-}
{-# LANGUAGE RankNTypes         #-}
{-# LANGUAGE Safe #-}

--------------------------------------------------------------------------------
-- |
-- A faster free applicative.
-- Based on <https://www.eyrie.org/~zednenem/2013/05/27/freeapp Dave Menendez's work>.
--------------------------------------------------------------------------------
module Control.Applicative.Free.Fast
  (
  -- * The Sequence of Effects
    ASeq(..)
  , reduceASeq
  , hoistASeq
  , traverseASeq
  , rebaseASeq
  -- * The Faster Free Applicative
  , Ap(..)
  , liftAp
  , retractAp
  , runAp
  , runAp_
  , hoistAp
  ) where

import           Control.Applicative
import           Data.Functor.Apply

-- | The free applicative is composed of a sequence of effects,
-- and a pure function to apply that sequence to.
-- The fast free applicative separates these from each other,
-- so that the sequence may be built up independently,
-- and so that 'fmap' can run in constant time by having immediate access to the pure function.
data ASeq f a where
  ANil :: ASeq f ()
  ACons :: f a -> ASeq f u -> ASeq f (a,u)

-- | Interprets the sequence of effects using the semantics for
--   `pure` and `<*>` given by the Applicative instance for 'f'.
reduceASeq :: Applicative f => ASeq f u -> f u
reduceASeq :: forall (f :: * -> *) u. Applicative f => ASeq f u -> f u
reduceASeq ASeq f u
ANil         = forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
reduceASeq (ACons f a
x ASeq f u
xs) = (,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f a
x forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) u. Applicative f => ASeq f u -> f u
reduceASeq ASeq f u
xs

-- | Given a natural transformation from @f@ to @g@ this gives a natural transformation from @ASeq f@ to @ASeq g@.
hoistASeq :: (forall x. f x -> g x) -> ASeq f a -> ASeq g a
hoistASeq :: forall (f :: * -> *) (g :: * -> *) a.
(forall x. f x -> g x) -> ASeq f a -> ASeq g a
hoistASeq forall x. f x -> g x
_ ASeq f a
ANil = forall (f :: * -> *). ASeq f ()
ANil
hoistASeq forall x. f x -> g x
u (ACons f a
x ASeq f u
xs) = forall (f :: * -> *) a u. f a -> ASeq f u -> ASeq f (a, u)
ACons (forall x. f x -> g x
u f a
x) (forall x. f x -> g x
u forall (f :: * -> *) (g :: * -> *) a.
(forall x. f x -> g x) -> ASeq f a -> ASeq g a
`hoistASeq` ASeq f u
xs)

-- | Traverse a sequence with resepect to its interpretation type 'f'.
traverseASeq :: Applicative h => (forall x. f x -> h (g x)) -> ASeq f a -> h (ASeq g a)
traverseASeq :: forall (h :: * -> *) (f :: * -> *) (g :: * -> *) a.
Applicative h =>
(forall x. f x -> h (g x)) -> ASeq f a -> h (ASeq g a)
traverseASeq forall x. f x -> h (g x)
_ ASeq f a
ANil      = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall (f :: * -> *). ASeq f ()
ANil
traverseASeq forall x. f x -> h (g x)
f (ACons f a
x ASeq f u
xs) = forall (f :: * -> *) a u. f a -> ASeq f u -> ASeq f (a, u)
ACons forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall x. f x -> h (g x)
f f a
x forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (h :: * -> *) (f :: * -> *) (g :: * -> *) a.
Applicative h =>
(forall x. f x -> h (g x)) -> ASeq f a -> h (ASeq g a)
traverseASeq forall x. f x -> h (g x)
f ASeq f u
xs

-- | It may not be obvious, but this essentially acts like ++,
-- traversing the first sequence and creating a new one by appending the second sequence.
-- The difference is that this also has to modify the return functions and that the return type depends on the input types.
--
-- See the source of 'hoistAp' as an example usage.
rebaseASeq :: ASeq f u -> (forall x. (x -> y) -> ASeq f x -> z) ->
  (v -> u -> y) -> ASeq f v -> z
rebaseASeq :: forall (f :: * -> *) u y z v.
ASeq f u
-> (forall x. (x -> y) -> ASeq f x -> z)
-> (v -> u -> y)
-> ASeq f v
-> z
rebaseASeq ASeq f u
ANil         forall x. (x -> y) -> ASeq f x -> z
k v -> u -> y
f = forall x. (x -> y) -> ASeq f x -> z
k (\v
v -> v -> u -> y
f v
v ())
rebaseASeq (ACons f a
x ASeq f u
xs) forall x. (x -> y) -> ASeq f x -> z
k v -> u -> y
f =
  forall (f :: * -> *) u y z v.
ASeq f u
-> (forall x. (x -> y) -> ASeq f x -> z)
-> (v -> u -> y)
-> ASeq f v
-> z
rebaseASeq ASeq f u
xs (\x -> a -> y
g ASeq f x
s -> forall x. (x -> y) -> ASeq f x -> z
k (\(a
a,x
u) -> x -> a -> y
g x
u a
a) (forall (f :: * -> *) a u. f a -> ASeq f u -> ASeq f (a, u)
ACons f a
x ASeq f x
s))
    (\v
v u
u a
a -> v -> u -> y
f v
v (a
a,u
u))


-- | The faster free 'Applicative'.
newtype Ap f a = Ap
  { forall (f :: * -> *) a.
Ap f a
-> forall u y z.
   (forall x. (x -> y) -> ASeq f x -> z)
   -> (u -> a -> y) -> ASeq f u -> z
unAp :: forall u y z.
    (forall x. (x -> y) -> ASeq f x -> z) ->
    (u -> a -> y) -> ASeq f u -> z }

-- | Given a natural transformation from @f@ to @g@, this gives a canonical monoidal natural transformation from @'Ap' f@ to @g@.
--
-- prop> runAp t == retractApp . hoistApp t
runAp :: Applicative g => (forall x. f x -> g x) -> Ap f a -> g a
runAp :: forall (g :: * -> *) (f :: * -> *) a.
Applicative g =>
(forall x. f x -> g x) -> Ap f a -> g a
runAp forall x. f x -> g x
u = forall (f :: * -> *) a. Applicative f => Ap f a -> f a
retractAp forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) (g :: * -> *) a.
(forall x. f x -> g x) -> Ap f a -> Ap g a
hoistAp forall x. f x -> g x
u

-- | Perform a monoidal analysis over free applicative value.
--
-- Example:
--
-- @
-- count :: Ap f a -> Int
-- count = getSum . runAp_ (\\_ -> Sum 1)
-- @
runAp_ :: Monoid m => (forall a. f a -> m) -> Ap f b -> m
runAp_ :: forall m (f :: * -> *) b.
Monoid m =>
(forall a. f a -> m) -> Ap f b -> m
runAp_ forall a. f a -> m
f = forall {k} a (b :: k). Const a b -> a
getConst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (g :: * -> *) (f :: * -> *) a.
Applicative g =>
(forall x. f x -> g x) -> Ap f a -> g a
runAp (forall {k} a (b :: k). a -> Const a b
Const forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. f a -> m
f)

instance Functor (Ap f) where
  fmap :: forall a b. (a -> b) -> Ap f a -> Ap f b
fmap a -> b
g Ap f a
x = forall (f :: * -> *) a.
(forall u y z.
 (forall x. (x -> y) -> ASeq f x -> z)
 -> (u -> a -> y) -> ASeq f u -> z)
-> Ap f a
Ap (\forall x. (x -> y) -> ASeq f x -> z
k u -> b -> y
f -> forall (f :: * -> *) a.
Ap f a
-> forall u y z.
   (forall x. (x -> y) -> ASeq f x -> z)
   -> (u -> a -> y) -> ASeq f u -> z
unAp Ap f a
x forall x. (x -> y) -> ASeq f x -> z
k (\u
s -> u -> b -> y
f u
s forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b
g))

instance Apply (Ap f) where
  <.> :: forall a b. Ap f (a -> b) -> Ap f a -> Ap f b
(<.>) = forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
(<*>)

instance Applicative (Ap f) where
  pure :: forall a. a -> Ap f a
pure a
a = forall (f :: * -> *) a.
(forall u y z.
 (forall x. (x -> y) -> ASeq f x -> z)
 -> (u -> a -> y) -> ASeq f u -> z)
-> Ap f a
Ap (\forall x. (x -> y) -> ASeq f x -> z
k u -> a -> y
f -> forall x. (x -> y) -> ASeq f x -> z
k (u -> a -> y
`f` a
a))
  Ap f (a -> b)
x <*> :: forall a b. Ap f (a -> b) -> Ap f a -> Ap f b
<*> Ap f a
y = forall (f :: * -> *) a.
(forall u y z.
 (forall x. (x -> y) -> ASeq f x -> z)
 -> (u -> a -> y) -> ASeq f u -> z)
-> Ap f a
Ap (\forall x. (x -> y) -> ASeq f x -> z
k u -> b -> y
f -> forall (f :: * -> *) a.
Ap f a
-> forall u y z.
   (forall x. (x -> y) -> ASeq f x -> z)
   -> (u -> a -> y) -> ASeq f u -> z
unAp Ap f a
y (forall (f :: * -> *) a.
Ap f a
-> forall u y z.
   (forall x. (x -> y) -> ASeq f x -> z)
   -> (u -> a -> y) -> ASeq f u -> z
unAp Ap f (a -> b)
x forall x. (x -> y) -> ASeq f x -> z
k) (\u
s a
a a -> b
g -> u -> b -> y
f u
s (a -> b
g a
a)))

-- | A version of 'lift' that can be used with just a 'Functor' for @f@.
liftAp :: f a -> Ap f a
liftAp :: forall (f :: * -> *) a. f a -> Ap f a
liftAp f a
a = forall (f :: * -> *) a.
(forall u y z.
 (forall x. (x -> y) -> ASeq f x -> z)
 -> (u -> a -> y) -> ASeq f u -> z)
-> Ap f a
Ap (\forall x. (x -> y) -> ASeq f x -> z
k u -> a -> y
f ASeq f u
s -> forall x. (x -> y) -> ASeq f x -> z
k (\(a
a',u
s') -> u -> a -> y
f u
s' a
a') (forall (f :: * -> *) a u. f a -> ASeq f u -> ASeq f (a, u)
ACons f a
a ASeq f u
s))
{-# INLINE liftAp #-}

-- | Given a natural transformation from @f@ to @g@ this gives a monoidal natural transformation from @Ap f@ to @Ap g@.
hoistAp :: (forall x. f x -> g x) -> Ap f a -> Ap g a
hoistAp :: forall (f :: * -> *) (g :: * -> *) a.
(forall x. f x -> g x) -> Ap f a -> Ap g a
hoistAp forall x. f x -> g x
g Ap f a
x = forall (f :: * -> *) a.
(forall u y z.
 (forall x. (x -> y) -> ASeq f x -> z)
 -> (u -> a -> y) -> ASeq f u -> z)
-> Ap f a
Ap (\forall x. (x -> y) -> ASeq g x -> z
k u -> a -> y
f ASeq g u
s ->
  forall (f :: * -> *) a.
Ap f a
-> forall u y z.
   (forall x. (x -> y) -> ASeq f x -> z)
   -> (u -> a -> y) -> ASeq f u -> z
unAp Ap f a
x
    (\x -> a
f' ASeq f x
s' ->
      forall (f :: * -> *) u y z v.
ASeq f u
-> (forall x. (x -> y) -> ASeq f x -> z)
-> (v -> u -> y)
-> ASeq f v
-> z
rebaseASeq (forall (f :: * -> *) (g :: * -> *) a.
(forall x. f x -> g x) -> ASeq f a -> ASeq g a
hoistASeq forall x. f x -> g x
g ASeq f x
s') forall x. (x -> y) -> ASeq g x -> z
k
        (\u
v x
u -> u -> a -> y
f u
v (x -> a
f' x
u)) ASeq g u
s)
    (forall a b. a -> b -> a
const forall a. a -> a
id)
    forall (f :: * -> *). ASeq f ()
ANil)

-- | Interprets the free applicative functor over f using the semantics for
--   `pure` and `<*>` given by the Applicative instance for f.
--
--   prop> retractApp == runAp id
retractAp :: Applicative f => Ap f a -> f a
retractAp :: forall (f :: * -> *) a. Applicative f => Ap f a -> f a
retractAp Ap f a
x = forall (f :: * -> *) a.
Ap f a
-> forall u y z.
   (forall x. (x -> y) -> ASeq f x -> z)
   -> (u -> a -> y) -> ASeq f u -> z
unAp Ap f a
x (\x -> a
f ASeq f x
s -> x -> a
f forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (f :: * -> *) u. Applicative f => ASeq f u -> f u
reduceASeq ASeq f x
s) (\() -> forall a. a -> a
id) forall (f :: * -> *). ASeq f ()
ANil