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

A global state can be hidden by an automaton by making it an internal state.

This module is based on the /strict/ state monad 'Control.Monad.Trans.State.Strict',
so when combining it with other modules such as @mtl@'s,
the strict version has to be included, i.e. 'Control.Monad.State.Strict'
instead of 'Control.Monad.State' or 'Control.Monad.State.Lazy'.
-}
module Data.Automaton.Trans.State (
  module Control.Monad.Trans.State.Strict,
  stateS,
  runStateS,
  runStateS_,
  runStateS__,
)
where

-- base
import Control.Arrow (arr, (>>>))
import Data.Tuple (swap)

-- transformers
import Control.Monad.Trans.State.Strict

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

-- * 'State' 'Automaton' running and wrapping

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

The original automaton is interpreted to take a state as input and return the updated state as output.

This is the opposite of 'runStateS'.
-}
stateS :: (Functor m, Monad m) => Automaton m (s, a) (s, b) -> Automaton (StateT s m) a b
stateS :: forall (m :: Type -> Type) s a b.
(Functor m, Monad m) =>
Automaton m (s, a) (s, b) -> Automaton (StateT s m) a b
stateS = (forall s.
 ((s, a) -> m (Result s (s, b))) -> a -> StateT s m (Result s b))
-> Automaton m (s, a) (s, b) -> Automaton (StateT s 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.
  ((s, a) -> m (Result s (s, b))) -> a -> StateT s m (Result s b))
 -> Automaton m (s, a) (s, b) -> Automaton (StateT s m) a b)
-> (forall s.
    ((s, a) -> m (Result s (s, b))) -> a -> StateT s m (Result s b))
-> Automaton m (s, a) (s, b)
-> Automaton (StateT s m) a b
forall a b. (a -> b) -> a -> b
$ \(s, a) -> m (Result s (s, b))
f a
a -> (s -> m (Result s b, s)) -> StateT s m (Result s b)
forall s (m :: Type -> Type) a. (s -> m (a, s)) -> StateT s m a
StateT ((s -> m (Result s b, s)) -> StateT s m (Result s b))
-> (s -> m (Result s b, s)) -> StateT s m (Result s b)
forall a b. (a -> b) -> a -> b
$ \s
s ->
  (\(Result s
s' (s
s, b
b)) -> (s -> b -> Result s b
forall s a. s -> a -> Result s a
Result s
s' b
b, s
s))
    (Result s (s, b) -> (Result s b, s))
-> m (Result s (s, b)) -> m (Result s b, s)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> (s, a) -> m (Result s (s, b))
f (s
s, a
a)

{- | Make the state transition in 'StateT' explicit as 'Automaton' inputs and outputs.

This is the opposite of 'stateS'.
-}
runStateS :: (Functor m, Monad m) => Automaton (StateT s m) a b -> Automaton m (s, a) (s, b)
runStateS :: forall (m :: Type -> Type) s a b.
(Functor m, Monad m) =>
Automaton (StateT s m) a b -> Automaton m (s, a) (s, b)
runStateS = (forall s.
 (a -> StateT s m (Result s b)) -> (s, a) -> m (Result s (s, b)))
-> Automaton (StateT s m) a b -> Automaton m (s, a) (s, 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 -> StateT s m (Result s b)) -> (s, a) -> m (Result s (s, b)))
 -> Automaton (StateT s m) a b -> Automaton m (s, a) (s, b))
-> (forall s.
    (a -> StateT s m (Result s b)) -> (s, a) -> m (Result s (s, b)))
-> Automaton (StateT s m) a b
-> Automaton m (s, a) (s, b)
forall a b. (a -> b) -> a -> b
$ \a -> StateT s m (Result s b)
f (s
s, a
a) ->
  (\(Result s
s' b
b, s
s) -> s -> (s, b) -> Result s (s, b)
forall s a. s -> a -> Result s a
Result s
s' (s
s, b
b))
    ((Result s b, s) -> Result s (s, b))
-> m (Result s b, s) -> m (Result s (s, b))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> StateT s m (Result s b) -> s -> m (Result s b, s)
forall s (m :: Type -> Type) a. StateT s m a -> s -> m (a, s)
runStateT (a -> StateT s m (Result s b)
f a
a) s
s

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

The current state is output on every step.
-}
runStateS_ ::
  (Functor m, Monad m) =>
  -- | An automaton with a global state effect
  Automaton (StateT s m) a b ->
  -- | The initial global state
  s ->
  Automaton m a (s, b)
runStateS_ :: forall (m :: Type -> Type) s a b.
(Functor m, Monad m) =>
Automaton (StateT s m) a b -> s -> Automaton m a (s, b)
runStateS_ Automaton (StateT s m) a b
automaton s
s =
  s -> Automaton m (a, s) ((s, b), s) -> Automaton m a (s, b)
forall (m :: Type -> Type) c a b.
Functor m =>
c -> Automaton m (a, c) (b, c) -> Automaton m a b
feedback s
s (Automaton m (a, s) ((s, b), s) -> Automaton m a (s, b))
-> Automaton m (a, s) ((s, b), s) -> Automaton m a (s, b)
forall a b. (a -> b) -> a -> b
$
    ((a, s) -> (s, a)) -> Automaton m (a, s) (s, a)
forall b c. (b -> c) -> Automaton m b c
forall (a :: Type -> Type -> Type) b c.
Arrow a =>
(b -> c) -> a b c
arr (a, s) -> (s, a)
forall a b. (a, b) -> (b, a)
swap Automaton m (a, s) (s, a)
-> Automaton m (s, a) ((s, b), s) -> Automaton m (a, s) ((s, b), s)
forall {k} (cat :: k -> k -> Type) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> Automaton (StateT s m) a b -> Automaton m (s, a) (s, b)
forall (m :: Type -> Type) s a b.
(Functor m, Monad m) =>
Automaton (StateT s m) a b -> Automaton m (s, a) (s, b)
runStateS Automaton (StateT s m) a b
automaton Automaton m (s, a) (s, b)
-> Automaton m (s, b) ((s, b), s) -> Automaton m (s, a) ((s, b), s)
forall {k} (cat :: k -> k -> Type) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> ((s, b) -> ((s, b), s)) -> Automaton m (s, b) ((s, b), s)
forall b c. (b -> c) -> Automaton m b c
forall (a :: Type -> Type -> Type) b c.
Arrow a =>
(b -> c) -> a b c
arr (\(s
s', b
b) -> ((s
s', b
b), s
s'))

-- | Like 'runStateS_', but don't output the current state.
runStateS__ :: (Functor m, Monad m) => Automaton (StateT s m) a b -> s -> Automaton m a b
runStateS__ :: forall (m :: Type -> Type) s a b.
(Functor m, Monad m) =>
Automaton (StateT s m) a b -> s -> Automaton m a b
runStateS__ Automaton (StateT s m) a b
automaton s
s = Automaton (StateT s m) a b -> s -> Automaton m a (s, b)
forall (m :: Type -> Type) s a b.
(Functor m, Monad m) =>
Automaton (StateT s m) a b -> s -> Automaton m a (s, b)
runStateS_ Automaton (StateT s m) a b
automaton s
s Automaton m a (s, b) -> Automaton m (s, 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
>>> ((s, b) -> b) -> Automaton m (s, 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 (s, b) -> b
forall a b. (a, b) -> b
snd