{-# LANGUAGE CPP, MultiParamTypeClasses, FlexibleInstances #-}
module Data.Lens.Light.State
  ( access
  , (~=)
  , (!=)
  , (%=)
  , (!%=)
  , zoom
  , MonadStateT
  )
  where

import Control.Monad.State.Class
import qualified Control.Monad.State.Strict as Strict
import qualified Control.Monad.State as Lazy
import Control.Monad.Trans
import Data.Lens.Light.Core

-- | Get the value of a lens into state
access :: MonadState a m => Lens a b -> m b
access :: Lens a b -> m b
access Lens a b
l = (a -> b) -> m b
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (Lens a b -> a -> b
forall a b. Lens a b -> a -> b
getL Lens a b
l)

-- | Set a value using a lens into state
(~=) :: MonadState a m => Lens a b -> b -> m ()
Lens a b
l ~= :: Lens a b -> b -> m ()
~= b
b = (a -> a) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((a -> a) -> m ()) -> (a -> a) -> m ()
forall a b. (a -> b) -> a -> b
$ Lens a b -> b -> a -> a
forall a b. Lens a b -> b -> a -> a
setL Lens a b
l b
b

-- | Set a value using a lens into state. Forces both the value and the
-- whole state.
(!=) :: MonadState a m => Lens a b -> b -> m ()
Lens a b
l != :: Lens a b -> b -> m ()
!= b
b = (a -> a) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify' ((a -> a) -> m ()) -> (a -> a) -> m ()
forall a b. (a -> b) -> a -> b
$ Lens a b -> b -> a -> a
forall a b. Lens a b -> b -> a -> a
setL Lens a b
l (b -> a -> a) -> b -> a -> a
forall a b. (a -> b) -> a -> b
$! b
b

#if !MIN_VERSION_mtl(2,2,0)
-- Copied from mtl-2.2.0.1
modify' :: MonadState s m => (s -> s) -> m ()
modify' f = state (\s -> let s' = f s in s' `seq` ((), s'))
#endif

infixr 4 ~=, !=

-- | Infix modification of a value through a lens into state
(%=) :: MonadState a m => Lens a b -> (b -> b) -> m ()
Lens a b
l %= :: Lens a b -> (b -> b) -> m ()
%= b -> b
f = (a -> a) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((a -> a) -> m ()) -> (a -> a) -> m ()
forall a b. (a -> b) -> a -> b
$ Lens a b -> (b -> b) -> a -> a
forall a b. Lens a b -> (b -> b) -> a -> a
modL Lens a b
l b -> b
f

-- | Infix modification of a value through a lens into state. Forces both
-- the function application and the whole state.
(!%=) :: MonadState a m => Lens a b -> (b -> b) -> m ()
Lens a b
l !%= :: Lens a b -> (b -> b) -> m ()
!%= b -> b
f = (a -> a) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify' ((a -> a) -> m ()) -> (a -> a) -> m ()
forall a b. (a -> b) -> a -> b
$ Lens a b -> (b -> b) -> a -> a
forall a b. Lens a b -> (b -> b) -> a -> a
modL' Lens a b
l b -> b
f

infixr 4 %=, !%=

-- | The purpose of this class is to abstract the difference between the
-- lazy and strict state monads, so that 'zoom' can work with either of
-- them.
class MonadStateT t where
  runStateT :: t s m a -> s -> m (a, s)

instance MonadStateT Strict.StateT where runStateT :: StateT s m a -> s -> m (a, s)
runStateT = StateT s m a -> s -> m (a, s)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
Strict.runStateT
instance MonadStateT Lazy.StateT where runStateT :: StateT s m a -> s -> m (a, s)
runStateT = StateT s m a -> s -> m (a, s)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
Lazy.runStateT

-- | Run a stateful computation with a smaller state inside another
-- computation with a bigger state.
zoom
  :: ( MonadStateT stateT
     , MonadState s (stateT s m)
     , MonadTrans (stateT s)
     , Monad m
     )
  => Lens s s'
  -> stateT s' m a
  -> stateT s m a
zoom :: Lens s s' -> stateT s' m a -> stateT s m a
zoom Lens s s'
l stateT s' m a
a = do
  s
s <- stateT s m s
forall s (m :: * -> *). MonadState s m => m s
get
  (a
r, s'
s') <- m (a, s') -> stateT s m (a, s')
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m (a, s') -> stateT s m (a, s'))
-> m (a, s') -> stateT s m (a, s')
forall a b. (a -> b) -> a -> b
$ stateT s' m a -> s' -> m (a, s')
forall (t :: * -> (* -> *) -> * -> *) s (m :: * -> *) a.
MonadStateT t =>
t s m a -> s -> m (a, s)
runStateT stateT s' m a
a (s
s s -> Lens s s' -> s'
forall b c. b -> Lens b c -> c
^. Lens s s'
l)
  (s -> s) -> stateT s m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify' ((s -> s) -> stateT s m ()) -> (s -> s) -> stateT s m ()
forall a b. (a -> b) -> a -> b
$ Lens s s' -> s' -> s -> s
forall a b. Lens a b -> b -> a -> a
setL Lens s s'
l s'
s'
  a -> stateT s m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
r