{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE UndecidableInstances #-}

-- | See 'EGraph'.
module Overeasy.EGraph
  ( EClassId (..)
  , ENodeId (..)
  , EAnalysis
  , noAnalysis
  , EClassInfo (..)
  , EGraph
  , WorkItem
  , WorkList
  , ClassReplacements
  , MergeResult (..)
  , egClassSource
  , egNodeSource
  , egEquivFind
  , egClassMap
  , egNodeAssoc
  , egHashCons
  , egClassSize
  , egNodeSize
  , egFindNode
  , egFindTerm
  , egClassInfo
  , egNew
  , egClasses
  , egCanonicalize
  , egCanonicalizePartial
  , egAddTerm
  , egMerge
  , egMergeMany
  , egReanalyzeSubset
  , egReanalyze
  ) where

import Control.DeepSeq (NFData)
import Control.Monad (unless, void)
import Control.Monad.State.Strict (State, gets, modify', state)
import Data.Foldable (foldl', toList)
import Data.Functor.Foldable (project)
import Data.Hashable (Hashable)
import Data.List.NonEmpty (NonEmpty (..))
import Data.Maybe (maybeToList)
import Data.Semigroup (sconcat)
import Data.Sequence (Seq (..))
import qualified Data.Sequence as Seq
import GHC.Generics (Generic)
import IntLike.Map (IntLikeMap (..))
import qualified IntLike.Map as ILM
import IntLike.Set (IntLikeSet (..))
import qualified IntLike.Set as ILS
import Overeasy.Assoc (Assoc, AssocInsertRes (..), assocCompactInc, assocFwd, assocInsertInc, assocLookupByValue,
                       assocMember, assocMembers, assocNew, assocPartialLookupByKey, assocRemoveAllInc, assocSingleton,
                       assocUnion)
import Overeasy.EquivFind (EquivFind (..), EquivMergeSetsRes (..), efAddInc, efCanonicalize, efCanonicalizePartial,
                           efClosure, efCompactInc, efFindRoot, efLookupRoot, efMergeSetsInc, efNew, efRootsSize,
                           efSubset)
import Overeasy.Source (Source, sourceAddInc, sourceNew)
import Overeasy.Util (Changed (..), RecursiveWhole, foldWholeM, stateFold)

-- | An opaque class id.
-- Constructor exported for coercibility.
-- Num instance for literals only.
newtype EClassId = EClassId { EClassId -> Int
unEClassId :: Int }
  deriving stock (Int -> EClassId -> ShowS
[EClassId] -> ShowS
EClassId -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [EClassId] -> ShowS
$cshowList :: [EClassId] -> ShowS
show :: EClassId -> String
$cshow :: EClassId -> String
showsPrec :: Int -> EClassId -> ShowS
$cshowsPrec :: Int -> EClassId -> ShowS
Show)
  deriving newtype (EClassId -> EClassId -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: EClassId -> EClassId -> Bool
$c/= :: EClassId -> EClassId -> Bool
== :: EClassId -> EClassId -> Bool
$c== :: EClassId -> EClassId -> Bool
Eq, Eq EClassId
EClassId -> EClassId -> Bool
EClassId -> EClassId -> Ordering
EClassId -> EClassId -> EClassId
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: EClassId -> EClassId -> EClassId
$cmin :: EClassId -> EClassId -> EClassId
max :: EClassId -> EClassId -> EClassId
$cmax :: EClassId -> EClassId -> EClassId
>= :: EClassId -> EClassId -> Bool
$c>= :: EClassId -> EClassId -> Bool
> :: EClassId -> EClassId -> Bool
$c> :: EClassId -> EClassId -> Bool
<= :: EClassId -> EClassId -> Bool
$c<= :: EClassId -> EClassId -> Bool
< :: EClassId -> EClassId -> Bool
$c< :: EClassId -> EClassId -> Bool
compare :: EClassId -> EClassId -> Ordering
$ccompare :: EClassId -> EClassId -> Ordering
Ord, Int -> EClassId
EClassId -> Int
EClassId -> [EClassId]
EClassId -> EClassId
EClassId -> EClassId -> [EClassId]
EClassId -> EClassId -> EClassId -> [EClassId]
forall a.
(a -> a)
-> (a -> a)
-> (Int -> a)
-> (a -> Int)
-> (a -> [a])
-> (a -> a -> [a])
-> (a -> a -> [a])
-> (a -> a -> a -> [a])
-> Enum a
enumFromThenTo :: EClassId -> EClassId -> EClassId -> [EClassId]
$cenumFromThenTo :: EClassId -> EClassId -> EClassId -> [EClassId]
enumFromTo :: EClassId -> EClassId -> [EClassId]
$cenumFromTo :: EClassId -> EClassId -> [EClassId]
enumFromThen :: EClassId -> EClassId -> [EClassId]
$cenumFromThen :: EClassId -> EClassId -> [EClassId]
enumFrom :: EClassId -> [EClassId]
$cenumFrom :: EClassId -> [EClassId]
fromEnum :: EClassId -> Int
$cfromEnum :: EClassId -> Int
toEnum :: Int -> EClassId
$ctoEnum :: Int -> EClassId
pred :: EClassId -> EClassId
$cpred :: EClassId -> EClassId
succ :: EClassId -> EClassId
$csucc :: EClassId -> EClassId
Enum, Eq EClassId
Int -> EClassId -> Int
EClassId -> Int
forall a. Eq a -> (Int -> a -> Int) -> (a -> Int) -> Hashable a
hash :: EClassId -> Int
$chash :: EClassId -> Int
hashWithSalt :: Int -> EClassId -> Int
$chashWithSalt :: Int -> EClassId -> Int
Hashable, EClassId -> ()
forall a. (a -> ()) -> NFData a
rnf :: EClassId -> ()
$crnf :: EClassId -> ()
NFData, Integer -> EClassId
EClassId -> EClassId
EClassId -> EClassId -> EClassId
forall a.
(a -> a -> a)
-> (a -> a -> a)
-> (a -> a -> a)
-> (a -> a)
-> (a -> a)
-> (a -> a)
-> (Integer -> a)
-> Num a
fromInteger :: Integer -> EClassId
$cfromInteger :: Integer -> EClassId
signum :: EClassId -> EClassId
$csignum :: EClassId -> EClassId
abs :: EClassId -> EClassId
$cabs :: EClassId -> EClassId
negate :: EClassId -> EClassId
$cnegate :: EClassId -> EClassId
* :: EClassId -> EClassId -> EClassId
$c* :: EClassId -> EClassId -> EClassId
- :: EClassId -> EClassId -> EClassId
$c- :: EClassId -> EClassId -> EClassId
+ :: EClassId -> EClassId -> EClassId
$c+ :: EClassId -> EClassId -> EClassId
Num)

-- | An opaque node id
-- Constructor exported for coercibility.
-- Num instance for literals only.
newtype ENodeId = ENodeId { ENodeId -> Int
unENodeId :: Int }
  deriving stock (Int -> ENodeId -> ShowS
[ENodeId] -> ShowS
ENodeId -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ENodeId] -> ShowS
$cshowList :: [ENodeId] -> ShowS
show :: ENodeId -> String
$cshow :: ENodeId -> String
showsPrec :: Int -> ENodeId -> ShowS
$cshowsPrec :: Int -> ENodeId -> ShowS
Show)
  deriving newtype (ENodeId -> ENodeId -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: ENodeId -> ENodeId -> Bool
$c/= :: ENodeId -> ENodeId -> Bool
== :: ENodeId -> ENodeId -> Bool
$c== :: ENodeId -> ENodeId -> Bool
Eq, Eq ENodeId
ENodeId -> ENodeId -> Bool
ENodeId -> ENodeId -> Ordering
ENodeId -> ENodeId -> ENodeId
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: ENodeId -> ENodeId -> ENodeId
$cmin :: ENodeId -> ENodeId -> ENodeId
max :: ENodeId -> ENodeId -> ENodeId
$cmax :: ENodeId -> ENodeId -> ENodeId
>= :: ENodeId -> ENodeId -> Bool
$c>= :: ENodeId -> ENodeId -> Bool
> :: ENodeId -> ENodeId -> Bool
$c> :: ENodeId -> ENodeId -> Bool
<= :: ENodeId -> ENodeId -> Bool
$c<= :: ENodeId -> ENodeId -> Bool
< :: ENodeId -> ENodeId -> Bool
$c< :: ENodeId -> ENodeId -> Bool
compare :: ENodeId -> ENodeId -> Ordering
$ccompare :: ENodeId -> ENodeId -> Ordering
Ord, Int -> ENodeId
ENodeId -> Int
ENodeId -> [ENodeId]
ENodeId -> ENodeId
ENodeId -> ENodeId -> [ENodeId]
ENodeId -> ENodeId -> ENodeId -> [ENodeId]
forall a.
(a -> a)
-> (a -> a)
-> (Int -> a)
-> (a -> Int)
-> (a -> [a])
-> (a -> a -> [a])
-> (a -> a -> [a])
-> (a -> a -> a -> [a])
-> Enum a
enumFromThenTo :: ENodeId -> ENodeId -> ENodeId -> [ENodeId]
$cenumFromThenTo :: ENodeId -> ENodeId -> ENodeId -> [ENodeId]
enumFromTo :: ENodeId -> ENodeId -> [ENodeId]
$cenumFromTo :: ENodeId -> ENodeId -> [ENodeId]
enumFromThen :: ENodeId -> ENodeId -> [ENodeId]
$cenumFromThen :: ENodeId -> ENodeId -> [ENodeId]
enumFrom :: ENodeId -> [ENodeId]
$cenumFrom :: ENodeId -> [ENodeId]
fromEnum :: ENodeId -> Int
$cfromEnum :: ENodeId -> Int
toEnum :: Int -> ENodeId
$ctoEnum :: Int -> ENodeId
pred :: ENodeId -> ENodeId
$cpred :: ENodeId -> ENodeId
succ :: ENodeId -> ENodeId
$csucc :: ENodeId -> ENodeId
Enum, Eq ENodeId
Int -> ENodeId -> Int
ENodeId -> Int
forall a. Eq a -> (Int -> a -> Int) -> (a -> Int) -> Hashable a
hash :: ENodeId -> Int
$chash :: ENodeId -> Int
hashWithSalt :: Int -> ENodeId -> Int
$chashWithSalt :: Int -> ENodeId -> Int
Hashable, ENodeId -> ()
forall a. (a -> ()) -> NFData a
rnf :: ENodeId -> ()
$crnf :: ENodeId -> ()
NFData, Integer -> ENodeId
ENodeId -> ENodeId
ENodeId -> ENodeId -> ENodeId
forall a.
(a -> a -> a)
-> (a -> a -> a)
-> (a -> a -> a)
-> (a -> a)
-> (a -> a)
-> (a -> a)
-> (Integer -> a)
-> Num a
fromInteger :: Integer -> ENodeId
$cfromInteger :: Integer -> ENodeId
signum :: ENodeId -> ENodeId
$csignum :: ENodeId -> ENodeId
abs :: ENodeId -> ENodeId
$cabs :: ENodeId -> ENodeId
negate :: ENodeId -> ENodeId
$cnegate :: ENodeId -> ENodeId
* :: ENodeId -> ENodeId -> ENodeId
$c* :: ENodeId -> ENodeId -> ENodeId
- :: ENodeId -> ENodeId -> ENodeId
$c- :: ENodeId -> ENodeId -> ENodeId
+ :: ENodeId -> ENodeId -> ENodeId
$c+ :: ENodeId -> ENodeId -> ENodeId
Num)

-- | The definition of an 'EGraph' analysis.
-- 'd' must be a join semilattice.
-- This function must be monotonic.
type EAnalysis d f = f d -> d

-- | A disabled analysis
noAnalysis :: EAnalysis () f
noAnalysis :: forall (f :: * -> *). EAnalysis () f
noAnalysis = forall a b. a -> b -> a
const ()

-- private
-- An internal triple of node, class, and data
data ENodeTriple d = ENodeTriple
  { forall d. ENodeTriple d -> ENodeId
entNode :: !ENodeId
  , forall d. ENodeTriple d -> EClassId
entClass :: !EClassId
  , forall d. ENodeTriple d -> d
entData :: !d
  } deriving stock (ENodeTriple d -> ENodeTriple d -> Bool
forall d. Eq d => ENodeTriple d -> ENodeTriple d -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: ENodeTriple d -> ENodeTriple d -> Bool
$c/= :: forall d. Eq d => ENodeTriple d -> ENodeTriple d -> Bool
== :: ENodeTriple d -> ENodeTriple d -> Bool
$c== :: forall d. Eq d => ENodeTriple d -> ENodeTriple d -> Bool
Eq, Int -> ENodeTriple d -> ShowS
forall d. Show d => Int -> ENodeTriple d -> ShowS
forall d. Show d => [ENodeTriple d] -> ShowS
forall d. Show d => ENodeTriple d -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ENodeTriple d] -> ShowS
$cshowList :: forall d. Show d => [ENodeTriple d] -> ShowS
show :: ENodeTriple d -> String
$cshow :: forall d. Show d => ENodeTriple d -> String
showsPrec :: Int -> ENodeTriple d -> ShowS
$cshowsPrec :: forall d. Show d => Int -> ENodeTriple d -> ShowS
Show, forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall d x. Rep (ENodeTriple d) x -> ENodeTriple d
forall d x. ENodeTriple d -> Rep (ENodeTriple d) x
$cto :: forall d x. Rep (ENodeTriple d) x -> ENodeTriple d
$cfrom :: forall d x. ENodeTriple d -> Rep (ENodeTriple d) x
Generic)
    deriving anyclass (forall d. NFData d => ENodeTriple d -> ()
forall a. (a -> ()) -> NFData a
rnf :: ENodeTriple d -> ()
$crnf :: forall d. NFData d => ENodeTriple d -> ()
NFData)

-- | Info stored for every class: analysis data, class members (nodes), and parent nodes.
data EClassInfo d f = EClassInfo
  { forall d (f :: * -> *). EClassInfo d f -> d
eciData :: !d
  , forall d (f :: * -> *). EClassInfo d f -> Assoc ENodeId (f ())
eciNodes :: !(Assoc ENodeId (f ()))
  , forall d (f :: * -> *). EClassInfo d f -> IntLikeSet ENodeId
eciParents :: !(IntLikeSet ENodeId)
  } deriving stock (forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall d (f :: * -> *) x. Rep (EClassInfo d f) x -> EClassInfo d f
forall d (f :: * -> *) x. EClassInfo d f -> Rep (EClassInfo d f) x
$cto :: forall d (f :: * -> *) x. Rep (EClassInfo d f) x -> EClassInfo d f
$cfrom :: forall d (f :: * -> *) x. EClassInfo d f -> Rep (EClassInfo d f) x
Generic)

deriving stock instance (Eq d, Eq (f ())) => Eq (EClassInfo d f)
deriving stock instance (Show d, Show (f ())) => Show (EClassInfo d f)
deriving anyclass instance (NFData d, NFData (f ())) => NFData (EClassInfo d f)

-- | A set of class ids to merge
type WorkItem = IntLikeSet EClassId

-- | A sequences of groups of class ids to mrege
type WorkList = Seq WorkItem

-- | An invertible multimap of new root class to the sets of old classes it subsumes
-- Can be used to externally recanonicalize any structures that reference class ids
-- after merges.
type ClassReplacements = EquivFind EClassId

-- | Merging classes can result in a few outcomes:
data MergeResult a =
    MergeResultUnchanged
  -- ^ All classes already merged, no change
  | MergeResultMissing !WorkItem
  -- ^ Some classes missing, returns first problematic worklist entry
  -- (note that not all classes in worklist item will be missing,
  -- only at least one from the set)
  | MergeResultChanged !a
  -- ^ Some classes merged, returns root map or merged class id
  deriving stock (MergeResult a -> MergeResult a -> Bool
forall a. Eq a => MergeResult a -> MergeResult a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: MergeResult a -> MergeResult a -> Bool
$c/= :: forall a. Eq a => MergeResult a -> MergeResult a -> Bool
== :: MergeResult a -> MergeResult a -> Bool
$c== :: forall a. Eq a => MergeResult a -> MergeResult a -> Bool
Eq, Int -> MergeResult a -> ShowS
forall a. Show a => Int -> MergeResult a -> ShowS
forall a. Show a => [MergeResult a] -> ShowS
forall a. Show a => MergeResult a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [MergeResult a] -> ShowS
$cshowList :: forall a. Show a => [MergeResult a] -> ShowS
show :: MergeResult a -> String
$cshow :: forall a. Show a => MergeResult a -> String
showsPrec :: Int -> MergeResult a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> MergeResult a -> ShowS
Show, forall a b. a -> MergeResult b -> MergeResult a
forall a b. (a -> b) -> MergeResult a -> MergeResult b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> MergeResult b -> MergeResult a
$c<$ :: forall a b. a -> MergeResult b -> MergeResult a
fmap :: forall a b. (a -> b) -> MergeResult a -> MergeResult b
$cfmap :: forall a b. (a -> b) -> MergeResult a -> MergeResult b
Functor, forall a. Eq a => a -> MergeResult a -> Bool
forall a. Num a => MergeResult a -> a
forall a. Ord a => MergeResult a -> a
forall m. Monoid m => MergeResult m -> m
forall a. MergeResult a -> Bool
forall a. MergeResult a -> Int
forall a. MergeResult a -> [a]
forall a. (a -> a -> a) -> MergeResult a -> a
forall m a. Monoid m => (a -> m) -> MergeResult a -> m
forall b a. (b -> a -> b) -> b -> MergeResult a -> b
forall a b. (a -> b -> b) -> b -> MergeResult a -> b
forall (t :: * -> *).
(forall m. Monoid m => t m -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. t a -> [a])
-> (forall a. t a -> Bool)
-> (forall a. t a -> Int)
-> (forall a. Eq a => a -> t a -> Bool)
-> (forall a. Ord a => t a -> a)
-> (forall a. Ord a => t a -> a)
-> (forall a. Num a => t a -> a)
-> (forall a. Num a => t a -> a)
-> Foldable t
product :: forall a. Num a => MergeResult a -> a
$cproduct :: forall a. Num a => MergeResult a -> a
sum :: forall a. Num a => MergeResult a -> a
$csum :: forall a. Num a => MergeResult a -> a
minimum :: forall a. Ord a => MergeResult a -> a
$cminimum :: forall a. Ord a => MergeResult a -> a
maximum :: forall a. Ord a => MergeResult a -> a
$cmaximum :: forall a. Ord a => MergeResult a -> a
elem :: forall a. Eq a => a -> MergeResult a -> Bool
$celem :: forall a. Eq a => a -> MergeResult a -> Bool
length :: forall a. MergeResult a -> Int
$clength :: forall a. MergeResult a -> Int
null :: forall a. MergeResult a -> Bool
$cnull :: forall a. MergeResult a -> Bool
toList :: forall a. MergeResult a -> [a]
$ctoList :: forall a. MergeResult a -> [a]
foldl1 :: forall a. (a -> a -> a) -> MergeResult a -> a
$cfoldl1 :: forall a. (a -> a -> a) -> MergeResult a -> a
foldr1 :: forall a. (a -> a -> a) -> MergeResult a -> a
$cfoldr1 :: forall a. (a -> a -> a) -> MergeResult a -> a
foldl' :: forall b a. (b -> a -> b) -> b -> MergeResult a -> b
$cfoldl' :: forall b a. (b -> a -> b) -> b -> MergeResult a -> b
foldl :: forall b a. (b -> a -> b) -> b -> MergeResult a -> b
$cfoldl :: forall b a. (b -> a -> b) -> b -> MergeResult a -> b
foldr' :: forall a b. (a -> b -> b) -> b -> MergeResult a -> b
$cfoldr' :: forall a b. (a -> b -> b) -> b -> MergeResult a -> b
foldr :: forall a b. (a -> b -> b) -> b -> MergeResult a -> b
$cfoldr :: forall a b. (a -> b -> b) -> b -> MergeResult a -> b
foldMap' :: forall m a. Monoid m => (a -> m) -> MergeResult a -> m
$cfoldMap' :: forall m a. Monoid m => (a -> m) -> MergeResult a -> m
foldMap :: forall m a. Monoid m => (a -> m) -> MergeResult a -> m
$cfoldMap :: forall m a. Monoid m => (a -> m) -> MergeResult a -> m
fold :: forall m. Monoid m => MergeResult m -> m
$cfold :: forall m. Monoid m => MergeResult m -> m
Foldable, Functor MergeResult
Foldable MergeResult
forall (t :: * -> *).
Functor t
-> Foldable t
-> (forall (f :: * -> *) a b.
    Applicative f =>
    (a -> f b) -> t a -> f (t b))
-> (forall (f :: * -> *) a. Applicative f => t (f a) -> f (t a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> t a -> m (t b))
-> (forall (m :: * -> *) a. Monad m => t (m a) -> m (t a))
-> Traversable t
forall (m :: * -> *) a.
Monad m =>
MergeResult (m a) -> m (MergeResult a)
forall (f :: * -> *) a.
Applicative f =>
MergeResult (f a) -> f (MergeResult a)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> MergeResult a -> m (MergeResult b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> MergeResult a -> f (MergeResult b)
sequence :: forall (m :: * -> *) a.
Monad m =>
MergeResult (m a) -> m (MergeResult a)
$csequence :: forall (m :: * -> *) a.
Monad m =>
MergeResult (m a) -> m (MergeResult a)
mapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> MergeResult a -> m (MergeResult b)
$cmapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> MergeResult a -> m (MergeResult b)
sequenceA :: forall (f :: * -> *) a.
Applicative f =>
MergeResult (f a) -> f (MergeResult a)
$csequenceA :: forall (f :: * -> *) a.
Applicative f =>
MergeResult (f a) -> f (MergeResult a)
traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> MergeResult a -> f (MergeResult b)
$ctraverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> MergeResult a -> f (MergeResult b)
Traversable)

-- | An E-Graph implementation
data EGraph d f = EGraph
  { forall d (f :: * -> *). EGraph d f -> Source EClassId
egClassSource :: !(Source EClassId)
  -- ^ Id source for classes
  , forall d (f :: * -> *). EGraph d f -> Source ENodeId
egNodeSource :: !(Source ENodeId)
  -- ^ Id source for nodes
  , forall d (f :: * -> *). EGraph d f -> EquivFind EClassId
egEquivFind :: !(EquivFind EClassId)
  -- ^ Class equivalences
  , forall d (f :: * -> *).
EGraph d f -> IntLikeMap EClassId (EClassInfo d f)
egClassMap :: !(IntLikeMap EClassId (EClassInfo d f))
  -- ^ Map of class to info
  -- Invariant: Only contains root classes.
  , forall d (f :: * -> *). EGraph d f -> Assoc ENodeId (f EClassId)
egNodeAssoc :: !(Assoc ENodeId (f EClassId))
  -- ^ Assoc of node id to node structure
  -- Invariant: only contains canonical structures (with root classes).
  , forall d (f :: * -> *). EGraph d f -> IntLikeMap ENodeId EClassId
egHashCons :: !(IntLikeMap ENodeId EClassId)
  -- ^ Map of node to class
  -- Invariant: only contains root classes.
  } deriving stock (forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall d (f :: * -> *) x. Rep (EGraph d f) x -> EGraph d f
forall d (f :: * -> *) x. EGraph d f -> Rep (EGraph d f) x
$cto :: forall d (f :: * -> *) x. Rep (EGraph d f) x -> EGraph d f
$cfrom :: forall d (f :: * -> *) x. EGraph d f -> Rep (EGraph d f) x
Generic)

deriving stock instance (Eq d, Eq (f EClassId), Eq (f ())) => Eq (EGraph d f)
deriving stock instance (Show d, Show (f EClassId), Show (f ())) => Show (EGraph d f)
deriving anyclass instance (NFData d, NFData (f EClassId), NFData (f ())) => NFData (EGraph d f)

-- | Number of equivalent classes in the 'EGraph' (see 'ufSize')
egClassSize :: EGraph d f -> Int
egClassSize :: forall d (f :: * -> *). EGraph d f -> Int
egClassSize = forall x. EquivFind x -> Int
efRootsSize forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall d (f :: * -> *). EGraph d f -> EquivFind EClassId
egEquivFind

-- | Number of nodes in the 'EGraph'
egNodeSize :: EGraph d f -> Int
egNodeSize :: forall d (f :: * -> *). EGraph d f -> Int
egNodeSize = forall x a. IntLikeMap x a -> Int
ILM.size forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall d (f :: * -> *). EGraph d f -> IntLikeMap ENodeId EClassId
egHashCons

-- | Lookup info for the given 'EClass'
egClassInfo :: EClassId -> EGraph d f -> Maybe (EClassInfo d f)
egClassInfo :: forall d (f :: * -> *).
EClassId -> EGraph d f -> Maybe (EClassInfo d f)
egClassInfo EClassId
c = forall x a. Coercible x Int => x -> IntLikeMap x a -> Maybe a
ILM.lookup EClassId
c forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall d (f :: * -> *).
EGraph d f -> IntLikeMap EClassId (EClassInfo d f)
egClassMap

-- | Find the class of the given node, if it exists.
-- Note that you may have to canonicalize first to find it!
egFindNode :: (Eq (f EClassId), Hashable (f EClassId)) => f EClassId -> EGraph d f -> Maybe EClassId
egFindNode :: forall (f :: * -> *) d.
(Eq (f EClassId), Hashable (f EClassId)) =>
f EClassId -> EGraph d f -> Maybe EClassId
egFindNode f EClassId
fc EGraph d f
eg = do
  ENodeId
n <- forall a x. (Eq a, Hashable a) => a -> Assoc x a -> Maybe x
assocLookupByValue f EClassId
fc (forall d (f :: * -> *). EGraph d f -> Assoc ENodeId (f EClassId)
egNodeAssoc EGraph d f
eg)
  forall x a. Coercible x Int => x -> IntLikeMap x a -> Maybe a
ILM.lookup ENodeId
n (forall d (f :: * -> *). EGraph d f -> IntLikeMap ENodeId EClassId
egHashCons EGraph d f
eg)

-- | Find the class of the given term, if it exists
egFindTerm :: (RecursiveWhole t f, Traversable f, Eq (f EClassId), Hashable (f EClassId)) => t -> EGraph d f -> Maybe EClassId
egFindTerm :: forall t (f :: * -> *) d.
(RecursiveWhole t f, Traversable f, Eq (f EClassId),
 Hashable (f EClassId)) =>
t -> EGraph d f -> Maybe EClassId
egFindTerm t
t EGraph d f
eg = forall t (f :: * -> *) (m :: * -> *) a.
(RecursiveWhole t f, Traversable f, Monad m) =>
(f a -> m a) -> t -> m a
foldWholeM (forall (f :: * -> *) d.
(Eq (f EClassId), Hashable (f EClassId)) =>
f EClassId -> EGraph d f -> Maybe EClassId
`egFindNode` EGraph d f
eg) t
t

-- | Creates a new 'EGraph'
egNew :: EGraph d f
egNew :: forall d (f :: * -> *). EGraph d f
egNew = forall d (f :: * -> *).
Source EClassId
-> Source ENodeId
-> EquivFind EClassId
-> IntLikeMap EClassId (EClassInfo d f)
-> Assoc ENodeId (f EClassId)
-> IntLikeMap ENodeId EClassId
-> EGraph d f
EGraph (forall x. Coercible x Int => x -> Source x
sourceNew (Int -> EClassId
EClassId Int
0)) (forall x. Coercible x Int => x -> Source x
sourceNew (Int -> ENodeId
ENodeId Int
0)) forall x. EquivFind x
efNew forall x a. IntLikeMap x a
ILM.empty forall x a. Assoc x a
assocNew forall x a. IntLikeMap x a
ILM.empty

-- | Yields all root classes
egClasses :: State (EGraph d f) (IntLikeSet EClassId)
egClasses :: forall d (f :: * -> *). State (EGraph d f) WorkItem
egClasses = forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (forall x a. IntLikeMap x a -> IntLikeSet x
ILM.keysSet forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall d (f :: * -> *).
EGraph d f -> IntLikeMap EClassId (EClassInfo d f)
egClassMap)

-- | Find the canonical form of a node.
-- If any classes are missing, the first missing is returned.
egCanonicalize :: Traversable f => f EClassId -> State (EGraph d f) (Either EClassId (f EClassId))
egCanonicalize :: forall (f :: * -> *) d.
Traversable f =>
f EClassId -> State (EGraph d f) (Either EClassId (f EClassId))
egCanonicalize f EClassId
fc = forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (forall (f :: * -> *) x.
(Traversable f, Coercible x Int) =>
f x -> EquivFind x -> Either x (f x)
efCanonicalize f EClassId
fc forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall d (f :: * -> *). EGraph d f -> EquivFind EClassId
egEquivFind)

-- | Find the canonical form of a node.
-- If any classes are missing, simply skip them.
egCanonicalizePartial :: Traversable f => f EClassId -> State (EGraph d f) (f EClassId)
egCanonicalizePartial :: forall (f :: * -> *) d.
Traversable f =>
f EClassId -> State (EGraph d f) (f EClassId)
egCanonicalizePartial f EClassId
fc = forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (forall (f :: * -> *) x.
(Functor f, Coercible x Int) =>
f x -> EquivFind x -> f x
efCanonicalizePartial f EClassId
fc forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall d (f :: * -> *). EGraph d f -> EquivFind EClassId
egEquivFind)

-- private
-- Variant of canonicalization used in rebuilding - updates node assocs and returns
-- 1. New canonical node id (could be the same)
-- 2. Maybe a set of node ids to delete (no longer canonical)
egCanonicalizeInternal :: (Traversable f, Eq (f EClassId), Hashable (f EClassId)) => ENodeId -> State (EGraph d f) (ENodeId, Maybe (IntLikeSet ENodeId))
egCanonicalizeInternal :: forall (f :: * -> *) d.
(Traversable f, Eq (f EClassId), Hashable (f EClassId)) =>
ENodeId -> State (EGraph d f) (ENodeId, Maybe (IntLikeSet ENodeId))
egCanonicalizeInternal ENodeId
x = forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state forall a b. (a -> b) -> a -> b
$ \EGraph d f
eg ->
  let ef :: EquivFind EClassId
ef = forall d (f :: * -> *). EGraph d f -> EquivFind EClassId
egEquivFind EGraph d f
eg
      assoc :: Assoc ENodeId (f EClassId)
assoc = forall d (f :: * -> *). EGraph d f -> Assoc ENodeId (f EClassId)
egNodeAssoc EGraph d f
eg
      node :: f EClassId
node = forall x a. Coercible x Int => x -> Assoc x a -> a
assocPartialLookupByKey ENodeId
x Assoc ENodeId (f EClassId)
assoc
      fz :: f EClassId
fz = forall (f :: * -> *) x.
(Functor f, Coercible x Int) =>
f x -> EquivFind x -> f x
efCanonicalizePartial f EClassId
node EquivFind EClassId
ef
      ((ENodeId
y, AssocInsertRes ENodeId
res), Assoc ENodeId (f EClassId)
assoc') = forall x a.
(Coercible x Int, Ord x, Eq a, Hashable a) =>
x -> a -> Assoc x a -> ((x, AssocInsertRes x), Assoc x a)
assocInsertInc ENodeId
x f EClassId
fz Assoc ENodeId (f EClassId)
assoc
  in case AssocInsertRes ENodeId
res of
    AssocInsertRes ENodeId
AssocInsertResUnchanged -> ((ENodeId
y, forall a. Maybe a
Nothing), EGraph d f
eg)
    AssocInsertResMerged IntLikeSet ENodeId
toDelete ->
      ((ENodeId
y, forall a. a -> Maybe a
Just IntLikeSet ENodeId
toDelete), EGraph d f
eg { egNodeAssoc :: Assoc ENodeId (f EClassId)
egNodeAssoc = Assoc ENodeId (f EClassId)
assoc' })
    AssocInsertRes ENodeId
_ -> ((ENodeId
y, forall a. Maybe a
Nothing), EGraph d f
eg { egNodeAssoc :: Assoc ENodeId (f EClassId)
egNodeAssoc = Assoc ENodeId (f EClassId)
assoc' })

-- private
data AddNodeRes d = AddNodeRes !Changed !(Seq (ENodeTriple d))
  deriving stock (AddNodeRes d -> AddNodeRes d -> Bool
forall d. Eq d => AddNodeRes d -> AddNodeRes d -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: AddNodeRes d -> AddNodeRes d -> Bool
$c/= :: forall d. Eq d => AddNodeRes d -> AddNodeRes d -> Bool
== :: AddNodeRes d -> AddNodeRes d -> Bool
$c== :: forall d. Eq d => AddNodeRes d -> AddNodeRes d -> Bool
Eq, Int -> AddNodeRes d -> ShowS
forall d. Show d => Int -> AddNodeRes d -> ShowS
forall d. Show d => [AddNodeRes d] -> ShowS
forall d. Show d => AddNodeRes d -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [AddNodeRes d] -> ShowS
$cshowList :: forall d. Show d => [AddNodeRes d] -> ShowS
show :: AddNodeRes d -> String
$cshow :: forall d. Show d => AddNodeRes d -> String
showsPrec :: Int -> AddNodeRes d -> ShowS
$cshowsPrec :: forall d. Show d => Int -> AddNodeRes d -> ShowS
Show, forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall d x. Rep (AddNodeRes d) x -> AddNodeRes d
forall d x. AddNodeRes d -> Rep (AddNodeRes d) x
$cto :: forall d x. Rep (AddNodeRes d) x -> AddNodeRes d
$cfrom :: forall d x. AddNodeRes d -> Rep (AddNodeRes d) x
Generic)
  deriving anyclass (forall d. NFData d => AddNodeRes d -> ()
forall a. (a -> ()) -> NFData a
rnf :: AddNodeRes d -> ()
$crnf :: forall d. NFData d => AddNodeRes d -> ()
NFData)

instance Semigroup (AddNodeRes d) where
  AddNodeRes Changed
c1 Seq (ENodeTriple d)
p1 <> :: AddNodeRes d -> AddNodeRes d -> AddNodeRes d
<> AddNodeRes Changed
c2 Seq (ENodeTriple d)
p2 = forall d. Changed -> Seq (ENodeTriple d) -> AddNodeRes d
AddNodeRes (Changed
c1 forall a. Semigroup a => a -> a -> a
<> Changed
c2) (Seq (ENodeTriple d)
p1 forall a. Semigroup a => a -> a -> a
<> Seq (ENodeTriple d)
p2)

instance Monoid (AddNodeRes d) where
  mempty :: AddNodeRes d
mempty = forall d. Changed -> Seq (ENodeTriple d) -> AddNodeRes d
AddNodeRes Changed
ChangedNo forall a. Seq a
Seq.empty
  mappend :: AddNodeRes d -> AddNodeRes d -> AddNodeRes d
mappend = forall a. Semigroup a => a -> a -> a
(<>)

-- private
egAddNodeSub :: (Functor f, Eq (f EClassId), Hashable (f EClassId), Hashable (f ())) => EAnalysis d f -> f (ENodeTriple d) -> State (EGraph d f) (Changed, ENodeTriple d)
egAddNodeSub :: forall (f :: * -> *) d.
(Functor f, Eq (f EClassId), Hashable (f EClassId),
 Hashable (f ())) =>
EAnalysis d f
-> f (ENodeTriple d) -> State (EGraph d f) (Changed, ENodeTriple d)
egAddNodeSub EAnalysis d f
ana f (ENodeTriple d)
fcd = do
  let fc :: f EClassId
fc = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall d. ENodeTriple d -> EClassId
entClass f (ENodeTriple d)
fcd
  -- important: node should already be canonicalized!
  -- first lookup the node in the assoc to ensure uniqueness
  Maybe ENodeId
mayNodeId <- forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (forall a x. (Eq a, Hashable a) => a -> Assoc x a -> Maybe x
assocLookupByValue f EClassId
fc forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall d (f :: * -> *). EGraph d f -> Assoc ENodeId (f EClassId)
egNodeAssoc)
  case Maybe ENodeId
mayNodeId of
    Just ENodeId
n -> do
      EClassId
x <- forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (forall x a. Coercible x Int => x -> IntLikeMap x a -> a
ILM.partialLookup ENodeId
n forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall d (f :: * -> *). EGraph d f -> IntLikeMap ENodeId EClassId
egHashCons)
      EClassInfo d f
eci <- forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (forall x a. Coercible x Int => x -> IntLikeMap x a -> a
ILM.partialLookup EClassId
x forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall d (f :: * -> *).
EGraph d f -> IntLikeMap EClassId (EClassInfo d f)
egClassMap)
      let d :: d
d = forall d (f :: * -> *). EClassInfo d f -> d
eciData EClassInfo d f
eci
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (Changed
ChangedNo, forall d. ENodeId -> EClassId -> d -> ENodeTriple d
ENodeTriple ENodeId
n EClassId
x d
d)
    Maybe ENodeId
Nothing -> forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state forall a b. (a -> b) -> a -> b
$ \EGraph d f
eg ->
      -- node does not exist; get new node and class ids
      let (ENodeId
n, Source ENodeId
nodeSource') = forall x. Coercible x Int => Source x -> (x, Source x)
sourceAddInc (forall d (f :: * -> *). EGraph d f -> Source ENodeId
egNodeSource EGraph d f
eg)
          (EClassId
x, Source EClassId
classSource') = forall x. Coercible x Int => Source x -> (x, Source x)
sourceAddInc (forall d (f :: * -> *). EGraph d f -> Source EClassId
egClassSource EGraph d f
eg)
          -- add it to the uf (can discard return value since it's a new id, will be the same)
          (EquivAddRes EClassId
_, EquivFind EClassId
ef') = forall x.
Coercible x Int =>
x -> EquivFind x -> (EquivAddRes x, EquivFind x)
efAddInc EClassId
x (forall d (f :: * -> *). EGraph d f -> EquivFind EClassId
egEquivFind EGraph d f
eg)
          -- add it to the assoc (ignore and partial by construction)
          ((ENodeId, AssocInsertRes ENodeId)
_, Assoc ENodeId (f EClassId)
assoc') = forall x a.
(Coercible x Int, Ord x, Eq a, Hashable a) =>
x -> a -> Assoc x a -> ((x, AssocInsertRes x), Assoc x a)
assocInsertInc ENodeId
n f EClassId
fc (forall d (f :: * -> *). EGraph d f -> Assoc ENodeId (f EClassId)
egNodeAssoc EGraph d f
eg)
          -- insert into the hashcons
          hc' :: IntLikeMap ENodeId EClassId
hc' = forall x a.
Coercible x Int =>
x -> a -> IntLikeMap x a -> IntLikeMap x a
ILM.insert ENodeId
n EClassId
x (forall d (f :: * -> *). EGraph d f -> IntLikeMap ENodeId EClassId
egHashCons EGraph d f
eg)
          -- analyze the node and put that info in the class map
          d :: d
d = EAnalysis d f
ana (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall d. ENodeTriple d -> d
entData f (ENodeTriple d)
fcd)
          eci :: EClassInfo d f
eci = forall d (f :: * -> *).
d -> Assoc ENodeId (f ()) -> IntLikeSet ENodeId -> EClassInfo d f
EClassInfo d
d (forall x a. (Coercible x Int, Hashable a) => x -> a -> Assoc x a
assocSingleton ENodeId
n (forall (f :: * -> *) a. Functor f => f a -> f ()
void f (ENodeTriple d)
fcd)) forall x. IntLikeSet x
ILS.empty
          classMap' :: IntLikeMap EClassId (EClassInfo d f)
classMap' = forall x a.
Coercible x Int =>
x -> a -> IntLikeMap x a -> IntLikeMap x a
ILM.insert EClassId
x EClassInfo d f
eci (forall d (f :: * -> *).
EGraph d f -> IntLikeMap EClassId (EClassInfo d f)
egClassMap EGraph d f
eg)
          eg' :: EGraph d f
eg' = EGraph d f
eg { egNodeSource :: Source ENodeId
egNodeSource = Source ENodeId
nodeSource', egClassSource :: Source EClassId
egClassSource = Source EClassId
classSource', egEquivFind :: EquivFind EClassId
egEquivFind = EquivFind EClassId
ef', egNodeAssoc :: Assoc ENodeId (f EClassId)
egNodeAssoc = Assoc ENodeId (f EClassId)
assoc', egHashCons :: IntLikeMap ENodeId EClassId
egHashCons = IntLikeMap ENodeId EClassId
hc', egClassMap :: IntLikeMap EClassId (EClassInfo d f)
egClassMap = IntLikeMap EClassId (EClassInfo d f)
classMap' }
      in ((Changed
ChangedYes, forall d. ENodeId -> EClassId -> d -> ENodeTriple d
ENodeTriple ENodeId
n EClassId
x d
d), EGraph d f
eg')

-- private
egAddTermSub :: (RecursiveWhole t f, Traversable f, Eq (f EClassId), Hashable (f EClassId), Hashable (f ())) => EAnalysis d f -> t -> State (EGraph d f) (AddNodeRes d, ENodeTriple d)
egAddTermSub :: forall t (f :: * -> *) d.
(RecursiveWhole t f, Traversable f, Eq (f EClassId),
 Hashable (f EClassId), Hashable (f ())) =>
EAnalysis d f
-> t -> State (EGraph d f) (AddNodeRes d, ENodeTriple d)
egAddTermSub EAnalysis d f
ana = t -> StateT (EGraph d f) Identity (AddNodeRes d, ENodeTriple d)
go where
  go :: t -> StateT (EGraph d f) Identity (AddNodeRes d, ENodeTriple d)
go t
t = do
    -- unwrap to work with the functor layer
    let ft :: Base t t
ft = forall t. Recursive t => t -> Base t t
project t
t
    -- add all child nodes
    f (AddNodeRes d, ENodeTriple d)
frx <- forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse t -> StateT (EGraph d f) Identity (AddNodeRes d, ENodeTriple d)
go Base t t
ft
    -- collect info generated from child nodes and leave pure structure
    let (AddNodeRes Changed
changed1 Seq (ENodeTriple d)
children, f (ENodeTriple d)
fx) = forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
sequenceA f (AddNodeRes d, ENodeTriple d)
frx
    -- now fx should be canonicalized by construction
    -- add the node to get its node and class ids
    (Changed
changed2, z :: ENodeTriple d
z@(ENodeTriple ENodeId
n EClassId
_ d
_)) <- forall (f :: * -> *) d.
(Functor f, Eq (f EClassId), Hashable (f EClassId),
 Hashable (f ())) =>
EAnalysis d f
-> f (ENodeTriple d) -> State (EGraph d f) (Changed, ENodeTriple d)
egAddNodeSub EAnalysis d f
ana f (ENodeTriple d)
fx
    -- now update all its children to add this as a parent
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (forall a. Seq a -> Bool
Seq.null Seq (ENodeTriple d)
children) forall a b. (a -> b) -> a -> b
$
      forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify' forall a b. (a -> b) -> a -> b
$ \EGraph d f
eg ->
        -- Add node to class parents (unless it's a self parent)
        let cm' :: IntLikeMap EClassId (EClassInfo d f)
cm' = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\IntLikeMap EClassId (EClassInfo d f)
cm (ENodeTriple ENodeId
_ EClassId
c d
_) -> forall x a.
Coercible x Int =>
(a -> a) -> x -> IntLikeMap x a -> IntLikeMap x a
ILM.adjust (\EClassInfo d f
v -> if forall x a. Coercible x Int => x -> Assoc x a -> Bool
assocMember ENodeId
n (forall d (f :: * -> *). EClassInfo d f -> Assoc ENodeId (f ())
eciNodes EClassInfo d f
v) then EClassInfo d f
v else EClassInfo d f
v { eciParents :: IntLikeSet ENodeId
eciParents = forall x. Coercible x Int => x -> IntLikeSet x -> IntLikeSet x
ILS.insert ENodeId
n (forall d (f :: * -> *). EClassInfo d f -> IntLikeSet ENodeId
eciParents EClassInfo d f
v) }) EClassId
c IntLikeMap EClassId (EClassInfo d f)
cm) (forall d (f :: * -> *).
EGraph d f -> IntLikeMap EClassId (EClassInfo d f)
egClassMap EGraph d f
eg) Seq (ENodeTriple d)
children
        in EGraph d f
eg { egClassMap :: IntLikeMap EClassId (EClassInfo d f)
egClassMap = IntLikeMap EClassId (EClassInfo d f)
cm' }
    forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall d. Changed -> Seq (ENodeTriple d) -> AddNodeRes d
AddNodeRes (Changed
changed1 forall a. Semigroup a => a -> a -> a
<> Changed
changed2) (forall a. a -> Seq a
Seq.singleton ENodeTriple d
z), ENodeTriple d
z)

-- | Adds a term (recursively) to the graph. If already in the graph, returns 'ChangedNo' and existing class id. Otherwise
-- returns 'ChangedYes' and a new class id.
egAddTerm :: (RecursiveWhole t f, Traversable f, Eq (f EClassId), Hashable (f EClassId), Hashable (f ())) => EAnalysis d f -> t -> State (EGraph d f) (Changed, EClassId)
egAddTerm :: forall t (f :: * -> *) d.
(RecursiveWhole t f, Traversable f, Eq (f EClassId),
 Hashable (f EClassId), Hashable (f ())) =>
EAnalysis d f -> t -> State (EGraph d f) (Changed, EClassId)
egAddTerm EAnalysis d f
ana t
t = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(AddNodeRes Changed
c Seq (ENodeTriple d)
_, ENodeTriple ENodeId
_ EClassId
x d
_) -> (Changed
c, EClassId
x)) (forall t (f :: * -> *) d.
(RecursiveWhole t f, Traversable f, Eq (f EClassId),
 Hashable (f EClassId), Hashable (f ())) =>
EAnalysis d f
-> t -> State (EGraph d f) (AddNodeRes d, ENodeTriple d)
egAddTermSub EAnalysis d f
ana t
t)

-- | Merges two classes:
-- Returns 'Nothing' if the classes are not found or if they're already equal.
-- Otherwise returns the class remapping.
-- Note that it's MUCH more efficient to accumulate a 'WorkList' and use 'egMergeMany'.
egMerge :: (Semigroup d, Traversable f, Eq (f EClassId), Hashable (f EClassId), Eq (f ()), Hashable (f ()))
  => EClassId -> EClassId -> State (EGraph d f) (MergeResult EClassId)
egMerge :: forall d (f :: * -> *).
(Semigroup d, Traversable f, Eq (f EClassId),
 Hashable (f EClassId), Eq (f ()), Hashable (f ())) =>
EClassId -> EClassId -> State (EGraph d f) (MergeResult EClassId)
egMerge EClassId
i EClassId
j = do
  MergeResult (EquivFind EClassId, WorkItem)
mr <- forall d (f :: * -> *).
(Semigroup d, Traversable f, Eq (f EClassId),
 Hashable (f EClassId), Eq (f ()), Hashable (f ())) =>
WorkList
-> State (EGraph d f) (MergeResult (EquivFind EClassId, WorkItem))
egMergeMany (forall a. a -> Seq a
Seq.singleton (forall x. Coercible x Int => [x] -> IntLikeSet x
ILS.fromList [EClassId
i, EClassId
j]))
  -- We're guaranteed to have one and only one root in the map, so this won't fail
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. [a] -> a
head forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall x a. Coercible x Int => IntLikeMap x a -> [(x, a)]
ILM.toList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall x. EquivFind x -> IntLikeMap x (IntLikeSet x)
efFwd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) MergeResult (EquivFind EClassId, WorkItem)
mr)

-- private
data BuildWorkResult a =
    BuildWorkResultUnchanged
  | BuildWorkResultMissing !WorkItem
  | BuildWorkResultOk !a

-- private
egBuildWorkItem :: WorkItem -> State (EGraph d f) (BuildWorkResult WorkItem)
egBuildWorkItem :: forall d (f :: * -> *).
WorkItem -> State (EGraph d f) (BuildWorkResult WorkItem)
egBuildWorkItem WorkItem
cs = do
  Maybe [EClassId]
mayRoots <- forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\EquivFind EClassId
ef -> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (forall x. Coercible x Int => x -> EquivFind x -> Maybe x
`efFindRoot` EquivFind EClassId
ef) (forall x. Coercible x Int => IntLikeSet x -> [x]
ILS.toList WorkItem
cs)) (forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets forall d (f :: * -> *). EGraph d f -> EquivFind EClassId
egEquivFind)
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$! case Maybe [EClassId]
mayRoots of
    Maybe [EClassId]
Nothing -> forall a. WorkItem -> BuildWorkResult a
BuildWorkResultMissing WorkItem
cs
    Just [EClassId]
roots ->
      let rootsSet :: WorkItem
rootsSet = forall x. Coercible x Int => [x] -> IntLikeSet x
ILS.fromList [EClassId]
roots
      in if forall x. IntLikeSet x -> Int
ILS.size WorkItem
rootsSet forall a. Ord a => a -> a -> Bool
< Int
2
        then forall a. BuildWorkResult a
BuildWorkResultUnchanged
        else forall a. a -> BuildWorkResult a
BuildWorkResultOk WorkItem
rootsSet

-- private
egBuildWorklist :: WorkList -> State (EGraph d f) (BuildWorkResult WorkList)
egBuildWorklist :: forall d (f :: * -> *).
WorkList -> State (EGraph d f) (BuildWorkResult WorkList)
egBuildWorklist = forall {d} {f :: * -> *}.
WorkList
-> WorkList
-> StateT (EGraph d f) Identity (BuildWorkResult WorkList)
go forall a. Seq a
Empty where
  go :: WorkList
-> WorkList
-> StateT (EGraph d f) Identity (BuildWorkResult WorkList)
go !WorkList
acc = \case
    WorkList
Empty ->
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$! if forall a. Seq a -> Bool
Seq.null WorkList
acc
        then forall a. BuildWorkResult a
BuildWorkResultUnchanged
        else forall a. a -> BuildWorkResult a
BuildWorkResultOk WorkList
acc
    WorkItem
cs :<| WorkList
wl' -> do
      BuildWorkResult WorkItem
rcs <- forall d (f :: * -> *).
WorkItem -> State (EGraph d f) (BuildWorkResult WorkItem)
egBuildWorkItem WorkItem
cs
      case BuildWorkResult WorkItem
rcs of
        BuildWorkResult WorkItem
BuildWorkResultUnchanged -> WorkList
-> WorkList
-> StateT (EGraph d f) Identity (BuildWorkResult WorkList)
go WorkList
acc WorkList
wl'
        BuildWorkResultMissing WorkItem
cs' -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. WorkItem -> BuildWorkResult a
BuildWorkResultMissing WorkItem
cs')
        BuildWorkResultOk WorkItem
cs' -> WorkList
-> WorkList
-> StateT (EGraph d f) Identity (BuildWorkResult WorkList)
go (WorkList
acc forall a. Seq a -> a -> Seq a
:|> WorkItem
cs') WorkList
wl'

-- | Merges many sets of classes.
-- Returns 'Nothing' if the classes are not found or if they're already equal.
-- Otherwise returns the class remapping (equiv map of root to set of leaf classes).
-- It is important to note that the leaf classes in the returned mapping have been
-- REMOVED from the egraph, so they cannot be used to lookup classes in the future.
-- Therefore, if you have any class ids stored externally, you'll want to (partially)
-- canonicalize with the returned mapping.
-- Also note that the analysis of a given class is going to be an UNDER-APPROXIMATION
-- of the true analysis value, because per-node analyses are not recomputed.
egMergeMany :: (Semigroup d, Traversable f, Eq (f EClassId), Hashable (f EClassId), Eq (f ()), Hashable (f ()))
  => WorkList -> State (EGraph d f) (MergeResult (ClassReplacements, IntLikeSet EClassId))
egMergeMany :: forall d (f :: * -> *).
(Semigroup d, Traversable f, Eq (f EClassId),
 Hashable (f EClassId), Eq (f ()), Hashable (f ())) =>
WorkList
-> State (EGraph d f) (MergeResult (EquivFind EClassId, WorkItem))
egMergeMany WorkList
wl0 = do
  BuildWorkResult WorkList
br <- forall d (f :: * -> *).
WorkList -> State (EGraph d f) (BuildWorkResult WorkList)
egBuildWorklist WorkList
wl0
  case BuildWorkResult WorkList
br of
    BuildWorkResult WorkList
BuildWorkResultUnchanged -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. MergeResult a
MergeResultUnchanged
    BuildWorkResultMissing WorkItem
cs -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. WorkItem -> MergeResult a
MergeResultMissing WorkItem
cs)
    BuildWorkResultOk WorkList
wl1 -> forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. a -> MergeResult a
MergeResultChanged (forall d (f :: * -> *).
(Semigroup d, Traversable f, Eq (f EClassId),
 Hashable (f EClassId), Eq (f ()), Hashable (f ())) =>
WorkList -> State (EGraph d f) (EquivFind EClassId, WorkItem)
egRebuild WorkList
wl1)

-- private
-- Folds over items in worklist to merge, returning:
-- 1. map of old class -> new class for changed classes only
-- 2. closure of remapped classes (includes roots)
egRebuildMerge :: WorkList -> State (EGraph d f) (IntLikeMap EClassId EClassId, IntLikeSet EClassId)
egRebuildMerge :: forall d (f :: * -> *).
WorkList
-> State (EGraph d f) (IntLikeMap EClassId EClassId, WorkItem)
egRebuildMerge WorkList
wl = StateT
  (EGraph d f) Identity (IntLikeMap EClassId EClassId, WorkItem)
finalRes where
  finalRes :: StateT
  (EGraph d f) Identity (IntLikeMap EClassId EClassId, WorkItem)
finalRes = forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state forall a b. (a -> b) -> a -> b
$ \EGraph d f
eg ->
    let ef :: EquivFind EClassId
ef = forall d (f :: * -> *). EGraph d f -> EquivFind EClassId
egEquivFind EGraph d f
eg
    in case forall x.
Coercible x Int =>
[IntLikeSet x] -> EquivFind x -> EquivMergeSetsRes x
efMergeSetsInc (forall (t :: * -> *) a. Foldable t => t a -> [a]
toList WorkList
wl) EquivFind EClassId
ef of
      EquivMergeSetsResChanged WorkItem
roots WorkItem
classRemapSet EquivFind EClassId
ef' ->
        let classRemap :: IntLikeMap EClassId EClassId
classRemap = forall x a. Coercible x Int => [(x, a)] -> IntLikeMap x a
ILM.fromList (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\EClassId
c -> (EClassId
c, forall x. Coercible x Int => x -> EquivFind x -> x
efLookupRoot EClassId
c EquivFind EClassId
ef')) (forall x. Coercible x Int => IntLikeSet x -> [x]
ILS.toList WorkItem
classRemapSet))
            closure :: WorkItem
closure = forall x. Coercible x Int => [x] -> EquivFind x -> IntLikeSet x
efClosure (forall x. Coercible x Int => IntLikeSet x -> [x]
ILS.toList WorkItem
roots) EquivFind EClassId
ef'
        in ((IntLikeMap EClassId EClassId
classRemap, WorkItem
closure), EGraph d f
eg { egEquivFind :: EquivFind EClassId
egEquivFind = EquivFind EClassId
ef' })
      EquivMergeSetsRes EClassId
_ -> ((forall x a. IntLikeMap x a
ILM.empty, forall x. IntLikeSet x
ILS.empty), EGraph d f
eg)

-- private
-- Loop through nodes of all changed classes and update the hashcons to point to new classes
egRebuildHashCons :: IntLikeMap EClassId EClassId -> State (EGraph d f) ()
egRebuildHashCons :: forall d (f :: * -> *).
IntLikeMap EClassId EClassId -> State (EGraph d f) ()
egRebuildHashCons IntLikeMap EClassId EClassId
classRemap =
  forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify' (\EGraph d f
eg -> let hc' :: IntLikeMap ENodeId EClassId
hc' = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (forall {x} {d} {f :: * -> *} {b}.
Coercible x Int =>
IntLikeMap x (EClassInfo d f)
-> IntLikeMap ENodeId b -> (x, b) -> IntLikeMap ENodeId b
go (forall d (f :: * -> *).
EGraph d f -> IntLikeMap EClassId (EClassInfo d f)
egClassMap EGraph d f
eg)) (forall d (f :: * -> *). EGraph d f -> IntLikeMap ENodeId EClassId
egHashCons EGraph d f
eg) (forall x a. Coercible x Int => IntLikeMap x a -> [(x, a)]
ILM.toList IntLikeMap EClassId EClassId
classRemap) in EGraph d f
eg { egHashCons :: IntLikeMap ENodeId EClassId
egHashCons = IntLikeMap ENodeId EClassId
hc' }) where
  go :: IntLikeMap x (EClassInfo d f)
-> IntLikeMap ENodeId b -> (x, b) -> IntLikeMap ENodeId b
go IntLikeMap x (EClassInfo d f)
cm IntLikeMap ENodeId b
hc (x
oldClassId, b
newClassId) =
    let eci :: EClassInfo d f
eci = forall x a. Coercible x Int => x -> IntLikeMap x a -> a
ILM.partialLookup x
oldClassId IntLikeMap x (EClassInfo d f)
cm
        nodes :: Assoc ENodeId (f ())
nodes = forall d (f :: * -> *). EClassInfo d f -> Assoc ENodeId (f ())
eciNodes EClassInfo d f
eci
    in forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (forall a b c. (a -> b -> c) -> b -> a -> c
flip (forall x a.
Coercible x Int =>
x -> a -> IntLikeMap x a -> IntLikeMap x a
`ILM.insert` b
newClassId)) IntLikeMap ENodeId b
hc (forall x a. Coercible x Int => Assoc x a -> [x]
assocMembers Assoc ENodeId (f ())
nodes)

-- private
-- For each touched class, recanonicalize all its nodes
-- Return pair of
-- 1. Set of parent class ids that can observe changes (i.e. need recanonicalization/reanalysis)
-- 2. Worklist of induced parent equalities found by recanonicalization
egRebuildAssoc :: (Traversable f, Eq (f EClassId), Hashable (f EClassId)) => IntLikeMap ENodeId EClassId -> IntLikeMap EClassId EClassId -> IntLikeSet EClassId -> State (EGraph d f) (IntLikeSet EClassId, WorkList)
egRebuildAssoc :: forall (f :: * -> *) d.
(Traversable f, Eq (f EClassId), Hashable (f EClassId)) =>
IntLikeMap ENodeId EClassId
-> IntLikeMap EClassId EClassId
-> WorkItem
-> State (EGraph d f) (WorkItem, WorkList)
egRebuildAssoc IntLikeMap ENodeId EClassId
origHc IntLikeMap EClassId EClassId
classRemap WorkItem
touchedClasses = do
  IntLikeMap ENodeId EClassId
hc <- forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets forall d (f :: * -> *). EGraph d f -> IntLikeMap ENodeId EClassId
egHashCons
  IntLikeMap EClassId (EClassInfo d f)
cm <- forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets forall d (f :: * -> *).
EGraph d f -> IntLikeMap EClassId (EClassInfo d f)
egClassMap
  -- For each class that we're going to merge
  forall (t :: * -> *) b a s.
Foldable t =>
b -> t a -> (b -> a -> State s b) -> State s b
stateFold (forall x. IntLikeSet x
ILS.empty, forall a. Seq a
Empty) (forall x. Coercible x Int => IntLikeSet x -> [x]
ILS.toList WorkItem
touchedClasses) forall a b. (a -> b) -> a -> b
$ \(WorkItem
ps, WorkList
parentWl) EClassId
c -> do
    -- Get the class info
    let eci :: EClassInfo d f
eci = forall x a. Coercible x Int => x -> IntLikeMap x a -> a
ILM.partialLookup EClassId
c IntLikeMap EClassId (EClassInfo d f)
cm
    -- For each node in the class
    (Bool
finalChanged, WorkList
finalParentWl) <- forall (t :: * -> *) b a s.
Foldable t =>
b -> t a -> (b -> a -> State s b) -> State s b
stateFold (Bool
False, WorkList
parentWl) (forall x a. Coercible x Int => Assoc x a -> [x]
assocMembers (forall d (f :: * -> *). EClassInfo d f -> Assoc ENodeId (f ())
eciNodes EClassInfo d f
eci)) forall a b. (a -> b) -> a -> b
$ \(Bool
changed', WorkList
parentWl') ENodeId
n -> do
      -- Canonicalize it and add to the node map
      (ENodeId
newN, Maybe (IntLikeSet ENodeId)
mayEquivNs) <- forall (f :: * -> *) d.
(Traversable f, Eq (f EClassId), Hashable (f EClassId)) =>
ENodeId -> State (EGraph d f) (ENodeId, Maybe (IntLikeSet ENodeId))
egCanonicalizeInternal ENodeId
n
      case Maybe (IntLikeSet ENodeId)
mayEquivNs of
        Maybe (IntLikeSet ENodeId)
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool
changed', WorkList
parentWl')
        Just IntLikeSet ENodeId
equivNs ->
          let allNs :: IntLikeSet ENodeId
allNs = forall x. Coercible x Int => x -> IntLikeSet x -> IntLikeSet x
ILS.insert ENodeId
newN IntLikeSet ENodeId
equivNs
              allEquivClasses :: WorkItem
allEquivClasses = forall x y.
(Coercible x Int, Coercible y Int) =>
(x -> y) -> IntLikeSet x -> IntLikeSet y
ILS.map (forall x a. Coercible x Int => x -> IntLikeMap x a -> a
`ILM.partialLookup` IntLikeMap ENodeId EClassId
hc) IntLikeSet ENodeId
allNs
          in if forall x. IntLikeSet x -> Int
ILS.size WorkItem
allEquivClasses forall a. Ord a => a -> a -> Bool
> Int
1
            then forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool
True, WorkList
parentWl' forall a. Seq a -> a -> Seq a
:|> WorkItem
allEquivClasses)
            else forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool
changed', WorkList
parentWl')
    -- Emit observing parents if:
    --   1. class has changed
    --   2. any nodes have changed during canonicalization
    -- Note that we look up parents in the ORIGINAL hashcons because those are the ones that have the nodes pointing to this
    let emitParents :: Bool
emitParents = Bool
finalChanged Bool -> Bool -> Bool
|| forall x a. Coercible x Int => x -> IntLikeMap x a -> Bool
ILM.member EClassId
c IntLikeMap EClassId EClassId
classRemap
        addlParents :: WorkItem
addlParents = forall x y.
(Coercible x Int, Coercible y Int) =>
(x -> y) -> IntLikeSet x -> IntLikeSet y
ILS.map (forall x a. Coercible x Int => x -> IntLikeMap x a -> a
`ILM.partialLookup` IntLikeMap ENodeId EClassId
origHc) (forall d (f :: * -> *). EClassInfo d f -> IntLikeSet ENodeId
eciParents EClassInfo d f
eci)
        ps' :: WorkItem
ps' = if Bool
emitParents then forall x. IntLikeSet x -> IntLikeSet x -> IntLikeSet x
ILS.union WorkItem
addlParents WorkItem
ps else WorkItem
ps
    forall (f :: * -> *) a. Applicative f => a -> f a
pure (WorkItem
ps', WorkList
finalParentWl)

-- private
-- One round of rebuilding
egRebuildNodeRound :: (Traversable f, Eq (f EClassId), Hashable (f EClassId)) => IntLikeMap ENodeId EClassId -> WorkList -> IntLikeSet EClassId -> State (EGraph d f) (IntLikeSet EClassId, WorkList, IntLikeSet EClassId)
egRebuildNodeRound :: forall (f :: * -> *) d.
(Traversable f, Eq (f EClassId), Hashable (f EClassId)) =>
IntLikeMap ENodeId EClassId
-> WorkList
-> WorkItem
-> State (EGraph d f) (WorkItem, WorkList, WorkItem)
egRebuildNodeRound IntLikeMap ENodeId EClassId
origHc WorkList
wl WorkItem
parents = do
  -- First merge all classes together and get merged class sets
  (IntLikeMap EClassId EClassId
classRemap, WorkItem
classClosure) <- forall d (f :: * -> *).
WorkList
-> State (EGraph d f) (IntLikeMap EClassId EClassId, WorkItem)
egRebuildMerge WorkList
wl
  -- Now update the hashcons so node ids point to merged classes
  forall d (f :: * -> *).
IntLikeMap EClassId EClassId -> State (EGraph d f) ()
egRebuildHashCons IntLikeMap EClassId EClassId
classRemap
  -- Track all classes touched here
  let touchedClasses :: WorkItem
touchedClasses = forall x. IntLikeSet x -> IntLikeSet x -> IntLikeSet x
ILS.union WorkItem
parents WorkItem
classClosure
  -- Traverse all touched classes and canonicalize their nodes,
  -- recording the mapping from old -> new
  -- Also track parents that can observe changes to this class
  (WorkItem
candParents, WorkList
parentWl) <- forall (f :: * -> *) d.
(Traversable f, Eq (f EClassId), Hashable (f EClassId)) =>
IntLikeMap ENodeId EClassId
-> IntLikeMap EClassId EClassId
-> WorkItem
-> State (EGraph d f) (WorkItem, WorkList)
egRebuildAssoc IntLikeMap ENodeId EClassId
origHc IntLikeMap EClassId EClassId
classRemap WorkItem
touchedClasses
  -- (We ignore parents that we have just now rebuilt)
  let finalParents :: WorkItem
finalParents = forall x. IntLikeSet x -> IntLikeSet x -> IntLikeSet x
ILS.difference WorkItem
candParents WorkItem
touchedClasses
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (WorkItem
touchedClasses, WorkList
parentWl, WorkItem
finalParents)

-- private
-- Rebuild just the class info corresponding to 'newClass'
egRebuildClassSingle :: (Semigroup d, Eq (f ()), Hashable (f ())) => EClassId -> IntLikeSet EClassId -> IntLikeMap EClassId (EClassInfo d f) -> IntLikeMap EClassId (EClassInfo d f)
egRebuildClassSingle :: forall d (f :: * -> *).
(Semigroup d, Eq (f ()), Hashable (f ())) =>
EClassId
-> WorkItem
-> IntLikeMap EClassId (EClassInfo d f)
-> IntLikeMap EClassId (EClassInfo d f)
egRebuildClassSingle EClassId
newClass WorkItem
oldClasses IntLikeMap EClassId (EClassInfo d f)
initCm =
  let EClassInfo d
rootData Assoc ENodeId (f ())
rootNodes IntLikeSet ENodeId
rootParents = forall x a. Coercible x Int => x -> IntLikeMap x a -> a
ILM.partialLookup EClassId
newClass IntLikeMap EClassId (EClassInfo d f)
initCm
      finalData :: d
finalData = forall a. Semigroup a => NonEmpty a -> a
sconcat (d
rootData forall a. a -> [a] -> NonEmpty a
:| forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\EClassId
c -> forall d (f :: * -> *). EClassInfo d f -> d
eciData (forall x a. Coercible x Int => x -> IntLikeMap x a -> a
ILM.partialLookup EClassId
c IntLikeMap EClassId (EClassInfo d f)
initCm)) (forall x. Coercible x Int => IntLikeSet x -> [x]
ILS.toList WorkItem
oldClasses))
      -- keep dead self nodes here. will be dropped in compact
      finalNodes :: Assoc ENodeId (f ())
finalNodes = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Assoc ENodeId (f ())
s EClassId
c -> forall x a.
(Coercible x Int, Ord x, Eq a, Hashable a) =>
Assoc x a -> Assoc x a -> Assoc x a
assocUnion Assoc ENodeId (f ())
s (forall d (f :: * -> *). EClassInfo d f -> Assoc ENodeId (f ())
eciNodes (forall x a. Coercible x Int => x -> IntLikeMap x a -> a
ILM.partialLookup EClassId
c IntLikeMap EClassId (EClassInfo d f)
initCm))) Assoc ENodeId (f ())
rootNodes (forall x. Coercible x Int => IntLikeSet x -> [x]
ILS.toList WorkItem
oldClasses)
      -- keep dead parent nodes here, just exclude self nodes. will be dropped in compact
      lookupParents :: EClassId -> IntLikeSet ENodeId
lookupParents EClassId
c = forall d (f :: * -> *). EClassInfo d f -> IntLikeSet ENodeId
eciParents (forall x a. Coercible x Int => x -> IntLikeMap x a -> a
ILM.partialLookup EClassId
c IntLikeMap EClassId (EClassInfo d f)
initCm)
      candParents :: IntLikeSet ENodeId
candParents = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\IntLikeSet ENodeId
s EClassId
c -> forall x. IntLikeSet x -> IntLikeSet x -> IntLikeSet x
ILS.union IntLikeSet ENodeId
s (EClassId -> IntLikeSet ENodeId
lookupParents EClassId
c)) IntLikeSet ENodeId
rootParents (forall x. Coercible x Int => IntLikeSet x -> [x]
ILS.toList WorkItem
oldClasses)
      finalParents :: IntLikeSet ENodeId
finalParents = forall x. IntLikeSet x -> IntLikeSet x -> IntLikeSet x
ILS.difference IntLikeSet ENodeId
candParents (forall x. Coercible x Int => [x] -> IntLikeSet x
ILS.fromList (forall x a. Coercible x Int => Assoc x a -> [x]
assocMembers Assoc ENodeId (f ())
finalNodes))
      finalInfo :: EClassInfo d f
finalInfo = forall d (f :: * -> *).
d -> Assoc ENodeId (f ()) -> IntLikeSet ENodeId -> EClassInfo d f
EClassInfo d
finalData Assoc ENodeId (f ())
finalNodes IntLikeSet ENodeId
finalParents
      finalCm :: IntLikeMap EClassId (EClassInfo d f)
finalCm = forall x a.
Coercible x Int =>
x -> a -> IntLikeMap x a -> IntLikeMap x a
ILM.insert EClassId
newClass EClassInfo d f
finalInfo IntLikeMap EClassId (EClassInfo d f)
initCm
  in IntLikeMap EClassId (EClassInfo d f)
finalCm

-- private
-- Rebuilds the classmap: merges old class infos into root class infos
-- Returns list of modified root classes
egRebuildClassMap :: (Semigroup d, Eq (f ()), Hashable (f ())) => IntLikeSet EClassId -> State (EGraph d f) ClassReplacements
egRebuildClassMap :: forall d (f :: * -> *).
(Semigroup d, Eq (f ()), Hashable (f ())) =>
WorkItem -> State (EGraph d f) (EquivFind EClassId)
egRebuildClassMap WorkItem
touchedClasses = forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state forall a b. (a -> b) -> a -> b
$ \EGraph d f
eg ->
  let ef :: EquivFind EClassId
ef = forall d (f :: * -> *). EGraph d f -> EquivFind EClassId
egEquivFind EGraph d f
eg
      -- Find roots corresponding to all touched classes
      roots :: WorkItem
roots = forall x y.
(Coercible x Int, Coercible y Int) =>
(x -> y) -> IntLikeSet x -> IntLikeSet y
ILS.map (forall x. Coercible x Int => x -> EquivFind x -> x
`efLookupRoot` EquivFind EClassId
ef) WorkItem
touchedClasses
      -- Prepare a replacement map for external consumers that just contains changed classes
      classReplacements :: EquivFind EClassId
classReplacements = forall x. Coercible x Int => [x] -> EquivFind x -> EquivFind x
efSubset (forall x. Coercible x Int => IntLikeSet x -> [x]
ILS.toList WorkItem
roots) EquivFind EClassId
ef
      -- Rebuild the class map
      cm' :: IntLikeMap EClassId (EClassInfo d f)
cm' = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\IntLikeMap EClassId (EClassInfo d f)
cm (EClassId
r, WorkItem
vs) -> forall d (f :: * -> *).
(Semigroup d, Eq (f ()), Hashable (f ())) =>
EClassId
-> WorkItem
-> IntLikeMap EClassId (EClassInfo d f)
-> IntLikeMap EClassId (EClassInfo d f)
egRebuildClassSingle EClassId
r WorkItem
vs IntLikeMap EClassId (EClassInfo d f)
cm) (forall d (f :: * -> *).
EGraph d f -> IntLikeMap EClassId (EClassInfo d f)
egClassMap EGraph d f
eg) (forall x a. Coercible x Int => IntLikeMap x a -> [(x, a)]
ILM.toList (forall x. EquivFind x -> IntLikeMap x (IntLikeSet x)
efFwd EquivFind EClassId
classReplacements))
  in (EquivFind EClassId
classReplacements, EGraph d f
eg { egClassMap :: IntLikeMap EClassId (EClassInfo d f)
egClassMap = IntLikeMap EClassId (EClassInfo d f)
cm' })

-- private
-- Rebuilds the 'EGraph' - merges classes as requested in the worklist and recanonicalizes.
-- This may take several rounds as changes propagate "upward" to parents.
-- Returns
-- 1. class remapping (roots -> removed classes)
-- 2. touched root classes
egRebuild :: (Semigroup d, Traversable f, Eq (f EClassId), Hashable (f EClassId), Eq (f ()), Hashable (f ())) => WorkList -> State (EGraph d f) (ClassReplacements, IntLikeSet EClassId)
egRebuild :: forall d (f :: * -> *).
(Semigroup d, Traversable f, Eq (f EClassId),
 Hashable (f EClassId), Eq (f ()), Hashable (f ())) =>
WorkList -> State (EGraph d f) (EquivFind EClassId, WorkItem)
egRebuild WorkList
wl0 = StateT (EGraph d f) Identity (EquivFind EClassId, WorkItem)
goRec where
  goRec :: StateT (EGraph d f) Identity (EquivFind EClassId, WorkItem)
goRec = do
    -- Note the existing hashcons
    IntLikeMap ENodeId EClassId
origHc <- forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets forall d (f :: * -> *). EGraph d f -> IntLikeMap ENodeId EClassId
egHashCons
    -- Merge and induce equivalences
    -- We track "touched classes" to know which to later rebuild in the classmap
    WorkItem
tc <- forall {f :: * -> *} {d}.
(Traversable f, Hashable (f EClassId)) =>
IntLikeMap ENodeId EClassId
-> WorkItem
-> WorkList
-> WorkItem
-> StateT (EGraph d f) Identity WorkItem
goNodeRounds IntLikeMap ENodeId EClassId
origHc forall x. IntLikeSet x
ILS.empty WorkList
wl0 forall x. IntLikeSet x
ILS.empty
    -- Compute final "touched roots"
    EquivFind EClassId
ef <- forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets forall d (f :: * -> *). EGraph d f -> EquivFind EClassId
egEquivFind
    let tr :: WorkItem
tr = forall x. Coercible x Int => [x] -> IntLikeSet x
ILS.fromList [EClassId
y | EClassId
x <- forall x. Coercible x Int => IntLikeSet x -> [x]
ILS.toList WorkItem
tc, EClassId
y <- forall a. Maybe a -> [a]
maybeToList (forall x. Coercible x Int => x -> EquivFind x -> Maybe x
efFindRoot EClassId
x EquivFind EClassId
ef)]
    -- Now everything is merged, so rewrite the changed parts of the classmap
    EquivFind EClassId
rm <- forall d (f :: * -> *).
(Semigroup d, Eq (f ()), Hashable (f ())) =>
WorkItem -> State (EGraph d f) (EquivFind EClassId)
egRebuildClassMap WorkItem
tc
    -- Finally, cleanup all "dead" classes and nodes
    forall (f :: * -> *) d.
(Foldable f, Eq (f ()), Hashable (f ())) =>
EquivFind EClassId -> State (EGraph d f) ()
egCompact EquivFind EClassId
rm
    -- And return the final class remapping and touched roots
    forall (f :: * -> *) a. Applicative f => a -> f a
pure (EquivFind EClassId
rm, WorkItem
tr)
  goNodeRounds :: IntLikeMap ENodeId EClassId
-> WorkItem
-> WorkList
-> WorkItem
-> StateT (EGraph d f) Identity WorkItem
goNodeRounds !IntLikeMap ENodeId EClassId
origHc !WorkItem
tc !WorkList
wl !WorkItem
parents =
    if forall (t :: * -> *) a. Foldable t => t a -> Bool
null WorkList
wl Bool -> Bool -> Bool
&& forall x. IntLikeSet x -> Bool
ILS.null WorkItem
parents
      then forall (f :: * -> *) a. Applicative f => a -> f a
pure WorkItem
tc
      else do
        (WorkItem
newTc, WorkList
newWl, WorkItem
newParents) <- forall (f :: * -> *) d.
(Traversable f, Eq (f EClassId), Hashable (f EClassId)) =>
IntLikeMap ENodeId EClassId
-> WorkList
-> WorkItem
-> State (EGraph d f) (WorkItem, WorkList, WorkItem)
egRebuildNodeRound IntLikeMap ENodeId EClassId
origHc WorkList
wl WorkItem
parents
        let mergedTc :: WorkItem
mergedTc = forall x. IntLikeSet x -> IntLikeSet x -> IntLikeSet x
ILS.union WorkItem
newTc WorkItem
tc
        IntLikeMap ENodeId EClassId
-> WorkItem
-> WorkList
-> WorkItem
-> StateT (EGraph d f) Identity WorkItem
goNodeRounds IntLikeMap ENodeId EClassId
origHc WorkItem
mergedTc WorkList
newWl WorkItem
newParents

-- private
-- Replace parent nodes with correct (remapped) ones
egCompactParentClass :: IntLikeMap ENodeId ENodeId -> EClassInfo d f -> EClassInfo d f
egCompactParentClass :: forall d (f :: * -> *).
IntLikeMap ENodeId ENodeId -> EClassInfo d f -> EClassInfo d f
egCompactParentClass IntLikeMap ENodeId ENodeId
nodeReplacements (EClassInfo d
dat Assoc ENodeId (f ())
nodes IntLikeSet ENodeId
parents) =
  forall d (f :: * -> *).
d -> Assoc ENodeId (f ()) -> IntLikeSet ENodeId -> EClassInfo d f
EClassInfo d
dat Assoc ENodeId (f ())
nodes (forall x y.
(Coercible x Int, Coercible y Int) =>
(x -> y) -> IntLikeSet x -> IntLikeSet y
ILS.map (\ENodeId
n -> forall x a. Coercible x Int => a -> x -> IntLikeMap x a -> a
ILM.findWithDefault ENodeId
n ENodeId
n IntLikeMap ENodeId ENodeId
nodeReplacements) IntLikeSet ENodeId
parents)

-- private
-- Remove dead nodes from given class info
egCompactSelfClass :: (Eq (f ()), Hashable (f ())) => IntLikeMap ENodeId ENodeId -> EClassInfo d f -> EClassInfo d f
egCompactSelfClass :: forall (f :: * -> *) d.
(Eq (f ()), Hashable (f ())) =>
IntLikeMap ENodeId ENodeId -> EClassInfo d f -> EClassInfo d f
egCompactSelfClass IntLikeMap ENodeId ENodeId
nodeReplacements (EClassInfo d
dat Assoc ENodeId (f ())
nodes IntLikeSet ENodeId
parents) =
  forall d (f :: * -> *).
d -> Assoc ENodeId (f ()) -> IntLikeSet ENodeId -> EClassInfo d f
EClassInfo d
dat (forall x a.
(Coercible x Int, Eq a, Hashable a) =>
[x] -> Assoc x a -> Assoc x a
assocRemoveAllInc (forall x a. Coercible x Int => IntLikeMap x a -> [x]
ILM.keys IntLikeMap ENodeId ENodeId
nodeReplacements) Assoc ENodeId (f ())
nodes) IntLikeSet ENodeId
parents

-- private
-- Find all classes that have dead nodes
findDeadNodeParentClasses :: Foldable f => Assoc ENodeId (f EClassId) -> [ENodeId] -> IntLikeSet EClassId
findDeadNodeParentClasses :: forall (f :: * -> *).
Foldable f =>
Assoc ENodeId (f EClassId) -> [ENodeId] -> WorkItem
findDeadNodeParentClasses Assoc ENodeId (f EClassId)
assoc = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' WorkItem -> ENodeId -> WorkItem
go forall x. IntLikeSet x
ILS.empty where
  go :: WorkItem -> ENodeId -> WorkItem
go WorkItem
s ENodeId
n = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall x. Coercible x Int => x -> IntLikeSet x -> IntLikeSet x
ILS.insert) WorkItem
s (forall x a. Coercible x Int => x -> Assoc x a -> a
assocPartialLookupByKey ENodeId
n Assoc ENodeId (f EClassId)
assoc)

-- private
-- Remove all dead nodes and classes from the graph
egCompactInc :: (Foldable f, Eq (f ()), Hashable (f ())) => ClassReplacements -> EGraph d f -> EGraph d f
egCompactInc :: forall (f :: * -> *) d.
(Foldable f, Eq (f ()), Hashable (f ())) =>
EquivFind EClassId -> EGraph d f -> EGraph d f
egCompactInc EquivFind EClassId
rm EGraph d f
eg =
  let ef :: EquivFind EClassId
ef = forall d (f :: * -> *). EGraph d f -> EquivFind EClassId
egEquivFind EGraph d f
eg
      assoc :: Assoc ENodeId (f EClassId)
assoc = forall d (f :: * -> *). EGraph d f -> Assoc ENodeId (f EClassId)
egNodeAssoc EGraph d f
eg
      hc :: IntLikeMap ENodeId EClassId
hc = forall d (f :: * -> *). EGraph d f -> IntLikeMap ENodeId EClassId
egHashCons EGraph d f
eg
      cm :: IntLikeMap EClassId (EClassInfo d f)
cm = forall d (f :: * -> *).
EGraph d f -> IntLikeMap EClassId (EClassInfo d f)
egClassMap EGraph d f
eg
      deadClasses :: WorkItem
deadClasses = forall x a. IntLikeMap x a -> IntLikeSet x
ILM.keysSet (forall x. EquivFind x -> IntLikeMap x x
efBwd EquivFind EClassId
rm)
      -- remove dead nodes from assoc
      (IntLikeMap ENodeId ENodeId
nodeReplacements, Assoc ENodeId (f EClassId)
assoc') = forall x a.
Coercible x Int =>
Assoc x a -> (IntLikeMap x x, Assoc x a)
assocCompactInc Assoc ENodeId (f EClassId)
assoc
      -- select all live classes that are parents of dead nodes
      deadNodeParentClasses :: WorkItem
deadNodeParentClasses = forall (f :: * -> *).
Foldable f =>
Assoc ENodeId (f EClassId) -> [ENodeId] -> WorkItem
findDeadNodeParentClasses Assoc ENodeId (f EClassId)
assoc (forall x a. Coercible x Int => IntLikeMap x a -> [x]
ILM.keys IntLikeMap ENodeId ENodeId
nodeReplacements)
      -- select all live classes that contain dead nodes
      deadNodeSelfClasses :: WorkItem
deadNodeSelfClasses = forall x. Coercible x Int => [x] -> IntLikeSet x
ILS.fromList (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall x a. Coercible x Int => x -> IntLikeMap x a -> a
`ILM.partialLookup` IntLikeMap ENodeId EClassId
hc) (forall x a. Coercible x Int => IntLikeMap x a -> [x]
ILM.keys IntLikeMap ENodeId ENodeId
nodeReplacements))
      -- remove dead classes from hashcons
      hc' :: IntLikeMap ENodeId EClassId
hc' = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall x a.
Coercible x Int =>
x -> IntLikeMap x a -> IntLikeMap x a
ILM.delete) IntLikeMap ENodeId EClassId
hc (forall x a. Coercible x Int => IntLikeMap x a -> [x]
ILM.keys IntLikeMap ENodeId ENodeId
nodeReplacements)
      -- remove dead classes from unionfind
      (IntLikeMap EClassId WorkItem
_, EquivFind EClassId
ef') = forall x.
Coercible x Int =>
EquivFind x -> (IntLikeMap x (IntLikeSet x), EquivFind x)
efCompactInc EquivFind EClassId
ef
      -- remove dead classes from classmap
      cm' :: IntLikeMap EClassId (EClassInfo d f)
cm' = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall x a.
Coercible x Int =>
x -> IntLikeMap x a -> IntLikeMap x a
ILM.delete) IntLikeMap EClassId (EClassInfo d f)
cm (forall x. Coercible x Int => IntLikeSet x -> [x]
ILS.toList WorkItem
deadClasses)
      -- rewrite dead parent nodes in classmap
      cm'' :: IntLikeMap EClassId (EClassInfo d f)
cm'' = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (forall a b c. (a -> b -> c) -> b -> a -> c
flip (forall x a.
Coercible x Int =>
(a -> a) -> x -> IntLikeMap x a -> IntLikeMap x a
ILM.adjust (forall d (f :: * -> *).
IntLikeMap ENodeId ENodeId -> EClassInfo d f -> EClassInfo d f
egCompactParentClass IntLikeMap ENodeId ENodeId
nodeReplacements))) IntLikeMap EClassId (EClassInfo d f)
cm' (forall x. Coercible x Int => IntLikeSet x -> [x]
ILS.toList WorkItem
deadNodeParentClasses)
      -- rewrite dead self nodes in classmap
      cm''' :: IntLikeMap EClassId (EClassInfo d f)
cm''' = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (forall a b c. (a -> b -> c) -> b -> a -> c
flip (forall x a.
Coercible x Int =>
(a -> a) -> x -> IntLikeMap x a -> IntLikeMap x a
ILM.adjust (forall (f :: * -> *) d.
(Eq (f ()), Hashable (f ())) =>
IntLikeMap ENodeId ENodeId -> EClassInfo d f -> EClassInfo d f
egCompactSelfClass IntLikeMap ENodeId ENodeId
nodeReplacements))) IntLikeMap EClassId (EClassInfo d f)
cm'' (forall x. Coercible x Int => IntLikeSet x -> [x]
ILS.toList WorkItem
deadNodeSelfClasses)
  in EGraph d f
eg { egEquivFind :: EquivFind EClassId
egEquivFind = EquivFind EClassId
ef', egNodeAssoc :: Assoc ENodeId (f EClassId)
egNodeAssoc = Assoc ENodeId (f EClassId)
assoc', egClassMap :: IntLikeMap EClassId (EClassInfo d f)
egClassMap = IntLikeMap EClassId (EClassInfo d f)
cm''', egHashCons :: IntLikeMap ENodeId EClassId
egHashCons = IntLikeMap ENodeId EClassId
hc' }

-- private
egCompact :: (Foldable f, Eq (f ()), Hashable (f ())) => ClassReplacements -> State (EGraph d f) ()
egCompact :: forall (f :: * -> *) d.
(Foldable f, Eq (f ()), Hashable (f ())) =>
EquivFind EClassId -> State (EGraph d f) ()
egCompact = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify' forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) d.
(Foldable f, Eq (f ()), Hashable (f ())) =>
EquivFind EClassId -> EGraph d f -> EGraph d f
egCompactInc

-- | Reanalyze a subset of classes - touched roots from merging is sufficient to ensure
-- complete reanalysis. (Note this is implemented in a simplistic way, just taking the
-- fixed point of rounds of analysis. The number of rounds can be no more than the size
-- of the given set.)
-- It may be necessary to call this because merging may leave class analyses in an
-- under-approximating state. This method gives you the true analysis by fixed point.
egReanalyzeSubset :: (Eq d, Semigroup d, Functor f) => EAnalysis d f -> IntLikeSet EClassId -> State (EGraph d f) ()
egReanalyzeSubset :: forall d (f :: * -> *).
(Eq d, Semigroup d, Functor f) =>
EAnalysis d f -> WorkItem -> State (EGraph d f) ()
egReanalyzeSubset EAnalysis d f
ana WorkItem
tr = StateT (EGraph d f) Identity ()
goStart where
  goStart :: StateT (EGraph d f) Identity ()
goStart = do
    IntLikeMap EClassId (EClassInfo d f)
cm <- forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets forall d (f :: * -> *).
EGraph d f -> IntLikeMap EClassId (EClassInfo d f)
egClassMap
    let am :: IntLikeMap EClassId d
am = forall a b x. (a -> b) -> IntLikeMap x a -> IntLikeMap x b
ILM.map forall d (f :: * -> *). EClassInfo d f -> d
eciData IntLikeMap EClassId (EClassInfo d f)
cm
    IntLikeMap EClassId d -> StateT (EGraph d f) Identity ()
goRec IntLikeMap EClassId d
am
  goRec :: IntLikeMap EClassId d -> StateT (EGraph d f) Identity ()
goRec IntLikeMap EClassId d
am0 = do
    IntLikeMap EClassId (EClassInfo d f)
cm <- forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets forall d (f :: * -> *).
EGraph d f -> IntLikeMap EClassId (EClassInfo d f)
egClassMap
    Assoc ENodeId (f EClassId)
assoc <- forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets forall d (f :: * -> *). EGraph d f -> Assoc ENodeId (f EClassId)
egNodeAssoc
    let fwd :: IntLikeMap ENodeId (f EClassId)
fwd = forall x a. Assoc x a -> IntLikeMap x a
assocFwd Assoc ENodeId (f EClassId)
assoc
    let onNode :: ENodeId -> d
onNode ENodeId
n =
          let fc :: f EClassId
fc = forall x a. Coercible x Int => x -> IntLikeMap x a -> a
ILM.partialLookup ENodeId
n IntLikeMap ENodeId (f EClassId)
fwd
              fd :: f d
fd = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall x a. Coercible x Int => x -> IntLikeMap x a -> a
`ILM.partialLookup` IntLikeMap EClassId d
am0) f EClassId
fc
          in EAnalysis d f
ana f d
fd
    let calcClass :: EClassId -> d
calcClass EClassId
cr =
          let nodes :: Assoc ENodeId (f ())
nodes = forall d (f :: * -> *). EClassInfo d f -> Assoc ENodeId (f ())
eciNodes (forall x a. Coercible x Int => x -> IntLikeMap x a -> a
ILM.partialLookup EClassId
cr IntLikeMap EClassId (EClassInfo d f)
cm)
          in case forall x a. Coercible x Int => IntLikeMap x a -> [x]
ILM.keys (forall x a. Assoc x a -> IntLikeMap x a
assocFwd Assoc ENodeId (f ())
nodes) of
            ENodeId
n0:[ENodeId]
ns ->
              let d0 :: d
d0 = ENodeId -> d
onNode ENodeId
n0
              in forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\d
d ENodeId
n -> d
d forall a. Semigroup a => a -> a -> a
<> ENodeId -> d
onNode ENodeId
n) d
d0 [ENodeId]
ns
            [] -> forall a. HasCallStack => String -> a
error String
"impossible"
    let onClassRoot :: (Bool, IntLikeMap EClassId d)
-> EClassId -> (Bool, IntLikeMap EClassId d)
onClassRoot p :: (Bool, IntLikeMap EClassId d)
p@(Bool
_, IntLikeMap EClassId d
amx) EClassId
cr =
          let d0 :: d
d0 = forall x a. Coercible x Int => x -> IntLikeMap x a -> a
ILM.partialLookup EClassId
cr IntLikeMap EClassId d
am0
              d1 :: d
d1 = EClassId -> d
calcClass EClassId
cr
          in if d
d0 forall a. Eq a => a -> a -> Bool
== d
d1 then (Bool, IntLikeMap EClassId d)
p else (Bool
True, forall x a.
Coercible x Int =>
x -> a -> IntLikeMap x a -> IntLikeMap x a
ILM.insert EClassId
cr d
d1 IntLikeMap EClassId d
amx)
    let (Bool
changed, IntLikeMap EClassId d
am1) = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (Bool, IntLikeMap EClassId d)
-> EClassId -> (Bool, IntLikeMap EClassId d)
onClassRoot (Bool
False, IntLikeMap EClassId d
am0) (forall x. Coercible x Int => IntLikeSet x -> [x]
ILS.toList WorkItem
tr)
    if Bool
changed
      then IntLikeMap EClassId d -> StateT (EGraph d f) Identity ()
goRec IntLikeMap EClassId d
am1
      else forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify' forall a b. (a -> b) -> a -> b
$ \EGraph d f
eg ->
        let cm0 :: IntLikeMap EClassId (EClassInfo d f)
cm0 = forall d (f :: * -> *).
EGraph d f -> IntLikeMap EClassId (EClassInfo d f)
egClassMap EGraph d f
eg
            cm1 :: IntLikeMap EClassId (EClassInfo d f)
cm1 = forall x a b.
Coercible x Int =>
(x -> a -> b) -> IntLikeMap x a -> IntLikeMap x b
ILM.mapWithKey (\EClassId
i EClassInfo d f
c -> EClassInfo d f
c { eciData :: d
eciData = forall x a. Coercible x Int => x -> IntLikeMap x a -> a
ILM.partialLookup EClassId
i IntLikeMap EClassId d
am1 } ) IntLikeMap EClassId (EClassInfo d f)
cm0
        in EGraph d f
eg { egClassMap :: IntLikeMap EClassId (EClassInfo d f)
egClassMap = IntLikeMap EClassId (EClassInfo d f)
cm1 }

-- | Reanalyze all classes in the graph.
egReanalyze :: (Eq d, Semigroup d, Functor f) => EAnalysis d f -> State (EGraph d f) ()
egReanalyze :: forall d (f :: * -> *).
(Eq d, Semigroup d, Functor f) =>
EAnalysis d f -> State (EGraph d f) ()
egReanalyze EAnalysis d f
ana = forall d (f :: * -> *). State (EGraph d f) WorkItem
egClasses forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall d (f :: * -> *).
(Eq d, Semigroup d, Functor f) =>
EAnalysis d f -> WorkItem -> State (EGraph d f) ()
egReanalyzeSubset EAnalysis d f
ana