{-# LANGUAGE TupleSections #-}

-- | Extend a monad with a modifiable environment
module Mini.Transformers.StateT (
  -- * Type
  StateT (
    StateT
  ),

  -- * Runner
  runStateT,

  -- * Operations
  get,
  modify,
  put,
) where

import Control.Applicative (
  Alternative (
    empty,
    (<|>)
  ),
 )
import Control.Monad (
  ap,
  liftM,
  (>=>),
 )
import Mini.Transformers.Class (
  MonadTrans (
    lift
  ),
 )

{-
 - Type
 -}

-- | A transformer with state /s/, inner monad /m/, return /a/
newtype StateT s m a = StateT
  { forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT :: s -> m (a, s)
  -- ^ Unwrap a 'StateT' computation with an initial state
  }

instance (Monad m) => Functor (StateT s m) where
  fmap :: forall a b. (a -> b) -> StateT s m a -> StateT s m b
fmap = (a -> b) -> StateT s m a -> StateT s m b
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM

instance (Monad m) => Applicative (StateT s m) where
  pure :: forall a. a -> StateT s m a
pure a
a = (s -> m (a, s)) -> StateT s m a
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StateT ((s -> m (a, s)) -> StateT s m a)
-> (s -> m (a, s)) -> StateT s m a
forall a b. (a -> b) -> a -> b
$ (a, s) -> m (a, s)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((a, s) -> m (a, s)) -> (s -> (a, s)) -> s -> m (a, s)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a
a,)
  <*> :: forall a b. StateT s m (a -> b) -> StateT s m a -> StateT s m b
(<*>) = StateT s m (a -> b) -> StateT s m a -> StateT s m b
forall (m :: * -> *) a b. Monad m => m (a -> b) -> m a -> m b
ap

instance (Monad m, Alternative m) => Alternative (StateT s m) where
  empty :: forall a. StateT s m a
empty = (s -> m (a, s)) -> StateT s m a
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StateT ((s -> m (a, s)) -> StateT s m a)
-> (s -> m (a, s)) -> StateT s m a
forall a b. (a -> b) -> a -> b
$ m (a, s) -> s -> m (a, s)
forall a b. a -> b -> a
const m (a, s)
forall a. m a
forall (f :: * -> *) a. Alternative f => f a
empty
  StateT s m a
m <|> :: forall a. StateT s m a -> StateT s m a -> StateT s m a
<|> StateT s m a
n = (s -> m (a, s)) -> StateT s m a
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StateT ((s -> m (a, s)) -> StateT s m a)
-> (s -> m (a, s)) -> StateT s m a
forall a b. (a -> b) -> a -> b
$ \s
s -> StateT s m a -> s -> m (a, s)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT StateT s m a
m s
s m (a, s) -> m (a, s) -> m (a, s)
forall a. m a -> m a -> m a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> StateT s m a -> s -> m (a, s)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT StateT s m a
n s
s

instance (Monad m) => Monad (StateT s m) where
  StateT s m a
m >>= :: forall a b. StateT s m a -> (a -> StateT s m b) -> StateT s m b
>>= a -> StateT s m b
k = (s -> m (b, s)) -> StateT s m b
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StateT ((s -> m (b, s)) -> StateT s m b)
-> (s -> m (b, s)) -> StateT s m b
forall a b. (a -> b) -> a -> b
$ StateT s m a -> s -> m (a, s)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT StateT s m a
m (s -> m (a, s)) -> ((a, s) -> m (b, s)) -> s -> m (b, s)
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> (\(a
a, s
s) -> StateT s m b -> s -> m (b, s)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT (a -> StateT s m b
k a
a) s
s)

instance MonadTrans (StateT s) where
  lift :: forall (m :: * -> *) a. Monad m => m a -> StateT s m a
lift m a
m = (s -> m (a, s)) -> StateT s m a
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StateT ((s -> m (a, s)) -> StateT s m a)
-> (s -> m (a, s)) -> StateT s m a
forall a b. (a -> b) -> a -> b
$ \s
s -> (,s
s) (a -> (a, s)) -> m a -> m (a, s)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m a
m

{-
 - Operations
 -}

-- | Fetch the current state
get :: (Monad m) => StateT s m s
get :: forall (m :: * -> *) s. Monad m => StateT s m s
get = (s -> m (s, s)) -> StateT s m s
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StateT ((s -> m (s, s)) -> StateT s m s)
-> (s -> m (s, s)) -> StateT s m s
forall a b. (a -> b) -> a -> b
$ \s
s -> (s, s) -> m (s, s)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (s
s, s
s)

-- | Update the current state with an operation
modify :: (Monad m) => (s -> s) -> StateT s m ()
modify :: forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
modify s -> s
f = (s -> m ((), s)) -> StateT s m ()
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StateT ((s -> m ((), s)) -> StateT s m ())
-> (s -> m ((), s)) -> StateT s m ()
forall a b. (a -> b) -> a -> b
$ ((), s) -> m ((), s)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (((), s) -> m ((), s)) -> (s -> ((), s)) -> s -> m ((), s)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((),) (s -> ((), s)) -> (s -> s) -> s -> ((), s)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. s -> s
f

-- | Overwrite the current state with a value
put :: (Monad m) => s -> StateT s m ()
put :: forall (m :: * -> *) s. Monad m => s -> StateT s m ()
put = (s -> m ((), s)) -> StateT s m ()
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StateT ((s -> m ((), s)) -> StateT s m ())
-> (s -> s -> m ((), s)) -> s -> StateT s m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. m ((), s) -> s -> m ((), s)
forall a b. a -> b -> a
const (m ((), s) -> s -> m ((), s))
-> (s -> m ((), s)) -> s -> s -> m ((), s)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((), s) -> m ((), s)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (((), s) -> m ((), s)) -> (s -> ((), s)) -> s -> m ((), s)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((),)