{- | Handle a global 'AccumT' layer in an 'Automaton'.

A global accumulation state can be hidden by an automaton by making it an internal state.
-}
module Data.Automaton.Trans.Accum (
  module Control.Monad.Trans.Accum,
  accumS,
  runAccumS,
  runAccumS_,
  runAccumS__,
)
where

-- base
import Control.Arrow (arr, returnA, (>>>))

-- transformers
import Control.Monad.Trans.Accum

-- automaton
import Data.Automaton (Automaton, feedback, withAutomaton)
import Data.Stream.Result (Result (..))

{- | Convert from explicit states to the 'AccumT' monad transformer.

The original automaton is interpreted to take the current accumulated state as input and return the log to be appended as output.

This is the opposite of 'runAccumS'.
-}
accumS :: (Functor m, Monad m) => Automaton m (w, a) (w, b) -> Automaton (AccumT w m) a b
accumS :: forall (m :: Type -> Type) w a b.
(Functor m, Monad m) =>
Automaton m (w, a) (w, b) -> Automaton (AccumT w m) a b
accumS = (forall s.
 ((w, a) -> m (Result s (w, b))) -> a -> AccumT w m (Result s b))
-> Automaton m (w, a) (w, b) -> Automaton (AccumT w m) a b
forall (m1 :: Type -> Type) (m2 :: Type -> Type) a1 b1 a2 b2.
(Functor m1, Functor m2) =>
(forall s. (a1 -> m1 (Result s b1)) -> a2 -> m2 (Result s b2))
-> Automaton m1 a1 b1 -> Automaton m2 a2 b2
withAutomaton ((forall s.
  ((w, a) -> m (Result s (w, b))) -> a -> AccumT w m (Result s b))
 -> Automaton m (w, a) (w, b) -> Automaton (AccumT w m) a b)
-> (forall s.
    ((w, a) -> m (Result s (w, b))) -> a -> AccumT w m (Result s b))
-> Automaton m (w, a) (w, b)
-> Automaton (AccumT w m) a b
forall a b. (a -> b) -> a -> b
$ \(w, a) -> m (Result s (w, b))
f a
a -> (w -> m (Result s b, w)) -> AccumT w m (Result s b)
forall w (m :: Type -> Type) a. (w -> m (a, w)) -> AccumT w m a
AccumT ((w -> m (Result s b, w)) -> AccumT w m (Result s b))
-> (w -> m (Result s b, w)) -> AccumT w m (Result s b)
forall a b. (a -> b) -> a -> b
$ \w
w ->
  (\(Result s
s (w
w', b
b)) -> (s -> b -> Result s b
forall s a. s -> a -> Result s a
Result s
s b
b, w
w'))
    (Result s (w, b) -> (Result s b, w))
-> m (Result s (w, b)) -> m (Result s b, w)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> (w, a) -> m (Result s (w, b))
f (w
w, a
a)

{- | Make the accumulation transition in 'AccumT' explicit as 'Automaton' inputs and outputs.

This is the opposite of 'accumS'.
-}
runAccumS :: (Functor m, Monad m) => Automaton (AccumT w m) a b -> Automaton m (w, a) (w, b)
runAccumS :: forall (m :: Type -> Type) w a b.
(Functor m, Monad m) =>
Automaton (AccumT w m) a b -> Automaton m (w, a) (w, b)
runAccumS = (forall s.
 (a -> AccumT w m (Result s b)) -> (w, a) -> m (Result s (w, b)))
-> Automaton (AccumT w m) a b -> Automaton m (w, a) (w, b)
forall (m1 :: Type -> Type) (m2 :: Type -> Type) a1 b1 a2 b2.
(Functor m1, Functor m2) =>
(forall s. (a1 -> m1 (Result s b1)) -> a2 -> m2 (Result s b2))
-> Automaton m1 a1 b1 -> Automaton m2 a2 b2
withAutomaton ((forall s.
  (a -> AccumT w m (Result s b)) -> (w, a) -> m (Result s (w, b)))
 -> Automaton (AccumT w m) a b -> Automaton m (w, a) (w, b))
-> (forall s.
    (a -> AccumT w m (Result s b)) -> (w, a) -> m (Result s (w, b)))
-> Automaton (AccumT w m) a b
-> Automaton m (w, a) (w, b)
forall a b. (a -> b) -> a -> b
$ \a -> AccumT w m (Result s b)
f (w
w, a
a) ->
  (\(Result s
s b
b, w
w') -> s -> (w, b) -> Result s (w, b)
forall s a. s -> a -> Result s a
Result s
s (w
w', b
b))
    ((Result s b, w) -> Result s (w, b))
-> m (Result s b, w) -> m (Result s (w, b))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AccumT w m (Result s b) -> w -> m (Result s b, w)
forall w (m :: Type -> Type) a. AccumT w m a -> w -> m (a, w)
runAccumT (a -> AccumT w m (Result s b)
f a
a) w
w

{- | Convert global accumulation state to internal state of an 'Automaton'.

The current state is output on every step.
-}
runAccumS_ ::
  (Functor m, Monoid w, Monad m) =>
  -- | An automaton with a global accumulation state effect
  Automaton (AccumT w m) a b ->
  Automaton m a (w, b)
runAccumS_ :: forall (m :: Type -> Type) w a b.
(Functor m, Monoid w, Monad m) =>
Automaton (AccumT w m) a b -> Automaton m a (w, b)
runAccumS_ Automaton (AccumT w m) a b
automaton = w -> Automaton m (a, w) ((w, b), w) -> Automaton m a (w, b)
forall (m :: Type -> Type) c a b.
Functor m =>
c -> Automaton m (a, c) (b, c) -> Automaton m a b
feedback w
forall a. Monoid a => a
mempty (Automaton m (a, w) ((w, b), w) -> Automaton m a (w, b))
-> Automaton m (a, w) ((w, b), w) -> Automaton m a (w, b)
forall a b. (a -> b) -> a -> b
$ proc (a
a, w
wState) -> do
  (w
wAdd, b
b) <- Automaton (AccumT w m) a b -> Automaton m (w, a) (w, b)
forall (m :: Type -> Type) w a b.
(Functor m, Monad m) =>
Automaton (AccumT w m) a b -> Automaton m (w, a) (w, b)
runAccumS Automaton (AccumT w m) a b
automaton -< (w
wState, a
a)
  let wState' :: w
wState' = w
wState w -> w -> w
forall a. Semigroup a => a -> a -> a
<> w
wAdd
  Automaton m ((w, b), w) ((w, b), w)
forall (a :: Type -> Type -> Type) b. Arrow a => a b b
returnA -< ((w
wState', b
b), w
wState')

-- | Like 'runAccumS_', but don't output the current accum.
runAccumS__ :: (Functor m, Monoid w, Monad m) => Automaton (AccumT w m) a b -> Automaton m a b
runAccumS__ :: forall (m :: Type -> Type) w a b.
(Functor m, Monoid w, Monad m) =>
Automaton (AccumT w m) a b -> Automaton m a b
runAccumS__ Automaton (AccumT w m) a b
automaton = Automaton (AccumT w m) a b -> Automaton m a (w, b)
forall (m :: Type -> Type) w a b.
(Functor m, Monoid w, Monad m) =>
Automaton (AccumT w m) a b -> Automaton m a (w, b)
runAccumS_ Automaton (AccumT w m) a b
automaton Automaton m a (w, b) -> Automaton m (w, b) b -> Automaton m a b
forall {k} (cat :: k -> k -> Type) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> ((w, b) -> b) -> Automaton m (w, b) b
forall b c. (b -> c) -> Automaton m b c
forall (a :: Type -> Type -> Type) b c.
Arrow a =>
(b -> c) -> a b c
arr (w, b) -> b
forall a b. (a, b) -> b
snd