-- | A cut-down implementation of lenses, with names taken from
--   Edward Kmett's lens package.

module Agda.Utils.Lens
  ( module Agda.Utils.Lens
  , (<&>) -- reexported from Agda.Utils.Functor
  ) where

import Control.Applicative ( Const(Const), getConst )
import Control.Monad.State
import Control.Monad.Reader
import Control.Monad.Writer

import Data.Map (Map)
import qualified Data.Map as Map

import Data.Functor.Identity

import Agda.Utils.Functor ((<&>))

-- * Type-preserving lenses.

-- | Van Laarhoven style homogeneous lenses.
--   Mnemoic: "Lens inner outer".
type Lens' i o = forall f. Functor f => (i -> f i) -> o -> f o

type LensGet i o = o -> i
type LensSet i o = i -> o -> o
type LensMap i o = (i -> i) -> o -> o

-- * Some simple lenses.

lFst :: Lens' a (a, b)
lFst :: forall a b. Lens' a (a, b)
lFst a -> f a
f (a
x, b
y) = (, b
y) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f a
f a
x

lSnd :: Lens' b (a, b)
lSnd :: forall b a. Lens' b (a, b)
lSnd b -> f b
f (a
x, b
y) = (a
x,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> b -> f b
f b
y

-- * Elementary lens operations.

infixl 8 ^.
-- | Get inner part @i@ of structure @o@ as designated by @Lens' i o@.
(^.) :: o -> Lens' i o -> i
o
o ^. :: forall o i. o -> Lens' i o -> i
^. Lens' i o
l = forall {k} a (b :: k). Const a b -> a
getConst forall a b. (a -> b) -> a -> b
$ Lens' i o
l forall {k} a (b :: k). a -> Const a b
Const o
o

-- | Set inner part @i@ of structure @o@ as designated by @Lens' i o@.
set :: Lens' i o -> LensSet i o
set :: forall i o. Lens' i o -> LensSet i o
set Lens' i o
l = forall i o. Lens' i o -> LensMap i o
over Lens' i o
l forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. a -> b -> a
const

-- | Modify inner part @i@ of structure @o@ using a function @i -> i@.
over :: Lens' i o -> LensMap i o
over :: forall i o. Lens' i o -> LensMap i o
over Lens' i o
l i -> i
f o
o = forall a. Identity a -> a
runIdentity forall a b. (a -> b) -> a -> b
$ Lens' i o
l (forall a. a -> Identity a
Identity forall b c a. (b -> c) -> (a -> b) -> a -> c
. i -> i
f) o
o


-- * State accessors and modifiers using 'StateT'.

-- | Focus on a part of the state for a stateful computation.
focus :: Monad m => Lens' i o -> StateT i m a -> StateT o m a
focus :: forall (m :: * -> *) i o a.
Monad m =>
Lens' i o -> StateT i m a -> StateT o m a
focus Lens' i o
l StateT i m a
m = forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StateT forall a b. (a -> b) -> a -> b
$ \ o
o -> do
  (a
a, i
i) <- forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT StateT i m a
m (o
o forall o i. o -> Lens' i o -> i
^. Lens' i o
l)
  forall (m :: * -> *) a. Monad m => a -> m a
return (a
a, forall i o. Lens' i o -> LensSet i o
set Lens' i o
l i
i o
o)

-- * State accessors and modifiers using 'MonadState'.

-- | Read a part of the state.
use :: MonadState o m => Lens' i o -> m i
use :: forall o (m :: * -> *) i. MonadState o m => Lens' i o -> m i
use Lens' i o
l = do !i
x <- forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (forall o i. o -> Lens' i o -> i
^. Lens' i o
l)
           forall (m :: * -> *) a. Monad m => a -> m a
return i
x

infix 4 .=
-- | Write a part of the state.
(.=) :: MonadState o m => Lens' i o -> i -> m ()
Lens' i o
l .= :: forall o (m :: * -> *) i. MonadState o m => Lens' i o -> i -> m ()
.= i
i = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ forall i o. Lens' i o -> LensSet i o
set Lens' i o
l i
i

infix 4 %=
-- | Modify a part of the state.
(%=) :: MonadState o m => Lens' i o -> (i -> i) -> m ()
Lens' i o
l %= :: forall o (m :: * -> *) i.
MonadState o m =>
Lens' i o -> (i -> i) -> m ()
%= i -> i
f = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ forall i o. Lens' i o -> LensMap i o
over Lens' i o
l i -> i
f

infix 4 %==
-- | Modify a part of the state monadically.
(%==) :: MonadState o m => Lens' i o -> (i -> m i) -> m ()
Lens' i o
l %== :: forall o (m :: * -> *) i.
MonadState o m =>
Lens' i o -> (i -> m i) -> m ()
%== i -> m i
f = forall s (m :: * -> *). MonadState s m => s -> m ()
put forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Lens' i o
l i -> m i
f forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall s (m :: * -> *). MonadState s m => m s
get

infix 4 %%=
-- | Modify a part of the state monadically, and return some result.
(%%=) :: MonadState o m => Lens' i o -> (i -> m (i, r)) -> m r
Lens' i o
l %%= :: forall o (m :: * -> *) i r.
MonadState o m =>
Lens' i o -> (i -> m (i, r)) -> m r
%%= i -> m (i, r)
f = do
  o
o <- forall s (m :: * -> *). MonadState s m => m s
get
  (o
o', r
r) <- forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT forall a b. (a -> b) -> a -> b
$ Lens' i o
l (forall w (m :: * -> *) a. m (a, w) -> WriterT w m a
WriterT forall b c a. (b -> c) -> (a -> b) -> a -> c
. i -> m (i, r)
f) o
o
  forall s (m :: * -> *). MonadState s m => s -> m ()
put o
o'
  forall (m :: * -> *) a. Monad m => a -> m a
return r
r

-- | Modify a part of the state locally.
locallyState :: MonadState o m => Lens' i o -> (i -> i) -> m r -> m r
locallyState :: forall o (m :: * -> *) i r.
MonadState o m =>
Lens' i o -> (i -> i) -> m r -> m r
locallyState Lens' i o
l i -> i
f m r
k = do
  i
old <- forall o (m :: * -> *) i. MonadState o m => Lens' i o -> m i
use Lens' i o
l
  Lens' i o
l forall o (m :: * -> *) i.
MonadState o m =>
Lens' i o -> (i -> i) -> m ()
%= i -> i
f
  r
x <- m r
k
  Lens' i o
l forall o (m :: * -> *) i. MonadState o m => Lens' i o -> i -> m ()
.= i
old
  forall (m :: * -> *) a. Monad m => a -> m a
return r
x

-- * Read-only state accessors and modifiers.

-- | Ask for part of read-only state.
view :: MonadReader o m => Lens' i o -> m i
view :: forall o (m :: * -> *) i. MonadReader o m => Lens' i o -> m i
view Lens' i o
l = forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (forall o i. o -> Lens' i o -> i
^. Lens' i o
l)

-- | Modify a part of the state in a subcomputation.
locally :: MonadReader o m => Lens' i o -> (i -> i) -> m a -> m a
locally :: forall o (m :: * -> *) i a.
MonadReader o m =>
Lens' i o -> (i -> i) -> m a -> m a
locally Lens' i o
l = forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall i o. Lens' i o -> LensMap i o
over Lens' i o
l

locally' :: ((o -> o) -> m a -> m a) -> Lens' i o -> (i -> i) -> m a -> m a
locally' :: forall o (m :: * -> *) a i.
((o -> o) -> m a -> m a) -> Lens' i o -> (i -> i) -> m a -> m a
locally' (o -> o) -> m a -> m a
local Lens' i o
l = (o -> o) -> m a -> m a
local forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall i o. Lens' i o -> LensMap i o
over Lens' i o
l

key :: Ord k => k -> Lens' (Maybe v) (Map k v)
key :: forall k v. Ord k => k -> Lens' (Maybe v) (Map k v)
key k
k Maybe v -> f (Maybe v)
f Map k v
m = Maybe v -> f (Maybe v)
f (forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
k Map k v
m) forall (m :: * -> *) a b. Functor m => m a -> (a -> b) -> m b
<&> \ Maybe v
v -> forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
Map.alter (forall a b. a -> b -> a
const Maybe v
v) k
k Map k v
m