-- | Lenses allow you to use fields of the state of a state monad as if they were variables in an imperative language.
-- 'use' 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.Family.State.Strict
  ( zoom
  , use, uses
  , (%=)
  , assign, (.=)
  , (%%=)
  , (<~)
-- * Compound Assignments
  , (+=), (-=), (*=)
  , (//=)
  , (&&=), (||=)
  , (<>=)
-- * Types
  , Zooming
-- * Re-exports
  , LensLike, LensLike'
  , FoldLike, Constant
  , ASetter, ASetter', Identity
  , StateT, Writer
  , Monoid
  ) where

import Data.Monoid (Monoid, mappend)
import Data.Tuple (swap)
import Control.Monad (liftM)
import Control.Monad.Trans.Writer.Lazy (Writer, writer, runWriter)
import Control.Monad.Trans.State.Strict (StateT(..), state, get, modify)
import Lens.Family ( LensLike, LensLike'
                   , FoldLike, Constant
                   , ASetter, ASetter', Identity
                   , view, views, (%~)
                   )
import Lens.Family.State.Zoom (Zooming(..))

{- all these Monad constraints could be weakened to Functor or Applicative constraints -}

zoom :: Monad m => LensLike' (Zooming m c) a b -> StateT b m c -> StateT a m c
-- ^ @
-- zoom :: Monad m => Lens' a b -> StateT b m c -> StateT a m c
-- @
--
-- 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.
--
-- @
-- zoom :: (Monoid c, Moand m) => Traversal' a b -> StateT b m c -> StateT a m c
-- @
--
-- Run the \"subroutine\" on each element of the traversal in turn and 'mconcat' all the results together.
--
-- @
-- zoom :: Monad m => Traversal' a b -> StateT b m () -> StateT a m ()
-- @
--
-- Run the \"subroutine\" on each element the traversal in turn.
zoom l m = StateT $ unZooming . l (Zooming . (runStateT m))

use :: Monad m => FoldLike b a a' b b' -> StateT a m b
-- ^ @
-- use :: Monad m => Getter a a' b b' -> StateT a m b
-- @
--
-- Retrieve a field of the state
--
-- @
-- use :: (Monoid b, Monad m) => Fold a a' b b' -> StateT a m b
-- @
--
-- Retrieve a monoidal summary of all the referenced fields from the state
use l = view l `liftM` get

uses :: Monad m => FoldLike r a a' b b' -> (b -> r) -> StateT a m r
-- ^ @
-- uses :: (Monoid r, Monad m) => Fold a a' b b' -> (b -> r) -> StateT a m r
-- @
--
-- Retrieve all the referenced fields from the state and foldMap the results together with @f :: b -> r@.
--
-- @
-- uses :: Monad m => Getter a a' b b' -> (b -> r) -> StateT a m r
-- @
--
-- Retrieve a field of the state and pass it through the function @f :: b -> r@.
--
-- @uses l f = f \<$> use l@
uses l f = views l f `liftM` get

infix 4 %=

-- | Modify a field of the state.
(%=) :: Monad m => ASetter a a b b' -> (b -> b') -> StateT a m ()
l %= f = modify (l %~ f)

infix 4 .=

-- | Set a field of the state.
(.=) :: Monad m => ASetter a a b b' -> b' -> StateT a m ()
l .= v = l %= const v

-- | Set a field of the state.
assign :: Monad m => ASetter a a b b' -> b' -> StateT a m ()
assign = (.=)

infixr 2 <~

-- | Set a field of the state using the result of executing a stateful command.
(<~) :: Monad m => ASetter a a b b' -> StateT a m b' -> StateT a m ()
l <~ v = assign l =<< v

infix 4 %%=

(%%=) :: Monad m => LensLike (Writer c) a a b b' -> (b -> (c, b')) -> StateT a m c
-- ^ @
-- (%%=) :: Monad m => Lens a a b b' -> (b -> (c, b')) -> StateT a m c
-- @
--
-- Modify a field of the state while returning another value.
--
-- @
-- (%%=) :: (Monad m, Monoid c) => Traversal a a b b' -> (b -> (c, b')) -> StateT a m c
-- @
--
-- Modify each field of the state and return the 'mconcat' of the other values.
l %%= f = state (swap . runWriter . l (writer . swap . f))

infixr 4 +=, -=, *=

(+=), (-=), (*=) :: (Monad m, Num b) => ASetter' a b -> b -> StateT a m ()
f += b = f %= (+ b)
f -= b = f %= subtract b
f *= b = f %= (* b)

infixr 4 //=

(//=) :: (Monad m, Fractional b) => ASetter' a b -> b -> StateT a m ()
f //= b = f %= (/ b)

infixr 4 &&=, ||=

(&&=), (||=) :: Monad m => ASetter' a Bool -> Bool -> StateT a m ()
f &&= b = f %= (&& b)
f ||= b = f %= (|| b)

infixr 4 <>=

-- | Monoidally append a value to all referenced fields of the state.
(<>=) :: (Monoid o, Monad m) => ASetter' a o -> o -> StateT a m ()
f <>= b = f %= (`mappend` b)