{-# 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 :: forall a (m :: * -> *) b. MonadState a m => Lens a b -> m b
access Lens a b
l = forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (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 ~= :: forall a (m :: * -> *) b. MonadState a m => Lens a b -> b -> m ()
~= b
b = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ 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 != :: forall a (m :: * -> *) b. MonadState a m => Lens a b -> b -> m ()
!= b
b = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify' forall a b. (a -> b) -> a -> b
$ forall a b. Lens a b -> b -> a -> a
setL Lens a b
l 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 %= :: forall a (m :: * -> *) b.
MonadState a m =>
Lens a b -> (b -> b) -> m ()
%= b -> b
f = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ 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 !%= :: forall a (m :: * -> *) b.
MonadState a m =>
Lens a b -> (b -> b) -> m ()
!%= b -> b
f = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify' forall a b. (a -> b) -> a -> b
$ 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 :: forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT = forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
Strict.runStateT
instance MonadStateT Lazy.StateT where runStateT :: forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT = 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 :: forall (stateT :: * -> (* -> *) -> * -> *) s (m :: * -> *) s' a.
(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'
l stateT s' m a
a = do
  s
s <- forall s (m :: * -> *). MonadState s m => m s
get
  (a
r, s'
s') <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall (t :: * -> (* -> *) -> * -> *) s (m :: * -> *) a.
MonadStateT t =>
t s m a -> s -> m (a, s)
runStateT stateT s' m a
a (s
s forall b c. b -> Lens b c -> c
^. Lens s s'
l)
  forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify' forall a b. (a -> b) -> a -> b
$ forall a b. Lens a b -> b -> a -> a
setL Lens s s'
l s'
s'
  forall (m :: * -> *) a. Monad m => a -> m a
return a
r