{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE Rank2Types #-}
{-|
  Hand-rolled lenses on e-graphs and e-classes which come in quite handy and
  are heavily used in 'Data.Equality.Graph'.
 -}
module Data.Equality.Graph.Lens where

import qualified Data.IntMap.Strict as IM
import qualified Data.Set as S

import Data.Functor.Identity
import Data.Functor.Const

import Data.Equality.Graph.Classes.Id
import Data.Equality.Graph.Nodes
import Data.Equality.Graph.Classes
import Data.Equality.Analysis
import {-# SOURCE #-} Data.Equality.Graph (EGraph(..), Memo, find)

-- | A 'Lens'' as defined in other lenses libraries
type Lens' s a = forall f. Functor f => (a -> f a) -> (s -> f s)


-- outdated comment for "getClass":
--
-- Get an e-class from an e-graph given its e-class id
--
-- Returns the canonical id of the class and the class itself
--
-- We'll find its canonical representation and then get it from the e-classes map
--
-- Invariant: The e-class exists.

-- | Lens for the e-class with the given id in an e-graph
--
-- Calls 'error' when the e-class doesn't exist
_class :: ClassId -> Lens' (EGraph l) (EClass l)
_class :: forall (l :: * -> *). ClassId -> Lens' (EGraph l) (EClass l)
_class ClassId
i EClass l -> f (EClass l)
afa EGraph l
s =
    let canon_id :: ClassId
canon_id = ClassId -> EGraph l -> ClassId
forall (l :: * -> *). ClassId -> EGraph l -> ClassId
find ClassId
i EGraph l
s
     in (\EClass l
c' -> EGraph l
s { classes :: ClassIdMap (EClass l)
classes = ClassId
-> EClass l -> ClassIdMap (EClass l) -> ClassIdMap (EClass l)
forall a. ClassId -> a -> IntMap a -> IntMap a
IM.insert ClassId
canon_id EClass l
c' (EGraph l -> ClassIdMap (EClass l)
forall (l :: * -> *). EGraph l -> ClassIdMap (EClass l)
classes EGraph l
s) }) (EClass l -> EGraph l) -> f (EClass l) -> f (EGraph l)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> EClass l -> f (EClass l)
afa (EGraph l -> ClassIdMap (EClass l)
forall (l :: * -> *). EGraph l -> ClassIdMap (EClass l)
classes EGraph l
s ClassIdMap (EClass l) -> ClassId -> EClass l
forall a. IntMap a -> ClassId -> a
IM.! ClassId
canon_id)
{-# INLINE _class #-}

-- | Lens for the 'Memo' of e-nodes in an e-graph
_memo :: Lens' (EGraph l) (Memo l)
_memo :: forall (l :: * -> *) (f :: * -> *).
Functor f =>
(Memo l -> f (Memo l)) -> EGraph l -> f (EGraph l)
_memo Memo l -> f (Memo l)
afa EGraph l
egr = (\Memo l
m1 -> EGraph l
egr {memo :: Memo l
memo = Memo l
m1}) (Memo l -> EGraph l) -> f (Memo l) -> f (EGraph l)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Memo l -> f (Memo l)
afa (EGraph l -> Memo l
forall (l :: * -> *). EGraph l -> Memo l
memo EGraph l
egr)
{-# INLINE _memo #-}

-- | Lens for the map of existing classes by id in an e-graph
_classes :: Lens' (EGraph l) (ClassIdMap (EClass l))
_classes :: forall (l :: * -> *) (f :: * -> *).
Functor f =>
(ClassIdMap (EClass l) -> f (ClassIdMap (EClass l)))
-> EGraph l -> f (EGraph l)
_classes ClassIdMap (EClass l) -> f (ClassIdMap (EClass l))
afa EGraph l
egr = (\ClassIdMap (EClass l)
m1 -> EGraph l
egr {classes :: ClassIdMap (EClass l)
classes = ClassIdMap (EClass l)
m1}) (ClassIdMap (EClass l) -> EGraph l)
-> f (ClassIdMap (EClass l)) -> f (EGraph l)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ClassIdMap (EClass l) -> f (ClassIdMap (EClass l))
afa (EGraph l -> ClassIdMap (EClass l)
forall (l :: * -> *). EGraph l -> ClassIdMap (EClass l)
classes EGraph l
egr)
{-# INLINE _classes #-}

-- | Lens for the 'Domain' of an e-class
_data :: Lens' (EClass l) (Domain l)
_data :: forall (l :: * -> *) (f :: * -> *).
Functor f =>
(Domain l -> f (Domain l)) -> EClass l -> f (EClass l)
_data Domain l -> f (Domain l)
afa EClass{ClassId
Set (ENode l)
NodeMap l ClassId
Domain l
eClassId :: ClassId
eClassNodes :: Set (ENode l)
eClassData :: Domain l
eClassParents :: NodeMap l ClassId
eClassId :: forall (l :: * -> *). EClass l -> ClassId
eClassNodes :: forall (l :: * -> *). EClass l -> Set (ENode l)
eClassData :: forall (l :: * -> *). EClass l -> Domain l
eClassParents :: forall (l :: * -> *). EClass l -> NodeMap l ClassId
..} = (\Domain l
d1 -> ClassId
-> Set (ENode l) -> Domain l -> NodeMap l ClassId -> EClass l
forall (l :: * -> *).
ClassId
-> Set (ENode l) -> Domain l -> NodeMap l ClassId -> EClass l
EClass ClassId
eClassId Set (ENode l)
eClassNodes Domain l
d1 NodeMap l ClassId
eClassParents) (Domain l -> EClass l) -> f (Domain l) -> f (EClass l)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Domain l -> f (Domain l)
afa Domain l
eClassData
{-# INLINE _data #-}

-- | Lens for the parent e-classes of an e-class
_parents :: Lens' (EClass l) (NodeMap l ClassId)
_parents :: forall (l :: * -> *) (f :: * -> *).
Functor f =>
(NodeMap l ClassId -> f (NodeMap l ClassId))
-> EClass l -> f (EClass l)
_parents NodeMap l ClassId -> f (NodeMap l ClassId)
afa EClass{ClassId
Set (ENode l)
NodeMap l ClassId
Domain l
eClassId :: forall (l :: * -> *). EClass l -> ClassId
eClassNodes :: forall (l :: * -> *). EClass l -> Set (ENode l)
eClassData :: forall (l :: * -> *). EClass l -> Domain l
eClassParents :: forall (l :: * -> *). EClass l -> NodeMap l ClassId
eClassId :: ClassId
eClassNodes :: Set (ENode l)
eClassData :: Domain l
eClassParents :: NodeMap l ClassId
..} = ClassId
-> Set (ENode l) -> Domain l -> NodeMap l ClassId -> EClass l
forall (l :: * -> *).
ClassId
-> Set (ENode l) -> Domain l -> NodeMap l ClassId -> EClass l
EClass ClassId
eClassId Set (ENode l)
eClassNodes Domain l
eClassData (NodeMap l ClassId -> EClass l)
-> f (NodeMap l ClassId) -> f (EClass l)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> NodeMap l ClassId -> f (NodeMap l ClassId)
afa NodeMap l ClassId
eClassParents
{-# INLINE _parents #-}

-- | Lens for the e-nodes in an e-class
_nodes :: Lens' (EClass l) (S.Set (ENode l))
_nodes :: forall (l :: * -> *) (f :: * -> *).
Functor f =>
(Set (ENode l) -> f (Set (ENode l))) -> EClass l -> f (EClass l)
_nodes Set (ENode l) -> f (Set (ENode l))
afa EClass{ClassId
Set (ENode l)
NodeMap l ClassId
Domain l
eClassId :: forall (l :: * -> *). EClass l -> ClassId
eClassNodes :: forall (l :: * -> *). EClass l -> Set (ENode l)
eClassData :: forall (l :: * -> *). EClass l -> Domain l
eClassParents :: forall (l :: * -> *). EClass l -> NodeMap l ClassId
eClassId :: ClassId
eClassNodes :: Set (ENode l)
eClassData :: Domain l
eClassParents :: NodeMap l ClassId
..} = (\Set (ENode l)
ns -> ClassId
-> Set (ENode l) -> Domain l -> NodeMap l ClassId -> EClass l
forall (l :: * -> *).
ClassId
-> Set (ENode l) -> Domain l -> NodeMap l ClassId -> EClass l
EClass ClassId
eClassId Set (ENode l)
ns Domain l
eClassData NodeMap l ClassId
eClassParents) (Set (ENode l) -> EClass l) -> f (Set (ENode l)) -> f (EClass l)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set (ENode l) -> f (Set (ENode l))
afa Set (ENode l)
eClassNodes
{-# INLINE _nodes #-}

-- | Like @'view'@ but with the arguments flipped
(^.) :: s -> Lens' s a -> a
^. :: forall s a. s -> Lens' s a -> a
(^.) s
s Lens' s a
ln = Lens' s a -> s -> a
forall s a. Lens' s a -> s -> a
view (a -> f a) -> s -> f s
Lens' s a
ln s
s
infixl 8 ^.
{-# INLINE (^.) #-}

-- | Synonym for @'set'@
(.~) :: Lens' s a -> a -> (s -> s)
.~ :: forall s a. Lens' s a -> a -> s -> s
(.~) = Lens' s a -> a -> s -> s
forall s a. Lens' s a -> a -> s -> s
set
infixr 4 .~
{-# INLINE (.~) #-}

-- | Synonym for @'over'@
(%~) :: Lens' s a -> (a -> a) -> (s -> s)
%~ :: forall s a. Lens' s a -> (a -> a) -> s -> s
(%~) = Lens' s a -> (a -> a) -> s -> s
forall s a. Lens' s a -> (a -> a) -> s -> s
over
infixr 4 %~
{-# INLINE (%~) #-}

-- | Applies a getter to a value
view :: Lens' s a -> (s -> a)
view :: forall s a. Lens' s a -> s -> a
view Lens' s a
ln = Const a s -> a
forall {k} a (b :: k). Const a b -> a
getConst (Const a s -> a) -> (s -> Const a s) -> s -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Const a a) -> s -> Const a s
Lens' s a
ln a -> Const a a
forall {k} a (b :: k). a -> Const a b
Const
{-# INLINE view #-}

-- | Applies a setter to a value
set :: Lens' s a -> a -> (s -> s)
set :: forall s a. Lens' s a -> a -> s -> s
set Lens' s a
ln a
x = Lens' s a -> (a -> a) -> s -> s
forall s a. Lens' s a -> (a -> a) -> s -> s
over (a -> f a) -> s -> f s
Lens' s a
ln (a -> a -> a
forall a b. a -> b -> a
const a
x)
{-# INLINE set #-}

-- | Applies a function to the target
over :: Lens' s a -> (a -> a) -> (s -> s)
over :: forall s a. Lens' s a -> (a -> a) -> s -> s
over Lens' s a
ln a -> a
f = Identity s -> s
forall a. Identity a -> a
runIdentity (Identity s -> s) -> (s -> Identity s) -> s -> s
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Identity a) -> s -> Identity s
Lens' s a
ln (a -> Identity a
forall a. a -> Identity a
Identity (a -> Identity a) -> (a -> a) -> a -> Identity a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a
f)
{-# INLINE over #-}