{-# OPTIONS_GHC -Wunused-imports #-}

-- | 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.Set (Set)
import qualified Data.Set as Set

import Data.Functor.Identity

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

-- * Type-preserving lenses.

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

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

-- * Some simple lenses.

lFst :: Lens' (a, b) a
lFst :: forall a b (f :: * -> *).
Functor f =>
(a -> f a) -> (a, b) -> f (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' (a, b) b
lSnd :: forall a b (f :: * -> *).
Functor f =>
(b -> f b) -> (a, b) -> f (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' o i@.
(^.) :: o -> Lens' o i -> i
o
o ^. :: forall o i. o -> Lens' o i -> i
^. Lens' o i
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' o i
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' o i@.
set :: Lens' o i -> LensSet o i
set :: forall o i. Lens' o i -> LensSet o i
set Lens' o i
l = Lens' o i -> LensMap o i
forall o i. Lens' o i -> LensMap o i
over (i -> f i) -> o -> f o
Lens' o i
l LensMap o i -> (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' o i -> LensMap o i
over :: forall o i. Lens' o i -> LensMap o i
over Lens' o i
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' o i
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

-- | Build a lens out of an isomorphism.
iso :: (o -> i) -> (i -> o) -> Lens' o i
iso :: forall o i. (o -> i) -> (i -> o) -> Lens' o i
iso o -> i
get i -> o
set i -> f i
f = (i -> o) -> f i -> f o
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap i -> o
set (f i -> f o) -> (o -> f i) -> o -> f o
forall b c a. (b -> c) -> (a -> b) -> a -> c
. i -> f i
f (i -> f i) -> (o -> i) -> o -> f i
forall b c a. (b -> c) -> (a -> b) -> a -> c
. o -> i
get

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

-- | Focus on a part of the state for a stateful computation.
focus :: Monad m => Lens' o i -> StateT i m a -> StateT o m a
focus :: forall (m :: * -> *) o i a.
Monad m =>
Lens' o i -> StateT i m a -> StateT o m a
focus Lens' o i
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' o i -> i
forall o i. o -> Lens' o i -> i
^. (i -> f i) -> o -> f o
Lens' o i
l)
  (a, o) -> m (a, o)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (a
a, Lens' o i -> LensSet o i
forall o i. Lens' o i -> LensSet o i
set (i -> f i) -> o -> f o
Lens' o i
l i
i o
o)

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

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

infix 4 .=
-- | Write a part of the state.
(.=) :: MonadState o m => Lens' o i -> i -> m ()
Lens' o i
l .= :: forall o (m :: * -> *) i. MonadState o m => Lens' o i -> 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' o i -> LensSet o i
forall o i. Lens' o i -> LensSet o i
set (i -> f i) -> o -> f o
Lens' o i
l i
i

infix 4 %=
-- | Modify a part of the state.
(%=) :: MonadState o m => Lens' o i -> (i -> i) -> m ()
Lens' o i
l %= :: forall o (m :: * -> *) i.
MonadState o m =>
Lens' o i -> (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' o i -> LensMap o i
forall o i. Lens' o i -> LensMap o i
over (i -> f i) -> o -> f o
Lens' o i
l i -> i
f

infix 4 %==
-- | Modify a part of the state monadically.
(%==) :: MonadState o m => Lens' o i -> (i -> m i) -> m ()
Lens' o i
l %== :: forall o (m :: * -> *) i.
MonadState o m =>
Lens' o i -> (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' o i
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' o i -> (i -> m (i, r)) -> m r
Lens' o i
l %%= :: forall o (m :: * -> *) i r.
MonadState o m =>
Lens' o i -> (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' o i
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 a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return r
r

-- | Modify a part of the state locally.
locallyState :: MonadState o m => Lens' o i -> (i -> i) -> m r -> m r
locallyState :: forall o (m :: * -> *) i r.
MonadState o m =>
Lens' o i -> (i -> i) -> m r -> m r
locallyState Lens' o i
l i -> i
f m r
k = do
  i
old <- Lens' o i -> m i
forall o (m :: * -> *) i. MonadState o m => Lens' o i -> m i
use (i -> f i) -> o -> f o
Lens' o i
l
  (i -> f i) -> o -> f o
Lens' o i
l Lens' o i -> (i -> i) -> m ()
forall o (m :: * -> *) i.
MonadState o m =>
Lens' o i -> (i -> i) -> m ()
%= i -> i
f
  r
x <- m r
k
  (i -> f i) -> o -> f o
Lens' o i
l Lens' o i -> i -> m ()
forall o (m :: * -> *) i. MonadState o m => Lens' o i -> i -> m ()
.= i
old
  r -> m r
forall a. a -> m a
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' o i -> m i
view :: forall o (m :: * -> *) i. MonadReader o m => Lens' o i -> m i
view Lens' o i
l = (o -> i) -> m i
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (o -> Lens' o i -> i
forall o i. o -> Lens' o i -> i
^. (i -> f i) -> o -> f o
Lens' o i
l)

-- | Modify a part of the state in a subcomputation.
locally :: MonadReader o m => Lens' o i -> (i -> i) -> m a -> m a
locally :: forall o (m :: * -> *) i a.
MonadReader o m =>
Lens' o i -> (i -> i) -> m a -> m a
locally Lens' o i
l = (o -> o) -> m a -> m a
forall a. (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' o i -> (i -> i) -> o -> o
forall o i. Lens' o i -> LensMap o i
over (i -> f i) -> o -> f o
Lens' o i
l

locally' :: ((o -> o) -> m a -> m a) -> Lens' o i -> (i -> i) -> m a -> m a
locally' :: forall o (m :: * -> *) a i.
((o -> o) -> m a -> m a) -> Lens' o i -> (i -> i) -> m a -> m a
locally' (o -> o) -> m a -> m a
local Lens' o i
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' o i -> (i -> i) -> o -> o
forall o i. Lens' o i -> LensMap o i
over (i -> f i) -> o -> f o
Lens' o i
l

-- * Lenses for collections

-- | Access a map value at a given key.
key :: Ord k => k -> Lens' (Map k v) (Maybe v)
key :: forall k v. Ord k => k -> Lens' (Map k v) (Maybe 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 (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f 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

-- | Focus on given element in a set.
contains :: Ord k => k -> Lens' (Set k) Bool
contains :: forall k. Ord k => k -> Lens' (Set k) Bool
contains k
k Bool -> f Bool
f Set k
s = Bool -> f Bool
f (k -> Set k -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member k
k Set k
s) f Bool -> (Bool -> Set k) -> f (Set k)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \case
  Bool
True  -> k -> Set k -> Set k
forall a. Ord a => a -> Set a -> Set a
Set.insert k
k Set k
s
  Bool
False -> k -> Set k -> Set k
forall a. Ord a => a -> Set a -> Set a
Set.delete k
k Set k
s