module Data.Record.Label (
    Getter, Setter, Modifier
  , Label (..)
  , lmod
  , (%), comp
  , getM, setM, modM
  , bothM
  , enterM, enterMT
  , withM, localM
  , list
  , module Data.Record.Label.TH
  ) where

import Control.Monad.State
import Data.Record.Label.TH

type Getter   a b = a -> b
type Setter   a b = b -> a -> a
type Modifier a b = (b -> b) -> a -> a

data Label a b = Label {
    lget :: Getter a b
  , lset :: Setter a b
  }

lmod :: Label a b -> Modifier a b
lmod l f a = lset l (f (lget l a)) a

infixr 8 %

(%) :: Label t a -> Label b t -> Label b a
a % b = Label (lget a . lget b) (lmod b . lset a)

-- Apply custom `parser' and 'printer' function.

comp :: (b -> c) -> (c -> b) -> Label t b -> Label t c
comp f g (Label a b) = Label (f . a) (\v -> b $ g v)

-- Extend the state monad with support for labels.

getM :: MonadState s m => Label s b -> m b
getM = gets . lget

setM :: MonadState s m => Label s b -> b -> m ()
setM l = modify . lset l

modM :: MonadState s m => Label s b -> (b -> b) -> m ()
modM l = modify . lmod l

-- Run a state computation for a sub element updating this part of the state afterwards.

enterM :: MonadState s m => Label s b -> State b a -> m a
enterM l c = do
  b <- getM l
  let (a, s) = runState c b
  setM l s
  return a

enterMT :: (MonadState s (t m), MonadTrans t, Monad m) => Label s b -> StateT b m a -> t m a
enterMT l c = do
  b <- getM l
  (a, s) <- lift $ runStateT c b
  setM l s
  return a

bothM :: MonadState s m => Label s b -> State b a -> m (b, a)
bothM parent cmp = do
  p <- getM parent
  c <- enterM parent cmp
  return (p, c)

localM :: MonadState s m => Label s b -> m c -> m c
localM l comp = do
  k <- getM l
  c <- comp
  setM l k
  return c

withM :: MonadState s m => Label s b -> State b a -> m c -> m c
withM l c d = localM l (enterM l c >> d)

-- Lift list indexing to a label.

list :: Int -> Label [a] a
list i = Label (!! i) (\v a -> take i a ++ [v] ++ drop (i+1) a)