{-# LANGUAGE Rank2Types #-}
-- | Lenses allow you to use fields of the state of a state monad as if they were variables in an imperative language.
-- 'access' is used to retrieve the value of a variable, and '~=' and '%=' allow you to set and modify a variable.
-- C-style compound assignments are also provided.
module Lens.Family2.State.Strict
  ( focus
  , access
  , (%=)
  , (~=)
  , (%%=)
  -- * Compound Assignments
  , (+=), (-=), (*=)
  , (//=)
  , (&&=), (||=)
  , (<>=)
  ) where

import Data.Monoid (Monoid, mappend)
import Control.Monad (liftM)
import Control.Monad.State.Strict (MonadState, StateT(..), get, modify, state)
import Lens.Family (Getter, Setter, (^.), (%~))
import qualified Lens.Family.State.Strict as LFS
import Lens.Family2.Stock (Lens)

-- | Lift a stateful operation on a field to a stateful operation on the whole state.
-- This is a good way to call a \"subroutine\" that only needs access to part of the state.
focus :: Monad m => Lens a b -> StateT b m c -> StateT a m c
focus l = LFS.focus l

-- | Retrieve a field of the state
access :: MonadState a m => Getter a b -> m b
access l = (^. l) `liftM` get

infix 4 %=

-- | Modify a field of the state
(%=) :: MonadState a m => Setter a b -> (b -> b) -> m ()
l %= f = modify (l %~ f)

infix 4 ~=

-- | Set a field of the state
(~=) :: MonadState a m => Setter a b -> b -> m ()
l ~= v = l %= const v

infix 4 %%=

-- | Modify a field of the state while returning another value
(%%=) :: MonadState a m => Lens a b -> (b -> (c, b)) -> m c
l %%= f = state (l f)

infixr 4 +=, -=, *=

(+=), (-=), (*=) :: (MonadState a m, Num b) => Setter a b -> b -> m ()
f += b = f %= (+ b)
f -= b = f %= subtract b
f *= b = f %= (* b)

infixr 4 //=

(//=) :: (MonadState a m, Fractional b) => Setter a b -> b -> m ()
f //= b = f %= (/ b)

infixr 4 &&=, ||=

(&&=), (||=) :: MonadState a m => Setter a Bool -> Bool -> m ()
f &&= b = f %= (&& b)
f ||= b = f %= (|| b)

infixr 4 <>=

(<>=) :: (Monoid o, MonadState a m) => Setter a o -> o -> m ()
f <>= b = f %= (`mappend` b)