{-# LANGUAGE
  BangPatterns,
  GADTs,
  KindSignatures,
  RankNTypes,
  ScopedTypeVariables,
  StandaloneKindSignatures,
  TypeOperators #-}

-- | = State as an algebraic effect
--
-- The 'runState' handler calls each continuation exactly once.
-- It is compatible with single-shot continuations.
module Bluefin.Algae.State
  ( -- * Operations
    State(..)
  , get
  , put
  , putL
  , modify
  , modifyL

    -- * Handlers
  , runState
  , evalState
  , execState
  ) where

import Data.Kind (Type)
import Bluefin.Eff (Eff, type (:&), type (:>))
import Bluefin.Algae

-- | The state effect.
data State (s :: Type) :: AEffect where
  -- | Get the current state.
  Get :: State s s
  -- | Put a new state.
  Put :: s -> State s ()

-- | Get the current state. Call the 'Get' operation.
get :: z :> zz => Handler (State s) z -> Eff zz s
get :: forall (z :: Effects) (zz :: Effects) s.
(z :> zz) =>
Handler (State s) z -> Eff zz s
get Handler (State s) z
h = Handler (State s) z -> State s s -> Eff zz s
forall (s :: Effects) (ss :: Effects) (f :: AEffect) a.
(s :> ss) =>
Handler f s -> f a -> Eff ss a
call Handler (State s) z
h State s s
forall s. State s s
Get

-- | Put a new state. Call the 'Put' operation.
--
-- This function is strict.
put :: z :> zz => Handler (State s) z -> s -> Eff zz ()
put :: forall (z :: Effects) (zz :: Effects) s.
(z :> zz) =>
Handler (State s) z -> s -> Eff zz ()
put Handler (State s) z
h !s
s = Handler (State s) z -> State s () -> Eff zz ()
forall (s :: Effects) (ss :: Effects) (f :: AEffect) a.
(s :> ss) =>
Handler f s -> f a -> Eff ss a
call Handler (State s) z
h (s -> State s ()
forall s. s -> State s ()
Put s
s)

-- | Lazy variant of 'put'.
putL :: z :> zz => Handler (State s) z -> s -> Eff zz ()
putL :: forall (z :: Effects) (zz :: Effects) s.
(z :> zz) =>
Handler (State s) z -> s -> Eff zz ()
putL Handler (State s) z
h s
s = Handler (State s) z -> State s () -> Eff zz ()
forall (s :: Effects) (ss :: Effects) (f :: AEffect) a.
(s :> ss) =>
Handler f s -> f a -> Eff ss a
call Handler (State s) z
h (s -> State s ()
forall s. s -> State s ()
Put s
s)

-- | Modify the state.
--
-- This function is strict in the modified state.
modify :: z :> zz => Handler (State s) z -> (s -> s) -> Eff zz ()
modify :: forall (z :: Effects) (zz :: Effects) s.
(z :> zz) =>
Handler (State s) z -> (s -> s) -> Eff zz ()
modify Handler (State s) z
h s -> s
f = Handler (State s) z -> Eff zz s
forall (z :: Effects) (zz :: Effects) s.
(z :> zz) =>
Handler (State s) z -> Eff zz s
get Handler (State s) z
h Eff zz s -> (s -> Eff zz ()) -> Eff zz ()
forall a b. Eff zz a -> (a -> Eff zz b) -> Eff zz b
forall (m :: AEffect) a b. Monad m => m a -> (a -> m b) -> m b
>>= Handler (State s) z -> s -> Eff zz ()
forall (z :: Effects) (zz :: Effects) s.
(z :> zz) =>
Handler (State s) z -> s -> Eff zz ()
put Handler (State s) z
h (s -> Eff zz ()) -> (s -> s) -> s -> Eff zz ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. s -> s
f

-- | Lazy variant of 'modify'.
modifyL :: z :> zz => Handler (State s) z -> (s -> s) -> Eff zz ()
modifyL :: forall (z :: Effects) (zz :: Effects) s.
(z :> zz) =>
Handler (State s) z -> (s -> s) -> Eff zz ()
modifyL Handler (State s) z
h s -> s
f = Handler (State s) z -> Eff zz s
forall (z :: Effects) (zz :: Effects) s.
(z :> zz) =>
Handler (State s) z -> Eff zz s
get Handler (State s) z
h Eff zz s -> (s -> Eff zz ()) -> Eff zz ()
forall a b. Eff zz a -> (a -> Eff zz b) -> Eff zz b
forall (m :: AEffect) a b. Monad m => m a -> (a -> m b) -> m b
>>= Handler (State s) z -> s -> Eff zz ()
forall (z :: Effects) (zz :: Effects) s.
(z :> zz) =>
Handler (State s) z -> s -> Eff zz ()
putL Handler (State s) z
h (s -> Eff zz ()) -> (s -> s) -> s -> Eff zz ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. s -> s
f

-- | Run a stateful computation from the given starting state.
runState ::
  s ->  -- ^ Initial state
  (forall z. Handler (State s) z -> Eff (z :& zz) a) ->  -- ^ Stateful computation
  Eff zz (a, s)
runState :: forall s (zz :: Effects) a.
s
-> (forall (z :: Effects). Handler (State s) z -> Eff (z :& zz) a)
-> Eff zz (a, s)
runState = (a -> s -> (a, s))
-> s
-> (forall (z :: Effects). Handler (State s) z -> Eff (z :& zz) a)
-> Eff zz (a, s)
forall s a r (zz :: Effects).
(a -> s -> r)
-> s
-> (forall (z :: Effects). Handler (State s) z -> Eff (z :& zz) a)
-> Eff zz r
runStateWith (,)

-- | Variant of 'runState' that returns only the result value.
evalState ::
  s ->  -- ^ Initial state
  (forall z. Handler (State s) z -> Eff (z :& zz) a) ->  -- ^ Stateful computation
  Eff zz a
evalState :: forall s (zz :: Effects) a.
s
-> (forall (z :: Effects). Handler (State s) z -> Eff (z :& zz) a)
-> Eff zz a
evalState = (a -> s -> a)
-> s
-> (forall (z :: Effects). Handler (State s) z -> Eff (z :& zz) a)
-> Eff zz a
forall s a r (zz :: Effects).
(a -> s -> r)
-> s
-> (forall (z :: Effects). Handler (State s) z -> Eff (z :& zz) a)
-> Eff zz r
runStateWith a -> s -> a
forall a b. a -> b -> a
const

-- | Variant of 'runState' that returns only the final state.
execState ::
  s ->  -- ^ Initial state
  (forall z. Handler (State s) z -> Eff (z :& zz) a) ->  -- ^ Stateful computation
  Eff zz s
execState :: forall s (zz :: Effects) a.
s
-> (forall (z :: Effects). Handler (State s) z -> Eff (z :& zz) a)
-> Eff zz s
execState = (a -> s -> s)
-> s
-> (forall (z :: Effects). Handler (State s) z -> Eff (z :& zz) a)
-> Eff zz s
forall s a r (zz :: Effects).
(a -> s -> r)
-> s
-> (forall (z :: Effects). Handler (State s) z -> Eff (z :& zz) a)
-> Eff zz r
runStateWith ((s -> s) -> a -> s -> s
forall a b. a -> b -> a
const s -> s
forall a. a -> a
id)

runStateWith :: forall s a r zz.
  (a -> s -> r) ->  -- ^ Combine the result and final state.
  s ->  -- ^ Initial state
  (forall z. Handler (State s) z -> Eff (z :& zz) a) ->  -- ^ Stateful computation
  Eff zz r
runStateWith :: forall s a r (zz :: Effects).
(a -> s -> r)
-> s
-> (forall (z :: Effects). Handler (State s) z -> Eff (z :& zz) a)
-> Eff zz r
runStateWith a -> s -> r
finish s
s0 forall (z :: Effects). Handler (State s) z -> Eff (z :& zz) a
f = s -> Eff zz (s -> Eff zz r) -> Eff zz r
unwrap s
s0 (HandlerBody (State s) zz (s -> Eff zz r)
-> ScopedEff (State s) zz (s -> Eff zz r) -> Eff zz (s -> Eff zz r)
forall (f :: AEffect) (ss :: Effects) a.
HandlerBody f ss a -> ScopedEff f ss a -> Eff ss a
handle State s x
-> (x -> Eff zz (s -> Eff zz r)) -> Eff zz (s -> Eff zz r)
HandlerBody (State s) zz (s -> Eff zz r)
stateHandler (Eff (s :& zz) a -> Eff (s :& zz) (s -> Eff zz r)
forall (z :: Effects).
Eff (z :& zz) a -> Eff (z :& zz) (s -> Eff zz r)
wrap (Eff (s :& zz) a -> Eff (s :& zz) (s -> Eff zz r))
-> (Handler (State s) s -> Eff (s :& zz) a)
-> Handler (State s) s
-> Eff (s :& zz) (s -> Eff zz r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Handler (State s) s -> Eff (s :& zz) a
forall (z :: Effects). Handler (State s) z -> Eff (z :& zz) a
f))
  where
    stateHandler :: HandlerBody (State s) zz (s -> Eff zz r)
    stateHandler :: HandlerBody (State s) zz (s -> Eff zz r)
stateHandler State s x
Get x -> Eff zz (s -> Eff zz r)
k = (s -> Eff zz r) -> Eff zz (s -> Eff zz r)
forall a. a -> Eff zz a
forall (f :: AEffect) a. Applicative f => a -> f a
pure (\s
s -> x -> Eff zz (s -> Eff zz r)
k s
x
s Eff zz (s -> Eff zz r) -> ((s -> Eff zz r) -> Eff zz r) -> Eff zz r
forall a b. Eff zz a -> (a -> Eff zz b) -> Eff zz b
forall (m :: AEffect) a b. Monad m => m a -> (a -> m b) -> m b
>>= \s -> Eff zz r
k1 -> s -> Eff zz r
k1 s
s)
    stateHandler (Put s
s) x -> Eff zz (s -> Eff zz r)
k = (s -> Eff zz r) -> Eff zz (s -> Eff zz r)
forall a. a -> Eff zz a
forall (f :: AEffect) a. Applicative f => a -> f a
pure (\s
_ -> x -> Eff zz (s -> Eff zz r)
k () Eff zz (s -> Eff zz r) -> ((s -> Eff zz r) -> Eff zz r) -> Eff zz r
forall a b. Eff zz a -> (a -> Eff zz b) -> Eff zz b
forall (m :: AEffect) a b. Monad m => m a -> (a -> m b) -> m b
>>= \s -> Eff zz r
k1 -> s -> Eff zz r
k1 s
s)

    wrap :: Eff (z :& zz) a -> Eff (z :& zz) (s -> Eff zz r)
    wrap :: forall (z :: Effects).
Eff (z :& zz) a -> Eff (z :& zz) (s -> Eff zz r)
wrap = (a -> s -> Eff zz r)
-> Eff (z :& zz) a -> Eff (z :& zz) (s -> Eff zz r)
forall a b. (a -> b) -> Eff (z :& zz) a -> Eff (z :& zz) b
forall (f :: AEffect) a b. Functor f => (a -> b) -> f a -> f b
fmap (\a
a s
s -> r -> Eff zz r
forall a. a -> Eff zz a
forall (f :: AEffect) a. Applicative f => a -> f a
pure (a -> s -> r
finish a
a s
s))

    unwrap :: s -> Eff zz (s -> Eff zz r) -> Eff zz r
    unwrap :: s -> Eff zz (s -> Eff zz r) -> Eff zz r
unwrap s
s Eff zz (s -> Eff zz r)
m = Eff zz (s -> Eff zz r)
m Eff zz (s -> Eff zz r) -> ((s -> Eff zz r) -> Eff zz r) -> Eff zz r
forall a b. Eff zz a -> (a -> Eff zz b) -> Eff zz b
forall (m :: AEffect) a b. Monad m => m a -> (a -> m b) -> m b
>>= \s -> Eff zz r
k -> s -> Eff zz r
k s
s