-- | 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.Family.State.Lazy 
  ( focus
  , access
  , (%=)
  , (~=)
  , (%%=)
  -- * Compound Assignments
  , (+=), (-=), (*=)
  , (//=)
  , (&&=), (||=)
  -- * Types
  , Focusing
  ) where

import Control.Monad (liftM)
import Control.Monad.Trans.State.Lazy (StateT(..), get, modify)
import Lens.Family (Getter, Setter, (^.), (^%=))
import Lens.Family.Stock (Lens)

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

newtype Focusing m c a = Focusing { unFocusing :: m (c, a) }
instance Monad m => Functor (Focusing m c) where
  fmap f (Focusing m) = Focusing (liftM (fmap f) m)

-- | 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 (Focusing m c) a b -> StateT b m c -> StateT a m c
focus l m = StateT $ unFocusing . l (Focusing . (runStateT m))

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

infix 4 %=

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

infix 4 ~=

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

infix 4 %%=

-- | Modify a field of the state while returning another value
(%%=) :: Monad m => Lens (Focusing m c) a b -> (b -> (c, b)) -> StateT a m c
l %%= f = focus l (StateT (return . f))

infixr 4 +=, -=, *=

(+=), (-=), (*=) :: (Monad m, Num b) => Setter 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) => Setter a b -> b -> StateT a m ()
f //= b = f %= (/ b)

infixr 4 &&=, ||=

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