{-# LANGUAGE MagicHash, UnboxedTuples, Strict #-}
{-# OPTIONS_GHC -Wno-redundant-bang-patterns #-}

module Agda.Utils.StrictState where

import GHC.Exts (oneShot)

newtype State s a = State {forall s a. State s a -> s -> (# a, s #)
runState# :: s -> (# a, s #)}

instance Functor (State s) where
  {-# INLINE fmap #-}
  fmap :: forall a b. (a -> b) -> State s a -> State s b
fmap a -> b
f (State s -> (# a, s #)
g) = (s -> (# b, s #)) -> State s b
forall s a. (s -> (# a, s #)) -> State s a
State ((s -> (# b, s #)) -> s -> (# b, s #)
forall a b. (a -> b) -> a -> b
oneShot \s
s -> case s -> (# a, s #)
g s
s of
    (# a
a, !s
s #) -> let b :: b
b = a -> b
f a
a in (# b
b, s
s #))
  {-# INLINE (<$) #-}
  <$ :: forall a b. a -> State s b -> State s a
(<$) a
a State s b
m = (\b
_ -> a
a) (b -> a) -> State s b -> State s a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> State s b
m

instance Applicative (State s) where
  {-# INLINE pure #-}
  pure :: forall a. a -> State s a
pure a
a = (s -> (# a, s #)) -> State s a
forall s a. (s -> (# a, s #)) -> State s a
State ((s -> (# a, s #)) -> s -> (# a, s #)
forall a b. (a -> b) -> a -> b
oneShot \s
s -> (# a
a, s
s #))
  {-# INLINE (<*>) #-}
  State s -> (# a -> b, s #)
mf <*> :: forall a b. State s (a -> b) -> State s a -> State s b
<*> State s -> (# a, s #)
ma = (s -> (# b, s #)) -> State s b
forall s a. (s -> (# a, s #)) -> State s a
State ((s -> (# b, s #)) -> s -> (# b, s #)
forall a b. (a -> b) -> a -> b
oneShot \s
s -> case s -> (# a -> b, s #)
mf s
s of
    (# a -> b
f, s
s #) -> case s -> (# a, s #)
ma s
s of
      (# a
a, !s
s #) -> let b :: b
b = a -> b
f a
a in (# b
b, s
s #))
  {-# INLINE (*>) #-}
  State s -> (# a, s #)
ma *> :: forall a b. State s a -> State s b -> State s b
*> State s -> (# b, s #)
mb = (s -> (# b, s #)) -> State s b
forall s a. (s -> (# a, s #)) -> State s a
State ((s -> (# b, s #)) -> s -> (# b, s #)
forall a b. (a -> b) -> a -> b
oneShot \s
s -> case s -> (# a, s #)
ma s
s of
    (# a
_, s
s #) -> s -> (# b, s #)
mb s
s)
  {-# INLINE (<*) #-}
  State s -> (# a, s #)
ma <* :: forall a b. State s a -> State s b -> State s a
<* State s -> (# b, s #)
mb = (s -> (# a, s #)) -> State s a
forall s a. (s -> (# a, s #)) -> State s a
State ((s -> (# a, s #)) -> s -> (# a, s #)
forall a b. (a -> b) -> a -> b
oneShot \s
s -> case s -> (# a, s #)
ma s
s of
    (# a
a, s
s #) -> case s -> (# b, s #)
mb s
s of
      (# b
_, s
s #) -> (# a
a, s
s #))

instance Monad (State s) where
  {-# INLINE return #-}
  return :: forall a. a -> State s a
return = a -> State s a
forall a. a -> State s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
  {-# INLINE (>>=) #-}
  State s -> (# a, s #)
ma >>= :: forall a b. State s a -> (a -> State s b) -> State s b
>>= a -> State s b
f = (s -> (# b, s #)) -> State s b
forall s a. (s -> (# a, s #)) -> State s a
State ((s -> (# b, s #)) -> s -> (# b, s #)
forall a b. (a -> b) -> a -> b
oneShot \s
s -> case s -> (# a, s #)
ma s
s of
    (# a
a, s
s #) -> State s b -> s -> (# b, s #)
forall s a. State s a -> s -> (# a, s #)
runState# (a -> State s b
f a
a) s
s)
  {-# INLINE (>>) #-}
  >> :: forall a b. State s a -> State s b -> State s b
(>>) = State s a -> State s b -> State s b
forall a b. State s a -> State s b -> State s b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
(*>)

{-# INLINE put #-}
put :: s -> State s ()
put :: forall s. s -> State s ()
put s
s = (s -> (# (), s #)) -> State s ()
forall s a. (s -> (# a, s #)) -> State s a
State \s
_ -> (# (), s
s #)

{-# INLINE get #-}
get :: State s s
get :: forall s. State s s
get = (s -> (# s, s #)) -> State s s
forall s a. (s -> (# a, s #)) -> State s a
State \s
s -> (# s
s, s
s #)

{-# INLINE gets #-}
gets :: (s -> a) -> State s a
gets :: forall s a. (s -> a) -> State s a
gets s -> a
f = s -> a
f (s -> a) -> State s s -> State s a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> State s s
forall s. State s s
get

{-# INLINE modify #-}
modify :: (s -> s) -> State s ()
modify :: forall s. (s -> s) -> State s ()
modify s -> s
f = (s -> (# (), s #)) -> State s ()
forall s a. (s -> (# a, s #)) -> State s a
State \s
s -> let s' :: s
s' = s -> s
f s
s in (# (), s
s' #)

{-# INLINE execState #-}
execState :: State s a -> s -> s
execState :: forall s a. State s a -> s -> s
execState (State s -> (# a, s #)
f) s
s = case s -> (# a, s #)
f s
s of
  (# a
_, s
s #) -> s
s

{-# INLINE runState #-}
runState :: State s a -> s -> (a, s)
runState :: forall s a. State s a -> s -> (a, s)
runState (State s -> (# a, s #)
f) s
s = case s -> (# a, s #)
f s
s of
  (# !a
a, !s
s #) -> (a
a, s
s)

{-# INLINE evalState #-}
evalState :: State s a -> s -> a
evalState :: forall s a. State s a -> s -> a
evalState (State s -> (# a, s #)
f) s
s = case s -> (# a, s #)
f s
s of
  (# a
a, s
_ #) -> a
a