{-# OPTIONS_HADDOCK hide #-}
{-# OPTIONS -Wno-orphans #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE LinearTypes #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE RebindableSyntax #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}

module Control.Functor.Linear.Internal.State
  ( StateT(..)
  , State
  , state
  , get, put, gets
  , modify
  , replace
  , runStateT, runState
  , mapStateT, mapState
  , execStateT, execState
  , withStateT, withState
  ) where

import Prelude.Linear.Internal
import Data.Unrestricted.Internal.Consumable
import Data.Unrestricted.Internal.Dupable
import Control.Functor.Linear.Internal.MonadTrans
import Control.Functor.Linear.Internal.Class
import Control.Functor.Linear.Internal.Instances ( Data(..) )
import qualified Data.Functor.Linear.Internal.Functor as Data
import qualified Data.Functor.Linear.Internal.Applicative as Data
import qualified Control.Monad.Trans.State.Strict as NonLinear
import qualified Control.Monad as NonLinear ()
import Data.Functor.Identity


-- # 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)

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

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)