hegg 0.1.0.0 → 0.2.0.0

raw patch · 23 files changed

+678/−325 lines, 23 filesdep +tasty-benchdep ~basedep ~containersdep ~deriving-compatPVP ok

version bump matches the API change (PVP)

Dependencies added: tasty-bench

Dependency ranges changed: base, containers, deriving-compat, hegg, tasty, tasty-hunit, tasty-quickcheck

API changes (from Hackage documentation)

Files

• CHANGELOG.md +13/−2
• README.md +233/−0
• hegg.cabal +33/−11
• src/Data/Equality/Analysis.hs +4/−2
• src/Data/Equality/Extraction.hs +27/−26
• src/Data/Equality/Graph.hs +52/−74
• src/Data/Equality/Graph.hs-boot +0/−22
• src/Data/Equality/Graph/Classes.hs +4/−2
• src/Data/Equality/Graph/Internal.hs +41/−0
• src/Data/Equality/Graph/Internal.hs-boot +9/−0
• src/Data/Equality/Graph/Lens.hs +13/−8
• src/Data/Equality/Graph/Nodes.hs +9/−15
• src/Data/Equality/Graph/ReprUnionFind.hs +4/−7
• src/Data/Equality/Matching.hs +4/−5
• src/Data/Equality/Matching/Database.hs +23/−44
• src/Data/Equality/Saturation.hs +103/−90
• src/Data/Equality/Saturation/Scheduler.hs +0/−1
• src/Data/Equality/Utils/IntToIntMap.hs +0/−1
• src/Data/Equality/Utils/SizedList.hs +66/−0
• test/Bench.hs +19/−0
• test/Invariants.hs +9/−8
• test/SimpleSym.hs +1/−1
• test/Sym.hs +11/−6

@@ -1,5 +1,16 @@-# Revision history for hsym+# Revision history for hegg -## 0.1.0.0 -- YYYY-mm-dd+## 0.2.0.0 -- 2022-09-19++* Expose `runEqualitySaturation` to run equality saturation on existing e-graphs+    whole instead of focusing on individual expressions+* (Very) significant performance improvements!+* Make `CostFunction` polymorphic over the `Cost` type, requiring that type+    to instance `Ord`+* Make e-graph abstract. The internal structure can still be modified through+    the available lenses in `Data.Equality.Graph.Lens`+* Fix a bug related to `NodeMap`'s size.++## 0.1.0.0 -- 2022-08-25  * First version. Released on an unsuspecting world.

@@ -0,0 +1,233 @@+## hegg++Fast equality saturation in Haskell++Based on [*egg: Fast and Extensible Equality Saturation*](https://arxiv.org/pdf/2004.03082.pdf), [*Relational E-matching*](https://arxiv.org/pdf/2108.02290.pdf) and the [rust implementation](https://github.com/egraphs-good/egg).++### Equality Saturation and E-graphs++Suggested material on equality saturation and e-graphs for beginners+* (tutorial) https://docs.rs/egg/latest/egg/tutorials/_01_background/index.html+* (5m video) https://www.youtube.com/watch?v=ap29SzDAzP0++## Equality saturation in Haskell++To get a feel for how we can use `hegg` and do equality saturation in Haskell,+we'll write a simple numeric *symbolic* manipulation library that can simplify expressions+according to a set of rewrite rules by leveraging equality saturation.++If you've never heard of symbolic mathematics you might get some intuition from+reading [Let’s Program a Calculus+Student](https://iagoleal.com/posts/calculus-symbolic/) first.++### Syntax++We'll start by defining the abstract syntax tree for our simple symbolic expressions:+```hs+data SymExpr = Const Double+             | Symbol String+             | SymExpr :+: SymExpr+             | SymExpr :*: SymExpr+             | SymExpr :/: SymExpr+infix 6 :+:+infix 7 :*:, :/:++e1 :: SymExpr+e1 = (Symbol "x" :*: Const 2) :/: (Const 2) -- (x*2)/2+```++You might notice that `(x*2)/2` is the same as just `x`. Our goal is to get+equality saturation to do that for us.++Our second step is to instance `Language` for our `SymExpr`++### Language++`Language` is the required constraint on *expressions* that are to be+represented in e-graph and on which equality saturation can be run:++```hs+class (Analysis l, Traversable l, Ord1 l) => Language l+```++To declare a `Language` we must write the "base functor" of `SymExpr` +(i.e. use a type parameter where the recursion points used to be in the original `SymExpr`),+then instance `Traversable`, `Ord1`, and write an `Analysis` instance for it (see next section).++```hs+data SymExpr a = Const Double+               | Symbol String+               | a :+: a+               | a :*: a+               | a :/: a+               deriving (Functor, Foldable, Traversable)+infix 6 :+:+infix 7 :*:, :/:+```++Suggested reading on defining recursive data types in their parametrized+version: [Introduction To Recursion+Schemes](https://blog.sumtypeofway.com/posts/introduction-to-recursion-schemes.html)++If we now wanted to represent an expression, we'd write it in its+fixed-point form++```hs+e1 :: Fix SymExpr+e1 = Fix (Fix (Fix (Symbol "x") :*: Fix (Const 2)) :/: (Fix (Const 2))) -- (x*2)/2+```++We've already automagically derived `Functor`, `Foldable` and `Traversable`+instances, and can use the following template haskell functions from `derive-compat` to derive `Ord1`.+```hs+deriveEq1   ''SymExpr+deriveOrd1  ''SymExpr+```++Then, we define an `Analysis` for our `SymExpr`.++### Analysis++E-class analysis is first described in [*egg: Fast and Extensible Equality+Saturation*](https://arxiv.org/pdf/2004.03082.pdf) as a way to make equality+saturation more *extensible*.++With it, we can attach *analysis data* from a semilattice to each e-class. More+can be read about e-class analysis in the [`Data.Equality.Analsysis`]() module and+in the paper.++We could easily define constant folding (`2+2` being simplified to `4`) through+an `Analysis` instance, but for the sake of simplicity we'll simply define the+analysis data as `()` and always ignore it.++```hs+instance Analysis SymExpr where+  type Domain SymExpr = ()+  makeA _ _ = ()+  joinA _ _ = ()+```++### Language, again++With this setup, we can now express that `SymExpr` forms a `Language` which we+can represent and manipulate in an e-graph by simply instancing it (there are no+additional functions to define).+```hs+instance Language SymExpr+```++### Equality saturation++Equality saturation is defined as the function+```hs+equalitySaturation :: forall l. Language l+                   => Fix l             -- ^ Expression to run equality saturation on+                   -> [Rewrite l]       -- ^ List of rewrite rules+                   -> CostFunction l    -- ^ Cost function to extract the best equivalent representation+                   -> (Fix l, EGraph l) -- ^ Best equivalent expression and resulting e-graph+```++To recap, our goal is to reach `x` starting from `(x*2)/2` by means of equality+saturation.++We already have a starting expression, so we're missing a list of rewrite rules+(`[Rewrite l]`) and a cost function (`CostFunction`).++### Cost function++Picking up the easy one first:+```hs+type CostFunction l cost = l cost -> cost+```++A cost function is used to attribute a cost to representations in the e-graph and to extract the best one.+The first type parameter `l` is the language we're going to attribute a cost to, and+the second type parameter `cost` is the type with which we will model cost. For+the cost function to be valid, `cost` must instance `Ord`.++We'll say `Const`s and `Symbol`s are the cheapest and then in increasing cost we+have `:+:`, `:*:` and `:/:`, and model cost with the `Int` type.+```hs+cost :: CostFunction SymExpr Int+cost = \case+  Const  x -> 1+  Symbol x -> 1+  c1 :+: c2 -> c1 + c2 + 2+  c1 :*: c2 -> c1 + c2 + 3+  c1 :/: c2 -> c1 + c2 + 4+```++### Rewrite rules++Rewrite rules are transformations applied to matching expressions represented in+an e-graph.++We can write simple rewrite rules and conditional rewrite rules, but we'll only look at the simple ones.++A simple rewrite is formed of its left hand side and right hand side. When the+left hand side is matched in the e-graph, the right hand side is added to the+e-class where the left hand side was found.+```hs+data Rewrite lang = Pattern lang := Pattern lang          -- Simple rewrite rule+                  | Rewrite lang :| RewriteCondition lang -- Conditional rewrite rule+```++A `Pattern` is basically an expression that might contain variables and which can be matched against actual expressions.+```hs+data Pattern lang+    = NonVariablePattern (lang (Pattern lang))+    | VariablePattern Var+```+A patterns is defined by its non-variable and variable parts, and can be+constructed directly or using the helper function `pat` and using+`OverloadedStrings` for the variables, where `pat` is just a synonym for+`NonVariablePattern` and a string literal `"abc"` is turned into a `Pattern`+constructed with `VariablePattern`.++We can then write the following very specific set of rewrite rules to simplify+our simple symbolic expressions.+```hs+rewrites :: [Rewrite SymExpr]+rewrites =+  [ pat (pat ("a" :*: "b") :/: "c") := pat ("a" :*: pat ("b" :/: "c"))+  , pat ("x" :/: "x")               := pat (Const 1)+  , pat ("x" :*: (pat (Const 1)))   := "x"+  ]+```+### Equality saturation, again++We can now run equality saturation on our expression!++```hs+let expr = fst (equalitySaturation e1 rewrites cost)+```+And upon printing we'd see `expr = Symbol "x"`!++This was a first introduction which skipped over some details but that tried to+walk through fundamental concepts for using e-graphs and equality saturation+with this library.++The final code for this tutorial is available under `test/SimpleSym.hs`++A more complicated symbolic rewrite system which simplifies some derivatives and+integrals was written for the testsuite. It can be found at `test/Sym.hs`.++This library could also be used not only for equality-saturation but also for+the equality-graphs and other equality-things (such as e-matching) available.+For example, using just the e-graphs from `Data.Equality.Graph` to improve GHC's+pattern match checker (https://gitlab.haskell.org/ghc/ghc/-/issues/19272).++## Profiling++Notes on profiling for development.++For producing the info table, ghc-options must include `-finfo-table-map+-fdistinct-constructor-tables`++```+cabal run --enable-profiling hegg-test -- +RTS -p -s -hi -l-agu+ghc-prof-flamegraph hegg-test.prof+eventlog2html hegg-test.eventlog+open hegg-test.svg+open hegg-test.eventlog.html+```

@@ -1,24 +1,28 @@ cabal-version:      2.4 name:               hegg-version:            0.1.0.0+version:            0.2.0.0 Tested-With:        GHC ==9.4.1 || ==9.2.2 || ==9.0.2 || ==8.10.7 synopsis:           Fast equality saturation in Haskell  description:        Fast equality saturation and equality graphs based on "egg:                     Fast and Extensible Equality Saturation" and "Relational E-matching".                     .-                    This package provides e-graphs (see 'Data.Equality.Graph'),+                    This package provides e-graphs (see "Data.Equality.Graph"),                     a data structure which efficiently represents a congruence-                    relation over many expressions+                    relation over many expressions.                     .+                    For a monadic interface to e-graphs check out+                    "Data.Equality.Graph.Monad" (home to the convenient+                    function 'represent').+                    .                     Secondly, it provides functions for doing equality-                    saturation (see 'Data.Equality.Saturation'), an+                    saturation (see "Data.Equality.Saturation"), an                     optimization/term-rewriting technique that applies rewrite                     rules non-destructively to an expression represented in an                     e-graph until saturation, and then extracts the best                     representation.                     .-                    Equality matching (see 'Data.Equality.Matching') is done as+                    Equality matching (see "Data.Equality.Matching") is done as                     described in "Relational E-Matching"                     .                     For a walkthrough of writing a simple symbolic@@ -35,6 +39,7 @@ copyright:          Copyright (C) 2022 Rodrigo Mesquita category:           Data extra-source-files: CHANGELOG.md+                    README.md  source-repository head     type: git@@ -58,6 +63,7 @@                       -- -dsuppress-var-kinds      exposed-modules:  Data.Equality.Graph,+                      Data.Equality.Graph.Internal,                       Data.Equality.Graph.ReprUnionFind,                       Data.Equality.Graph.Classes,                       Data.Equality.Graph.Classes.Id,@@ -73,7 +79,8 @@                       Data.Equality.Analysis,                       Data.Equality.Saturation.Scheduler,                       Data.Equality.Saturation.Rewrites,-                      Data.Equality.Utils+                      Data.Equality.Utils,+                      Data.Equality.Utils.SizedList     if impl(ghc >= 9.2)         exposed-modules: Data.Equality.Utils.IntToIntMap @@ -93,22 +100,37 @@     default-language: Haskell2010  test-suite hegg-test-    ghc-options:      -threaded -Wall+    ghc-options:      -Wall                       -- -finfo-table-map -fdistinct-constructor-tables-                      -- -threaded+                      -threaded     default-language: Haskell2010     type:             exitcode-stdio-1.0     hs-source-dirs:   test     main-is:          Test.hs     other-modules:    Invariants, Sym, Lambda, SimpleSym     other-extensions: OverloadedStrings-    build-depends:    base >= 4.4 && < 5,-                      hegg >= 0.1 && < 0.2,-                      containers >= 0.4 && < 0.7,+    build-depends:    base,+                      hegg,+                      containers,                       deriving-compat  >= 0.6 && < 0.7,                       tasty            >= 1.4 && < 1.5,                       tasty-hunit      >= 0.10 && < 0.11,                       tasty-quickcheck >= 0.10 && < 0.11++benchmark hegg-bench+  default-language: Haskell2010+  hs-source-dirs: test+  other-modules:  Invariants, Sym, Lambda, SimpleSym+  main-is:        Bench.hs+  type:           exitcode-stdio-1.0+  build-depends:  base, hegg,+                  containers,+                  deriving-compat,+                  tasty,+                  tasty-hunit,+                  tasty-quickcheck,+                  tasty-bench >= 0.2  && < 0.4+  ghc-options:    -with-rtsopts=-A32m -threaded  Flag vizdot     Description: Compile 'Data.Equality.Graph.Dot' module to visualize e-graphs

@@ -22,13 +22,15 @@ -} module Data.Equality.Analysis where +import Data.Kind (Type)+ import Data.Equality.Graph.Classes.Id import Data.Equality.Graph.Nodes -import {-# SOURCE #-} Data.Equality.Graph (EGraph)+import {-# SOURCE #-} Data.Equality.Graph.Internal (EGraph)  -- | The e-class analysis defined for a language @l@.-class Eq (Domain l) => Analysis l where+class Eq (Domain l) => Analysis (l :: Type -> Type) where      -- | Domain of data stored in e-class according to e-class analysis     type Domain l

@@ -21,7 +21,6 @@    -- * Cost   , CostFunction-  , Cost   , depthCost   ) where @@ -30,6 +29,7 @@  import Data.Equality.Utils import Data.Equality.Graph+import Data.Equality.Graph.Lens  -- vvvv and necessarily all the best sub-expressions from children equilalence classes @@ -45,18 +45,19 @@ -- @ -- -- For a real example you might want to check out the source code of 'Data.Equality.Saturation.equalitySaturation''-extractBest :: forall lang. Language lang-            => EGraph lang       -- ^ The e-graph out of which we are extracting an expression-            -> CostFunction lang -- ^ The cost function to define /best/-            -> ClassId           -- ^ The e-class from which we'll extract the expression-            -> Fix lang          -- ^ The resulting /best/ expression, in its fixed point form.-extractBest g@EGraph{classes = eclasses'} cost (flip find g -> i) = +extractBest :: forall lang cost+             . (Language lang, Ord cost)+            => EGraph lang            -- ^ The e-graph out of which we are extracting an expression+            -> CostFunction lang cost -- ^ The cost function to define /best/+            -> ClassId                -- ^ The e-class from which we'll extract the expression+            -> Fix lang               -- ^ The resulting /best/ expression, in its fixed point form.+extractBest egr cost (flip find egr -> i) =       -- Use `egg`s strategy of find costs for all possible classes and then just     -- picking up the best from the target e-class.  In practice this shouldn't     -- find the cost of unused nodes because the "topmost" e-class will be the     -- target, and all sub-classes must be calculated?-    let allCosts = findCosts eclasses' mempty+    let allCosts = findCosts (egr^._classes) mempty       in case findBest i allCosts of         Just (CostWithExpr (_,n)) -> n@@ -65,15 +66,14 @@   where      -- | Find the lowest cost of all e-classes in an e-graph in an extraction-    findCosts :: ClassIdMap (EClass lang) -> ClassIdMap (CostWithExpr lang) -> ClassIdMap (CostWithExpr lang)+    findCosts :: ClassIdMap (EClass lang) -> ClassIdMap (CostWithExpr lang cost) -> ClassIdMap (CostWithExpr lang cost)     findCosts eclasses current =        let (modified, updated) = IM.foldlWithKey f (False, current) eclasses            {-# INLINE f #-}-          f :: (Bool, ClassIdMap (CostWithExpr lang)) -> Int -> EClass lang -> (Bool, ClassIdMap (CostWithExpr lang))-          f = \acc@(_, beingUpdated) i' (EClass _ nodes _ _) ->-+          f :: (Bool, ClassIdMap (CostWithExpr lang cost)) -> Int -> EClass lang -> (Bool, ClassIdMap (CostWithExpr lang cost))+          f = \acc@(_, beingUpdated) i' EClass{eClassNodes = nodes} ->                 let                     currentCost = IM.lookup i' beingUpdated @@ -104,20 +104,24 @@     -- For a node to have a cost, all its (canonical) sub-classes have a cost and     -- an associated better expression. We return the constructed best expression     -- with its cost-    nodeTotalCost :: Traversable lang => ClassIdMap (CostWithExpr lang) -> ENode lang -> Maybe (CostWithExpr lang)+    nodeTotalCost :: Traversable lang => ClassIdMap (CostWithExpr lang cost) -> ENode lang -> Maybe (CostWithExpr lang cost)     nodeTotalCost m (Node n) = do-        expr <- traverse ((`IM.lookup` m) . flip find g) n+        expr <- traverse ((`IM.lookup` m) . flip find egr) n         return $ CostWithExpr (cost ((fst . unCWE) <$> expr), (Fix $ (snd . unCWE) <$> expr))     {-# INLINE nodeTotalCost #-}--{-# SCC extractBest #-}+{-# INLINABLE extractBest #-}  -- | A cost function is used to attribute a cost to representations in the -- e-graph and to extract the best one. --+-- The cost function is polymorphic over the type used for the cost, however+-- @cost@ must instance 'Ord' in order for the defined 'CostFunction' to+-- fulfill its purpose. That's why we have an @Ord cost@ constraint in+-- 'Data.Equality.Saturation.equalitySaturation' and 'extractBest'+-- -- === Example -- @--- symCost :: Expr Cost -> Cost+-- symCost :: Expr Int -> Int -- symCost = \case --     BinOp Integral e1 e2 -> e1 + e2 + 20000 --     BinOp Diff e1 e2 -> e1 + e2 + 500@@ -126,29 +130,26 @@ --     Sym _ -> 1 --     Const _ -> 1 -- @-type CostFunction l = l Cost -> Cost---- | 'Cost' is simply an integer-type Cost = Int+type CostFunction l cost = l cost -> cost  -- | Simple cost function: the deeper the expression, the bigger the cost-depthCost :: Language l => CostFunction l+depthCost :: Language l => CostFunction l Int depthCost = (+1) . sum {-# INLINE depthCost #-}  -- | Find the current best node and its cost in an equivalence class given only the class and the current extraction -- This is not necessarily the best node in the e-graph, only the best in the current extraction state-findBest :: ClassId -> ClassIdMap (CostWithExpr lang) -> Maybe (CostWithExpr lang)+findBest :: ClassId -> ClassIdMap (CostWithExpr lang a) -> Maybe (CostWithExpr lang a) findBest i = IM.lookup i {-# INLINE findBest #-} -newtype CostWithExpr lang = CostWithExpr { unCWE :: (Cost, Fix lang) }+newtype CostWithExpr lang a = CostWithExpr { unCWE :: (a, Fix lang) } -instance Eq (CostWithExpr lang) where+instance Eq a => Eq (CostWithExpr lang a) where   (==) (CostWithExpr (a,_)) (CostWithExpr (b,_)) = a == b   {-# INLINE (==) #-} -instance Ord (CostWithExpr lang) where+instance Ord a => Ord (CostWithExpr lang a) where   compare (CostWithExpr (a,_)) (CostWithExpr (b,_)) = a `compare` b   {-# INLINE compare #-} 

@@ -3,7 +3,6 @@ {-# LANGUAGE TupleSections #-} -- {-# LANGUAGE ApplicativeDo #-} {-# LANGUAGE BlockArguments #-}-{-# LANGUAGE UndecidableInstances #-} -- tmp show {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-}@@ -15,9 +14,7 @@ module Data.Equality.Graph     (       -- * Definition of e-graph-      EGraph(..)--    , Memo, Worklist+      EGraph        -- * Functions on e-graphs     , emptyEGraph@@ -38,12 +35,15 @@ -- import GHC.Conc  import Data.Function--import Data.Functor.Classes+import Data.Bifunctor+import Data.Containers.ListUtils  import qualified Data.IntMap.Strict as IM import qualified Data.Set    as S +import Data.Equality.Utils.SizedList++import Data.Equality.Graph.Internal import Data.Equality.Graph.ReprUnionFind import Data.Equality.Graph.Classes import Data.Equality.Graph.Nodes@@ -51,50 +51,23 @@ import Data.Equality.Language import Data.Equality.Graph.Lens --- | E-graph representing terms of language @l@.------ Intuitively, an e-graph is a set of equivalence classes (e-classes). Each e-class is a--- set of e-nodes representing equivalent terms from a given language, and an e-node is a function--- symbol paired with a list of children e-classes.-data EGraph l = EGraph-    { unionFind :: !ReprUnionFind           -- ^ Union find like structure to find canonical representation of an e-class id-    , classes   :: !(ClassIdMap (EClass l)) -- ^ Map canonical e-class ids to their e-classes-    , memo      :: !(Memo l)                -- ^ Hashcons maps all canonical e-nodes to their e-class ids-    , worklist  :: !(Worklist l)            -- ^ Worklist of e-class ids that need to be upward merged-    , analysisWorklist :: !(Worklist l)     -- ^ Like 'worklist' but for analysis repairing-    }---- | The hashcons 𝐻  is a map from e-nodes to e-class ids-type Memo l = NodeMap l ClassId---- | Maintained worklist of e-class ids that need to be “upward merged”-type Worklist l = NodeMap l ClassId- -- ROMES:TODO: join things built in paralell? -- instance Ord1 l => Semigroup (EGraph l) where --     (<>) eg1 eg2 = undefined -- not so easy -- instance Ord1 l => Monoid (EGraph l) where --     mempty = EGraph emptyUF mempty mempty mempty -instance (Show (Domain l), Show1 l) => Show (EGraph l) where-    show (EGraph a b c d e) =-        "UnionFind: " <> show a <>-            "\n\nE-Classes: " <> show b <>-                "\n\nHashcons: " <> show c <>-                    "\n\nWorklist: " <> show d <>-                        "\n\nAnalWorklist: " <> show e - -- | Add an e-node to the e-graph -- -- If the e-node is already represented in this e-graph, the class-id of the -- class it's already represented in will be returned. add :: forall l. Language l => ENode l -> EGraph l -> (ClassId, EGraph l) add uncanon_e egr =-    let !new_en = {-# SCC "-2" #-} canonicalize uncanon_e egr+    let !new_en = canonicalize uncanon_e egr -     in case {-# SCC "-1" #-} lookupNM new_en (memo egr) of-      Just canon_enode_id -> {-# SCC "0" #-} (find canon_enode_id egr, egr)+     in case lookupNM new_en (memo egr) of+      Just canon_enode_id -> (find canon_enode_id egr, egr)       Nothing ->          let@@ -111,9 +84,9 @@             -- to the e-class parents the new e-node and its e-class id             --             -- And add new e-class to existing e-classes-            new_parents      = insertNM new_en new_eclass_id-            new_classes      = {-# SCC "2" #-} IM.insert new_eclass_id new_eclass $-                                    foldr  (IM.adjust (_parents %~ new_parents))+            new_parents      = ((new_eclass_id, new_en) |:)+            new_classes      = IM.insert new_eclass_id new_eclass $+                                    foldr  (IM.adjust ((_parents %~ new_parents)))                                            (classes egr)                                            (unNode new_en) @@ -142,24 +115,24 @@             -- something else?             --             -- So in the end, we do need to addToWorklist to get correct results-            new_worklist     = {-# SCC "4" #-} insertNM new_en new_eclass_id (worklist egr)+            new_worklist     = (new_eclass_id, new_en):(worklist egr)              -- Add the e-node's e-class id at the e-node's id-            new_memo         = {-# SCC "5" #-} insertNM new_en new_eclass_id (memo egr)+            new_memo         = insertNM new_en new_eclass_id (memo egr)           in ( new_eclass_id              , egr { unionFind = new_uf                   , classes   = new_classes                   , worklist  = new_worklist-                  , memo     = new_memo+                  , memo      = new_memo                   }                    -- Modify created node according to analysis-                  & {-# SCC "6" #-} modifyA new_eclass_id+                  & modifyA new_eclass_id              )-{-# SCC add #-}+{-# INLINABLE add #-}  -- | Merge 2 e-classes by id merge :: forall l. Language l => ClassId -> ClassId -> EGraph l -> (ClassId, EGraph l)@@ -180,7 +153,7 @@             -- Leader is the class with more parents            (leader, leader_class, sub, sub_class) =-               if (sizeNM (class_a^._parents)) < (sizeNM (class_b^._parents))+               if sizeSL (class_a^._parents) < sizeSL (class_b^._parents)                   then (b', class_b, a', class_a) -- b is leader                   else (a', class_a, b', class_b) -- a is leader @@ -189,8 +162,8 @@             -- Update leader class with all e-nodes and parents from the            -- subsumed class-           updatedLeader = leader_class & _parents %~ (<> sub_class^._parents)-                                        & _nodes   %~ (<> sub_class^._nodes)+           updatedLeader = leader_class & _parents %~ (sub_class^._parents <>)+                                        & _nodes   %~ (sub_class^._nodes <>)                                         & _data    .~ new_data            new_data = joinA @l (leader_class^._data) (sub_class^._data) @@ -201,18 +174,18 @@            -- Add all subsumed parents to worklist We can do this instead of            -- adding the new e-class itself to the worklist because it would end            -- up adding its parents anyway-           new_worklist = sub_class^._parents <> (worklist egr0)+           new_worklist = toListSL (sub_class^._parents) <> (worklist egr0)             -- If the new_data is different from the classes, the parents of the            -- class whose data is different from the merged must be put on the            -- analysisWorklist            new_analysis_worklist =+             (if new_data /= (sub_class^._data)+                then toListSL (sub_class^._parents)+                else mempty) <>              (if new_data /= (leader_class^._data)-                then leader_class^._parents+                then toListSL (leader_class^._parents)                 else mempty) <>-             (if new_data /= (sub_class^._data)-                 then sub_class^._parents-                 else mempty) <>              (analysisWorklist egr0)             -- ROMES:TODO: The code that makes the -1 * cos test pass when some other things are tweaked@@ -231,10 +204,10 @@              & modifyA new_id          in (new_id, new_egr)-{-# SCC merge #-}+{-# INLINEABLE merge #-}              --- | The rebuild operation processes the e-graph's current 'Worklist',+-- | The rebuild operation processes the e-graph's current worklist, -- restoring the invariants of deduplication and congruence. Rebuilding is -- similar to other approaches in how it restores congruence; but it uniquely -- allows the client to choose when to restore invariants in the context of a@@ -244,47 +217,52 @@   -- empty worklists   -- repair deduplicated e-classes   let-    egr'  = foldrWithKeyNM' repair (EGraph uf cls mm mempty mempty) wl-    egr'' = foldrWithKeyNM' repairAnal egr' awl+    emptiedEgr = (EGraph uf cls mm mempty mempty)++    wl'   = nubOrd $ bimap (`find` emptiedEgr) (`canonicalize` emptiedEgr) <$> wl+    egr'  = foldr repair emptiedEgr wl'++    awl'  = nubIntOn fst $ first (`find` egr') <$> awl+    egr'' = foldr repairAnal egr' awl'   in   -- Loop until worklist is completely empty   if null (worklist egr'') && null (analysisWorklist egr'')      then egr''-     else rebuild egr''--{-# SCC rebuild #-}+     else rebuild egr'' -- ROMES:TODO: Doesn't seem to be needed at all in the testsuite.+{-# INLINEABLE rebuild #-}  -- ROMES:TODO: find repair_id could be shared between repair and repairAnal?  -- | Repair a single worklist entry.-repair :: forall l. Language l => ENode l -> ClassId -> EGraph l -> EGraph l-repair node repair_id egr =+repair :: forall l. Language l => (ClassId, ENode l) -> EGraph l -> EGraph l+repair (repair_id, node) egr = -   case insertLookupNM (node `canonicalize` egr) (find repair_id egr) (deleteNM node $ memo egr) of-- TODO: I seem to really need it. Is find needed? (they don't use it)+   -- TODO We're no longer deleting the uncanonicalized node, how much does it matter that the structure keeps growing? -      (Nothing, memo2) -> egr { memo = memo2 } -- Return new memo but delete uncanonicalized node+   case insertLookupNM node repair_id (memo egr) of -      (Just existing_class, memo2) -> snd (merge existing_class repair_id egr{memo = memo2})-{-# SCC repair #-}+      (Nothing, memo') -> egr { memo = memo' } -- new memo with inserted node +      (Just existing_class, memo') -> snd (merge existing_class repair_id egr{memo = memo'})+{-# INLINE repair #-}+ -- | Repair a single analysis-worklist entry.-repairAnal :: forall l. Language l => ENode l -> ClassId -> EGraph l -> EGraph l-repairAnal node repair_id egr =+repairAnal :: forall l. Language l => (ClassId, ENode l) -> EGraph l -> EGraph l+repairAnal (repair_id, node) egr =     let-        canon_id = find repair_id egr-        c        = egr^._class canon_id+        c        = (egr^._classes) IM.! repair_id         new_data = joinA @l (c^._data) (makeA node egr)     in     -- Take action if the new_data is different from the existing data     if c^._data /= new_data         -- Merge result is different from original class data, update class         -- with new_data-       then egr { analysisWorklist = c^._parents <> analysisWorklist egr+       then egr { analysisWorklist = toListSL (c^._parents) <> analysisWorklist egr                 }-                & _class canon_id._data .~ new_data-                & modifyA canon_id+                & _classes %~ (IM.adjust (_data .~ new_data) repair_id)+                & modifyA repair_id        else egr-{-# SCC repairAnal #-}+{-# INLINE repairAnal #-}  -- | Canonicalize an e-node --@@ -296,7 +274,7 @@ -- canonicalize(𝑓(𝑎,𝑏,𝑐,...)) = 𝑓((find 𝑎), (find 𝑏), (find 𝑐),...) canonicalize :: Functor l => ENode l -> EGraph l -> ENode l canonicalize (Node enode) eg = Node $ fmap (`find` eg) enode-{-# SCC canonicalize #-}+{-# INLINE canonicalize #-}  -- | Find the canonical representation of an e-class id in the e-graph -- Invariant: The e-class id always exists.

@@ -1,22 +0,0 @@-{-# LANGUAGE RoleAnnotations #-}-{-# LANGUAGE KindSignatures #-}-module Data.Equality.Graph where--import Data.Equality.Graph.Classes.Id-import Data.Equality.Graph.Nodes-import Data.Equality.Graph.ReprUnionFind-import {-# SOURCE #-} Data.Equality.Graph.Classes (EClass)--type role EGraph nominal-data EGraph l = EGraph-    { unionFind :: !ReprUnionFind-    , classes   :: !(ClassIdMap (EClass l))-    , memo      :: !(Memo l)-    , worklist  :: !(Worklist l)-    , analysisWorklist :: !(Worklist l)-    }--find :: ClassId -> EGraph l -> ClassId--type Memo l = NodeMap l ClassId-type Worklist l = NodeMap l ClassId

@@ -15,6 +15,8 @@ import Data.Equality.Graph.Classes.Id import Data.Equality.Graph.Nodes +import Data.Equality.Utils.SizedList+ import Data.Equality.Analysis  -- | An e-class (an equivalence class of terms) of a language @l@.@@ -27,9 +29,9 @@     { eClassId      :: {-# UNPACK #-} !ClassId -- ^ E-class identifier     , eClassNodes   :: !(S.Set (ENode l))      -- ^ E-nodes in this class     , eClassData    :: Domain l                -- ^ The analysis data associated with this eclass.-    , eClassParents :: !(NodeMap l ClassId)    -- ^ E-nodes which are parents of this e-class and their corresponding e-class ids. We found a mapping from nodes to e-class ids a better representation than @[(ENode l, ClassId)]@, and we get de-duplication built-in.+    , eClassParents :: !(SList (ClassId, ENode l))   -- ^ E-nodes which are parents of this e-class and their corresponding e-class ids.     }  instance (Show (Domain l), Show1 l) => Show (EClass l) where-    show (EClass a b d c) = "Id: " <> show a <> "\nNodes: " <> show b <> "\nParents: " <> show c <> "\nData: " <> show d+    show (EClass a b d (SList c _)) = "Id: " <> show a <> "\nNodes: " <> show b <> "\nParents: " <> show c <> "\nData: " <> show d 

@@ -0,0 +1,41 @@+{-# LANGUAGE UndecidableInstances #-} -- tmp show+{-# LANGUAGE FlexibleContexts #-}+{-# OPTIONS_HADDOCK hide #-}+{-|+ Non-abstract definition of e-graphs+ -}+module Data.Equality.Graph.Internal where++import Data.Functor.Classes++import Data.Equality.Graph.ReprUnionFind+import Data.Equality.Graph.Classes+import Data.Equality.Graph.Nodes+import Data.Equality.Analysis++-- | E-graph representing terms of language @l@.+--+-- Intuitively, an e-graph is a set of equivalence classes (e-classes). Each e-class is a+-- set of e-nodes representing equivalent terms from a given language, and an e-node is a function+-- symbol paired with a list of children e-classes.+data EGraph l = EGraph+    { unionFind :: !ReprUnionFind           -- ^ Union find like structure to find canonical representation of an e-class id+    , classes   :: !(ClassIdMap (EClass l)) -- ^ Map canonical e-class ids to their e-classes+    , memo      :: !(Memo l)                -- ^ Hashcons maps all canonical e-nodes to their e-class ids+    , worklist  :: !(Worklist l)            -- ^ Worklist of e-class ids that need to be upward merged+    , analysisWorklist :: !(Worklist l)     -- ^ Like 'worklist' but for analysis repairing+    }++-- | The hashcons 𝐻  is a map from e-nodes to e-class ids+type Memo l = NodeMap l ClassId++-- | Maintained worklist of e-class ids that need to be “upward merged”+type Worklist l = [(ClassId, ENode l)]++instance (Show (Domain l), Show1 l) => Show (EGraph l) where+    show (EGraph a b c d e) =+        "UnionFind: " <> show a <>+            "\n\nE-Classes: " <> show b <>+                "\n\nHashcons: " <> show c <>+                    "\n\nWorklist: " <> show d <>+                        "\n\nAnalWorklist: " <> show e

@@ -0,0 +1,9 @@+{-# LANGUAGE RoleAnnotations #-}+{-# LANGUAGE StandaloneKindSignatures #-}+module Data.Equality.Graph.Internal where++import Data.Kind++type EGraph :: (Type -> Type) -> Type+type role EGraph nominal+data EGraph l

@@ -1,8 +1,10 @@ {-# LANGUAGE RecordWildCards #-} {-# LANGUAGE Rank2Types #-} {-|-  Hand-rolled lenses on e-graphs and e-classes which come in quite handy and-  are heavily used in 'Data.Equality.Graph'.+  Hand-rolled lenses on e-graphs and e-classes which come in quite handy, are+  heavily used in 'Data.Equality.Graph', and are the only exported way of+  editing the structure of the e-graph. If you want to write some complex+  'Analysis' you'll probably need these.  -} module Data.Equality.Graph.Lens where @@ -12,11 +14,13 @@ import Data.Functor.Identity import Data.Functor.Const +import Data.Equality.Utils.SizedList+import Data.Equality.Graph.Internal import Data.Equality.Graph.Classes.Id import Data.Equality.Graph.Nodes import Data.Equality.Graph.Classes+import Data.Equality.Graph.ReprUnionFind import Data.Equality.Analysis-import {-# SOURCE #-} Data.Equality.Graph (EGraph(..), Memo, find)  -- | A 'Lens'' as defined in other lenses libraries type Lens' s a = forall f. Functor f => (a -> f a) -> (s -> f s)@@ -37,12 +41,13 @@ -- Calls 'error' when the e-class doesn't exist _class :: ClassId -> Lens' (EGraph l) (EClass l) _class i afa s =-    let canon_id = find i s+    let canon_id = findRepr i (unionFind s)      in (\c' -> s { classes = IM.insert canon_id c' (classes s) }) <$> afa (classes s IM.! canon_id) {-# INLINE _class #-} --- | Lens for the 'Memo' of e-nodes in an e-graph-_memo :: Lens' (EGraph l) (Memo l)+-- | Lens for the memo of e-nodes in an e-graph, that is, a mapping from+-- e-nodes to the e-class they're represented in+_memo :: Lens' (EGraph l) (NodeMap l ClassId) _memo afa egr = (\m1 -> egr {memo = m1}) <$> afa (memo egr) {-# INLINE _memo #-} @@ -57,8 +62,8 @@ {-# INLINE _data #-}  -- | Lens for the parent e-classes of an e-class-_parents :: Lens' (EClass l) (NodeMap l ClassId)-_parents afa EClass{..} = EClass eClassId eClassNodes eClassData <$> afa eClassParents+_parents :: Lens' (EClass l) (SList (ClassId, ENode l))+_parents afa EClass{..} = (\ps -> EClass eClassId eClassNodes eClassData ps) <$> afa eClassParents {-# INLINE _parents #-}  -- | Lens for the e-nodes in an e-class

@@ -37,7 +37,7 @@ -- | Get the children e-class ids of an e-node children :: Traversable l => ENode l -> [ClassId] children = toList . unNode-{-# SCC children #-}+{-# INLINE children #-}  -- * Operator @@ -48,7 +48,7 @@ -- | Get the operator (function symbol) of an e-node operator :: Traversable l => ENode l -> Operator l operator = Operator . void . unNode-{-# SCC operator #-}+{-# INLINE operator #-}  instance Eq1 l => (Eq (ENode l)) where     (==) (Node a) (Node b) = liftEq (==) a b@@ -75,20 +75,14 @@ -- * Node Map  -- | A mapping from e-nodes of @l@ to @a@-data NodeMap (l :: Type -> Type) a = NodeMap { unNodeMap :: !(M.Map (ENode l) a), sizeNodeMap :: {-# UNPACK #-} !Int }+newtype NodeMap (l :: Type -> Type) a = NodeMap { unNodeMap :: M.Map (ENode l) a } -- TODO: Investigate whether it would be worth it requiring a trie-map for the -- e-node definition. Probably it isn't better since e-nodes aren't recursive.-  deriving (Show, Functor, Foldable, Traversable)--instance (Eq1 l, Ord1 l) => Semigroup (NodeMap l a) where-  NodeMap m1 s1 <> NodeMap m2 s2 = NodeMap (m1 <> m2) (s1 + s2)--instance (Eq1 l, Ord1 l) => Monoid (NodeMap l a) where-  mempty = NodeMap mempty 0+  deriving (Show, Functor, Foldable, Traversable, Semigroup, Monoid)  -- | Insert a value given an e-node in a 'NodeMap' insertNM :: Ord1 l => ENode l -> a -> NodeMap l a -> NodeMap l a-insertNM e v (NodeMap m s) = NodeMap (M.insert e v m) (s+1)+insertNM e v (NodeMap m) = NodeMap (M.insert e v m) {-# INLINE insertNM #-}  -- | Lookup an e-node in a 'NodeMap'@@ -98,12 +92,12 @@  -- | Delete an e-node in a 'NodeMap' deleteNM :: Ord1 l => ENode l -> NodeMap l a -> NodeMap l a-deleteNM e (NodeMap m s) = NodeMap (M.delete e m) (s-1)+deleteNM e (NodeMap m) = NodeMap (M.delete e m) {-# INLINE deleteNM #-}  -- | Insert a value and lookup by e-node in a 'NodeMap' insertLookupNM :: Ord1 l => ENode l -> a -> NodeMap l a -> (Maybe a, NodeMap l a)-insertLookupNM e v (NodeMap m s) = second (flip NodeMap (s+1)) $ M.insertLookupWithKey (\_ a _ -> a) e v m+insertLookupNM e v (NodeMap m) = second NodeMap $ M.insertLookupWithKey (\_ a _ -> a) e v m {-# INLINE insertLookupNM #-}  -- | As 'Data.Map.foldlWithKeyNM'' but in a 'NodeMap'@@ -120,12 +114,12 @@ -- -- This operation takes constant time (__O(1)__) sizeNM :: NodeMap l a -> Int-sizeNM = sizeNodeMap+sizeNM = M.size . unNodeMap {-# INLINE sizeNM #-}  -- | As 'Data.Map.traverseWithKeyNM' but in a 'NodeMap' traverseWithKeyNM :: Applicative t => (ENode l -> a -> t b) -> NodeMap l a -> t (NodeMap l b) -traverseWithKeyNM f (NodeMap m s) = (`NodeMap` s) <$> M.traverseWithKey f m+traverseWithKeyNM f (NodeMap m) = NodeMap <$> M.traverseWithKey f m {-# INLINE traverseWithKeyNM #-}  -- Node Set

@@ -75,7 +75,6 @@ #else makeNewSet (RUF im si) = (si, RUF (IIM.insert si 0 im) (si + 1)) #endif-{-# SCC makeNewSet #-}  -- | Union operation of the union find. --@@ -94,21 +93,20 @@     -- represented_by_b = hc IM.! b     -- -- Overwrite previous id of b (which should be 0#) with new representative (a)     -- -- AND "rebuild" all nodes represented by b by making them represented directly by a-    -- new_im = {-# SCC "rebuild_im" #-} IIM.unliftedFoldr (\(I# x) -> IIM.insert x a#) (IIM.insert b# a# im) represented_by_b-    -- new_hc = {-# SCC "adjust_hc" #-} IM.adjust ((b:) . (represented_by_b <>)) a (IM.delete b hc)-{-# SCC unionSets #-}+    -- new_im = IIM.unliftedFoldr (\(I# x) -> IIM.insert x a#) (IIM.insert b# a# im) represented_by_b+    -- new_hc = IM.adjust ((b:) . (represented_by_b <>)) a (IM.delete b hc)  -- | Find the canonical representation of an e-class id findRepr :: ClassId -> ReprUnionFind          -> ClassId -- ^ The found canonical representation #if __GLASGOW_HASKELL__ >= 902 findRepr v@(I# v#) (RUF m s) =-  case {-# SCC "findRepr_TAKE" #-} m IIM.! v# of+  case m IIM.! v# of     0# -> v     x  -> findRepr (I# x) (RUF m s) #else findRepr v (RUF m s) =-  case {-# SCC "findRepr_TAKE" #-} m IIM.! v of+  case m IIM.! v of     0 -> v     x -> findRepr x (RUF m s) #endif@@ -121,7 +119,6 @@ -- -- When using the ad-hoc path compression in `unionSets`, the depth of -- recursion never even goes above 1!-{-# SCC findRepr #-}   -- {-# RULES

@@ -28,6 +28,7 @@ import qualified Data.IntSet as IS  import Data.Equality.Graph+import Data.Equality.Graph.Lens import Data.Equality.Matching.Database import Data.Equality.Matching.Pattern @@ -69,7 +70,7 @@  -- | Convert an e-graph into a database eGraphToDatabase :: Language l => EGraph l -> Database l-eGraphToDatabase EGraph{..} = foldrWithKeyNM' addENodeToDB (DB mempty) memo+eGraphToDatabase egr = foldrWithKeyNM' addENodeToDB (DB mempty) (egr^._memo)   where      -- Add an enode in an e-graph, given its class, to a database@@ -78,7 +79,6 @@         -- ROMES:TODO map find         -- Insert or create a relation R_f(i1,i2,...,in) for lang in which          DB $ M.alter (Just . populate (classid:children enode)) (operator enode) m-    {-# SCC addENodeToDB #-}      -- Populate or create a triemap given the population D_x (ClassIds)     -- Insert remaining ids population doesn't exist, recursively merge tries with remaining ids@@ -89,8 +89,7 @@     -- If trie map entry already exists, populate the existing map with the remaining ids     populate []     (Just it)              = it     populate (x:xs) (Just (MkIntTrie k m)) = MkIntTrie (x `IS.insert` k) $ IM.alter (Just . populate xs) x m-    {-# SCC populate #-}-{-# SCC eGraphToDatabase #-}+{-# INLINABLE eGraphToDatabase #-}   -- * Database related internals@@ -134,4 +133,4 @@         vars :: Foldable lang => Pattern lang -> [Var]         vars (VariablePattern x) = [x]         vars (NonVariablePattern p) = nubInt $ join $ map vars $ toList p-{-# SCC compileToQuery #-}+{-# INLINABLE compileToQuery #-}

@@ -123,35 +123,35 @@   where    genericJoin' :: [Atom l] -> [Var] -> [Subst]-   genericJoin' !atoms' = \case+   genericJoin' atoms' = \case -     [] -> map mempty atoms+     [] -> mempty <$> atoms' -     (!x):xs -> -       -- IS.foldl' (\acc x_in_D -> genericJoin' (substitute x x_in_D atoms') (map (IM.insert x x_in_D) substs) xs <> acc)-       --           mempty-       --           (domainX x atoms')-       IS.foldl'-         (\acc x_in_D ->-           map (\y -> let !y' = IM.insert x x_in_D y in y') -- TODO: A bit contrieved, perhaps better to avoid map ?-             -- Each valid sub-query assumed the x -> x_in_D substitution-             (genericJoin' (substitute x x_in_D atoms') xs)-               <> acc)-         mempty-         (domainX x atoms')-   {-# SCC genericJoin' #-}+     (!x):xs -> do -   atomsWithX :: Var -> [Atom l] -> [Atom l]-   atomsWithX x = filter (x `elemOfAtom`)-   {-# INLINE atomsWithX #-}+       x_in_D <- domainX x atoms' -   domainX :: Var -> [Atom l] -> IS.IntSet-   domainX x = intersectAtoms x d . atomsWithX x+       -- Each valid sub-query assumes x -> x_in_D substitution+       y <- genericJoin' (substitute x x_in_D atoms') xs++       return $! IM.insert x x_in_D y -- TODO: A bit contrieved, perhaps better to avoid map ?++   domainX :: Var -> [Atom l] -> [Int]+   domainX x = IS.toList . intersectAtoms x d . filter (x `elemOfAtom`)    {-# INLINE domainX #-} +   -- | Substitute all occurrences of 'Var' with given 'ClassId' in all given atoms.+   substitute :: Functor lang => Var -> ClassId -> [Atom lang] -> [Atom lang]+   substitute r i = map $ \case+      Atom x l -> Atom (if CVar r == x then CClassId i else x) $ fmap (\v -> if CVar r == v then CClassId i else v) l+ {-# INLINABLE genericJoin #-}-{-# SCC genericJoin #-} +-- | Returns True if 'Var' occurs in given 'Atom'+elemOfAtom :: (Functor lang, Foldable lang) => Var -> Atom lang -> Bool+elemOfAtom !x (Atom v l) = case v of+ CVar v' -> x == v'+ _ -> or $ fmap (\v' -> CVar x == v') l  -- ROMES:TODO: Batching? How? https://arxiv.org/pdf/2108.02290.pdf @@ -182,7 +182,7 @@         -- | Get the size of an atom         atomLength :: Foldable lang => Atom lang -> Int         atomLength (Atom _ l) = 1 + F.length l-        {-# SCC atomLength #-}+        {-# INLINE atomLength #-}          -- | Extract 'Var' from 'ClassIdOrVar'         toVar :: ClassIdOrVar -> Maybe Var@@ -190,23 +190,7 @@         toVar (CClassId _) = Nothing         {-# INLINE toVar #-} -{-# SCC orderedVarsInQuery #-}  ---- | Substitute all occurrences of 'Var' with given 'ClassId' in all given atoms.-substitute :: Functor lang => Var -> ClassId -> [Atom lang] -> [Atom lang]-substitute !r !i = map $ \case-   Atom x l -> Atom (if CVar r == x then CClassId i else x) $ fmap (\v -> if CVar r == v then CClassId i else v) l-{-# SCC substitute #-}---- | Returns True if 'Var' occurs in given 'Atom'-elemOfAtom :: (Functor lang, Foldable lang) => Var -> Atom lang -> Bool-elemOfAtom !x (Atom v l) = case v of-  CVar v' -> x == v'-  _ -> or $ fmap (\v' -> CVar x == v') l-{-# SCC elemOfAtom #-}-- -- ROMES:TODO Terrible name 'intersectAtoms'  -- | Given a database and a list of Atoms with an occurring var @x@, find@@ -231,8 +215,6 @@                      Just xs -> xs  intersectAtoms _ _ [] = error "can't intersect empty list of atoms?"-{-# INLINABLE intersectAtoms #-}-{-# SCC intersectAtoms #-}  -- | Find the matching ids that a variable can take given a list of variables -- and ids that must match the structure@@ -306,7 +288,7 @@                    Just _  -> k `IS.insert` acc                          ) mempty m           -- (3)-          -- else {-# SCC "intersect_new_OTHER_var" #-} IS.unions $ IM.elems $ IM.mapMaybeWithKey (\k ls -> intersectInTrie var ({-# SCC "putSubst" #-} IM.insert x k substs) ls xs) m+          -- else IS.unions $ IM.elems $ IM.mapMaybeWithKey (\k ls -> intersectInTrie var (IM.insert x k substs) ls xs) m           else IM.foldrWithKey (\k ls (!acc) ->             case intersectInTrie var (IM.insert x k substs) ls xs of                 Nothing -> acc@@ -318,6 +300,3 @@       isVarDifferentFrom _ (CClassId _) = False       isVarDifferentFrom x (CVar     y) = x /= y       {-# INLINE isVarDifferentFrom #-}--{-# INLINABLE intersectInTrie #-}-{-# SCC intersectInTrie #-}

@@ -23,7 +23,7 @@ module Data.Equality.Saturation     (       -- * Equality saturation-      equalitySaturation, equalitySaturation'+      equalitySaturation, equalitySaturation', runEqualitySaturation        -- * Re-exports for equality saturation @@ -33,7 +33,7 @@       -- ** Writing cost functions       --       -- | 'CostFunction' re-exported from 'Data.Equality.Extraction' since they are required to do equality saturation-    , CostFunction --, Cost, depthCost+    , CostFunction --, depthCost        -- ** Writing expressions       -- @@ -51,6 +51,8 @@ import Data.Proxy  import Data.Equality.Utils+import Data.Equality.Graph.Nodes+import Data.Equality.Graph.Lens import qualified Data.Equality.Graph as G import Data.Equality.Graph.Monad import Data.Equality.Language@@ -63,11 +65,12 @@ import Data.Equality.Saturation.Scheduler  -- | Equality saturation with defaults-equalitySaturation :: forall l. Language l-                   => Fix l             -- ^ Expression to run equality saturation on-                   -> [Rewrite l]       -- ^ List of rewrite rules-                   -> CostFunction l    -- ^ Cost function to extract the best equivalent representation-                   -> (Fix l, EGraph l) -- ^ Best equivalent expression and resulting e-graph+equalitySaturation :: forall l cost+                    . (Language l, Ord cost)+                   => Fix l               -- ^ Expression to run equality saturation on+                   -> [Rewrite l]         -- ^ List of rewrite rules+                   -> CostFunction l cost -- ^ Cost function to extract the best equivalent representation+                   -> (Fix l, EGraph l)   -- ^ Best equivalent expression and resulting e-graph equalitySaturation = equalitySaturation' (Proxy @BackoffScheduler)  @@ -75,113 +78,123 @@ -- extract the best equivalent expression according to the given cost function -- -- This variant takes all arguments instead of using defaults-equalitySaturation' :: forall l schd-                    . (Language l, Scheduler schd)-                    => Proxy schd        -- ^ Proxy for the scheduler to use-                    -> Fix l             -- ^ Expression to run equality saturation on-                    -> [Rewrite l]       -- ^ List of rewrite rules-                    -> CostFunction l    -- ^ Cost function to extract the best equivalent representation-                    -> (Fix l, EGraph l) -- ^ Best equivalent expression and resulting e-graph-equalitySaturation' _ expr rewrites cost = egraph $ do+equalitySaturation' :: forall l schd cost+                    . (Language l, Scheduler schd, Ord cost)+                    => Proxy schd          -- ^ Proxy for the scheduler to use+                    -> Fix l               -- ^ Expression to run equality saturation on+                    -> [Rewrite l]         -- ^ List of rewrite rules+                    -> CostFunction l cost -- ^ Cost function to extract the best equivalent representation+                    -> (Fix l, EGraph l)   -- ^ Best equivalent expression and resulting e-graph+equalitySaturation' proxy expr rewrites cost = egraph $ do      -- Represent expression as an e-graph     origClass <- represent expr      -- Run equality saturation (by applying non-destructively all rewrites)-    equalitySaturation'' 0 mempty -- Start at iteration 0+    runEqualitySaturation proxy rewrites      -- Extract best solution from the e-class of the original expression     gets $ \g -> extractBest g cost origClass+{-# INLINABLE equalitySaturation' #-} -      where -        -- Take map each rewrite rule to stats on its usage so we can do-        -- backoff scheduling. Each rewrite rule is assigned an integer-        -- (corresponding to its position in the list of rewrite rules)-        equalitySaturation'' :: Int -> IM.IntMap (Stat schd) -> EGraphM l ()-        equalitySaturation'' 30 _ = return () -- Stop after X iterations-        equalitySaturation'' i stats = do+-- | Run equality saturation on an e-graph by non-destructively applying all+-- given rewrite rules until saturation (using the given 'Scheduler')+runEqualitySaturation :: forall l schd+                       . (Language l, Scheduler schd)+                      => Proxy schd          -- ^ Proxy for the scheduler to use+                      -> [Rewrite l]         -- ^ List of rewrite rules+                      -> EGraphM l ()+runEqualitySaturation _ rewrites = runEqualitySaturation' 0 mempty where -- Start at iteration 0 -            egr@G.EGraph{ G.memo = beforeMemo, G.classes = beforeClasses } <- get+  -- Take map each rewrite rule to stats on its usage so we can do+  -- backoff scheduling. Each rewrite rule is assigned an integer+  -- (corresponding to its position in the list of rewrite rules)+  runEqualitySaturation' :: Int -> IM.IntMap (Stat schd) -> EGraphM l ()+  runEqualitySaturation' 30 _ = return () -- Stop after X iterations+  runEqualitySaturation' i stats = do -            let db = eGraphToDatabase egr+      egr <- get -            -- Read-only phase, invariants are preserved-            -- With backoff scheduler-            -- ROMES:TODO parMap with chunks-            let (!matches, newStats) = mconcat (fmap (matchWithScheduler db i stats) (zip [1..] rewrites))+      let (beforeMemo, beforeClasses) = (egr^._memo, egr^._classes)+          db = eGraphToDatabase egr -            -- Write-only phase, temporarily break invariants-            forM_ matches applyMatchesRhs+      -- Read-only phase, invariants are preserved+      -- With backoff scheduler+      -- ROMES:TODO parMap with chunks+      let (!matches, newStats) = mconcat (fmap (matchWithScheduler db i stats) (zip [1..] rewrites)) -            -- Restore the invariants once per iteration-            rebuild-            -            G.EGraph { G.memo = afterMemo, G.classes = afterClasses } <- get+      -- Write-only phase, temporarily break invariants+      forM_ matches applyMatchesRhs -            -- ROMES:TODO: Node limit...-            -- ROMES:TODO: Actual Timeout... not just iteration timeout-            -- ROMES:TODO Better saturation (see Runner)-            -- Apply rewrites until saturated or ROMES:TODO: timeout-            unless (G.sizeNM afterMemo == G.sizeNM beforeMemo-                      && IM.size afterClasses == IM.size beforeClasses)-                (equalitySaturation'' (i+1) newStats)+      -- Restore the invariants once per iteration+      rebuild+      +      (afterMemo, afterClasses) <- gets (\g -> (g^._memo, g^._classes)) -        matchWithScheduler :: Database l -> Int -> IM.IntMap (Stat schd) -> (Int, Rewrite l) -> ([(Rewrite l, Match)], IM.IntMap (Stat schd))-        matchWithScheduler db i stats = \case-            (rw_id, rw :| cnd) -> first (map (first (:| cnd))) $ matchWithScheduler db i stats (rw_id, rw)-            (rw_id, lhs := rhs) -> do-                case IM.lookup rw_id stats of-                  -- If it's banned until some iteration, don't match this rule-                  -- against anything.-                  Just s | isBanned @schd i s -> ([], stats)+      -- ROMES:TODO: Node limit...+      -- ROMES:TODO: Actual Timeout... not just iteration timeout+      -- ROMES:TODO Better saturation (see Runner)+      -- Apply rewrites until saturated or ROMES:TODO: timeout+      unless (G.sizeNM afterMemo == G.sizeNM beforeMemo+                && IM.size afterClasses == IM.size beforeClasses)+          (runEqualitySaturation' (i+1) newStats) -                  -- Otherwise, match and update stats-                  x -> do+  matchWithScheduler :: Database l -> Int -> IM.IntMap (Stat schd) -> (Int, Rewrite l) -> ([(Rewrite l, Match)], IM.IntMap (Stat schd))+  matchWithScheduler db i stats = \case+      (rw_id, rw :| cnd) -> first (map (first (:| cnd))) $ matchWithScheduler db i stats (rw_id, rw)+      (rw_id, lhs := rhs) -> do+          case IM.lookup rw_id stats of+            -- If it's banned until some iteration, don't match this rule+            -- against anything.+            Just s | isBanned @schd i s -> ([], stats) -                      -- Match pattern-                      let matches' = ematch db lhs -- Add rewrite to the e-match substitutions+            -- Otherwise, match and update stats+            x -> do -                      -- Backoff scheduler: update stats-                      let newStats = updateStats @schd i rw_id x stats matches'+                -- Match pattern+                let matches' = ematch db lhs -- Add rewrite to the e-match substitutions -                      (map (lhs := rhs,) matches', newStats)+                -- Backoff scheduler: update stats+                let newStats = updateStats @schd i rw_id x stats matches' -        applyMatchesRhs :: (Rewrite l, Match) -> EGraphM l ()-        applyMatchesRhs =-            \case-                (rw :| cond, m@(Match subst _)) -> do-                    -- If the rewrite condition is satisfied, applyMatchesRhs on the rewrite rule.-                    egr <- get-                    when (cond subst egr) $-                       applyMatchesRhs (rw, m)+                (map (lhs := rhs,) matches', newStats) -                (_ := VariablePattern v, Match subst eclass) -> do-                    -- rhs is equal to a variable, simply merge class where lhs-                    -- pattern was found (@eclass@) and the eclass the pattern-                    -- variable matched (@lookup v subst@)-                    case IM.lookup v subst of-                      Nothing -> error "impossible: couldn't find v in subst"-                      Just n  -> do-                          _ <- merge n eclass-                          return ()+  applyMatchesRhs :: (Rewrite l, Match) -> EGraphM l ()+  applyMatchesRhs =+      \case+          (rw :| cond, m@(Match subst _)) -> do+              -- If the rewrite condition is satisfied, applyMatchesRhs on the rewrite rule.+              egr <- get+              when (cond subst egr) $+                 applyMatchesRhs (rw, m) -                (_ := NonVariablePattern rhs, Match subst eclass) -> do-                    -- rhs is (at the top level) a non-variable pattern, so substitute-                    -- all pattern variables in the pattern and create a new e-node (and-                    -- e-class that represents it), then merge the e-class of the-                    -- substituted rhs with the class that matched the left hand side-                    eclass' <- reprPat subst rhs-                    _ <- merge eclass eclass'+          (_ := VariablePattern v, Match subst eclass) -> do+              -- rhs is equal to a variable, simply merge class where lhs+              -- pattern was found (@eclass@) and the eclass the pattern+              -- variable matched (@lookup v subst@)+              case IM.lookup v subst of+                Nothing -> error "impossible: couldn't find v in subst"+                Just n  -> do+                    _ <- merge n eclass                     return () -        -- | Represent a pattern in the e-graph a pattern given substitions-        reprPat :: Subst -> l (Pattern l) -> EGraphM l ClassId-        reprPat subst = add . G.Node <=< traverse \case-            VariablePattern v ->-                case IM.lookup v subst of-                    Nothing -> error "impossible: couldn't find v in subst?"-                    Just i  -> return i-            NonVariablePattern p -> reprPat subst p-{-# SCC equalitySaturation' #-}+          (_ := NonVariablePattern rhs, Match subst eclass) -> do+              -- rhs is (at the top level) a non-variable pattern, so substitute+              -- all pattern variables in the pattern and create a new e-node (and+              -- e-class that represents it), then merge the e-class of the+              -- substituted rhs with the class that matched the left hand side+              eclass' <- reprPat subst rhs+              _ <- merge eclass eclass'+              return () +  -- | Represent a pattern in the e-graph a pattern given substitions+  reprPat :: Subst -> l (Pattern l) -> EGraphM l ClassId+  reprPat subst = add . Node <=< traverse \case+      VariablePattern v ->+          case IM.lookup v subst of+              Nothing -> error "impossible: couldn't find v in subst?"+              Just i  -> return i+      NonVariablePattern p -> reprPat subst p++{-# INLINEABLE runEqualitySaturation #-}

@@ -79,7 +79,6 @@           updateBans = \case             Nothing -> Just (BSS (i + ban_length) 1)             Just (BSS _ n)  -> Just (BSS (i + ban_length) (n+1))-    {-# SCC updateStats #-}      isBanned i s = i < bannedUntil s 

@@ -78,7 +78,6 @@   | otherwise = find' k r find' k (Tip kx x) | isTrue# (k `eqWord#` kx) = x find' _ _ = error ("IntMap.!: key ___ is not an element of the map")-{-# SCC find' #-}  -- * Other stuff taken from IntMap 

@@ -0,0 +1,66 @@+{-# LANGUAGE DeriveTraversable #-}+{-# LANGUAGE TypeFamilies #-}+{-|+   Util: A list with a O(1) size function+ -}+module Data.Equality.Utils.SizedList where++import qualified Data.List+import GHC.Exts+import Data.Foldable++-- | A list with O(1) size access and O(1) conversion to normal list+data SList a = SList ![a] {-# UNPACK #-} !Int+  deriving Traversable++instance Semigroup (SList a) where+  (<>) (SList a i) (SList b j) = SList (a <> b) (i+j)+  {-# INLINE (<>) #-}++instance Monoid (SList a) where+  mempty = SList mempty 0+  {-# INLINE mempty #-}++instance Functor SList where+  fmap f (SList a i) = SList (fmap f a) i+  {-# INLINE fmap #-}++instance Foldable SList where+  fold       ( SList l _) = fold l+  foldMap f  ( SList l _) = foldMap f l+  foldMap' f ( SList l _) = foldMap' f l+  foldr f b  ( SList l _) = foldr f b l+  foldr' f b ( SList l _) = foldr' f b l+  foldl f b  ( SList l _) = foldl f b l+  foldl' f b ( SList l _) = foldl' f b l+  foldr1 f   ( SList l _) = foldr1 f l+  foldl1 f   ( SList l _) = foldl1 f l+  toList     ( SList l _) = l+  null       ( SList l _) = Data.List.null l+  length     ( SList _ i) = i+  elem x     ( SList l _) = x `elem` l+  maximum    ( SList l _) = maximum l+  minimum    ( SList l _) = minimum l+  sum        ( SList l _) = sum l+  product    ( SList l _) = product l++instance IsList (SList a) where+  type Item (SList a) = a+  fromList l          = SList l (length l)+  fromListN i l       = SList l i+  toList (SList l _)  = l++-- | Prepend an item to the list in O(1)+(|:) :: a -> SList a -> SList a+(|:) a (SList l i) = SList (a:l) (i+1)+{-# INLINE (|:) #-}++-- | Make a normal list from the sized list in O(1)+toListSL :: SList a -> [a]+toListSL (SList l _) = l+{-# INLINE toListSL #-}++-- | Get the size of the list in O(1)+sizeSL :: SList a -> Int+sizeSL (SList _ i) = i+{-# INLINE sizeSL #-}

@@ -0,0 +1,19 @@+{-# LANGUAGE OverloadedStrings #-}+import Test.Tasty.Bench++import Data.Equality.Utils+import Invariants+import Sym+import Lambda+import SimpleSym++tests :: [Benchmark]+tests = [ bgroup "Tests"+  [ symTests+  , lambdaTests+  , simpleSymTests+  , invariants+  ] ]++main :: IO ()+main = defaultMain tests

@@ -25,6 +25,7 @@ import qualified Data.IntMap.Strict as IM  import Data.Equality.Graph.Monad as GM+import Data.Equality.Graph.Lens import Data.Equality.Graph import Data.Equality.Analysis import Data.Equality.Extraction@@ -52,12 +53,12 @@ patFoldAllClasses :: forall l. (Language l, Num (Pattern l))                   => Fix l -> Integer -> Bool patFoldAllClasses expr i =-    case IM.toList $ classes eg of+    case IM.toList $ (eg^._classes) of         [_] -> True         _   -> False     where         eg :: EGraph l-        eg = snd $ equalitySaturation expr [VariablePattern 1:=fromInteger i] (error "Cost function shouldn't be used")+        eg = snd $ equalitySaturation expr [VariablePattern 1:=fromInteger i] (error "Cost function shouldn't be used" :: CostFunction l Int)  -- | Test 'compileToQuery'. --@@ -97,7 +98,7 @@     let         db = eGraphToDatabase eg         matches = S.fromList $ map matchClassId $ ematch db (VariablePattern v)-        eclasses = S.fromList $ map fst $ IM.toList $ classes eg+        eclasses = S.fromList $ map fst $ IM.toList (eg^._classes)     in         matches == eclasses  @@ -122,18 +123,18 @@ -- ROMES:TODO Should I rebuild it here? Then the property test is that after rebuilding ...HashConsInvariant hashConsInvariant :: forall l. Language l                   => EGraph l -> Bool-hashConsInvariant eg@EGraph{..} =-    all f (IM.toList classes)+hashConsInvariant eg =+    all f (IM.toList (eg^._classes))     where       -- e-node 𝑛 ∈ 𝑀 [𝑎] ⇐⇒ 𝐻 [canonicalize(𝑛)] = find(𝑎)-      f (i, EClass _ nodes _ _) = all g nodes+      f (i, EClass{eClassNodes=nodes}) = all g nodes         where-          g en = case lookupNM (canonicalize en eg) memo of+          g en = case lookupNM (canonicalize en eg) (eg^._memo) of             Nothing -> error "how can we not find canonical thing in map? :)" -- False             Just i' -> i' == find i eg   benchSaturate :: forall l. Language l-              => [Rewrite l] -> (l Cost -> Cost) -> Fix l -> Bool+              => [Rewrite l] -> (l Int -> Int) -> Fix l -> Bool benchSaturate rws cost expr =     equalitySaturation expr rws cost `seq` True 

@@ -38,7 +38,7 @@  instance Language SymExpr -cost :: CostFunction SymExpr+cost :: CostFunction SymExpr Int cost = \case   Const  _ -> 1   Symbol _ -> 1

@@ -20,6 +20,8 @@ import Data.Ord.Deriving import Text.Show.Deriving +import qualified Data.Foldable as F+ import Control.Applicative (liftA2) import Control.Monad (unless) @@ -76,7 +78,7 @@     (/) a b = Fix (BinOp Div a b)     fromRational = Fix . Const . fromRational -symCost :: Expr Cost -> Cost+symCost :: CostFunction Expr Int symCost = \case     BinOp Pow e1 e2 -> e1 + e2 + 6     BinOp Div e1 e2 -> e1 + e2 + 5@@ -110,11 +112,9 @@ instance Analysis Expr where     type Domain Expr = Maybe Double -    {-# SCC makeA #-}     makeA (Node e) egr = evalConstant ((\c -> egr^._class c._data) <$> e)      -- joinA = (<|>)-    {-# SCC joinA #-}     joinA ma mb = do         a <- ma         b <- mb@@ -124,7 +124,6 @@         !_ <- unless (a == b || (a == 0 && b == (-0)) || (a == (-0) && b == 0)) (error "Merged non-equal constants!")         return a -    {-# SCC modifyA #-}     modifyA i egr =         case egr ^._class i._data of           Nothing -> egr@@ -135,7 +134,7 @@             _     <- GM.merge i new_c              -- Prune all except leaf e-nodes-            modify (_class i._nodes %~ S.filter (null . children))+            modify (_class i._nodes %~ S.filter (F.null . unNode))   @@ -343,10 +342,16 @@     , testCase "i6" $         rewrite (_i (_ln "x") "x") @?= "x"*(_ln "x" + fromInteger(-1)) +    -- TODO: Require ability to fine tune parameters+    -- , testCase "diff_power_harder" $+    --     rewrite (_d "x" ((_pow "x" 3) - 7*(_pow "x" 2))) @?= "x"*(3*"x"-14)+     ] -_i :: Fix Expr -> Fix Expr -> Fix Expr+_i, _d, _pow :: Fix Expr -> Fix Expr -> Fix Expr _i a b = Fix (BinOp Integral a b)+_d a b = Fix (BinOp Diff a b)+_pow a b = Fix (BinOp Pow a b) _ln, _cos, _sin :: Fix Expr -> Fix Expr _ln a = Fix (UnOp Ln a) _cos a = Fix (UnOp Cos a)