{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE Rank2Types #-}
{-|
  Hand-rolled lenses on e-graphs and e-classes which come in quite handy, are
  heavily used in 'Data.Equality.Graph', and are the only exported way of
  editing the structure of the e-graph. If you want to write some complex
  'Analysis' you'll probably need these.
 -}
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.Utils.SizedList
import Data.Equality.Graph.Internal
import Data.Equality.Graph.Classes.Id
import Data.Equality.Graph.Nodes
import Data.Equality.Graph.Classes
import Data.Equality.Graph.ReprUnionFind
import Data.Equality.Analysis

-- | 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 -> ReprUnionFind -> ClassId
findRepr ClassId
i (forall (l :: * -> *). EGraph l -> ReprUnionFind
unionFind EGraph l
s)
     in (\EClass l
c' -> EGraph l
s { classes :: ClassIdMap (EClass l)
classes = forall a. ClassId -> a -> IntMap a -> IntMap a
IM.insert ClassId
canon_id EClass l
c' (forall (l :: * -> *). EGraph l -> ClassIdMap (EClass l)
classes EGraph l
s) }) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> EClass l -> f (EClass l)
afa (forall (l :: * -> *). EGraph l -> ClassIdMap (EClass l)
classes EGraph l
s forall a. IntMap a -> ClassId -> a
IM.! ClassId
canon_id)
{-# INLINE _class #-}

-- | Lens for the memo of e-nodes in an e-graph, that is, a mapping from
-- e-nodes to the e-class they're represented in
_memo :: Lens' (EGraph l) (NodeMap l ClassId)
_memo :: forall (l :: * -> *). Lens' (EGraph l) (NodeMap l ClassId)
_memo NodeMap l ClassId -> f (NodeMap l ClassId)
afa EGraph l
egr = (\NodeMap l ClassId
m1 -> EGraph l
egr {memo :: NodeMap l ClassId
memo = NodeMap l ClassId
m1}) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> NodeMap l ClassId -> f (NodeMap l ClassId)
afa (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 :: * -> *). Lens' (EGraph l) (ClassIdMap (EClass 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}) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ClassIdMap (EClass l) -> f (ClassIdMap (EClass l))
afa (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 :: * -> *). Lens' (EClass l) (Domain l)
_data Domain l -> f (Domain l)
afa EClass{ClassId
Set (ENode l)
Domain l
SList (ClassId, ENode l)
eClassParents :: forall (l :: * -> *). EClass l -> SList (ClassId, ENode l)
eClassData :: forall (l :: * -> *). EClass l -> Domain l
eClassNodes :: forall (l :: * -> *). EClass l -> Set (ENode l)
eClassId :: forall (l :: * -> *). EClass l -> ClassId
eClassParents :: SList (ClassId, ENode l)
eClassData :: Domain l
eClassNodes :: Set (ENode l)
eClassId :: ClassId
..} = (\Domain l
d1 -> forall (l :: * -> *).
ClassId
-> Set (ENode l)
-> Domain l
-> SList (ClassId, ENode l)
-> EClass l
EClass ClassId
eClassId Set (ENode l)
eClassNodes Domain l
d1 SList (ClassId, ENode l)
eClassParents) 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) (SList (ClassId, ENode l))
_parents :: forall (l :: * -> *). Lens' (EClass l) (SList (ClassId, ENode l))
_parents SList (ClassId, ENode l) -> f (SList (ClassId, ENode l))
afa EClass{ClassId
Set (ENode l)
Domain l
SList (ClassId, ENode l)
eClassParents :: SList (ClassId, ENode l)
eClassData :: Domain l
eClassNodes :: Set (ENode l)
eClassId :: ClassId
eClassParents :: forall (l :: * -> *). EClass l -> SList (ClassId, ENode l)
eClassData :: forall (l :: * -> *). EClass l -> Domain l
eClassNodes :: forall (l :: * -> *). EClass l -> Set (ENode l)
eClassId :: forall (l :: * -> *). EClass l -> ClassId
..} = (\SList (ClassId, ENode l)
ps -> forall (l :: * -> *).
ClassId
-> Set (ENode l)
-> Domain l
-> SList (ClassId, ENode l)
-> EClass l
EClass ClassId
eClassId Set (ENode l)
eClassNodes Domain l
eClassData SList (ClassId, ENode l)
ps) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SList (ClassId, ENode l) -> f (SList (ClassId, ENode l))
afa SList (ClassId, ENode l)
eClassParents
{-# INLINE _parents #-}

-- | Lens for the e-nodes in an e-class
_nodes :: Lens' (EClass l) (S.Set (ENode l))
_nodes :: forall (l :: * -> *). Lens' (EClass l) (Set (ENode l))
_nodes Set (ENode l) -> f (Set (ENode l))
afa EClass{ClassId
Set (ENode l)
Domain l
SList (ClassId, ENode l)
eClassParents :: SList (ClassId, ENode l)
eClassData :: Domain l
eClassNodes :: Set (ENode l)
eClassId :: ClassId
eClassParents :: forall (l :: * -> *). EClass l -> SList (ClassId, ENode l)
eClassData :: forall (l :: * -> *). EClass l -> Domain l
eClassNodes :: forall (l :: * -> *). EClass l -> Set (ENode l)
eClassId :: forall (l :: * -> *). EClass l -> ClassId
..} = (\Set (ENode l)
ns -> forall (l :: * -> *).
ClassId
-> Set (ENode l)
-> Domain l
-> SList (ClassId, ENode l)
-> EClass l
EClass ClassId
eClassId Set (ENode l)
ns Domain l
eClassData SList (ClassId, ENode l)
eClassParents) 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 = forall s a. Lens' s a -> s -> a
view 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
(.~) = 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
(%~) = 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 = forall {k} a (b :: k). Const a b -> a
getConst forall b c a. (b -> c) -> (a -> b) -> a -> c
. Lens' s a
ln 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 = forall s a. Lens' s a -> (a -> a) -> s -> s
over Lens' s a
ln (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 = forall a. Identity a -> a
runIdentity forall b c a. (b -> c) -> (a -> b) -> a -> c
. Lens' s a
ln (forall a. a -> Identity a
Identity forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a
f)
{-# INLINE over #-}