overeasy (empty) → 0.1.0

raw patch · 14 files changed

+2932/−0 lines, 14 filesdep +algebraic-graphsdep +basedep +containerssetup-changed

Dependencies added: algebraic-graphs, base, containers, deepseq, hashable, hedgehog, int-like, logict, mtl, overeasy, prop-unit, recursion-schemes, text, transformers, unfree, unordered-containers

Files

• LICENSE +30/−0
• README.md +15/−0
• Setup.hs +2/−0
• overeasy.cabal +136/−0
• src/Overeasy/Assoc.hs +254/−0
• src/Overeasy/EGraph.hs +568/−0
• src/Overeasy/EquivFind.hs +350/−0
• src/Overeasy/Matching.hs +348/−0
• src/Overeasy/Source.hs +52/−0
• src/Overeasy/Streams.hs +40/−0
• src/Overeasy/Util.hs +85/−0
• test/Main.hs +953/−0
• test/Test/Overeasy/Arith.hs +30/−0
• test/Test/Overeasy/BinTree.hs +69/−0

@@ -0,0 +1,30 @@+Copyright Eric Conlon (c) 2021++All rights reserved.++Redistribution and use in source and binary forms, with or without+modification, are permitted provided that the following conditions are met:++    * Redistributions of source code must retain the above copyright+      notice, this list of conditions and the following disclaimer.++    * Redistributions in binary form must reproduce the above+      copyright notice, this list of conditions and the following+      disclaimer in the documentation and/or other materials provided+      with the distribution.++    * Neither the name of Eric Conlon nor the names of other+      contributors may be used to endorse or promote products derived+      from this software without specific prior written permission.++THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS+"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT+LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR+A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT+OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT+LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

@@ -0,0 +1,15 @@+# overeasy++[![CircleCI](https://circleci.com/gh/ejconlon/overeasy/tree/master.svg?style=svg)](https://circleci.com/gh/ejconlon/overeasy/tree/master)++A purely functional E-Graph library++## How to build and run++This repo is setup so you never have to `cd` out of the root directory.++To run the Haskell programs, you need `stack` installed on your system. [This](https://docs.haskellstack.org/en/stable/README/) is an easy way to do so, but you can also check your package manager or use [ghcup](https://www.haskell.org/ghcup/).++Most of the interesting stuff is going to be run in the test suite. Run it with `make test`. `stack` will get the appropriate Haskell compiler and package dependencies, and it will compile the project before running the test suite.++If you have Docker installed and running (with about 4G of RAM allocated to it), you can run `make docker-test` to build and run tests in a container instead of installing `stack` locally.

@@ -0,0 +1,2 @@+import Distribution.Simple+main = defaultMain

@@ -0,0 +1,136 @@+cabal-version: 1.12++-- This file has been generated from package.yaml by hpack version 0.34.4.+--+-- see: https://github.com/sol/hpack++name:           overeasy+version:        0.1.0+synopsis:       A purely functional E-Graph library+description:    Please see the README on GitHub at <https://github.com/ejconlon/overeasy#readme>+category:       Data Structures+homepage:       https://github.com/ejconlon/overeasy#readme+bug-reports:    https://github.com/ejconlon/overeasy/issues+author:         Eric Conlon+maintainer:     ejconlon@gmail.com+copyright:      (c) 2021 Eric Conlon+license:        BSD3+license-file:   LICENSE+build-type:     Simple+extra-source-files:+    README.md++source-repository head+  type: git+  location: https://github.com/ejconlon/overeasy++library+  exposed-modules:+      Overeasy.Assoc+      Overeasy.EGraph+      Overeasy.EquivFind+      Overeasy.Matching+      Overeasy.Source+      Overeasy.Streams+      Overeasy.Util+  other-modules:+      Paths_overeasy+  hs-source-dirs:+      src+  default-extensions:+      BangPatterns+      ConstraintKinds+      DeriveFunctor+      DeriveFoldable+      DeriveGeneric+      DeriveTraversable+      DerivingStrategies+      DerivingVia+      FlexibleContexts+      FlexibleInstances+      FunctionalDependencies+      GADTs+      GeneralizedNewtypeDeriving+      LambdaCase+      KindSignatures+      MultiParamTypeClasses+      PatternSynonyms+      Rank2Types+      ScopedTypeVariables+      StandaloneDeriving+      TemplateHaskell+      TupleSections+      TypeApplications+      TypeOperators+      TypeFamilies+  ghc-options: -Wall -Wcompat -Widentities -Wincomplete-record-updates -Wincomplete-uni-patterns -Wpartial-fields -Wredundant-constraints -fno-warn-unused-top-binds -fwrite-ide-info -hiedir=.hie+  build-depends:+      algebraic-graphs >=0.5 && <0.7+    , base >=4.12 && <5+    , containers ==0.6.*+    , deepseq ==1.4.*+    , hashable >=1.3 && <1.5+    , int-like >=0.1.1 && <0.2+    , logict >=0.7 && <0.9+    , mtl ==2.2.*+    , recursion-schemes ==5.2.*+    , text >=1.2 && <2.1+    , transformers >=0.5 && <0.7+    , unfree ==0.1.*+    , unordered-containers ==0.2.*+  default-language: Haskell2010++test-suite overeasy-test+  type: exitcode-stdio-1.0+  main-is: Main.hs+  other-modules:+      Test.Overeasy.Arith+      Test.Overeasy.BinTree+      Paths_overeasy+  hs-source-dirs:+      test+  default-extensions:+      BangPatterns+      ConstraintKinds+      DeriveFunctor+      DeriveFoldable+      DeriveGeneric+      DeriveTraversable+      DerivingStrategies+      DerivingVia+      FlexibleContexts+      FlexibleInstances+      FunctionalDependencies+      GADTs+      GeneralizedNewtypeDeriving+      LambdaCase+      KindSignatures+      MultiParamTypeClasses+      PatternSynonyms+      Rank2Types+      ScopedTypeVariables+      StandaloneDeriving+      TemplateHaskell+      TupleSections+      TypeApplications+      TypeOperators+      TypeFamilies+  ghc-options: -Wall -Wcompat -Widentities -Wincomplete-record-updates -Wincomplete-uni-patterns -Wpartial-fields -Wredundant-constraints -fno-warn-unused-top-binds -fwrite-ide-info -hiedir=.hie -threaded -rtsopts -with-rtsopts=-N+  build-depends:+      algebraic-graphs >=0.5 && <0.7+    , base >=4.12 && <5+    , containers ==0.6.*+    , deepseq ==1.4.*+    , hashable >=1.3 && <1.5+    , hedgehog ==1.0.*+    , int-like >=0.1.1 && <0.2+    , logict >=0.7 && <0.9+    , mtl ==2.2.*+    , overeasy+    , prop-unit ==0.1.*+    , recursion-schemes ==5.2.*+    , text >=1.2 && <2.1+    , transformers >=0.5 && <0.7+    , unfree ==0.1.*+    , unordered-containers ==0.2.*+  default-language: Haskell2010

@@ -0,0 +1,254 @@+{-# LANGUAGE DeriveAnyClass #-}++-- | See 'Assoc'.+module Overeasy.Assoc+  ( Assoc+  , assocFwd+  , assocBwd+  , assocEquiv+  , assocSize+  , assocNew+  , assocSingleton+  , AssocInsertRes (..)+  , assocInsertInc+  , assocInsert+  , assocFromList+  , assocToList+  , assocMember+  , assocLookupByKey+  , assocPartialLookupByKey+  , assocLookupByValue+  , assocPartialLookupByValue+  , assocLookupRoot+  , assocRoots+  , assocLeaves+  , assocMembers+  , assocCanCompact+  , assocCompactInc+  , assocCompact+  , assocRemoveAllInc+  , assocRemoveAll+  , assocUnion+  , assocFootprint+  ) where++import Control.DeepSeq (NFData)+import Control.Monad.State.Strict (MonadState (..), State, modify')+import Data.Coerce (Coercible)+import Data.Foldable (foldl')+import Data.Hashable (Hashable)+import Data.HashMap.Strict (HashMap)+import qualified Data.HashMap.Strict as HashMap+import Data.Maybe (fromJust)+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.EquivFind (EquivAddRes (..), EquivFind, efAddInc, efBwd, efCanCompact, efCompactInc, efEquivs, efLeaves,+                           efLookupRoot, efMember, efMembers, efNew, efRemoveAllInc, efRoots, efSingleton,+                           efUnsafeAddLeafInc, efUnsafeMerge)++-- | Associates keys and values in such a way that inserting+-- duplicate values induces equivalences on their keys.+-- Invariant: fwd and bwd maps contain only root keys.+data Assoc x a = Assoc+  { assocFwd :: !(IntLikeMap x a)+  -- ^ Map from id to element+  , assocBwd :: !(HashMap a x)+  -- ^ Map from element to id+  , assocEquiv :: !(EquivFind x)+  -- ^ Equivalence classes of ids+  } deriving stock (Eq, Show, Generic)+    deriving anyclass (NFData)++-- | How many values are in the map?+assocSize :: Assoc x a -> Int+assocSize = ILM.size . assocFwd++-- | Creates an empty assoc+assocNew :: Assoc x a+assocNew = Assoc ILM.empty HashMap.empty efNew++-- | Creates a singleton assoc+assocSingleton :: (Coercible x Int, Hashable a) => x -> a -> Assoc x a+assocSingleton x a = Assoc (ILM.singleton x a) (HashMap.singleton a x) (efSingleton x)++-- | The result of inserting into the assoc, if you're interested.+data AssocInsertRes x =+    AssocInsertResUnchanged+  | AssocInsertResCreated+  | AssocInsertResUpdated+  | AssocInsertResMerged !(IntLikeSet x)+  deriving stock (Eq, Show)++-- | Insert into the assoc (raw version)+assocInsertInc :: (Coercible x Int, Ord x, Eq a, Hashable a) => x -> a -> Assoc x a -> ((x, AssocInsertRes x), Assoc x a)+assocInsertInc x a1 assoc@(Assoc fwd bwd equiv) = finalRes where+  finalRes =+    let (res, equiv') = efAddInc x equiv+    in case res of+      EquivAddResNewRoot -> insertRoot x equiv'+      EquivAddResAlreadyLeafOf z -> updateRoot z+      EquivAddResAlreadyRoot -> updateRoot x+  updateRoot w =+    -- w is existing root and is guaranteed to map to something+    let a0 = ILM.partialLookup w fwd+    in if a0 == a1+      -- the value has not changed, don't need to change assoc+      then ((w, AssocInsertResUnchanged), assoc)+      else+        -- value has changed, need to check if it's fresh+        case HashMap.lookup a1 bwd of+          -- never seen; insert and return+          Nothing ->+            let fwd' = ILM.insert w a1 fwd+                bwd' = HashMap.insert a1 w (HashMap.delete a0 bwd)+            in ((w, AssocInsertResUpdated), Assoc fwd' bwd' equiv)+          -- mapped to another set of nodes, merge+          Just v ->+            let (toKeep, toDelete, equiv') = efUnsafeMerge w v equiv+                res = AssocInsertResMerged toDelete+            in if toKeep == w+              -- w wins+              then+                let fwd' = ILM.insert w a1 (ILM.delete v fwd)+                    bwd' = HashMap.insert a1 w (HashMap.delete a0 bwd)+                in ((w, res), Assoc fwd' bwd' equiv')+              -- v wins+              else+                let fwd' = ILM.delete w fwd+                    bwd' = HashMap.delete a0 bwd+                in ((v, res), Assoc fwd' bwd' equiv')+  insertRoot w equiv' =+    -- w is new root that doesn't exist+    case HashMap.lookup a1 bwd of+      -- never seen; insert and return+      Nothing ->+        let fwd' = ILM.insert w a1 fwd+            bwd' = HashMap.insert a1 w bwd+        in ((w, AssocInsertResCreated), Assoc fwd' bwd' equiv')+      Just v ->+        let (toKeep, toDelete, equiv'') = efUnsafeMerge w v equiv'+            res = AssocInsertResMerged toDelete+        in if toKeep == w+          -- w wins+          then+            let fwd' = ILM.insert w a1 (ILM.delete v fwd)+                bwd' = HashMap.insert a1 w bwd+            in ((w, res), Assoc fwd' bwd' equiv'')+          -- v wins+          else+            let fwd' = ILM.delete w fwd+            in ((v, res), Assoc fwd' bwd equiv'')++-- | Insert into the assoc (the state version)+assocInsert :: (Coercible x Int, Ord x, Eq a, Hashable a) => x -> a -> State (Assoc x a) (x, AssocInsertRes x)+assocInsert x a = state (assocInsertInc x a)++-- | Build an assoc from a list of pairs+assocFromList :: (Coercible x Int, Ord x, Eq a, Hashable a) => [(x, a)] -> Assoc x a+assocFromList = foldl' (\assoc (x, a) -> snd (assocInsertInc x a assoc)) assocNew++-- | Turn an assoc into a list of pairs (NOTE - emits only root keys!)+assocToList :: Coercible x Int => Assoc x a -> [(x, a)]+assocToList = ILM.toList . assocFwd++-- | Is the given key in the assoc?+assocMember :: Coercible x Int => x -> Assoc x a -> Bool+assocMember x (Assoc _ _ equiv) = efMember x equiv++-- | Forward lookup+assocLookupByKey :: Coercible x Int => x -> Assoc x a -> Maybe a+assocLookupByKey x (Assoc fwd _ equiv) = ILM.lookup (efLookupRoot x equiv) fwd++-- | PARTIAL forward lookup+assocPartialLookupByKey :: Coercible x Int => x -> Assoc x a -> a+assocPartialLookupByKey x = fromJust . assocLookupByKey x++-- | Backward lookup+assocLookupByValue :: (Eq a, Hashable a) => a -> Assoc x a -> Maybe x+assocLookupByValue a = HashMap.lookup a . assocBwd++-- | PARTIAL backward lookup+assocPartialLookupByValue :: (Eq a, Hashable a) => a -> Assoc x a -> x+assocPartialLookupByValue a = flip (HashMap.!) a . assocBwd++-- | Finds the root for the given key (id if not found)+assocLookupRoot :: Coercible x Int => x -> Assoc x a -> x+assocLookupRoot x = efLookupRoot x . assocEquiv++-- | List all root (live, non-compactible) keys+assocRoots :: Coercible x Int => Assoc x a -> [x]+assocRoots = efRoots . assocEquiv++-- | List all leaf (dead, compactible) keys+assocLeaves :: Coercible x Int => Assoc x a -> [x]+assocLeaves = efLeaves . assocEquiv++-- | List all entries (root and leaf)+assocMembers :: Coercible x Int => Assoc x a -> [x]+assocMembers = efMembers . assocEquiv++-- | Are there dead keys in the equiv from 'assocInsert'?+assocCanCompact :: Assoc x a -> Bool+assocCanCompact = efCanCompact . assocEquiv++-- | Removes all dead keys in the equiv (raw version).+assocCompactInc :: Coercible x Int => Assoc x a -> (IntLikeMap x x, Assoc x a)+assocCompactInc assoc@(Assoc fwd bwd equiv) =+  let replacements = efBwd equiv+      assoc' =+        if ILM.null replacements+          then assoc+          else let (_, equiv') = efCompactInc equiv in Assoc fwd bwd equiv'+  in (replacements, assoc')++-- | Removes all dead keys in the equiv (state version).+-- Returns map of dead leaf node -> live root node+assocCompact :: Coercible x Int => State (Assoc x a) (IntLikeMap x x)+assocCompact = state assocCompactInc++-- | Removes the given keys from the assoc (raw version)+assocRemoveAllInc :: (Coercible x Int, Eq a, Hashable a) => [x] -> Assoc x a -> Assoc x a+assocRemoveAllInc xs (Assoc fwd0 bwd0 equiv0) = Assoc fwdFinal bwdFinal equivFinal where+  (remap, equivFinal) = efRemoveAllInc xs equiv0+  (fwdFinal, bwdFinal) = foldl' go (fwd0, bwd0) xs+  go tup@(fwd, bwd) x =+    case ILM.lookup x fwd of+      -- Leaf, ignore+      Nothing -> tup+      -- Root+      Just a ->+        case ILM.lookup x remap of+          -- Singleton root, delete+          Nothing ->+            let fwd' = ILM.delete x fwd+                bwd' = HashMap.delete a bwd+            in (fwd', bwd')+          -- Remapped root, rotate+          Just y ->+            let fwd' = ILM.delete x (ILM.insert y a fwd)+                bwd' = HashMap.insert a y bwd+            in (fwd', bwd')++-- | Removes the given keys from the assoc (state version).+-- Values will only be removed from the assoc if the key is a singleton root.+-- If a key is not found, it is simply ignored.+assocRemoveAll :: (Coercible x Int, Eq a, Hashable a) => [x] -> State (Assoc x a) ()+assocRemoveAll = modify' . assocRemoveAllInc++-- | Join two assocs (uses the first as the base)+assocUnion :: (Coercible x Int, Ord x, Eq a, Hashable a) => Assoc x a -> Assoc x a -> Assoc x a+assocUnion base (Assoc fwd _ equiv) = Assoc fwdFinal bwdFinal equivFinal where+  goRoots assocGo (x, a) = snd (assocInsertInc x a assocGo)+  goLeaves equivGo (leaf, oldRoot) = efUnsafeAddLeafInc oldRoot leaf equivGo+  Assoc fwdFinal bwdFinal equivMid = foldl' goRoots base (ILM.toList fwd)+  equivFinal = foldl' goLeaves equivMid (ILM.toList (efBwd equiv))++-- | Returns the footprint of the given value - all keys that map to it (root and leaf)+assocFootprint :: (Coercible x Int, Eq a, Hashable a) => a -> Assoc x a -> IntLikeSet x+assocFootprint a (Assoc _ bwd equiv) =+  case HashMap.lookup a bwd of+    Nothing -> ILS.empty+    Just r -> efEquivs r equiv

@@ -0,0 +1,568 @@+{-# 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 { unEClassId :: Int }+  deriving stock (Show)+  deriving newtype (Eq, Ord, Enum, Hashable, NFData, Num)++-- | An opaque node id+-- Constructor exported for coercibility.+-- Num instance for literals only.+newtype ENodeId = ENodeId { unENodeId :: Int }+  deriving stock (Show)+  deriving newtype (Eq, Ord, Enum, Hashable, NFData, 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 = const ()++-- private+-- An internal triple of node, class, and data+data ENodeTriple d = ENodeTriple+  { entNode :: !ENodeId+  , entClass :: !EClassId+  , entData :: !d+  } deriving stock (Eq, Show, Generic)+    deriving anyclass (NFData)++-- | Info stored for every class: analysis data, class members (nodes), and parent nodes.+data EClassInfo d f = EClassInfo+  { eciData :: !d+  , eciNodes :: !(Assoc ENodeId (f ()))+  , eciParents :: !(IntLikeSet ENodeId)+  } deriving stock (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 (Eq, Show, Functor, Foldable, Traversable)++-- | An E-Graph implementation+data EGraph d f = EGraph+  { egClassSource :: !(Source EClassId)+  -- ^ Id source for classes+  , egNodeSource :: !(Source ENodeId)+  -- ^ Id source for nodes+  , egEquivFind :: !(EquivFind EClassId)+  -- ^ Class equivalences+  , egClassMap :: !(IntLikeMap EClassId (EClassInfo d f))+  -- ^ Map of class to info+  -- Invariant: Only contains root classes.+  , egNodeAssoc :: !(Assoc ENodeId (f EClassId))+  -- ^ Assoc of node id to node structure+  -- Invariant: only contains canonical structures (with root classes).+  , egHashCons :: !(IntLikeMap ENodeId EClassId)+  -- ^ Map of node to class+  -- Invariant: only contains root classes.+  } deriving stock (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 = efRootsSize . egEquivFind++-- | Number of nodes in the 'EGraph'+egNodeSize :: EGraph d f -> Int+egNodeSize = ILM.size . egHashCons++-- | Lookup info for the given 'EClass'+egClassInfo :: EClassId -> EGraph d f -> Maybe (EClassInfo d f)+egClassInfo c = ILM.lookup c . 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 fc eg = do+  n <- assocLookupByValue fc (egNodeAssoc eg)+  ILM.lookup n (egHashCons 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 t eg = foldWholeM (`egFindNode` eg) t++-- | Creates a new 'EGraph'+egNew :: EGraph d f+egNew = EGraph (sourceNew (EClassId 0)) (sourceNew (ENodeId 0)) efNew ILM.empty assocNew ILM.empty++-- | Yields all root classes+egClasses :: State (EGraph d f) (IntLikeSet EClassId)+egClasses = gets (ILM.keysSet . 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 fc = gets (efCanonicalize fc . 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 fc = gets (efCanonicalizePartial fc . 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 x = state $ \eg ->+  let ef = egEquivFind eg+      assoc = egNodeAssoc eg+      node = assocPartialLookupByKey x assoc+      fz = efCanonicalizePartial node ef+      ((y, res), assoc') = assocInsertInc x fz assoc+  in case res of+    AssocInsertResUnchanged -> ((y, Nothing), eg)+    AssocInsertResMerged toDelete ->+      ((y, Just toDelete), eg { egNodeAssoc = assoc' })+    _ -> ((y, Nothing), eg { egNodeAssoc = assoc' })++-- private+data AddNodeRes d = AddNodeRes !Changed !(Seq (ENodeTriple d))+  deriving stock (Eq, Show, Generic)+  deriving anyclass (NFData)++instance Semigroup (AddNodeRes d) where+  AddNodeRes c1 p1 <> AddNodeRes c2 p2 = AddNodeRes (c1 <> c2) (p1 <> p2)++instance Monoid (AddNodeRes d) where+  mempty = AddNodeRes ChangedNo Seq.empty+  mappend = (<>)++-- 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 ana fcd = do+  let fc = fmap entClass fcd+  -- important: node should already be canonicalized!+  -- first lookup the node in the assoc to ensure uniqueness+  mayNodeId <- gets (assocLookupByValue fc . egNodeAssoc)+  case mayNodeId of+    Just n -> do+      x <- gets (ILM.partialLookup n . egHashCons)+      eci <- gets (ILM.partialLookup x . egClassMap)+      let d = eciData eci+      pure (ChangedNo, ENodeTriple n x d)+    Nothing -> state $ \eg ->+      -- node does not exist; get new node and class ids+      let (n, nodeSource') = sourceAddInc (egNodeSource eg)+          (x, classSource') = sourceAddInc (egClassSource eg)+          -- add it to the uf (can discard return value since it's a new id, will be the same)+          (_, ef') = efAddInc x (egEquivFind eg)+          -- add it to the assoc (ignore and partial by construction)+          (_, assoc') = assocInsertInc n fc (egNodeAssoc eg)+          -- insert into the hashcons+          hc' = ILM.insert n x (egHashCons eg)+          -- analyze the node and put that info in the class map+          d = ana (fmap entData fcd)+          eci = EClassInfo d (assocSingleton n (void fcd)) ILS.empty+          classMap' = ILM.insert x eci (egClassMap eg)+          eg' = eg { egNodeSource = nodeSource', egClassSource = classSource', egEquivFind = ef', egNodeAssoc = assoc', egHashCons = hc', egClassMap = classMap' }+      in ((ChangedYes, ENodeTriple n x d), 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 ana = go where+  go t = do+    -- unwrap to work with the functor layer+    let ft = project t+    -- add all child nodes+    frx <- traverse go ft+    -- collect info generated from child nodes and leave pure structure+    let (AddNodeRes changed1 children, fx) = sequenceA frx+    -- now fx should be canonicalized by construction+    -- add the node to get its node and class ids+    (changed2, z@(ENodeTriple n _ _)) <- egAddNodeSub ana fx+    -- now update all its children to add this as a parent+    unless (Seq.null children) $+      modify' $ \eg ->+        -- Add node to class parents (unless it's a self parent)+        let cm' = foldl' (\cm (ENodeTriple _ c _) -> ILM.adjust (\v -> if assocMember n (eciNodes v) then v else v { eciParents = ILS.insert n (eciParents v) }) c cm) (egClassMap eg) children+        in eg { egClassMap = cm' }+    pure (AddNodeRes (changed1 <> changed2) (Seq.singleton z), 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 ana t = fmap (\(AddNodeRes c _, ENodeTriple _ x _) -> (c, x)) (egAddTermSub ana 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 i j = do+  mr <- egMergeMany (Seq.singleton (ILS.fromList [i, j]))+  -- We're guaranteed to have one and only one root in the map, so this won't fail+  pure (fmap (fst . head . ILM.toList . efFwd . fst) mr)++-- private+data BuildWorkResult a =+    BuildWorkResultUnchanged+  | BuildWorkResultMissing !WorkItem+  | BuildWorkResultOk !a++-- private+egBuildWorkItem :: WorkItem -> State (EGraph d f) (BuildWorkResult WorkItem)+egBuildWorkItem cs = do+  mayRoots <- fmap (\ef -> traverse (`efFindRoot` ef) (ILS.toList cs)) (gets egEquivFind)+  pure $! case mayRoots of+    Nothing -> BuildWorkResultMissing cs+    Just roots ->+      let rootsSet = ILS.fromList roots+      in if ILS.size rootsSet < 2+        then BuildWorkResultUnchanged+        else BuildWorkResultOk rootsSet++-- private+egBuildWorklist :: WorkList -> State (EGraph d f) (BuildWorkResult WorkList)+egBuildWorklist = go Empty where+  go !acc = \case+    Empty ->+      pure $! if Seq.null acc+        then BuildWorkResultUnchanged+        else BuildWorkResultOk acc+    cs :<| wl' -> do+      rcs <- egBuildWorkItem cs+      case rcs of+        BuildWorkResultUnchanged -> go acc wl'+        BuildWorkResultMissing cs' -> pure (BuildWorkResultMissing cs')+        BuildWorkResultOk cs' -> go (acc :|> cs') 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 wl0 = do+  br <- egBuildWorklist wl0+  case br of+    BuildWorkResultUnchanged -> pure MergeResultUnchanged+    BuildWorkResultMissing cs -> pure (MergeResultMissing cs)+    BuildWorkResultOk wl1 -> fmap MergeResultChanged (egRebuild 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 wl = finalRes where+  finalRes = state $ \eg ->+    let ef = egEquivFind eg+    in case efMergeSetsInc (toList wl) ef of+      EquivMergeSetsResChanged roots classRemapSet ef' ->+        let classRemap = ILM.fromList (fmap (\c -> (c, efLookupRoot c ef')) (ILS.toList classRemapSet))+            closure = efClosure (ILS.toList roots) ef'+        in ((classRemap, closure), eg { egEquivFind = ef' })+      _ -> ((ILM.empty, ILS.empty), 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 classRemap =+  modify' (\eg -> let hc' = foldl' (go (egClassMap eg)) (egHashCons eg) (ILM.toList classRemap) in eg { egHashCons = hc' }) where+  go cm hc (oldClassId, newClassId) =+    let eci = ILM.partialLookup oldClassId cm+        nodes = eciNodes eci+    in foldl' (flip (`ILM.insert` newClassId)) hc (assocMembers 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 origHc classRemap touchedClasses = do+  hc <- gets egHashCons+  cm <- gets egClassMap+  -- For each class that we're going to merge+  stateFold (ILS.empty, Empty) (ILS.toList touchedClasses) $ \(ps, parentWl) c -> do+    -- Get the class info+    let eci = ILM.partialLookup c cm+    -- For each node in the class+    (finalChanged, finalParentWl) <- stateFold (False, parentWl) (assocMembers (eciNodes eci)) $ \(changed', parentWl') n -> do+      -- Canonicalize it and add to the node map+      (newN, mayEquivNs) <- egCanonicalizeInternal n+      case mayEquivNs of+        Nothing -> pure (changed', parentWl')+        Just equivNs ->+          let allNs = ILS.insert newN equivNs+              allEquivClasses = ILS.map (`ILM.partialLookup` hc) allNs+          in if ILS.size allEquivClasses > 1+            then pure (True, parentWl' :|> allEquivClasses)+            else pure (changed', 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 = finalChanged || ILM.member c classRemap+        addlParents = ILS.map (`ILM.partialLookup` origHc) (eciParents eci)+        ps' = if emitParents then ILS.union addlParents ps else ps+    pure (ps', 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 origHc wl parents = do+  -- First merge all classes together and get merged class sets+  (classRemap, classClosure) <- egRebuildMerge wl+  -- Now update the hashcons so node ids point to merged classes+  egRebuildHashCons classRemap+  -- Track all classes touched here+  let touchedClasses = ILS.union parents 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+  (candParents, parentWl) <- egRebuildAssoc origHc classRemap touchedClasses+  -- (We ignore parents that we have just now rebuilt)+  let finalParents = ILS.difference candParents touchedClasses+  pure (touchedClasses, parentWl, 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 newClass oldClasses initCm =+  let EClassInfo rootData rootNodes rootParents = ILM.partialLookup newClass initCm+      finalData = sconcat (rootData :| fmap (\c -> eciData (ILM.partialLookup c initCm)) (ILS.toList oldClasses))+      -- keep dead self nodes here. will be dropped in compact+      finalNodes = foldl' (\s c -> assocUnion s (eciNodes (ILM.partialLookup c initCm))) rootNodes (ILS.toList oldClasses)+      -- keep dead parent nodes here, just exclude self nodes. will be dropped in compact+      lookupParents c = eciParents (ILM.partialLookup c initCm)+      candParents = foldl' (\s c -> ILS.union s (lookupParents c)) rootParents (ILS.toList oldClasses)+      finalParents = ILS.difference candParents (ILS.fromList (assocMembers finalNodes))+      finalInfo = EClassInfo finalData finalNodes finalParents+      finalCm = ILM.insert newClass finalInfo initCm+  in 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 touchedClasses = state $ \eg ->+  let ef = egEquivFind eg+      -- Find roots corresponding to all touched classes+      roots = ILS.map (`efLookupRoot` ef) touchedClasses+      -- Prepare a replacement map for external consumers that just contains changed classes+      classReplacements = efSubset (ILS.toList roots) ef+      -- Rebuild the class map+      cm' = foldl' (\cm (r, vs) -> egRebuildClassSingle r vs cm) (egClassMap eg) (ILM.toList (efFwd classReplacements))+  in (classReplacements, eg { egClassMap = 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 wl0 = goRec where+  goRec = do+    -- Note the existing hashcons+    origHc <- gets egHashCons+    -- Merge and induce equivalences+    -- We track "touched classes" to know which to later rebuild in the classmap+    tc <- goNodeRounds origHc ILS.empty wl0 ILS.empty+    -- Compute final "touched roots"+    ef <- gets egEquivFind+    let tr = ILS.fromList [y | x <- ILS.toList tc, y <- maybeToList (efFindRoot x ef)]+    -- Now everything is merged, so rewrite the changed parts of the classmap+    rm <- egRebuildClassMap tc+    -- Finally, cleanup all "dead" classes and nodes+    egCompact rm+    -- And return the final class remapping and touched roots+    pure (rm, tr)+  goNodeRounds !origHc !tc !wl !parents =+    if null wl && ILS.null parents+      then pure tc+      else do+        (newTc, newWl, newParents) <- egRebuildNodeRound origHc wl parents+        let mergedTc = ILS.union newTc tc+        goNodeRounds origHc mergedTc newWl newParents++-- private+-- Replace parent nodes with correct (remapped) ones+egCompactParentClass :: IntLikeMap ENodeId ENodeId -> EClassInfo d f -> EClassInfo d f+egCompactParentClass nodeReplacements (EClassInfo dat nodes parents) =+  EClassInfo dat nodes (ILS.map (\n -> ILM.findWithDefault n n nodeReplacements) parents)++-- private+-- Remove dead nodes from given class info+egCompactSelfClass :: (Eq (f ()), Hashable (f ())) => IntLikeMap ENodeId ENodeId -> EClassInfo d f -> EClassInfo d f+egCompactSelfClass nodeReplacements (EClassInfo dat nodes parents) =+  EClassInfo dat (assocRemoveAllInc (ILM.keys nodeReplacements) nodes) parents++-- private+-- Find all classes that have dead nodes+findDeadNodeParentClasses :: Foldable f => Assoc ENodeId (f EClassId) -> [ENodeId] -> IntLikeSet EClassId+findDeadNodeParentClasses assoc = foldl' go ILS.empty where+  go s n = foldl' (flip ILS.insert) s (assocPartialLookupByKey n 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 rm eg =+  let ef = egEquivFind eg+      assoc = egNodeAssoc eg+      hc = egHashCons eg+      cm = egClassMap eg+      deadClasses = ILM.keysSet (efBwd rm)+      -- remove dead nodes from assoc+      (nodeReplacements, assoc') = assocCompactInc assoc+      -- select all live classes that are parents of dead nodes+      deadNodeParentClasses = findDeadNodeParentClasses assoc (ILM.keys nodeReplacements)+      -- select all live classes that contain dead nodes+      deadNodeSelfClasses = ILS.fromList (fmap (`ILM.partialLookup` hc) (ILM.keys nodeReplacements))+      -- remove dead classes from hashcons+      hc' = foldl' (flip ILM.delete) hc (ILM.keys nodeReplacements)+      -- remove dead classes from unionfind+      (_, ef') = efCompactInc ef+      -- remove dead classes from classmap+      cm' = foldl' (flip ILM.delete) cm (ILS.toList deadClasses)+      -- rewrite dead parent nodes in classmap+      cm'' = foldl' (flip (ILM.adjust (egCompactParentClass nodeReplacements))) cm' (ILS.toList deadNodeParentClasses)+      -- rewrite dead self nodes in classmap+      cm''' = foldl' (flip (ILM.adjust (egCompactSelfClass nodeReplacements))) cm'' (ILS.toList deadNodeSelfClasses)+  in eg { egEquivFind = ef', egNodeAssoc = assoc', egClassMap = cm''', egHashCons = hc' }++-- private+egCompact :: (Foldable f, Eq (f ()), Hashable (f ())) => ClassReplacements -> State (EGraph d f) ()+egCompact = modify' . 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 ana tr = goStart where+  goStart = do+    cm <- gets egClassMap+    let am = ILM.map eciData cm+    goRec am+  goRec am0 = do+    cm <- gets egClassMap+    assoc <- gets egNodeAssoc+    let fwd = assocFwd assoc+    let onNode n =+          let fc = ILM.partialLookup n fwd+              fd = fmap (`ILM.partialLookup` am0) fc+          in ana fd+    let calcClass cr =+          let nodes = eciNodes (ILM.partialLookup cr cm)+          in case ILM.keys (assocFwd nodes) of+            n0:ns ->+              let d0 = onNode n0+              in foldl' (\d n -> d <> onNode n) d0 ns+            [] -> error "impossible"+    let onClassRoot p@(_, amx) cr =+          let d0 = ILM.partialLookup cr am0+              d1 = calcClass cr+          in if d0 == d1 then p else (True, ILM.insert cr d1 amx)+    let (changed, am1) = foldl' onClassRoot (False, am0) (ILS.toList tr)+    if changed+      then goRec am1+      else modify' $ \eg ->+        let cm0 = egClassMap eg+            cm1 = ILM.mapWithKey (\i c -> c { eciData = ILM.partialLookup i am1 } ) cm0+        in eg { egClassMap = cm1 }++-- | Reanalyze all classes in the graph.+egReanalyze :: (Eq d, Semigroup d, Functor f) => EAnalysis d f -> State (EGraph d f) ()+egReanalyze ana = egClasses >>= egReanalyzeSubset ana

@@ -0,0 +1,350 @@+{-# LANGUAGE DeriveAnyClass #-}++-- | See 'EquivFind'.+module Overeasy.EquivFind+  ( EquivFind+  , efFwd+  , efBwd+  , efRootsSize+  , efLeavesSize+  , efTotalSize+  , efCanonicalize+  , efCanonicalizePartial+  , efNew+  , efSingleton+  , efMember+  , efRoots+  , efLeaves+  , efMembers+  , EquivAddRes (..)+  , efAddInc+  , efAdd+  , efEquivs+  , efClosure+  , efFindRoot+  , efFindLeaves+  , efSubset+  , efLookupRoot+  , efLookupLeaves+  , efFindAll+  , EquivMergeRes (..)+  , efUnsafeMerge+  , efMergeInc+  , efMerge+  , EquivMergeSetsRes (..)+  , efMergeSetsInc+  , efMergeSets+  , efCanCompact+  , efCompactInc+  , efCompact+  , efRemoveAllInc+  , efRemoveAll+  , efUnsafeAddLeafInc+  ) where++import Control.Applicative ((<|>))+import Control.DeepSeq (NFData)+import Control.Monad.State.Strict (State, state)+import Data.Coerce (Coercible)+import Data.Foldable (foldl')+import Data.Maybe (fromJust, fromMaybe)+import GHC.Generics (Generic)+import IntLike.Map (IntLikeMap)+import qualified IntLike.Map as ILM+import IntLike.Set (IntLikeSet)+import qualified IntLike.Set as ILS++-- | A "Union-Find" implementation using explicit equivalence classes.+-- Sure, the asympotics aren't as good, but we eventually have to construct these+-- classes, so we might as well just do it as we go!+data EquivFind x = EquivFind+  { efFwd :: !(IntLikeMap x (IntLikeSet x))+  -- ^ Map of root to equivalent leaves+  -- Invariant: Map keys are only roots+  -- Invariant: Sets only contain leaf keys (and not the root itself)+  , efBwd :: !(IntLikeMap x x)+  -- ^ Map of leaf to root+  -- Invariant: Map keys are only leaves, values are only roots+  } deriving stock (Eq, Show, Generic)+    deriving anyclass (NFData)++-- | Number of roots in the equiv.+efRootsSize :: EquivFind x -> Int+efRootsSize = ILM.size . efFwd++-- | Number of leaves in the equiv.+efLeavesSize :: EquivFind x -> Int+efLeavesSize = ILM.size . efBwd++-- | Total number of keys in the equiv.+efTotalSize :: EquivFind x -> Int+efTotalSize ef = efRootsSize ef + efLeavesSize ef++-- | Canonicalize the given expression functor by replacing leaves with roots.+-- If any elements are missing, the first is returned.+efCanonicalize :: (Traversable f, Coercible x Int) => f x -> EquivFind x -> Either x (f x)+efCanonicalize fx ef = traverse (\x -> maybe (Left x) pure (efFindRoot x ef)) fx++-- | Canonicalize the given expression functor by replacing leaves with roots.+-- If any elements are missing, they are simply skipped.+efCanonicalizePartial :: (Functor f, Coercible x Int) => f x -> EquivFind x -> f x+efCanonicalizePartial fx ef = fmap (`efLookupRoot` ef) fx++-- | Creates an empty equiv+efNew :: EquivFind x+efNew = EquivFind ILM.empty ILM.empty++-- | Creates a singleton equiv+efSingleton :: Coercible x Int => x -> EquivFind x+efSingleton x = EquivFind (ILM.singleton x ILS.empty) ILM.empty++-- private+allocMM :: Coercible x Int => x -> IntLikeMap x (IntLikeSet x) -> IntLikeMap x (IntLikeSet x)+allocMM = ILM.alter (<|> Just ILS.empty)++-- private+insertMM :: Coercible x Int => x -> x -> IntLikeMap x (IntLikeSet x) -> IntLikeMap x (IntLikeSet x)+insertMM x y = ILM.alter (\case { Nothing -> Just (ILS.singleton y); Just s -> Just (ILS.insert y s) }) x++-- | Result of adding something to the equiv, if you're interested.+data EquivAddRes x =+    EquivAddResAlreadyRoot+  | EquivAddResAlreadyLeafOf !x+  | EquivAddResNewRoot+  deriving stock (Eq, Show, Generic)+  deriving anyclass (NFData)++-- | Add the given key to the equiv (raw version).+efAddInc :: Coercible x Int => x -> EquivFind x -> (EquivAddRes x, EquivFind x)+efAddInc x ef@(EquivFind fwd bwd) =+  case ILM.lookup x bwd of+    Nothing ->+      if ILM.member x fwd+        then (EquivAddResAlreadyRoot, ef)+        else (EquivAddResNewRoot, EquivFind (ILM.insert x ILS.empty fwd) bwd)+    Just y -> (EquivAddResAlreadyLeafOf y, ef)++-- | Add the given key to the equiv (raw version).+efAdd :: Coercible x Int => x -> State (EquivFind x) (EquivAddRes x)+efAdd = state . efAddInc++-- | All keys equivalent to the given key in the equiv.+-- Always returns a set with the given key, even if it's not present.+efEquivs :: Coercible x Int => x -> EquivFind x -> IntLikeSet x+efEquivs x ef = let r = efLookupRoot x ef in ILS.insert r (efLookupLeaves r ef)++-- | Set of all keys equivalent to the given keys in the equiv.+efClosure :: Coercible x Int => [x] -> EquivFind x -> IntLikeSet x+efClosure xs ef = foldl' (\c x -> if ILS.member x c then c else ILS.union (efEquivs x ef) c) ILS.empty xs++-- | Find the root equivalent to the given key (if it exists).+efFindRoot :: Coercible x Int => x -> EquivFind x -> Maybe x+efFindRoot x ef = ILM.lookup x (efBwd ef) <|> if ILM.member x (efFwd ef) then Just x else Nothing++-- | Find the leaves equivalent to the given key (if they exist).+efFindLeaves :: Coercible x Int => x -> EquivFind x -> Maybe (IntLikeSet x)+efFindLeaves x ef = ILM.lookup x (efFwd ef)++-- | Returns an EquivFind subset representing the given list of keys.+efSubset :: Coercible x Int => [x] -> EquivFind x -> EquivFind x+efSubset xs0 ef0 = foldl' go efNew xs0 where+  go (EquivFind fwd1 bwd1) x =+    let r = efLookupRoot x ef0+        ls = efLookupLeaves r ef0+    in EquivFind (ILM.insert r ls fwd1) (foldl' (\b l -> ILM.insert l r b) bwd1 (ILS.toList ls))++-- | Like 'efFindRoot' but returns same key if not found - does not guarantee presence in map.+efLookupRoot :: Coercible x Int => x -> EquivFind x -> x+efLookupRoot x = fromMaybe x . ILM.lookup x . efBwd++-- | Like 'efFindLeaves' but returns empty set if not found - does not guarantee presence in map.+efLookupLeaves :: Coercible x Int => x -> EquivFind x -> IntLikeSet x+efLookupLeaves x = fromMaybe ILS.empty . ILM.lookup x . efFwd++-- | Returns the set of roots for the given set of keys, or an error with the first key+-- not found in the equiv.+efFindAll :: Coercible x Int => [x] -> EquivFind x -> Either x (IntLikeSet x)+efFindAll xs ef = go ILS.empty xs where+  go !acc = \case+    [] -> Right acc+    y:ys ->+      case efFindRoot y ef of+        Nothing -> Left y+        Just z -> go (ILS.insert z acc) ys++-- | Is the key in the equiv?+efMember :: Coercible x Int => x -> EquivFind x -> Bool+efMember x (EquivFind fwd bwd) = ILM.member x fwd || ILM.member x bwd++-- | List all roots in the equiv.+efRoots :: Coercible x Int => EquivFind x -> [x]+efRoots = ILM.keys . efFwd++-- | List all leaves in the equiv.+efLeaves :: Coercible x Int => EquivFind x -> [x]+efLeaves = ILM.keys . efBwd++-- | List all members (roots and leaves) in the equiv.+efMembers :: Coercible x Int => EquivFind x -> [x]+efMembers ef = efRoots ef ++ efLeaves ef++-- | The result of trying to merge two keys, if you care.+data EquivMergeRes x =+    EquivMergeResMissing !x+  | EquivMergeResUnchanged !x+  | EquivMergeResChanged !x !(IntLikeSet x) !(EquivFind x)+  deriving stock (Eq, Show, Generic)+  deriving anyclass (NFData)++-- | Don't even think about it, it's got unsafe in the name.+efUnsafeMerge :: (Coercible x Int, Ord x) => x -> x -> EquivFind x -> (x, IntLikeSet x, EquivFind x)+efUnsafeMerge ix jx (EquivFind fwd bwd) =+  let loKey = min ix jx+      hiKey = max ix jx+      hiSet = ILS.insert hiKey (ILM.partialLookup hiKey fwd)+      finalFwd = ILM.adjust (hiSet <>) loKey (ILM.delete hiKey fwd)+      finalBwd = foldl' (flip (`ILM.insert` loKey)) bwd (ILS.toList hiSet)+  in (loKey, hiSet, EquivFind finalFwd finalBwd)++-- | Merge two keys (raw version).+efMergeInc :: (Coercible x Int, Ord x) => x -> x -> EquivFind x -> EquivMergeRes x+efMergeInc i j ef =+  case efFindRoot i ef of+    Nothing -> EquivMergeResMissing i+    Just ix ->+      case efFindRoot j ef of+        Nothing -> EquivMergeResMissing j+        Just jx ->+          if ix == jx+            then EquivMergeResUnchanged ix+            else+              let (loKey, hiSet, ef') = efUnsafeMerge ix jx ef+              in EquivMergeResChanged loKey hiSet ef'++-- | Merge two keys (state version).+efMerge :: (Coercible x Int, Ord x) => x -> x -> State (EquivFind x) (Maybe (x, IntLikeSet x))+efMerge i j = state $ \ef ->+  case efMergeInc i j ef of+    EquivMergeResChanged loKey hiSet ef' -> (Just (loKey, hiSet), ef')+    _ -> (Nothing, ef)++-- | The result of trying to merge multiple keys, if you care.+data EquivMergeManyRes x =+    EquivMergeManyResEmpty+  | EquivMergeManyResEmbed !(EquivMergeRes x)+  deriving stock (Eq, Show, Generic)+  deriving anyclass (NFData)++-- | The result of trying to merge multiple sets of keys, if you care.+data EquivMergeSetsRes x =+    EquivMergeSetsResEmptySet+  | EquivMergeSetsResMissing !x+  | EquivMergeSetsResUnchanged !(IntLikeSet x)+  | EquivMergeSetsResChanged !(IntLikeSet x) !(IntLikeSet x) !(EquivFind x)+  deriving stock (Eq, Show, Generic)+  deriving anyclass (NFData)++-- | Merge sets of keys (raw version).+efMergeSetsInc :: Coercible x Int => [IntLikeSet x] -> EquivFind x -> EquivMergeSetsRes x+efMergeSetsInc css0 u0 = res where+  res =+    case css0 of+      [] -> EquivMergeSetsResUnchanged ILS.empty+      _ -> go ILS.empty ILS.empty u0 css0+  go !roots !classRemapSet ef@(EquivFind fwd bwd) css =+    case css of+      [] ->+        let finalRoots = ILS.map (`efLookupRoot` ef) roots+        in if ILS.null classRemapSet+          then EquivMergeSetsResUnchanged finalRoots+          else EquivMergeSetsResChanged finalRoots classRemapSet ef+      ds:dss ->+        case ILS.toList ds of+          [] -> go roots classRemapSet ef dss+          zs -> case efFindAll zs ef of+            Left x -> EquivMergeSetsResMissing x+            Right xs ->+              let (loKey, ys) = fromJust (ILS.minView xs)+                  newRoots = ILS.insert loKey roots+                  hiSet = ILS.unions (fmap (\y -> ILS.insert y (efLookupLeaves y ef)) (ILS.toList ys))+                  newClassRemapSet = ILS.union hiSet classRemapSet+                  newFwd = ILM.adjust (ILS.union hiSet) loKey (foldl' (flip ILM.delete) fwd (ILS.toList ys))+                  newBwd = foldl' (flip (`ILM.insert` loKey)) bwd (ILS.toList hiSet)+                  newU = EquivFind newFwd newBwd+              in go newRoots newClassRemapSet newU dss++-- | Merge sets of keys (state version).+efMergeSets :: Coercible x Int => [IntLikeSet x] -> State (EquivFind x) (Maybe (IntLikeSet x, IntLikeSet x))+efMergeSets css = state $ \ef ->+  case efMergeSetsInc css ef of+    EquivMergeSetsResChanged roots classRemapSet ef' -> (Just (roots, classRemapSet), ef')+    _ -> (Nothing, ef)++-- | Are they compactible keys?+efCanCompact :: EquivFind x -> Bool+efCanCompact = not . ILM.null . efBwd++-- | See 'efCompact' (this is the raw version).+efCompactInc :: Coercible x Int => EquivFind x -> (IntLikeMap x (IntLikeSet x), EquivFind x)+efCompactInc (EquivFind origFwd origBwd) = finalRes where+  finalRes =+    let (rootMap, fwd') = foldl' go (ILM.empty, origFwd) (ILM.elems origBwd)+    in (rootMap, EquivFind fwd' ILM.empty)+  go p@(rootMap, fwd) r =+    if ILM.member r rootMap+      then p+      else+        let xs = ILM.partialLookup r fwd+        in (ILM.insert r xs rootMap, if ILS.null xs then fwd else ILM.insert r ILS.empty fwd)++-- | Removes leaves and returns map of root to deleted leaf.+efCompact :: Coercible x Int => State (EquivFind x) (IntLikeMap x (IntLikeSet x))+efCompact = state efCompactInc++-- | See 'efRemoveAll' (this is the raw version).+efRemoveAllInc :: Coercible x Int => [x] -> EquivFind x -> (IntLikeMap x x, EquivFind x)+efRemoveAllInc xs (EquivFind fwd0 bwd0) = (remapFinal, EquivFind fwdFinal bwdFinal) where+  (fwdFinal, bwdFinal, remapFinal) = foldl' go (fwd0, bwd0, ILM.empty) xs+  go tup@(fwd, bwd, remap) x =+    case ILM.lookup x fwd of+      -- Key is not root+      Nothing -> case ILM.lookup x bwd of+        -- Key is missing, skip it+        Nothing -> tup+        -- Key is leaf, remove from both containers+        Just r ->+          let bwd' = ILM.delete x bwd+              fwd' = ILM.adjust (ILS.delete x) r fwd+          in (fwd', bwd', remap)+      -- Key is root+      Just leaves ->+        -- ensure the remapping is from ORIGINAL roots to new roots+        let origRoot = fromMaybe x (ILM.lookup x bwd0)+        in case ILS.minView leaves of+          -- Singleton root, remove from fwd and remap+          Nothing ->+            let fwd' = ILM.delete x fwd+                remap' = ILM.delete origRoot remap+            in (fwd', bwd, remap')+          -- Non-singleton root, rotate+          Just (y, rest) ->+            let fwd' = ILM.delete x (ILM.insert y rest fwd)+                bwd' = ILM.delete y (foldl' (\m l -> ILM.insert l y m) bwd (ILS.toList rest))+                remap' = ILM.insert origRoot y remap+            in (fwd', bwd', remap')++-- | Removes the given keys from the equiv map.+-- If a key is a leaf or singleton root, simply remove it.+-- If it is a root of a larger class, select the min leaf and make it root.+-- Returns a map of old roots to new roots (only those changed in the process -+-- possibly empty). If a key is not found, it is simply ignored.+efRemoveAll :: Coercible x Int => [x] -> State (EquivFind x) (IntLikeMap x x)+efRemoveAll = state . efRemoveAllInc++-- | Given root, add leaf. Requires that root be present in the map+-- and that leaf would be picked as a leaf. (Therefore, unsafe.)+-- Exposed for efficient merging.+efUnsafeAddLeafInc :: Coercible x Int => x -> x -> EquivFind x -> EquivFind x+efUnsafeAddLeafInc root leaf ef@(EquivFind fwd bwd) =+  let trueRoot = efLookupRoot root ef+  in EquivFind (ILM.adjust (ILS.insert leaf) trueRoot fwd) (ILM.insert leaf trueRoot bwd)

@@ -0,0 +1,348 @@+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE UndecidableInstances #-}++-- | Methods to match patterns in 'EGraph's (aka e-matching)+module Overeasy.Matching+  ( Pat+  , patVars+  , Subst+  , substVars+  , VarId (..)+  , PatGraphC+  , PatGraph (..)+  , patGraph+  , Match (..)+  , MatchPat (..)+  , MatchF (..)+  , MatchPatF (..)+  , matchVars+  , matchClasses+  , MatchSubst (..)+  , SolGraphC+  , SolGraph (..)+  , solGraph+  , SolStream+  , SolveC+  , solve+  , match+  ) where++import Control.Applicative (Alternative (..))+import Control.DeepSeq (NFData)+import Control.Monad (void)+import Control.Monad.Reader (asks)+import Control.Monad.State.Strict (MonadState (..), State, evalState, execState, gets, modify', runState)+import Data.Bifunctor (bimap)+import Data.Coerce (Coercible)+import Data.Foldable (fold, foldMap', foldl', for_, toList)+import Data.Functor.Foldable (Base, Corecursive (..), Recursive (..), cata)+import Data.Hashable (Hashable)+import Data.HashMap.Strict (HashMap)+import qualified Data.HashMap.Strict as HashMap+import Data.HashSet (HashSet)+import qualified Data.HashSet as HashSet+import IntLike.Map (IntLikeMap)+import qualified IntLike.Map as ILM+import IntLike.Set (IntLikeSet)+import qualified IntLike.Set as ILS+import Overeasy.Assoc (Assoc, assocBwd, assocEquiv, assocFootprint, assocFwd, assocInsertInc, assocLookupByValue,+                       assocNew)+import Overeasy.EGraph (EClassId (..), EGraph (egHashCons), ENodeId (..), eciNodes, egClassMap, egNodeAssoc)+import Overeasy.EquivFind (efLookupRoot)+import Overeasy.Source (Source, sourceAddInc, sourceNew)+import Overeasy.Streams (Stream, chooseWith, streamAll)+import Unfree (Free, FreeF (..))++-- | A pattern is exactly the free monad over the expression functor+-- It has spots for var names ('FreePure') and spots for structural+-- pieces ('FreeEmbed')+type Pat = Free++-- | The base functor of 'Pat'.+type PatF = FreeF++-- | The set of vars for a pattern.+patVars :: (Foldable f, Eq v, Hashable v) => Pat f v -> HashSet v+patVars = foldMap' HashSet.singleton++-- | A substitution.+type Subst c v = HashMap v c++-- | The set of vars for a substitution.+substVars :: Subst a v -> HashSet v+substVars = HashMap.keysSet++-- | A match is a pattern annotated with classes (or other data).+data Match c f v = Match+  { matchAnno :: !c+  , matchPat :: !(MatchPat c f v)+  } deriving stock (Functor, Foldable, Traversable)++deriving stock instance (Eq c, Eq v, Eq (f (Match c f v))) => Eq (Match c f v)+deriving stock instance (Show c, Show v, Show (f (Match c f v))) => Show (Match c f v)++-- | Tie the knot - the inner layer of a match.+data MatchPat c f v =+    MatchPatPure !v+  | MatchPatEmbed !(f (Match c f v))+  deriving stock (Functor, Foldable, Traversable)++deriving stock instance (Eq v, Eq (f (Match c f v))) => Eq (MatchPat c f v)+deriving stock instance (Show v, Show (f (Match c f v))) => Show (MatchPat c f v)++-- | The base functor of 'Match'+data MatchF c f v r = MatchF+  { matchClassF :: !c+  , matchPatF :: !(MatchPatF f v r)+  } deriving stock (Functor, Foldable, Traversable)++-- | Tie the knot - the inner part of 'MatchF'.+data MatchPatF f v r =+    MatchPatPureF !v+  | MatchPatEmbedF !(f r)+  deriving stock (Functor, Foldable, Traversable)++type instance Base (Match c f v) = MatchF c f v++instance Functor f => Recursive (Match c f v) where+  project (Match cl mp) = MatchF cl $ case mp of+     MatchPatPure v -> MatchPatPureF v+     MatchPatEmbed f -> MatchPatEmbedF f++instance Functor f => Corecursive (Match c f v) where+  embed (MatchF cl mpf) = Match cl $ case mpf of+     MatchPatPureF v -> MatchPatPure v+     MatchPatEmbedF f -> MatchPatEmbed f++-- | The set of vars in a match.+matchVars :: (Foldable f, Eq v, Hashable v) => Match c f v -> HashSet v+matchVars = foldMap' HashSet.singleton++-- | The set of classes in a match.+matchClasses :: (Coercible c Int, Functor f, Foldable f) => Match c f v -> IntLikeSet c+matchClasses = cata go where+  go (MatchF cl mpf) = ILS.insert cl $ case mpf of+    MatchPatPureF _ -> ILS.empty+    MatchPatEmbedF fc -> fold fc++-- | A apri of match and substitution.+data MatchSubst c f v = MatchSubst+  { msMatch :: !(Match c f v)+  , msSubst :: !(Subst c v)+  }++deriving stock instance (Eq c, Eq v, Eq (f (Match c f v))) => Eq (MatchSubst c f v)+deriving stock instance (Show c, Show v, Show (f (Match c f v))) => Show (MatchSubst c f v)++-- | An opaque var id+-- Constructor exported for coercibility+newtype VarId = VarId { unVarId :: Int }+  deriving stock (Show)+  deriving newtype (Eq, Ord, Enum, Hashable, NFData)++-- | A pattern graph - can be created once for each pattern and reused+-- for many iterations of search.+data PatGraph f v = PatGraph+  { pgRoot :: !VarId+  , pgNodes :: !(IntLikeMap VarId (PatF f v VarId))+  , pgVars :: !(HashMap v VarId)+  }++deriving stock instance (Eq v, Eq (f VarId)) => Eq (PatGraph f v)+deriving stock instance (Show v, Show (f VarId)) => Show (PatGraph f v)++-- | The set of constraints necessary to build a pattern graph.+type PatGraphC f v = (Traversable f, Eq v, Eq (f VarId), Hashable v, Hashable (f VarId))++data GraphState f v = GraphState+  { gsSrc :: !(Source VarId)+  , gsAssoc :: !(Assoc VarId (PatF f v VarId))+  }++emptyGraphState :: GraphState f v+emptyGraphState = GraphState (sourceNew (VarId 0)) assocNew++graphEnsurePart :: PatGraphC f v => PatF f v VarId -> State (GraphState f v) VarId+graphEnsurePart part = do+  mi <- gets (assocLookupByValue part . gsAssoc)+  case mi of+    Just i -> pure i+    Nothing -> state $ \st ->+      let (i, src') = sourceAddInc (gsSrc st)+          (_, assoc') = assocInsertInc i part (gsAssoc st)+      in (i, st { gsSrc = src', gsAssoc = assoc' })++graphEnsurePat :: PatGraphC f v => Pat f v -> State (GraphState f v) VarId+graphEnsurePat = cata go where+  go = \case+    FreePureF v -> graphEnsurePart (FreePureF v)+    FreeEmbedF fp -> do+      fi <- sequenceA fp+      graphEnsurePart (FreeEmbedF fi)++graphCanonicalize :: PatGraphC f v => GraphState f v -> IntLikeMap VarId (PatF f v VarId)+graphCanonicalize (GraphState _ assoc) =+  let fwd = assocFwd assoc+      equiv = assocEquiv assoc+  in fmap (fmap (`efLookupRoot` equiv)) fwd++-- | Builds a pattern graph from a pattern.+patGraph :: PatGraphC f v => Pat f v -> PatGraph f v+patGraph p =+  let (i, st) = runState (graphEnsurePat p) emptyGraphState+      m = graphCanonicalize st+      n = HashMap.fromList (ILM.toList m >>= \(j, x) -> case x of { FreePureF v -> [(v, j)]; _ -> [] })+  in PatGraph i m n++-- | A solution graph - must be created from an e-graph each merge/rebuild.+data SolGraph c f = SolGraph+  { sgByVar :: !(IntLikeMap VarId (IntLikeSet c))+  -- ^ Map of var -> classes.+  -- Contains all vars.+  -- If the inner map is empty, that means the pattern was not matched.+  -- The inner set will not be empty.+  , sgNodes :: !(HashMap (f c) c)+  -- ^ Map of node structures to classes.+  }++deriving stock instance (Eq c, Eq (f c)) => Eq (SolGraph c f)+deriving stock instance (Show c, Show (f c)) => Show (SolGraph c f)++-- | The set of constraints necessary to build a solution graph.+type SolGraphC f = (Functor f, Foldable f, Eq (f ()), Hashable (f ()))++-- | Builds a solution graph from an e-graph.+solGraph :: SolGraphC f => PatGraph f v -> EGraph d f -> SolGraph EClassId f+solGraph pg eg =+  -- For each class, use footprint of reverse node assoc to find set of node ids+  -- Start with just the embedded nodes+  let byVarEmbed = ILM.fromList $ ILM.toList (pgNodes pg) >>= \(i, pf) ->+        case pf of+          FreePureF _ -> []+          FreeEmbedF fi ->+            let fu = void fi+                cns = ILM.toList (egClassMap eg) >>= \(c, inf) ->+                  let ns = eciNodes inf+                      fp = assocFootprint fu ns+                  in [(c, fp) | not (ILS.null fp)]+            in [(i, bimap ILS.fromList mconcat (unzip cns))]+      byVar = genByVar byVarEmbed (pgNodes pg) (assocFwd (egNodeAssoc eg))+      hc = egHashCons eg+      nodes = fmap (`ILM.partialLookup` hc) (assocBwd (egNodeAssoc eg))+  in SolGraph byVar nodes++data Record =+    RecordPure !VarId !(IntLikeSet EClassId)+  | RecordEmbed+  deriving stock (Eq, Show)++type Records = [Record]++initRecords :: Foldable f => IntLikeMap VarId (PatF f v VarId) -> f VarId -> Records+initRecords nodes = fmap (\i -> case ILM.partialLookup i nodes of { FreePureF _ -> RecordPure i ILS.empty; _ -> RecordEmbed }) . toList++updateRecords :: Foldable f => Records -> f EClassId -> Records+updateRecords rs = zipWith (\r c -> case r of { RecordPure v cs -> RecordPure v (ILS.insert c cs); _ -> r } ) rs . toList++genByVar :: Foldable f => IntLikeMap VarId (IntLikeSet EClassId, IntLikeSet ENodeId) -> IntLikeMap VarId (PatF f v VarId) -> IntLikeMap ENodeId (f EClassId) -> IntLikeMap VarId (IntLikeSet EClassId)+genByVar byVarEmbed nodes fwd = execState (for_ (ILM.toList nodes) go) (fmap fst byVarEmbed) where+  go (i, pf) =+    case pf of+      FreePureF _ -> pure ()+      FreeEmbedF fi -> do+        -- We've gone through embedded patterns before so we're able+        -- to look up nodes for each pattern+        let (_, ns) = ILM.partialLookup i byVarEmbed+        -- For each node, update positionally what it could be+        let rs = foldl' (\rsx n -> let fc = ILM.partialLookup n fwd in updateRecords rsx fc) (initRecords nodes fi) (ILS.toList ns)+        -- Finally update the map; if missing set the positions as is, otherwise take intersection+        modify' $ \m -> foldl' (\mx r -> case r of {RecordPure j cs -> ILM.alter (Just . maybe cs (ILS.intersection cs)) j mx; _ -> mx}) m rs++data SolEnv c f v = SolEnv+  { sePatGraph :: !(PatGraph f v)+  , seSolGraph :: !(SolGraph c f)+  }++newtype SolState c = SolState+  { ssClasses :: IntLikeMap VarId c+  } deriving (Eq, Show)++-- | A stream of solutions. Can be demanded all at once, unconsed one at a time,+-- or interleaved.+type SolStream c f v z = Stream (SolEnv c f v) (SolState c) z++-- | The set of constraints necessary to search for solutions.+type SolveC c f v = (Traversable f, Coercible c Int, Eq v, Hashable v, Eq (f c), Hashable (f c))++constructMatch :: Traversable f => IntLikeMap VarId (PatF f v VarId) -> IntLikeMap VarId c -> VarId -> Match c f v+constructMatch nodes classes i0 = evalState (go i0) ILM.empty where+  go i = do+    cache <- get+    case ILM.lookup i cache of+      Just res -> pure res+      Nothing -> do+        let c = ILM.partialLookup i classes+        mp <- case ILM.partialLookup i nodes of+          FreePureF v -> pure $! MatchPatPure v+          FreeEmbedF f -> fmap MatchPatEmbed (traverse go f)+        pure $! Match c mp++constructSubst :: HashMap v VarId -> IntLikeMap VarId a -> Subst a v+constructSubst vars classes = fmap (`ILM.partialLookup` classes) vars++solveYield :: Traversable f => SolStream c f v (MatchSubst c f v)+solveYield = do+  pg <- asks sePatGraph+  classes <- gets ssClasses+  let mat = constructMatch (pgNodes pg) classes (pgRoot pg)+      subst = constructSubst (pgVars pg) classes+      ms = MatchSubst mat subst+  pure ms++-- | Produces a stream of solutions (e-matches).+solve :: SolveC c f v => SolStream c f v (MatchSubst c f v)+solve = do+  i <- asks (pgRoot . sePatGraph)+  void (solveRec i)+  solveYield++solveChoose :: SolveC c f v => VarId -> IntLikeSet c -> SolStream c f v c+solveChoose i cs = chooseWith (ILS.toList cs) (solveSet i)++solveSet :: VarId -> c -> SolStream c f v c+solveSet i c =+  c <$ modify' (\ss -> ss { ssClasses = ILM.insert i c (ssClasses ss) })++solveRec :: SolveC c f v => VarId -> SolStream c f v c+solveRec i = do+  ms <- gets (ILM.lookup i . ssClasses)+  case ms of+    -- Seen before, return solution+    Just s -> pure s+    -- Unseen+    Nothing -> do+      n <- asks (ILM.partialLookup i . pgNodes . sePatGraph)+      case n of+        -- Free var, choose a solution for each class in `sgByVar i`+        FreePureF _ -> do+          cs <- asks (ILM.partialLookup i . sgByVar . seSolGraph)+          solveChoose i cs+        -- Embedded functor, traverse and emit solution if present+        FreeEmbedF fi -> do+          fa <- traverse solveRec fi+          mc <- asks (HashMap.lookup fa . sgNodes . seSolGraph)+          case mc of+            Nothing -> empty+            Just c -> solveSet i c++-- | The easiest way to do e-matching: given a pattern and an e-graph, yield the list of matches.+-- Note that it might be more efficient to keep a 'PatGraph' on hand and uncons the matches+-- one by one.+match :: (PatGraphC f v, SolGraphC f, SolveC EClassId f v) => Pat f v -> EGraph d f -> [MatchSubst EClassId f v]+match p eg =+  let pg = patGraph p+      sg = solGraph pg eg+  in if any ILS.null (ILM.elems (sgByVar sg))+    -- If any var id has no patches, the pattern won't match, so don't try to solve+    then []+    else streamAll solve (SolEnv pg sg) (SolState ILM.empty)

@@ -0,0 +1,52 @@+{-# LANGUAGE DeriveAnyClass #-}++-- | See 'Source'.+module Overeasy.Source+  ( Source+  , sourceSize+  , sourceNew+  , sourceAddInc+  , sourceAdd+  , sourceSkipInc+  , sourceSkip+  , sourcePeek+  ) where++import Control.DeepSeq (NFData)+import Control.Monad.State.Strict (State, modify', state)+import Data.Coerce (Coercible, coerce)+import GHC.Generics (Generic)++-- | A source of unique ids+data Source x = Source+  { sourceSize :: !Int+  -- ^ How many ids have ever been created?+  , sourceNextId :: !Int+  } deriving stock (Eq, Show, Generic)+    deriving anyclass (NFData)++-- | Creates a new 'Source' from a starting id+sourceNew :: Coercible x Int => x -> Source x+sourceNew = Source 0 . coerce++-- | Generates the next id from the source (purely)+sourceAddInc :: Coercible x Int => Source x -> (x, Source x)+sourceAddInc (Source s x) = (coerce x, Source (s + 1) (x + 1))++-- | Generates the next id from the source (statefully)+sourceAdd :: Coercible x Int => State (Source x) x+sourceAdd = state sourceAddInc++-- | Skips past the given id (purely)+sourceSkipInc :: Coercible x Int => x -> Source x -> Source x+sourceSkipInc y (Source s x) =+  let z = coerce y+  in Source (s + 1) (max x (z + 1))++-- | Skips past the given id (statefully)+sourceSkip :: Coercible x Int => x -> State (Source x) ()+sourceSkip = modify' . sourceSkipInc++-- | Peek at the next id+sourcePeek :: Coercible x Int => Source x -> x+sourcePeek = coerce . sourceNextId

@@ -0,0 +1,40 @@+-- | Stuff for streaming search results.+module Overeasy.Streams+  ( chooseWith+  , choose+  , Stream+  , streamAll+  ) where++import Control.Applicative (Alternative (..))+import Control.Monad (MonadPlus)+import Control.Monad.Logic (LogicT, MonadLogic, observeAllT)+import Control.Monad.Reader (MonadReader (..), ReaderT (..))+import Control.Monad.State.Strict (MonadState (..), State, runState)++newtype M r s a = M { unM :: ReaderT r (State s) a }+  deriving newtype (+    Functor, Applicative, Monad,+    MonadReader r, MonadState s)++runM :: M r s a -> r -> s -> (a, s)+runM m r = runState (runReaderT (unM m) r)++-- | Choose one of many alteratives and process it with the given function.+chooseWith :: (Foldable f, Alternative m) => f a -> (a -> m b) -> m b+chooseWith fa f = foldr ((<|>) . f) empty fa++-- | Choose one of many alteratives.+choose :: (Foldable f, Alternative m) => f a -> m a+choose fa = chooseWith fa pure++-- | A stream of results. Just a wrapper around 'LogicT' to keep things tidy.+newtype Stream r s a = Stream { unStream :: LogicT (M r s) a }+  deriving newtype (+    Functor, Applicative, Monad,+    MonadReader r, MonadState s, MonadLogic,+    Alternative, MonadPlus, Semigroup, Monoid)++-- | Produces all results from the stream.+streamAll :: Stream r s a -> r -> s -> [a]+streamAll stream env st = fst (runM (observeAllT (unStream stream)) env st)

@@ -0,0 +1,85 @@+{-# LANGUAGE DeriveAnyClass #-}++-- | A grab bag of fun stuff.+module Overeasy.Util+  ( Whole+  , RecursiveWhole+  , foldWholeM+  , Changed (..)+  , stateFail+  , stateOption+  , stateFailChanged+  , stateFold+  ) where++import Control.DeepSeq (NFData)+import Control.Monad (foldM, forM_)+import Control.Monad.State.Strict (State, get, put)+import Data.Functor.Foldable (Base, Recursive (..))+import Data.Hashable (Hashable)+import GHC.Generics (Generic)++-- | Often 'f' is primary, not 't'. Relate them with this constraint.+type Whole t f = (f ~ Base t)++-- | Constraint for recursive structures+type RecursiveWhole t f = (Recursive t, Whole t f)++-- | Traverses a recursive structure+foldWholeM :: (RecursiveWhole t f, Traversable f, Monad m) => (f a -> m a) -> t -> m a+foldWholeM h = go where+  go t = do+    let ft = project t+    fa <- traverse go ft+    h fa++-- | A nicely-named 'Bool' for tracking state changes+data Changed = ChangedNo | ChangedYes+  deriving stock (Eq, Ord, Show, Generic)+  deriving anyclass (Hashable, NFData)++instance Semigroup Changed where+  c1 <> c2 =+    case c1 of+      ChangedYes -> ChangedYes+      _ -> c2++instance Monoid Changed where+  mempty = ChangedNo+  mappend = (<>)++-- | Embeds a function that may fail in a stateful context+stateFail :: (s -> Maybe (b, s)) -> State s (Maybe b)+stateFail f = do+  s <- get+  case f s of+    Nothing -> pure Nothing+    Just (b, s') -> put s' >> pure (Just b)++-- | Embeds a function that may fail in a stateful context+stateOption :: (s -> (b, Maybe s)) -> State s b+stateOption f = do+  s <- get+  let (b, ms) = f s+  forM_ ms put+  pure b++-- | Embeds a function that may fail in a stateful context with change tracking+stateFailChanged :: (s -> Maybe s) -> State s Changed+stateFailChanged f = do+  s <- get+  case f s of+    Nothing -> pure ChangedNo+    Just s' -> put s' >> pure ChangedYes++-- -- | Embeds a stateful action in a larger context+-- stateLens :: Lens' s a -> State a b -> State s b+-- stateLens l act = state $ \s ->+--   let (b, a') = runState act (view l s)+--       s' = set l a' s+--   in (b, s')++-- | 'foldM' specialized and flipped.+stateFold :: Foldable t => b -> t a -> (b -> a -> State s b) -> State s b+stateFold b as f = foldM f b as+{-# INLINE stateFold #-}

@@ -0,0 +1,953 @@+module Main (main) where++import Control.DeepSeq (NFData, force)+import Control.Exception (evaluate)+import Control.Monad (foldM, unless, when)+import Control.Monad.IO.Class (MonadIO, liftIO)+import Control.Monad.State.Strict (MonadState (..), State, StateT, evalState, evalStateT, execState, execStateT, gets,+                                   runState)+import Control.Monad.Trans (MonadTrans (..))+import Data.Bifunctor (bimap)+import Data.Char (chr, ord)+import Data.Coerce (coerce)+import Data.Foldable (for_)+import Data.Hashable (Hashable)+import qualified Data.HashMap.Strict as HashMap+import Data.List (delete)+import Data.Maybe (fromJust, isJust)+import Data.Semigroup (Max (..))+import qualified Data.Sequence as Seq+import Data.Traversable (for)+import qualified Hedgehog.Gen as Gen+import qualified Hedgehog.Range as Range+import qualified IntLike.Equiv as ILE+import qualified IntLike.Graph as ILG+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 (..), assocBwd, assocCanCompact, assocCompact, assocEquiv, assocFromList,+                       assocFwd, assocInsert, assocLeaves, assocMember, assocMembers, assocNew, assocPartialLookupByKey,+                       assocRoots, assocSize)+import Overeasy.EGraph (EAnalysis, EClassId (..), EClassInfo (..), EGraph (..), ENodeId (..), MergeResult (..),+                        egAddTerm, egCanonicalize, egClassSize, egFindTerm, egMerge, egMergeMany, egNew, egNodeSize,+                        noAnalysis)+import Overeasy.EquivFind (EquivFind (..), efAdd, efCanCompact, efCompact, efFindRoot, efLeaves, efLeavesSize, efMember,+                           efMembers, efMerge, efMergeSets, efNew, efRemoveAll, efRoots, efRootsSize, efTotalSize)+import Overeasy.Matching (Match (..), MatchPat (..), MatchSubst (..), Pat, match)+import Overeasy.Util (Changed (..))+import PropUnit (DependencyType (..), Gen, MonadTest, PropertyT, Range, TestLimit, TestTree, after, assert, forAll,+                 testGroup, testMain, testProp, testUnit, (/==), (===))+import Test.Overeasy.Arith (Arith (..), ArithF (ArithPlusF))+import Test.Overeasy.BinTree (BinTree, pattern BinTreeBranch, BinTreeF (..), pattern BinTreeLeaf)+import Unfree (pattern FreeEmbed, pattern FreePure)++fullyEvaluate :: (MonadIO m, NFData a) => a -> m a+fullyEvaluate = liftIO . evaluate . force++applyS :: Monad m => State s a -> StateT s m a+applyS = state . runState++testS :: Monad m => (s -> m a) -> StateT s m a+testS p = get >>= lift . p++applyTestS :: Monad m => State s a -> (a -> s -> m b) -> StateT s m b+applyTestS act check = do+  a <- applyS act+  s <- get+  lift (check a s)++foldS_ :: (Monad m, Foldable t) => s -> t a -> (a -> StateT s m ()) -> m s+foldS_ z as f = execStateT (for_ as f) z++runS :: Monad m => s -> StateT s m () -> m ()+runS = flip evalStateT++flipFoldM :: Monad m => b -> [a] -> (b -> a -> m b) -> m b+flipFoldM b as f = foldM f b as++newtype V = V { unV :: Int }+  deriving newtype (Eq, Ord, Hashable, NFData)++instance Show V where+  show = show . fromV++toV :: Char -> V+toV = V . ord++fromV :: V -> Char+fromV = chr . unV++setV :: String -> IntLikeSet V+setV = ILS.fromList . fmap toV++mapV :: [(Char, Char)] -> IntLikeMap V V+mapV = ILM.fromList . fmap (bimap toV toV)++multiMapV :: [(Char, String)] -> IntLikeMap V (IntLikeSet V)+multiMapV = ILM.fromList . fmap (bimap toV setV)++type EF = EquivFind V++testEfSimple :: TestTree+testEfSimple = testUnit "EF simple" $ runS efNew $ do+  testS $ \ef -> do+    efRootsSize ef === 0+    efLeavesSize ef === 0+    efTotalSize ef === 0+    efRoots ef === []+    efLeaves ef === []+    efMembers ef === []+    efMember (toV 'a') ef === False+    efMember (toV 'c') ef === False+    efFwd ef === ILM.empty+    efBwd ef === ILM.empty+  _ <- applyS (efAdd (toV 'a'))+  testS $ \ef -> do+    efRootsSize ef === 1+    efLeavesSize ef === 0+    efTotalSize ef === 1+    ILS.fromList (efRoots ef) === setV "a"+    ILS.fromList (efLeaves ef) === ILS.empty+    efFwd ef === multiMapV [('a', "")]+    efBwd ef === ILM.empty+  _ <- applyS (efAdd (toV 'b'))+  _ <- applyS (efAdd (toV 'c'))+  testS $ \ef -> do+    efRootsSize ef === 3+    efLeavesSize ef === 0+    efTotalSize ef === 3+    ILS.fromList (efRoots ef) === setV "abc"+    ILS.fromList (efLeaves ef) === ILS.empty+    efFwd ef === multiMapV [('a', ""), ('b', ""), ('c', "")]+    efBwd ef === ILM.empty+  applyTestS (efMerge (toV 'a') (toV 'c')) $ \res ef -> do+    res === Just (toV 'a', setV "c")+    efRootsSize ef === 2+    efLeavesSize ef === 1+    efTotalSize ef === 3+    ILS.fromList (efRoots ef) === setV "ab"+    ILS.fromList (efLeaves ef) === setV "c"+    efFwd ef === multiMapV [('a', "c"), ('b', "")]+    efBwd ef === mapV [('c', 'a')]+    efMembers ef === fmap toV ['a', 'b', 'c']+    efMember (toV 'a') ef === True+    efMember (toV 'c') ef === True+  applyTestS (efMerge (toV 'c') (toV 'a')) $ \res _ -> res === Nothing+  applyTestS (efMerge (toV 'b') (toV 'z')) $ \res _ -> res === Nothing++resetEf :: StateT EF (PropertyT IO) ()+resetEf = do+  put efNew+  _ <- applyS (efAdd (toV 'a'))+  _ <- applyS (efAdd (toV 'b'))+  _ <- applyS (efAdd (toV 'c'))+  _ <- applyS (efMerge (toV 'a') (toV 'c'))+  ef <- get+  efFwd ef === multiMapV [('a', "c"), ('b', "")]+  efBwd ef === mapV [('c', 'a')]++addExtraEf :: StateT EF (PropertyT IO) ()+addExtraEf = do+  _ <- applyS (efAdd (toV 'd'))+  _ <- applyS (efMerge (toV 'a') (toV 'd'))+  ef <- get+  efFwd ef === multiMapV [('a', "cd"), ('b', "")]+  efBwd ef === mapV [('c', 'a'), ('d', 'a')]++testEfRemove :: TestTree+testEfRemove = testUnit "EF remove" $ runS efNew $ do+  -- remove leav+  resetEf+  applyTestS (efRemoveAll [toV 'c']) $ \res ef -> do+    res === ILM.empty+    efFwd ef === multiMapV [('a', ""), ('b', "")]+    efBwd ef === mapV []+  -- remove singleton root+  resetEf+  applyTestS (efRemoveAll [toV 'b']) $ \res ef -> do+    res === ILM.empty+    efFwd ef === multiMapV [('a', "c")]+    efBwd ef === mapV [('c', 'a')]+  -- remove non-singleton root+  resetEf+  applyTestS (efRemoveAll [toV 'a']) $ \res ef -> do+    res === mapV [('a', 'c')]+    efFwd ef === multiMapV [('b', ""), ('c', "")]+    efBwd ef === mapV []+  -- remove all in class (root -> leaf order)+  resetEf+  applyTestS (efRemoveAll [toV 'a', toV 'c']) $ \res ef -> do+    res === ILM.empty+    efFwd ef === multiMapV [('b', "")]+    efBwd ef === mapV []+  -- remove all in class (leaf -> root order)+  resetEf+  applyTestS (efRemoveAll [toV 'c', toV 'a']) $ \res ef -> do+    res === ILM.empty+    efFwd ef === multiMapV [('b', "")]+    efBwd ef === mapV []+  -- remove with rotation and leaf+  resetEf+  addExtraEf+  applyTestS (efRemoveAll [toV 'a']) $ \res ef -> do+    res === mapV [('a', 'c')]+    efFwd ef === multiMapV [('c', "d"), ('b', "")]+    efBwd ef === mapV [('d', 'c')]+  -- remove with two rotations+  resetEf+  addExtraEf+  applyTestS (efRemoveAll [toV 'a', toV 'c']) $ \res ef -> do+    res === mapV [('a', 'd')]+    efFwd ef === multiMapV [('d', ""), ('b', "")]+    efBwd ef === mapV []+  -- remove with (leaf, rotation)+  resetEf+  addExtraEf+  applyTestS (efRemoveAll [toV 'c', toV 'a']) $ \res ef -> do+    res === mapV [('a', 'd')]+    efFwd ef === multiMapV [('d', ""), ('b', "")]+    efBwd ef === mapV []++testEfRec :: TestTree+testEfRec = testUnit "EF rec" $ runS efNew $ do+  _ <- applyS (efAdd (toV 'a'))+  _ <- applyS (efAdd (toV 'b'))+  _ <- applyS (efAdd (toV 'c'))+  applyTestS (efMerge (toV 'b') (toV 'c')) $ \res ef -> do+    res === Just (toV 'b', setV "c")+    efRootsSize ef === 2+    efLeavesSize ef === 1+    efTotalSize ef === 3+    ILS.fromList (efRoots ef) === setV "ab"+    ILS.fromList (efLeaves ef) === setV "c"+    efFwd ef === multiMapV [('a', ""), ('b', "c")]+    efBwd ef === mapV [('c', 'b')]+  applyTestS (efMerge (toV 'a') (toV 'c')) $ \res ef -> do+    res === Just (toV 'a', setV "bc")+    efRootsSize ef === 1+    efLeavesSize ef === 2+    efTotalSize ef === 3+    ILS.fromList (efRoots ef) === setV "a"+    ILS.fromList (efLeaves ef) === setV "bc"+    efFwd ef === multiMapV [('a', "bc")]+    efBwd ef === mapV [('b', 'a'), ('c', 'a')]++testEfMany :: TestTree+testEfMany = testUnit "EF many" $ runS efNew $ do+  _ <- applyS (efAdd (toV 'a'))+  _ <- applyS (efAdd (toV 'b'))+  _ <- applyS (efAdd (toV 'c'))+  _ <- applyS (efAdd (toV 'd'))+  _ <- applyS (efAdd (toV 'e'))+  applyTestS (efMergeSets [setV "cde"]) $ \res ef -> do+    res === Just (setV "c", setV "de")+    efRootsSize ef === 3+    efLeavesSize ef === 2+    efTotalSize ef === 5+    ILS.fromList (efRoots ef) === setV "abc"+    ILS.fromList (efLeaves ef) === setV "de"+    efFwd ef === multiMapV [('a', ""), ('b', ""), ('c', "de")]+    efBwd ef === mapV [('d', 'c'), ('e', 'c')]+  applyTestS (efMergeSets [setV "abd"]) $ \res ef -> do+    res === Just (setV "a", setV "bcde")+    efRootsSize ef === 1+    efLeavesSize ef === 4+    efTotalSize ef === 5+    ILS.fromList (efRoots ef) === setV "a"+    ILS.fromList (efLeaves ef) === setV "bcde"+    efFwd ef === multiMapV [('a', "bcde")]+    efBwd ef === mapV [('b', 'a'), ('c', 'a'), ('d', 'a'), ('e', 'a')]++testEfSets :: TestTree+testEfSets = testUnit "EF sets" $ runS efNew $ do+  _ <- applyS (efAdd (toV 'a'))+  _ <- applyS (efAdd (toV 'b'))+  _ <- applyS (efAdd (toV 'c'))+  _ <- applyS (efAdd (toV 'd'))+  _ <- applyS (efAdd (toV 'e'))+  applyTestS (efMergeSets [setV "cde", setV "abc"]) $ \res ef -> do+    res === Just (setV "a", setV "bcde")+    efRootsSize ef === 1+    efLeavesSize ef === 4+    efTotalSize ef === 5+    ILS.fromList (efRoots ef) === setV "a"+    efFwd ef === multiMapV [('a', "bcde")]++testEfCompact :: TestTree+testEfCompact = testUnit "EF compact" $ runS efNew $ do+  _ <- applyS (efAdd (toV 'a'))+  _ <- applyS (efAdd (toV 'b'))+  _ <- applyS (efAdd (toV 'c'))+  _ <- applyS (efAdd (toV 'd'))+  _ <- applyS (efAdd (toV 'e'))+  testS $ \ef -> assert (not (efCanCompact ef))+  applyTestS (efMergeSets [setV "cde"]) $ \res ef -> do+    res === Just (setV "c", setV "de")+    efFwd ef === multiMapV [('a', ""), ('b', ""), ('c', "de")]+    efBwd ef === mapV [('d', 'c'), ('e', 'c')]+    assert (efCanCompact ef)+  applyTestS efCompact $ \res ef -> do+    efFwd ef === multiMapV [('a', ""), ('b', ""), ('c', "")]+    efBwd ef === ILM.empty+    res === multiMapV [('c', "de")]+    assert (not (efCanCompact ef))++testEfUnit :: TestTree+testEfUnit = testGroup "EF unit" [testEfSimple, testEfRec, testEfMany, testEfSets, testEfCompact, testEfRemove]++genDistinctPairFromList :: Eq a => [a] -> Gen (a, a)+genDistinctPairFromList = \case+  xs@(_:_:_) -> do+    a <- Gen.element xs+    b <- Gen.element (delete a xs)+    pure (a, b)+  _ -> error "List needs more than two elements"++genListOfDistinctPairs :: Eq a => Range Int -> [a] -> Gen [(a, a)]+genListOfDistinctPairs nOpsRange vs =+  if length vs < 2+    then pure []+    else Gen.list nOpsRange (genDistinctPairFromList vs)++genV :: Int -> Gen V+genV maxElems =+  let minVal = ord 'a'+      maxVal = minVal + maxElems - 1+  in fmap V (Gen.int (Range.linear minVal maxVal))++genMembers :: Int -> Gen [V]+genMembers maxElems = do+  let nElemsRange = Range.linear 0 maxElems+      minVal = ord 'a'+  n <- Gen.int nElemsRange+  pure (fmap (\i -> V (minVal + i)) [0..n-1])++mkInitEf :: [V] -> EF+mkInitEf vs = execState (for_ vs efAdd) efNew++mkPairsMergedEf :: [(V, V)] -> EF -> EF+mkPairsMergedEf vvs = execState (for_ vvs (uncurry efMerge))++mkSetsMergedEf :: [(V, V)] -> EF -> EF+mkSetsMergedEf vvs = execState (for_ (fmap (\(x, y) -> [ILS.fromList [x, y]]) vvs) efMergeSets)++mkSingleMergedEf :: [(V, V)] -> EF -> EF+mkSingleMergedEf vvs = execState (efMergeSets (fmap (\(x, y) -> ILS.fromList [x, y]) vvs))++data MergeStrat = MergeStratPairs | MergeStratSets | MergeStratSingle+  deriving stock (Eq, Show, Enum, Bounded)++genMergeStrat :: Gen MergeStrat+genMergeStrat = Gen.enumBounded++testEfProp :: TestLimit -> TestTree+testEfProp lim = after AllSucceed "EF unit" $ testProp "EF prop" lim $ do+  let maxElems = 50+  -- generate elements+  memberList <- forAll (genMembers maxElems)+  let memberSet = ILS.fromList memberList+      nMembers = ILS.size memberSet+      allPairs = ILS.unorderedPairs memberSet+      nOpsRange = Range.linear 0 (nMembers * nMembers)+  let initEf = mkInitEf memberList+  -- assert that sizes indicate nothing is merged+  efRootsSize initEf === nMembers+  efLeavesSize initEf === 0+  efTotalSize initEf === nMembers+  -- assert that find indicates nothing is merged+  for_ allPairs $ \(a, b) -> flip evalStateT initEf $ do+    x <- applyS (gets (efFindRoot a))+    y <- applyS (gets (efFindRoot b))+    assert (isJust x)+    assert (isJust y)+    x /== y+  -- generate some pairs and merge them+  mergePairs <- forAll (genListOfDistinctPairs nOpsRange memberList)+  mergeStrat <- forAll genMergeStrat+  let mergedUf =+        case mergeStrat of+          MergeStratPairs -> mkPairsMergedEf mergePairs initEf+          MergeStratSets -> mkSetsMergedEf mergePairs initEf+          MergeStratSingle -> mkSingleMergedEf mergePairs initEf+  -- assert that total size is unchanged+  efTotalSize mergedUf === nMembers+  -- calculate components by graph reachability+  let components = ILG.undirectedComponents mergePairs+  -- assert that elements are equal or not according to component+  _ <- foldS_ mergedUf allPairs $ \(a, b) -> do+    x <- applyS (gets (efFindRoot a))+    y <- applyS (gets (efFindRoot b))+    let aComponent = ILE.lookupClass a components+        bComponent = ILE.lookupClass b components+    if isJust aComponent && aComponent == bComponent+      then x === y+      else x /== y+  pure ()++type AV = Assoc ENodeId V++-- | Asserts assoc is compact - should also check 'assertAssocInvariants'+assertAssocCompact :: (MonadTest m, Eq a, Hashable a, Show a) => Assoc ENodeId a -> m ()+assertAssocCompact av = do+  let fwd = assocFwd av+      bwd = assocBwd av+  -- Assert that the assoc has been rebuilt+  assert $ not (assocCanCompact av)+  -- Look at sizes to confirm that assoc could map 1-1+  ILM.size fwd === HashMap.size bwd+  -- Go through keys forward+  for_ (ILM.toList fwd) $ \(x, fc) -> do+    -- Assert is found in backward map AND maps back+    HashMap.lookup fc bwd === Just x+  -- Go through keys backward+  for_ (HashMap.toList bwd) $ \(fc, x) ->+    -- Assert is present in forward map AND maps back+    ILM.lookup x fwd === Just fc++-- | Asserts assoc is correctly structured (compact or not)+assertAssocInvariants :: (MonadTest m, Eq a, Hashable a) => Assoc ENodeId a -> m ()+assertAssocInvariants av = do+  let fwd = assocFwd av+      bwd = assocBwd av+      equiv = assocEquiv av+  -- First check that fwd and bwd are 1-1+  -- Go through keys forward+  for_ (ILM.toList fwd) $ \(_, fc) -> do+    -- Assert is found in backward map+    assert $ HashMap.member fc bwd+  -- Go through keys backward+  for_ (HashMap.toList bwd) $ \(_, x) ->+    -- Assert is present in forward map+    assert $ ILM.member x fwd+  -- Assert that fwd keys are exactly the equiv roots+  ILS.fromList (ILM.keys fwd) === ILS.fromList (efRoots equiv)++data AssocCase = AssocCase !String ![(Int, Char)] ![(Int, Char, Int, AssocInsertRes Int)] ![(Int, Char)]++allAssocCases :: [AssocCase]+allAssocCases =+  let start = [(0, 'a'), (1, 'b'), (2, 'c')]+  in [ AssocCase "base" start [] start+     , AssocCase "ident" start+        [(0, 'a', 0, AssocInsertResUnchanged)]+        start+     , AssocCase "superfluous" start+        [(4, 'a', 0, AssocInsertResMerged (ILS.singleton 4))]+        start+     , AssocCase "internal" start+        [(0, 'b', 0, AssocInsertResMerged (ILS.singleton 1))]+        [(0, 'b'), (2, 'c')]+     , AssocCase "external" start+        [(0, 'd', 0, AssocInsertResUpdated)]+        [(0, 'd'), (1, 'b'), (2, 'c')]+     , AssocCase "additional" start+        [(4, 'd', 4, AssocInsertResCreated)]+        [(0, 'a'), (1, 'b'), (2, 'c'), (4, 'd')]+     , AssocCase "chain fwd" start+        -- The singleton set in the second result is just the children (and self) of the clobbered node+        -- We don't have to lookup the old clobbered nodes for 1 bc when this is used everything will be merged+        [(0, 'b', 0, AssocInsertResMerged (ILS.singleton 1)), (1, 'c', 0, AssocInsertResMerged (ILS.singleton 2))]+        [(0, 'c')]+     , AssocCase "chain bwd" start+        -- The set in the second result is not a singleton here because it already had children+        [(1, 'c', 1, AssocInsertResMerged (ILS.singleton 2)), (0, 'c', 0, AssocInsertResMerged (ILS.fromList [1,2]))]+        [(0, 'c')]+     , AssocCase "chain self" start+        [(1, 'c', 1, AssocInsertResMerged (ILS.singleton 2)), (2, 'c', 1, AssocInsertResUnchanged)]+        [(0, 'a'), (1, 'c')]+     , AssocCase "chain change" start+        [(1, 'c', 1, AssocInsertResMerged (ILS.singleton 2)), (2, 'd', 1, AssocInsertResUpdated)]+        [(0, 'a'), (1, 'd')]+     , AssocCase "chain back id" start+        [(1, 'c', 1, AssocInsertResMerged (ILS.singleton 2)), (1, 'b', 1, AssocInsertResUpdated)]+        [(0, 'a'), (1, 'b')]+     , AssocCase "chain back del" start+        [(1, 'c', 1, AssocInsertResMerged (ILS.singleton 2)), (2, 'b', 1, AssocInsertResUpdated)]+        [(0, 'a'), (1, 'b')]+     , AssocCase "chain change rev" start+        [(2, 'd', 2, AssocInsertResUpdated), (1, 'c', 1, AssocInsertResUpdated)]+        [(0, 'a'), (1, 'c'), (2, 'd')]+     ]++mkAssoc :: [(Int, Char)] -> AV+mkAssoc rawPairs =+  let pairs = fmap (bimap ENodeId toV) rawPairs+  in assocFromList pairs++runAV :: Monad m => [(Int, Char)] -> StateT AV m () -> m ()+runAV = runS . mkAssoc++testAssocCase :: AssocCase -> TestTree+testAssocCase (AssocCase name start act end) = testUnit name $ runAV start $ do+  testS $ \av -> do+    assertAssocInvariants av+    assertAssocCompact av+    assocSize av === length start+  for_ act $ \(x, a, expectedY, expectedRes) -> do+    (actualY, actualRes) <- applyS (assocInsert (ENodeId x) (toV a))+    (actualY, actualRes) === coerce (expectedY, expectedRes)+    testS assertAssocInvariants+  _ <- applyS assocCompact+  testS $ \av -> do+    assertAssocInvariants av+    assertAssocCompact av+    assocSize av === length end+    let endAv = mkAssoc end+    assocFwd av === assocFwd endAv+    assocBwd av === assocBwd endAv++testAssocCases :: TestTree+testAssocCases = testGroup "Assoc case" (fmap testAssocCase allAssocCases)++testAssocUnit :: TestTree+testAssocUnit = testUnit "Assoc unit" $ do+  let a0 = assocNew :: AV+  assertAssocInvariants a0+  assertAssocCompact a0+  assocSize a0 === 0+  let aKey = ENodeId 0+      aVal = toV 'a'+      bKey = ENodeId 1+      bVal = toV 'b'+      cKey = ENodeId 2+      cVal = toV 'c'+  let members = [(aKey, aVal), (bKey, bVal), (cKey, cVal)]+  let a1 = execState (for_ members (uncurry assocInsert)) a0+  assertAssocInvariants a1+  assertAssocCompact a0+  assocSize a1 === 3+  assocRoots a1 === [aKey, bKey, cKey]+  assocLeaves a1 === []+  let (res, a2) = runState (assocInsert aKey bVal) a1+  res === (aKey, AssocInsertResMerged (ILS.singleton bKey))+  assertAssocInvariants a2+  assert $ assocCanCompact a2+  assocSize a2 === 2+  assocRoots a2 === [aKey, cKey]+  assocLeaves a2 === [bKey]+  let a3 = execState assocCompact a2+  assertAssocInvariants a3+  assertAssocCompact a3+  assocSize a3 === 2+  assocRoots a3 === [aKey, cKey]+  assocLeaves a3 === []++type EGA = EGraph () ArithF+type EGP = Pat ArithF String++testEgUnit :: TestTree+testEgUnit = after AllSucceed "Assoc unit" $ testUnit "EG unit" $ runS egNew $ do+  -- We're going to have our egraph track the equality `2 + 2 = 4`.+  -- We disable analysis+  let ana = noAnalysis+  -- Some simple terms:+  let termFour = ArithConst 4+      termTwo = ArithConst 2+      termPlus = ArithPlus termTwo termTwo+  -- And a simple pattern:+  let pat = FreeEmbed (ArithPlusF (FreePure "x") (FreePure "y")) :: EGP+  -- Test that the empty egraph is sane+  testS $ \eg -> do+    egClassSize eg === 0+    egNodeSize eg === 0+  -- Nothing is matched+  testS $ \eg ->+    match pat eg === []+  -- Add the term `4`+  cidFour <- applyTestS (egAddTerm ana termFour) $ \(c, x) eg -> do+    c === ChangedYes+    egFindTerm termFour eg === Just x+    egClassSize eg === 1+    egNodeSize eg === 1+    pure x+  -- Add the term `2`+  cidTwo <- applyTestS (egAddTerm ana termTwo) $ \(c, x) eg -> do+    c === ChangedYes+    x /== cidFour+    egFindTerm termTwo eg === Just x+    egClassSize eg === 2+    egNodeSize eg === 2+    pure x+  -- Add the term `4` again and assert things haven't changed+  applyTestS (egAddTerm ana termFour) $ \(c, x) eg -> do+    c === ChangedNo+    x === cidFour+    egFindTerm termFour eg === Just x+    egClassSize eg === 2+    egNodeSize eg === 2+  -- Still, nothing is matched+  testS $ \eg ->+    match pat eg === []+  -- Add the term `2 + 2`+  cidPlus <- applyTestS (egAddTerm ana termPlus) $ \(c, x) eg -> do+    c === ChangedYes+    x /== cidFour+    x /== cidTwo+    egFindTerm termPlus eg === Just x+    egClassSize eg === 3+    egNodeSize eg === 3+    pure x+  -- We now match `2 + 2`+  testS $ \eg ->+    match pat eg ===+      [ MatchSubst+          (Match cidPlus+            (MatchPatEmbed+              (ArithPlusF+                (Match cidTwo (MatchPatPure "x"))+                (Match cidTwo (MatchPatPure "y"))+              )+            )+          )+          (HashMap.fromList [("x", cidTwo), ("y", cidTwo)])+      ]+  -- Merge `4` and `4` and assert things haven't changed+  applyTestS (egMerge cidFour cidFour) $ \m _ -> do+    case m of+      MergeResultUnchanged -> pure ()+      _ -> fail "expected unchanged merge"+  -- Merge `2 + 2` and `4`+  applyTestS (egMerge cidPlus cidFour) $ \m eg -> do+    case m of+      MergeResultChanged _ -> pure ()+      _ -> fail "expected changed merge"+    egFindTerm termFour eg === Just cidFour+    egFindTerm termPlus eg === Just cidFour+    egFindTerm termTwo eg === Just cidTwo+  -- We still match `2 + 2`, but the class is different+  testS $ \eg ->+    match pat eg ===+      [ MatchSubst+          (Match cidFour+            (MatchPatEmbed+              (ArithPlusF+                (Match cidTwo (MatchPatPure "x"))+                (Match cidTwo (MatchPatPure "y"))+              )+            )+          )+          (HashMap.fromList [("x", cidTwo), ("y", cidTwo)])+      ]++type EGD = Max V+type EGF = BinTreeF V+type EGT = BinTree V+type EGV = EGraph EGD EGF++maxVAnalysis :: EAnalysis EGD EGF+maxVAnalysis = \case+  BinTreeLeafF v -> Max v+  BinTreeBranchF d1 d2 -> d1 <> d2++assertEgInvariants :: (MonadTest m, Traversable f, Eq (f EClassId), Hashable (f EClassId), Show (f EClassId)) => EGraph d f -> m ()+assertEgInvariants eg = do+  let assoc = egNodeAssoc eg+      hc = egHashCons eg+      bwd = assocBwd assoc+      rootNodes = ILS.fromList (assocRoots assoc)+      leafNodes = ILS.fromList (assocLeaves assoc)+      allNodes = ILS.union rootNodes leafNodes+      ef = egEquivFind eg+      rootClasses = ILS.fromList (efRoots ef)+      leafClasses = ILS.fromList (efLeaves ef)+      cm = egClassMap eg+      cmClasses = ILS.fromList (ILM.keys cm)+  -- Assert that root nodes and leaf nodes are disjoint+  ILS.intersection rootNodes leafNodes === ILS.empty+  -- Assert that root classes and leaf classes are disjoint+  ILS.intersection rootClasses leafClasses === ILS.empty+  -- Assert that the assoc is 1-1 etc+  assertAssocInvariants assoc+  -- Assert that the hashcons and assoc have equal key sets+  ILS.fromList (ILM.keys hc) === allNodes+  -- Assert that hashcons has exactly the same values as unionfind roots for all nodes+  for_ (ILM.elems hc) $ \c ->+    assert $ ILS.member c rootClasses+  -- Assert that classmap only contains unionfind roots+  cmClasses === rootClasses+  -- For every node, assert in the nodes of some class+  for_ (ILM.toList hc) $ \(n, c) -> do+    let nodes = eciNodes (ILM.partialLookup c cm)+    assert (assocMember n nodes)+  -- For every root, assert is in all parent classes+  for_ (ILM.toList hc) $ \(n, c) ->+    when (ILS.member n rootNodes) $ do+      -- for all children that are not of the node's own class+      let children = ILS.filter (/= c) (foldMap ILS.singleton (assocPartialLookupByKey n assoc))+      for_ (ILS.toList children) $ \y -> do+        -- look up child and assert in child's parents+        let parents = eciParents (ILM.partialLookup y cm)+        assert (ILS.member n parents)+  -- For every class+  cmNodes <- flipFoldM ILS.empty (ILM.toList cm) $ \accNodesSet (c, eci) -> do+    let nodes = eciNodes eci+        nodesSet = ILS.fromList (assocMembers nodes)+        parents = eciParents eci+    -- Assert that classmap node values are non-empty+    nodesSet /== ILS.empty+    -- Assert that classmap class has node values that are hashconsed to class+    for_ (ILS.toList nodesSet) $ \n -> do+      ILM.lookup n hc === Just c+    -- Assert that classmap class has NO parents that are hashconsed to class+    for_ (ILS.toList parents) $ \p ->+      ILM.lookup p hc /== Just c+    -- Assert we haven't seen these nodes before+    assert $ ILS.disjoint nodesSet accNodesSet+    -- Assert that the nodes and parents are disjoint+    assert $ ILS.disjoint nodesSet parents+    pure (ILS.union accNodesSet nodesSet)+  let hcNodes = ILS.fromList (ILM.keys hc)+  -- Assert hc keys are exactly the class nodes+  cmNodes === hcNodes+  -- Now test recanonicalization - we already know assoc fwd and bwd are 1-1+  for_ (HashMap.toList bwd) $ \(fc, _) ->+    let recanon = evalState (egCanonicalize fc) eg+    in recanon === Right fc++data EgRound = EgRound+  { egRoundTerms :: ![EGT]+  , egRoundSets :: ![[EGT]]+  , egRoundEqTests :: ![[EGT]]+  , egRoundNeqTests :: ![(EGT, EGT)]+  } deriving stock (Eq, Show)++data EgCase = EgCase+  { egCaseName :: !String+  , egCaseRounds :: ![EgRound]+  } deriving stock (Eq, Show)++allEgCases :: [EgCase]+allEgCases =+  let leafA = BinTreeLeaf (toV 'a')+      leafB = BinTreeLeaf (toV 'b')+      leafC = BinTreeLeaf (toV 'c')+      leafD = BinTreeLeaf (toV 'd')+      leafE = BinTreeLeaf (toV 'e')+      leafTerms = [leafA, leafB, leafC, leafD]+      parentAA = BinTreeBranch leafA leafA+      parentAB = BinTreeBranch leafA leafB+      parentAC = BinTreeBranch leafA leafC+      parentAD = BinTreeBranch leafA leafD+      parentBD = BinTreeBranch leafB leafD+      parentCA = BinTreeBranch leafC leafA+      simpleParentTerms = [parentAC, parentAD]+      complexParentTerms = [parentAC, parentBD]+      grandparentAAC = BinTreeBranch leafA parentAC+      grandparentAAD = BinTreeBranch leafA parentAD+      grandparentBAC = BinTreeBranch leafB parentAC+      grandparentEAD = BinTreeBranch leafE parentAD+      simpleGrandparentTerms = [grandparentAAC, grandparentAAD]+      complexGrandparentTerms = [grandparentBAC, grandparentEAD]+  in [ EgCase "simple"+          [ EgRound leafTerms [] [] [(leafA, leafB), (leafA, leafC), (leafB, leafC)]+          , EgRound [] [[leafA, leafB]] [[leafA, leafB]] [(leafA, leafC), (leafB, leafC)]+          ]+     , EgCase "transitive one round"+        [ EgRound leafTerms [] [] [(leafA, leafB), (leafA, leafC), (leafB, leafC), (leafA, leafD)]+        , EgRound [] [[leafA, leafB], [leafB, leafC]] [[leafA, leafB, leafC]] [(leafA, leafD)]+        ]+     , EgCase "transitive two round"+        [ EgRound leafTerms [] [] [(leafA, leafB), (leafA, leafC), (leafB, leafC), (leafA, leafD)]+        , EgRound [] [[leafA, leafB]] [[leafA, leafB]] [(leafA, leafC), (leafA, leafD)]+        , EgRound [] [[leafB, leafC]] [[leafA, leafB, leafC]] [(leafA, leafD)]+        ]+     , EgCase "simple parents"+        [ EgRound simpleParentTerms [] [] [(leafC, leafD), (parentAC, parentAD)]+        , EgRound [] [[leafC, leafD]] [[parentAC, parentAD]] []+        ]+     , EgCase "complex parents one round"+        [ EgRound complexParentTerms [] [] [(leafA, leafB), (leafC, leafD), (parentAC, parentBD)]+        , EgRound [] [[leafA, leafB], [leafC, leafD]] [[leafA, leafB], [leafC, leafD], [parentAC, parentBD]] []+        ]+     , EgCase "complex parents two round"+        [ EgRound complexParentTerms [] [] [(leafA, leafB), (leafC, leafD), (parentAC, parentBD)]+        , EgRound [] [[leafA, leafB]] [[leafA, leafB]] [(leafC, leafD), (parentAC, parentBD)]+        , EgRound [] [[leafC, leafD]] [[leafA, leafB], [leafC, leafD], [parentAC, parentBD]] []+        ]+     , EgCase "simple grandparents"+        [ EgRound simpleGrandparentTerms [] [] [(leafC, leafD), (parentAC, parentAD), (grandparentAAC, grandparentAAD)]+        , EgRound [] [[leafC, leafD]] [[leafC, leafD], [parentAC, parentAD], [grandparentAAC, grandparentAAD]] []+        ]+     , EgCase "complex grandparents bottom up"+        [ EgRound complexGrandparentTerms [] [] [(leafC, leafD), (leafB, leafE), (parentAC, parentAD), (grandparentBAC, grandparentEAD)]+        , EgRound [] [[leafC, leafD]] [[leafC, leafD], [parentAC, parentAD]] [(leafB, leafE), (grandparentBAC, grandparentEAD)]+        , EgRound [] [[leafB, leafE]] [[leafC, leafD], [leafB, leafE], [parentAC, parentAD], [grandparentBAC, grandparentEAD]] []+        ]+     , EgCase "complex grandparents top down"+        [ EgRound complexGrandparentTerms [] [] [(leafC, leafD), (leafB, leafE), (parentAC, parentAD), (grandparentBAC, grandparentEAD)]+        , EgRound [] [[leafB, leafE]] [[leafB, leafE]] [(leafC, leafD), (parentAC, parentAD), (grandparentBAC, grandparentEAD)]+        , EgRound [] [[leafC, leafD]] [[leafC, leafD], [leafB, leafE], [parentAC, parentAD], [grandparentBAC, grandparentEAD]] []+        ]+     , EgCase "connect"+        [ EgRound leafTerms [] [] [(leafA, leafB), (leafA, leafC), (leafB, leafC), (leafA, leafD)]+        , EgRound [] [[leafA, leafB]] [[leafA, leafB]] [(leafA, leafC), (leafA, leafD)]+        , EgRound [] [[leafC, leafD]] [[leafA, leafB], [leafC, leafD]] [(leafA, leafD)]+        , EgRound [] [[leafB, leafD]] [[leafA, leafB, leafC, leafD]] []+        ]+     , EgCase "mid grandparents"+        [ EgRound simpleGrandparentTerms [] [] [(leafC, leafD), (parentAC, parentAD), (grandparentAAC, grandparentAAD)]+        , EgRound [] [[parentAC, parentAD]] [[parentAC, parentAD], [grandparentAAC, grandparentAAD]] [(leafC, leafD)]+        , EgRound [] [[leafC, leafD]] [[leafC, leafD], [parentAC, parentAD], [grandparentAAC, grandparentAAD]] []+        ]+     , EgCase "unify node"+        [ EgRound [BinTreeBranch parentAC leafA, parentAA] [] [] [(parentAC, parentAA)]+        , EgRound [] [[leafA, leafC]] [[parentAC, parentAA]] []+        ]+     , EgCase "self parent"+        [ EgRound [BinTreeBranch parentAC leafB] [] [] []+        , EgRound [] [[parentAC, leafA]] [] []+        ]+     , EgCase "self parent again"+        [ EgRound [leafB, parentAA] [] [] []+        , EgRound [] [[leafB, leafA], [leafB, parentAA]] [] []+        ]+     , EgCase "dead add parent"+        [ EgRound [parentAC, leafB] [[leafA, parentAC], [leafA, leafC]] [[parentAC, leafC]] [(leafA, leafB)]+        , EgRound [parentCA] [] [[parentCA, parentAC]] []+        ]+     , EgCase "repro 1"+        [ EgRound [BinTreeBranch parentAA parentAB] [[leafB, leafA]] [] [(leafB, parentAA)]+        , EgRound [leafA] [[parentAA, leafA]] [[leafB, parentAA]] []+        ]+     ,  let grandparent = BinTreeBranch parentAA leafB+            greatGrandparent = BinTreeBranch grandparent leafA+        in EgCase "repro 2"+          [ EgRound [greatGrandparent, leafA, leafC] [[greatGrandparent, parentAA], [leafA, leafB]] [] []+          , EgRound [leafA, leafA] [[leafA, grandparent]] [] []+          , EgRound [leafA, leafA, leafA] [[parentAA, leafA]] [] []+          ]+     ]++testEgCase :: EgCase -> TestTree+testEgCase (EgCase name rounds) = kase where+  findMayTerm t = fmap (egFindTerm t) get+  findTerm t = fmap fromJust (findMayTerm t)+  findTerms ts = fmap ILS.fromList (for ts findTerm)+  assertTermFound t = findMayTerm t >>= \mi -> assert (isJust mi)+  assertTermsFound ts = for_ ts assertTermFound+  kase = testUnit name $ runS egNew $ do+    -- for each round+    for_ rounds $ \(EgRound start act endEq endNeq) -> do+      -- add initial terms and assert invariants hold+      applyS (for_ start (egAddTerm maxVAnalysis))+      -- liftIO (putStrLn "===== post add =====")+      -- testS $ liftIO . pPrint+      testS assertEgInvariants+      -- assert that all mentioned terms are in the egraph+      assertTermsFound start+      for_ act assertTermsFound+      for_ endEq assertTermsFound+      for_ endNeq $ \(x, y) -> do+        -- also assert that neq terms are not themselves equal+        x /== y+        assertTermFound x+        assertTermFound y+      -- merge sets of terms and rebuild+      applyS $ do+        sets <- for act findTerms+        mr <- egMergeMany (Seq.fromList sets)+        case mr of+          MergeResultMissing _ -> error "bad set"+          _ -> pure ()+      -- assert invariants hold+      testS assertEgInvariants+      -- find merged terms again and assert they are in same classes+      sets <- applyS $ for act findTerms+      for_ sets $ \set -> ILS.size set === 1+      -- find final eq terms and assert they are in same classes+      for_ endEq $ \ts -> do+        set <- applyS (findTerms ts)+        ILS.size set === 1+      -- find final neq terms and assert they are not in same class+      for_ endNeq $ \(x, y) -> do+        i <- applyS (findTerm x)+        j <- applyS (findTerm y)+        i /== j++testEgCases :: TestTree+testEgCases = testGroup "Eg case" $+  testEgCase <$> allEgCases++testEgNew :: TestTree+testEgNew = testUnit "EG new" $ do+  eg0 <- fullyEvaluate (egNew :: EGV)+  egNodeSize eg0 === 0+  egClassSize eg0 === 0+  assertEgInvariants eg0++genNodePairs :: Range Int -> EGV -> Gen [(EClassId, EClassId)]+genNodePairs nOpsRange eg = genListOfDistinctPairs nOpsRange (ILM.keys (egClassMap eg))++genSomeList :: [a] -> Gen [a]+genSomeList xs = go where+  go = Gen.recursive Gen.choice [Gen.constant [], fmap pure (Gen.element xs)] [Gen.subterm2 go go (++)]++genBinTree :: Gen a -> Gen (BinTree a)+genBinTree genA = genEither where+  genLeaf = fmap BinTreeLeaf genA+  genBranch = Gen.subterm2 genEither genEither BinTreeBranch+  genEither = Gen.recursive Gen.choice [genLeaf] [genBranch]++genBinTreeMembers :: Int -> Gen [BinTree V]+genBinTreeMembers maxElems = Gen.list (Range.linear 0 maxElems) (genBinTree (genV maxElems))++-- An alternative to 'genBinTreeMembers' that makes smaller trees+mkSimpleTreeLevels :: Int -> [BinTree V]+mkSimpleTreeLevels maxElems =+  let letters = take maxElems (['a'..'z'] ++ ['A'..'Z'])+      zeroLevel = fmap (BinTreeLeaf . toV) letters+      mkLevel y x = (BinTreeBranch <$> x <*> y) ++ (BinTreeBranch <$> y <*> x)+      mkLevels y = foldr (\x r -> mkLevel y x ++ r) (BinTreeBranch <$> y <*> y)+      oneLevel = mkLevels zeroLevel []+      twoLevel = mkLevels oneLevel [zeroLevel]+      anyLevel = zeroLevel ++ oneLevel ++ twoLevel+  in anyLevel++testEgProp :: TestLimit -> TestTree+testEgProp lim = after AllSucceed "EG unit" $ after AllSucceed "EG cases" $ testProp "EG prop" lim prop where+  maxElems = 10+  termGen = genBinTreeMembers maxElems+  -- Guarantee yourself small trees with this:+  -- termGen = genSomeList (mkSimpleTreeLevels maxElems)+  prop = do+    eg0 <- fullyEvaluate (egNew :: EGV)+    rounds <- forAll (Gen.element [1, 2, 3])+    body rounds eg0+  body (rounds :: Int) eg0 = do+    members <- forAll termGen+    let nMembers = length members+        nOpsRange = Range.linear 0 (nMembers * nMembers)+    eg1 <- fullyEvaluate (execState (for_ members (egAddTerm maxVAnalysis)) eg0)+    assertEgInvariants eg1+    pairs <- forAll (genNodePairs nOpsRange eg1)+    let merge = do+          mr <- egMergeMany (Seq.fromList (fmap (\(a, b) -> ILS.fromList [a, b]) pairs))+          case mr of+            MergeResultMissing _ -> error "bad set"+            _ -> pure ()+    eg2 <- fullyEvaluate (execState merge eg1)+    assertEgInvariants eg2+    unless (rounds == 1) (body (rounds - 1) eg2)++type M = IntLikeMap ENodeId Char++testILM :: TestTree+testILM = testUnit "ILM unit" $ do+  let mLeft = ILM.fromList [(ENodeId 0, 'a'), (ENodeId 1, 'b')] :: M+      mRight = ILM.fromList [(ENodeId 1, 'x'), (ENodeId 2, 'c')] :: M+      mMerged = ILM.fromList [(ENodeId 0, 'a'), (ENodeId 1, 'b'), (ENodeId 2, 'c')] :: M+  mLeft <> mRight === mMerged++main :: IO ()+main = testMain $ \lim -> testGroup "Overeasy"+    [ testILM+    , testEfUnit+    , testAssocUnit+    , testAssocCases+    , testEgUnit+    , testEgNew+    , testEgCases+    , testEfProp lim+    , testEgProp lim+    ]

@@ -0,0 +1,30 @@+{-# LANGUAGE DeriveAnyClass #-}+{-# LANGUAGE TemplateHaskell #-}++module Test.Overeasy.Arith+  ( ArithF (..)+  , Arith (..)+  ) where++import Control.DeepSeq (NFData)+import Data.Functor.Foldable.TH (makeBaseFunctor)+import Data.Hashable (Hashable)+import GHC.Generics (Generic)++data Arith =+    ArithPlus Arith Arith+  | ArithTimes Arith Arith+  | ArithShiftL Arith !Int+  | ArithShiftR Arith !Int+  | ArithConst !Int+  deriving stock (Eq, Show, Generic)+  deriving anyclass (Hashable, NFData)++-- Generates 'ArithF' and other recursion-schemes boilerplate+makeBaseFunctor ''Arith++deriving stock instance Eq a => Eq (ArithF a)+deriving stock instance Show a => Show (ArithF a)+deriving stock instance Generic (ArithF a)+deriving anyclass instance Hashable a => Hashable (ArithF a)+deriving anyclass instance NFData a => NFData (ArithF a)

@@ -0,0 +1,69 @@+{-# LANGUAGE DeriveAnyClass #-}++module Test.Overeasy.BinTree+  ( BinTreeF (..)+  , BinTree (..)+  , pattern BinTreeLeaf+  , pattern BinTreeBranch+  ) where++import Control.Applicative (liftA2)+import Control.DeepSeq (NFData)+import Data.Bifoldable (Bifoldable (..))+import Data.Bifunctor (Bifunctor (..))+import Data.Bitraversable (Bitraversable (..))+import Data.Functor.Foldable (Base, Corecursive (..), Recursive (..))+import Data.Hashable (Hashable)+import GHC.Generics (Generic)++data BinTreeF a r =+    BinTreeLeafF !a+  | BinTreeBranchF r r+  deriving stock (Eq, Ord, Show, Functor, Foldable, Traversable, Generic)+  deriving anyclass (Hashable, NFData)++instance Bifunctor BinTreeF where+  bimap f g = \case+    BinTreeLeafF a -> BinTreeLeafF (f a)+    BinTreeBranchF x y -> BinTreeBranchF (g x) (g y)++instance Bifoldable BinTreeF where+  bifoldr f g z = \case+    BinTreeLeafF a -> f a z+    BinTreeBranchF x y -> g x (g y z)++instance Bitraversable BinTreeF where+  bitraverse f g = \case+    BinTreeLeafF a -> fmap BinTreeLeafF (f a)+    BinTreeBranchF x y -> liftA2 BinTreeBranchF (g x) (g y)++newtype BinTree a = BinTree { unBinTree :: BinTreeF a (BinTree a) }+  deriving newtype (Eq, Ord, Show, Hashable, NFData)++pattern BinTreeLeaf :: a -> BinTree a+pattern BinTreeLeaf a = BinTree (BinTreeLeafF a)++pattern BinTreeBranch :: BinTree a -> BinTree a -> BinTree a+pattern BinTreeBranch a b = BinTree (BinTreeBranchF a b)++{-# COMPLETE BinTreeLeaf, BinTreeBranch #-}++instance Functor BinTree where+  fmap f = go where+    go = BinTree . bimap f go . unBinTree++instance Foldable BinTree where+  foldr f z0 x0 = go x0 z0 where+    go x z = bifoldr f go z (unBinTree x)++instance Traversable BinTree where+  traverse f = go where+    go = fmap BinTree . bitraverse f go . unBinTree++type instance Base (BinTree a) = BinTreeF a++instance Recursive (BinTree a) where+  project = unBinTree++instance Corecursive (BinTree a) where+  embed = BinTree