content: {-# OPTIONS -Wno-orphans #-}
  {-# LANGUAGE DerivingVia #-}
  {-# LANGUAGE LinearTypes #-}
  {-# LANGUAGE QuantifiedConstraints #-}
  {-# LANGUAGE RebindableSyntax #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TupleSections #-}
  {-# LANGUAGE NoImplicitPrelude #-}
  {-# OPTIONS_HADDOCK hide #-}
  
  module Control.Functor.Linear.Internal.State
    ( StateT (..),
      State,
      state,
      get,
      put,
      gets,
      modify,
      replace,
      runStateT,
      runState,
      mapStateT,
      mapState,
      evalStateT,
      evalState,
      execStateT,
      execState,
      withStateT,
      withState,
    )
  where
  
  import Control.Functor.Linear.Internal.Class
  import Control.Functor.Linear.Internal.Instances (Data (..))
  import Control.Functor.Linear.Internal.MonadTrans
  import qualified Control.Monad as NonLinear ()
  import qualified Control.Monad.Trans.State.Strict as NonLinear
  import Data.Functor.Identity
  import qualified Data.Functor.Linear.Internal.Applicative as Data
  import qualified Data.Functor.Linear.Internal.Functor as Data
  import qualified Data.Tuple.Linear as Linear
  import Data.Unrestricted.Linear.Internal.Consumable
  import Data.Unrestricted.Linear.Internal.Dupable
  import Prelude.Linear.Internal
  
  -- # StateT
  -------------------------------------------------------------------------------
  
  -- | A (strict) linear state monad transformer.
  newtype StateT s m a = StateT (s %1 -> m (a, s))
    deriving (Data.Applicative) via Data (StateT s m)
  
  -- We derive Data.Applicative and not Data.Functor since Data.Functor can use
  -- weaker constraints on m than Control.Functor, while
  -- Data.Applicative needs a Monad instance just like Control.Applicative.
  
  type State s = StateT s Identity
  
  get :: (Applicative m, Dupable s) => StateT s m s
  get = state dup
  
  put :: (Applicative m, Consumable s) => s %1 -> StateT s m ()
  put = Data.void . replace
  
  gets :: (Applicative m, Dupable s) => (s %1 -> a) %1 -> StateT s m a
  gets f = state ((\(s1, s2) -> (f s1, s2)) . dup)
  
  runStateT :: StateT s m a %1 -> s %1 -> m (a, s)
  runStateT (StateT f) = f
  
  state :: (Applicative m) => (s %1 -> (a, s)) %1 -> StateT s m a
  state f = StateT (pure . f)
  
  runState :: State s a %1 -> s %1 -> (a, s)
  runState f = runIdentity' . runStateT f
  
  mapStateT :: (m (a, s) %1 -> n (b, s)) %1 -> StateT s m a %1 -> StateT s n b
  mapStateT r (StateT f) = StateT (r . f)
  
  withStateT :: (s %1 -> s) %1 -> StateT s m a %1 -> StateT s m a
  withStateT r (StateT f) = StateT (f . r)
  
  execStateT :: (Functor m) => StateT s m () %1 -> s %1 -> m s
  execStateT f = fmap (\((), s) -> s) . (runStateT f)
  
  -- | Use with care!
  --   This consumes the final state, so might be costly at runtime.
  evalStateT :: (Functor m, Consumable s) => StateT s m a %1 -> s %1 -> m a
  evalStateT f = fmap Linear.fst . runStateT f
  
  mapState :: ((a, s) %1 -> (b, s)) %1 -> State s a %1 -> State s b
  mapState f = mapStateT (Identity . f . runIdentity')
  
  withState :: (s %1 -> s) %1 -> State s a %1 -> State s a
  withState = withStateT
  
  execState :: State s () %1 -> s %1 -> s
  execState f = runIdentity' . execStateT f
  
  -- | Use with care!
  --   This consumes the final state, so might be costly at runtime.
  evalState :: (Consumable s) => State s a %1 -> s %1 -> a
  evalState f = runIdentity' . evalStateT f
  
  modify :: (Applicative m) => (s %1 -> s) %1 -> StateT s m ()
  modify f = state $ \s -> ((), f s)
  
  -- TODO: add strict version of `modify`
  
  -- | @replace s@ will replace the current state with the new given state, and
  -- return the old state.
  replace :: (Applicative m) => s %1 -> StateT s m s
  replace s = state $ (\s' -> (s', s))
  
  -- # Instances of StateT
  -------------------------------------------------------------------------------
  
  instance (Functor m) => Functor (NonLinear.StateT s m) where
    fmap f (NonLinear.StateT x) = NonLinear.StateT $ \s -> fmap (\(a, s') -> (f a, s')) $ x s
  
  instance (Data.Functor m) => Data.Functor (StateT s m) where
    fmap f (StateT x) = StateT (\s -> Data.fmap (\(a, s') -> (f a, s')) (x s))
  
  instance (Functor m) => Functor (StateT s m) where
    fmap f (StateT x) = StateT (\s -> fmap (\(a, s') -> (f a, s')) (x s))
  
  instance (Monad m) => Applicative (StateT s m) where
    pure x = StateT (\s -> return (x, s))
    StateT mf <*> StateT mx = StateT $ \s -> do
      (f, s') <- mf s
      (x, s'') <- mx s'
      return (f x, s'')
  
  instance (Monad m) => Monad (StateT s m) where
    StateT mx >>= f = StateT $ \s -> do
      (x, s') <- mx s
      runStateT (f x) s'
  
  instance MonadTrans (StateT s) where
    lift x = StateT (\s -> fmap (,s) x)
  