-- | 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) (a -> (a, b)) -> f a -> f (a, b)
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,) (b -> (a, b)) -> f b -> f (a, b)
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 = Const i o -> i
forall {k} a (b :: k). Const a b -> a
getConst (Const i o -> i) -> Const i o -> i
forall a b. (a -> b) -> a -> b
$ (i -> Const i i) -> o -> Const i o
Lens' i o
l i -> Const i i
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 = Lens' i o -> LensMap i o
forall i o. Lens' i o -> LensMap i o
over Lens' i o
l LensMap i o -> (i -> i -> i) -> i -> o -> o
forall b c a. (b -> c) -> (a -> b) -> a -> c
. i -> i -> i
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 = Identity o -> o
forall a. Identity a -> a
runIdentity (Identity o -> o) -> Identity o -> o
forall a b. (a -> b) -> a -> b
$ (i -> Identity i) -> o -> Identity o
Lens' i o
l (i -> Identity i
forall a. a -> Identity a
Identity (i -> Identity i) -> (i -> i) -> i -> Identity i
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 = (o -> m (a, o)) -> StateT o m a
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StateT ((o -> m (a, o)) -> StateT o m a)
-> (o -> m (a, o)) -> StateT o m a
forall a b. (a -> b) -> a -> b
$ \ o
o -> do
  (a
a, i
i) <- StateT i m a -> i -> m (a, i)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT StateT i m a
m (o
o o -> Lens' i o -> i
forall o i. o -> Lens' i o -> i
^. Lens' i o
l)
  (a, o) -> m (a, o)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
a, Lens' i o -> LensSet i o
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 <- (o -> i) -> m i
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (o -> Lens' i o -> i
forall o i. o -> Lens' i o -> i
^. Lens' i o
l)
           i -> m i
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 = (o -> o) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((o -> o) -> m ()) -> (o -> o) -> m ()
forall a b. (a -> b) -> a -> b
$ Lens' i o -> LensSet i o
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 = (o -> o) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((o -> o) -> m ()) -> (o -> o) -> m ()
forall a b. (a -> b) -> a -> b
$ Lens' i o -> LensMap i o
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 = o -> m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (o -> m ()) -> m o -> m ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (i -> m i) -> o -> m o
Lens' i o
l i -> m i
f (o -> m o) -> m o -> m o
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m o
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 <- m o
forall s (m :: * -> *). MonadState s m => m s
get
  (o
o', r
r) <- WriterT r m o -> m (o, r)
forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (WriterT r m o -> m (o, r)) -> WriterT r m o -> m (o, r)
forall a b. (a -> b) -> a -> b
$ (i -> WriterT r m i) -> o -> WriterT r m o
Lens' i o
l (m (i, r) -> WriterT r m i
forall w (m :: * -> *) a. m (a, w) -> WriterT w m a
WriterT (m (i, r) -> WriterT r m i)
-> (i -> m (i, r)) -> i -> WriterT r m i
forall b c a. (b -> c) -> (a -> b) -> a -> c
. i -> m (i, r)
f) o
o
  o -> m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put o
o'
  r -> m r
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 <- Lens' i o -> m i
forall o (m :: * -> *) i. MonadState o m => Lens' i o -> m i
use Lens' i o
l
  Lens' i o
l Lens' i o -> (i -> i) -> m ()
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 Lens' i o -> i -> m ()
forall o (m :: * -> *) i. MonadState o m => Lens' i o -> i -> m ()
.= i
old
  r -> m r
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 = (o -> i) -> m i
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (o -> Lens' i o -> i
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 = (o -> o) -> m a -> m a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((o -> o) -> m a -> m a)
-> ((i -> i) -> o -> o) -> (i -> i) -> m a -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Lens' i o -> (i -> i) -> o -> o
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 ((o -> o) -> m a -> m a)
-> ((i -> i) -> o -> o) -> (i -> i) -> m a -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Lens' i o -> (i -> i) -> o -> o
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 (k -> Map k v -> Maybe v
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
k Map k v
m) f (Maybe v) -> (Maybe v -> Map k v) -> f (Map k v)
forall (m :: * -> *) a b. Functor m => m a -> (a -> b) -> m b
<&> \ Maybe v
v -> (Maybe v -> Maybe v) -> k -> Map k v -> Map k v
forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
Map.alter (Maybe v -> Maybe v -> Maybe v
forall a b. a -> b -> a
const Maybe v
v) k
k Map k v
m