hegg 0.4.0.0 → 0.5.0.0

raw patch · 19 files changed

+592/−162 lines, 19 filesdep +deepseqdep +textdep ~graphvizPVP ok

version bump matches the API change (PVP)

Dependencies added: deepseq, text

Dependency ranges changed: graphviz

API changes (from Hackage documentation)

Files

• CHANGELOG.md +21/−0
• README.md +15/−21
• hegg.cabal +21/−8
• src/Data/Equality/Analysis.hs +31/−28
• src/Data/Equality/Analysis/Monadic.hs +68/−0
• src/Data/Equality/Graph.hs +230/−49
• src/Data/Equality/Graph/Dot.hs +8/−10
• src/Data/Equality/Graph/Lens.hs +23/−6
• src/Data/Equality/Graph/Monad.hs +44/−1
• src/Data/Equality/Matching/Database.hs +4/−0
• src/Data/Equality/Saturation.hs +6/−6
• src/Data/Equality/Utils.hs +4/−1
• test/Bench.hs +31/−6
• test/Invariants.hs +0/−1
• test/Lambda.hs +10/−8
• test/SimpleSym.hs +8/−4
• test/Sym.hs +19/−13
• test/T3.hs +47/−0
• test/Test.hs +2/−0

@@ -2,6 +2,27 @@  ## Unreleased +## 0.5.0.0 -- 2023-10-31++* Change `'modifyA'` to instead operate over e-graphs, instead of being+    constrained to editing the e-class that prompted the modification.+    (Remember that the e-graph lenses in `'Data.Equality.Graph.Lens'` are the+    preferred way to edit the e-graph and the desired e-class (by id), and its+    data, etc...)++* Fix compilation of Data.Equality.Graph.Dot, the graphviz rendering backend+    (despite there being some usability bugs still) (by @BinderDavid)++* Dropped support for GHC 9.0 because of the QuantifiedConstraints bug (by @phadej)++* Add `AnalysisM`, a class for e-graph analysis that are only well-defined+    within a certain monadic context. Accordingly, we also add versions of the+    current e-graph transformation functions (such as `add` and `merge`) for+    analysis defined monadically (such as `addM` and `mergeM`).++* Add operation to create empty e-classes with explicit domain data+    (experimental, not sure whether this is something good to keep in the API)+ ## 0.4.0.0 -- 2023-06-24  * Make `Language` a constraint type synonym instead of a standalone empty class

@@ -117,7 +117,7 @@ e-node has a constant value or otherwise return `Nothing`:  ```hs-makeA :: SymExpr (Maybe Double) -> Maybe Int+makeA :: SymExpr (Maybe Double) -> Maybe Double makeA = \case   Const x -> Just x   Symbol _ -> Nothing@@ -144,19 +144,22 @@ Finally, `modifyA` describes how an e-class should (optionally) be modified according to the e-class data and what new language expressions are to be added to the e-class also w.r.t. the e-class data.-Its type is `EClass domain l -> (EClass domain l, [Fix l])`, where the argument-is the class to modify, the first element of the return tuple is the-(optionally) modified e-class and the second element is a list of the-expressions to represent and merge with this e-class.-For our example, if the e-class has a constant value associated to it, we want-to create a new e-class with that constant value and merge it to this e-class.+Its type is `ClassId -> EGraph domain l -> EGraph domain l`, where the first argument+is the id of the class to modify (the class which prompted the modification),+and then receives and returns an e-graph, in which the e-class has been+modified.  For our example, if the e-class has a constant value associated to+it, we want to create a new e-class with that constant value and merge it to+this e-class.  ```hs--- import Data.Equality.Graph.Lens ((^.), _data)-modifyA :: EClass (Maybe Double) SymExpr -> (EClass (Maybe Double) SymExpr, [Fix SymExpr])-modifyA c = case c^._data of-              Nothing -> (c, [])-              Just i  -> (c, [Fix (Const i)])+-- import Data.Equality.Graph.Lens ((^.), _class, _data)+modifyA :: ClassId -> EGraph (Maybe Double) SymExpr -> EGraph (Maybe Double) SymExpr+modifyA c egr+    = case egr ^._class c._data of+        Nothing -> egr+        Just i ->+          let (c', egr') = represent (Fix (Const i)) egr+           in snd $ merge c c' egr' ```  Modify is a bit trickier than the other methods, but it allows our e-graph to@@ -179,15 +182,6 @@ instance Analysis () lang where   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

@@ -1,7 +1,7 @@ cabal-version:      2.4 name:               hegg-version:            0.4.0.0-Tested-With:        GHC ==9.4.2 || ==9.2.2 || ==9.0.2 || ==8.10.7+version:            0.5.0.0+Tested-With:        GHC ==9.6.2 || ==9.4.5 || ==9.2.7 || ==8.10.7 synopsis:           Fast equality saturation in Haskell  description:        Fast equality saturation and equality graphs based on "egg:@@ -38,7 +38,7 @@ maintainer:         Rodrigo Mesquita <rodrigo.m.mesquita@gmail.com> copyright:          Copyright (C) 2022 Rodrigo Mesquita category:           Data-extra-source-files: CHANGELOG.md+extra-doc-files:    CHANGELOG.md                     README.md  source-repository head@@ -46,6 +46,13 @@     location: https://github.com/alt-romes/hegg  library+    other-extensions: BlockArguments StandaloneKindSignatures++    -- QuantifiedConstraints are quite broken with GHC-9.0+    -- Ord (Pattern l) instance triggers same issue like+    -- https://github.com/haskell/mtl/issues/138+    build-depends: base >=4.16  || <4.15+     ghc-options:      -Wall -Wcompat                        -- -fno-prof-auto@@ -77,6 +84,7 @@                       Data.Equality.Extraction,                       Data.Equality.Language,                       Data.Equality.Analysis,+                      Data.Equality.Analysis.Monadic,                       Data.Equality.Saturation.Scheduler,                       Data.Equality.Saturation.Rewrites,                       Data.Equality.Utils,@@ -91,11 +99,13 @@     -- other-modules:      -- LANGUAGE extensions used by modules in this package.-    build-depends:    base         >= 4.4 && < 5,+    -- base-4.13, because foldMap' is used.+    build-depends:    base         >= 4.13 && < 5,                       transformers >= 0.4 && < 0.7,                       containers   >= 0.4 && < 0.7     if flag(vizdot)-        build-depends: graphviz >= 2999.6 && < 2999.7+        build-depends: graphviz >= 2999.20 && < 2999.21,+                       text     hs-source-dirs:   src     default-language: Haskell2010 @@ -108,7 +118,7 @@     hs-source-dirs:   test     main-is:          Test.hs     other-modules:    Invariants, Sym, Lambda, SimpleSym,-                      T1, T2+                      T1, T2, T3      other-extensions: OverloadedStrings     build-depends:    base,@@ -129,8 +139,11 @@                   tasty,                   tasty-hunit,                   tasty-quickcheck,-                  tasty-bench >= 0.2  && < 0.4-  ghc-options:    -with-rtsopts=-A32m -threaded+                  tasty-bench >= 0.2 && < 0.4,+                  deepseq >= 1.4 && < 1.6+  ghc-options:    -with-rtsopts=-A32m+  if impl(ghc >= 8.6)+    ghc-options: -fproc-alignment=64  Flag vizdot     Description: Compile 'Data.Equality.Graph.Dot' module to visualize e-graphs

@@ -30,8 +30,10 @@ import Data.Kind (Type) import Control.Arrow ((***)) -import Data.Equality.Utils+import Data.Function ((&))+import Data.Equality.Graph.Lens import Data.Equality.Language+import Data.Equality.Graph.Internal (EGraph) import Data.Equality.Graph.Classes  -- | An e-class analysis with domain @domain@ defined for a language @l@.@@ -70,31 +72,29 @@     -- e-node to c. Modify should be idempotent if no other changes occur to     -- the e-class, i.e., modify(modify(c)) = modify(c)     ---    -- The return value of the modify function is both the modified class and-    -- the expressions (in their fixed-point form) to add to this class. We-    -- can't manually add them because not only would it skip some of the-    -- internal steps of representing + merging, but also because it's-    -- impossible to add any expression with depth > 0 without access to the-    -- e-graph (since we must represent every sub-expression in the e-graph-    -- first).-    ---    -- That's why we must return the modified class and the expressions to add-    -- to this class.-    --     -- === Example     --     -- Pruning an e-class with a constant value of all its nodes except for the     -- leaf values, and adding a constant value node     --     -- @-    --  -- Prune all except leaf e-nodes-    --  modifyA cl =-    --    case cl^._data of-    --      Nothing -> (cl, [])-    --      Just d -> ((_nodes %~ S.filter (F.null .unNode)) cl, [Fix (Const d)])+    -- modifyA cl eg0+    --   = case eg0^._class cl._data of+    --       Nothing -> eg0+    --       Just d  ->+    --             -- Add constant as e-node+    --         let (new_c,eg1) = represent (Fix (Const d)) eg0+    --             (rep, eg2)  = merge cl new_c eg1+    --             -- Prune all except leaf e-nodes+    --          in eg2 & _class rep._nodes %~ S.filter (F.null .unNode)     -- @-    modifyA :: EClass domain l -> (EClass domain l, [Fix l])-    modifyA c = (c, [])+    modifyA :: ClassId+            -- ^ Id of class @c@ whose new data @d_c@ triggered the modify call+            -> EGraph domain l+            -- ^ E-graph where class @c@ being modified exists+            -> EGraph domain l+            -- ^ E-graph resulting from the modification+    modifyA _ = id     {-# INLINE modifyA #-}  @@ -105,13 +105,14 @@   joinA = (<>)  --- This instance is not necessarily well behaved for any two analysis, so care+-- | This instance is not necessarily well behaved for any two analysis, so care -- must be taken when using it. -- -- A possible criterion is: -- -- For any two analysis, where 'modifyA' is called @m1@ and @m2@ respectively,--- this instance is well behaved if @m1@ and @m2@ commute.+-- this instance is well behaved if @m1@ and @m2@ commute, and the analysis+-- only change the e-class being modified. -- -- That is, if @m1@ and @m2@ satisfy the following law: -- @@@ -133,10 +134,12 @@   joinA :: (a,b) -> (a,b) -> (a,b)   joinA (x,y) = joinA @a @l x *** joinA @b @l y -  modifyA :: EClass (a, b) l -> (EClass (a, b) l, [Fix l])-  modifyA c =-    let (ca, la) = modifyA @a (c { eClassData = fst (eClassData c) })-        (cb, lb) = modifyA @b (c { eClassData = snd (eClassData c) })-     in ( EClass (eClassId c) (eClassNodes ca <> eClassNodes cb) (eClassData ca, eClassData cb) (eClassParents ca <> eClassParents cb)-        , la <> lb-        )+  modifyA :: ClassId -> EGraph (a, b) l -> EGraph (a, b) l+  modifyA c egr =+    let egra = modifyA @a c (egr & _classes._data %~ fst)+        egrb = modifyA @b c (egr & _classes._data %~ snd)+        ca = egra ^._class c+        cb = egrb ^._class c+     in+      egr &+        _class c .~ (EClass c (eClassNodes ca <> eClassNodes cb) (eClassData ca, eClassData cb) (eClassParents ca <> eClassParents cb))

@@ -0,0 +1,68 @@+{-# LANGUAGE AllowAmbiguousTypes #-} -- joinA+{-# LANGUAGE InstanceSigs #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE ImpredicativeTypes #-}+{-|++Like 'Data.Equality.Analysis' but for 'Analysis' that are only well-defined+within an (effectful) context. Mostly used with the monadic operations+'representM', 'addM', 'mergeM', and 'rebuildM'.++This effectful 'Analysis' could almost be trivially defined in terms of the+other, through a "contextful" domain and by means of the '_classes' 'Traversal'.++However, that would require an instance of 'Eq' for the monadic domain, which+is usually unnattainable.++Therefore, we do need this class for monadic 'Analysis'.++-}+module Data.Equality.Analysis.Monadic where+++import Data.Kind (Type)++import Data.Equality.Graph.Internal (EGraph)+import Data.Equality.Graph.Classes++-- | An e-class analysis with domain @domain@ defined for a language @l@, whose operations are only well-defined within some effectful context.+--+-- The @domain@ is the type of the domain of the e-class analysis, that is, the+-- type of the data stored in an e-class according to this e-class analysis+class (Monad m, Eq domain) => AnalysisM m domain (l :: Type -> Type) where++    -- | When a new e-node is added into a new, singleton e-class, construct a+    -- new value of the domain to be associated with the new e-class, by+    -- accessing the associated data of the node's children+    --+    -- The argument is the e-node term populated with its children data+    makeA :: l domain -> m domain++    -- | When e-classes c1 c2 are being merged into c, join d_c1 and+    -- d_c2 into a new value d_c to be associated with the new+    -- e-class c+    joinA :: domain -> domain -> m domain++    -- | Optionally modify the e-class c (based on d_c), typically by adding an+    -- e-node to c. Modify should be idempotent if no other changes occur to+    -- the e-class, i.e., modify(modify(c)) = modify(c)+    modifyA :: ClassId+            -- ^ Id of class @c@ whose new data @d_c@ triggered the modify call+            -> EGraph domain l+            -- ^ E-graph where class @c@ being modified exists+            -> m (EGraph domain l)+            -- ^ E-graph resulting from the modification+    modifyA _ = pure+    {-# INLINE modifyA #-}+++-- | The simplest analysis that defines the domain to be () and does nothing otherwise+instance Monad m => AnalysisM m () l where+  makeA _ = pure ()+  joinA _ _ = pure ()+

@@ -5,7 +5,8 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-}-{-# LANGUAGE NamedFieldPuns #-}+{-# LANGUAGE TupleSections #-}+{-# OPTIONS_GHC -ddump-to-file -ddump-simpl #-} {-|    An e-graph efficiently represents a congruence relation over many expressions. @@ -13,35 +14,45 @@  -} module Data.Equality.Graph     (-      -- * Definition of e-graph+      -- ** Definition of e-graph       EGraph -      -- * Functions on e-graphs-    , emptyEGraph--      -- ** Transformations+      -- ** E-graph transformations     , represent, add, merge, rebuild     -- , repair, repairAnal        -- ** Querying     , find, canonicalize +      -- ** Functions on e-graphs+    , emptyEGraph, newEClass++      -- * E-graph transformations for monadic analysis+      -- | These are the same operations over e-graphs as above but over a monad in which the analysis is defined.+      -- It is common to only have a valid 'Analysis' under a monadic context.+      -- In that case, these are the functions to use -- they are just like the+      -- non-monadic ones, but have require an 'Analysis' defined in a+      -- monadic context ('AnalysisM').+    , representM, addM, mergeM, rebuildM+       -- * Re-exports     , module Data.Equality.Graph.Classes     , module Data.Equality.Graph.Nodes     , module Data.Equality.Language     ) where --- ROMES:TODO: Is the E-Graph a Monad if the analysis data were the type arg? i.e. Monad (EGraph language)?+-- ROMES:TODO: Is the E-Graph a Monad if the analysis data were the type arg? i.e. instance Monad (EGraph language)?  -- import GHC.Conc import Prelude hiding (lookup)  import Data.Function+import Data.Foldable (foldlM) import Data.Bifunctor import Data.Containers.ListUtils  import Control.Monad+import Control.Monad.Trans.Class import Control.Monad.Trans.State import Control.Exception (assert) @@ -55,6 +66,7 @@ import Data.Equality.Graph.Classes import Data.Equality.Graph.Nodes import Data.Equality.Analysis+import qualified Data.Equality.Analysis.Monadic as AM import Data.Equality.Language import Data.Equality.Graph.Lens @@ -89,12 +101,7 @@             (new_eclass_id, new_uf) = makeNewSet (unionFind egr)              -- New singleton e-class stores the e-node and its analysis data-            -- which is modified according to analysis-            ---            -- The modification also produces a list of expressions to-            -- represent and merge with this class, which we'll do before-            -- returning from this function-            (new_eclass, added_nodes) = modifyA $ EClass new_eclass_id (S.singleton new_en) (makeA @a ((\i -> egr^._class i._data @a) <$> unNode new_en)) mempty+            new_eclass = EClass new_eclass_id (S.singleton new_en) (makeA @a ((\i -> egr^._class i._data @a) <$> unNode new_en)) mempty              -- TODO:Performance: All updates can be done to the map first? Parallelize?             --@@ -138,20 +145,14 @@             -- Add the e-node's e-class id at the e-node's id             new_memo         = insertNM new_en new_eclass_id (memo egr) -            -- So we have our almost final e-graph. We just need to represent-            -- and merge in it all expressions which resulted from 'modifyA'-            -- above-            egr1             = egr { unionFind = new_uf-                                   , classes   = new_classes-                                   , worklist  = new_worklist-                                   , memo      = new_memo-                                   }--            egr2             = foldr (representAndMerge new_eclass_id) egr1 added_nodes--          in ( new_eclass_id-            , egr2+            , egr { unionFind = new_uf+                  , classes   = new_classes+                  , worklist  = new_worklist+                  , memo      = new_memo+                  }+                  -- Modify created node according to analysis+                  & modifyA new_eclass_id             ) {-# INLINABLE add #-} @@ -184,11 +185,10 @@             -- Update leader class with all e-nodes and parents from the            -- subsumed class-           (updatedLeader, added_nodes) = leader_class-                                            & _parents %~ (sub_class^._parents <>)-                                            & _nodes   %~ (sub_class^._nodes <>)-                                            & _data    .~ new_data-                                            & modifyA+           updatedLeader = leader_class+                             & _parents %~ (sub_class^._parents <>)+                             & _nodes   %~ (sub_class^._nodes <>)+                             & _data    .~ new_data             new_data = joinA @a @l (leader_class^._data) (sub_class^._data) @@ -226,10 +226,10 @@              , worklist  = new_worklist              , analysisWorklist = new_analysis_worklist              }--           egr2 = foldr (representAndMerge leader) egr1 added_nodes+             -- Modify according to analysis+             & modifyA new_id -        in (new_id, egr2)+        in (new_id, egr1) {-# INLINEABLE merge #-}              @@ -276,22 +276,18 @@ repairAnal :: forall a l. (Analysis a l, Language l) => (ClassId, ENode l) -> EGraph a l -> EGraph a l repairAnal (repair_id, node) egr =     let-        c1                = egr^._class repair_id-        new_data          = joinA @a @l (c1^._data) (makeA @a ((\i -> egr^._class i^._data @a) <$> unNode node))-        (c2, added_nodes) = modifyA (c1 & _data .~ new_data)+        c                = egr^._class repair_id+        new_data          = joinA @a @l (c^._data) (makeA @a ((\i -> egr^._class i^._data @a) <$> unNode node))     in     -- Take action if the new_data is different from the existing data-    if c1^._data /= new_data+    if c^._data /= new_data         -- Merge result is different from original class data, update class         -- with new_data        then-        let-            new_classes = IM.insert repair_id c2 (classes egr)-            egr1 = egr { analysisWorklist = toListSL (c1^._parents) <> analysisWorklist egr-                       , classes = new_classes-                       }-            egr2 = foldr (representAndMerge repair_id) egr1 added_nodes-         in egr2+         egr { analysisWorklist = toListSL (c^._parents) <> analysisWorklist egr+             , classes = IM.adjust (_data .~ new_data) repair_id (classes egr)+             }+             & modifyA repair_id        else egr {-# INLINE repairAnal #-} @@ -319,7 +315,192 @@ emptyEGraph = EGraph emptyUF mempty mempty mempty mempty {-# INLINE emptyEGraph #-} --- | Represent an expression (in fix-point form) and merge it with the e-class with the given id-representAndMerge :: (Analysis a l, Language l) => ClassId -> Fix l -> EGraph a l -> EGraph a l-representAndMerge o f g = case represent f g of-                        (i, e) -> snd $ merge o i e+-- | Creates an empty e-class in an e-graph, with the explicitly given domain analysis data.+-- (That is, an e-class with no e-nodes)+newEClass :: (Language l) => a -> EGraph a l -> (ClassId, EGraph a l)+newEClass adata egr =+  let+    -- Make new equivalence class with a new id in the union-find+    (new_eclass_id, new_uf) = makeNewSet (unionFind egr)++    -- New empty e-class stores just the analysis data+    new_eclass = EClass new_eclass_id S.empty adata mempty+   in ( new_eclass_id+      , egr { unionFind = new_uf+            , classes   = IM.insert new_eclass_id new_eclass (classes egr)+            }+      )+{-# INLINE newEClass #-}++----- Monadic operations on e-graphs+-- Unfortunately, these cannot be defined in terms of the primary ones.+-- This could almost be done by defining the domain of the standard Analysis to+-- be (m a), for some Monad m, but this would require an instance Eq (m a),+-- which often won't exist.+--+-- Anyway, this allows us to have a better story for monadic analysis because+-- the type-class functions arguments don't need to be of the same monadic type+-- as the result (unlike if we were using a normal analysis with a monadic domain).+--+-- Be sure to update these functions if the above "canonical" versions are changed.++-- TODO: Move these to new module?++-- * E-graph operations for analysis defined monadically ('AM.AnalysisM')++-- | Like 'represent', but for a monadic analysis+representM :: forall a l m. (AM.AnalysisM m a l, Language l) => Fix l -> EGraph a l -> m (ClassId, EGraph a l)+representM = cata $ \l e -> do+  -- Canonical implementation is represent, this is just the monadic variant of it+  (l', e') <- (`runStateT` e) $ traverse (\f -> get >>= lift . f >>= StateT . const . pure) l+  addM (Node l') e'++-- | Like 'add', but for a monadic analysis+addM :: forall a l m. (AM.AnalysisM m a l, Language l) => ENode l -> EGraph a l -> m (ClassId, EGraph a l)+addM uncanon_e egr =+  -- Canonical implementation is add, this is just the monadic variant of it+    let !new_en = canonicalize uncanon_e egr++     in case lookupNM new_en (memo egr) of+      Just canon_enode_id -> pure (find canon_enode_id egr, egr)+      Nothing -> do+        let+            (new_eclass_id, new_uf) = makeNewSet (unionFind egr)++        new_data <- AM.makeA @m @a ((\i -> egr^._class i._data @a) <$> unNode new_en)++        let+            new_eclass   =  EClass new_eclass_id (S.singleton new_en) new_data mempty+            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)++            new_worklist = (new_eclass_id, new_en):worklist egr++            new_memo     = insertNM new_en new_eclass_id (memo egr)++        egr1 <- egr { unionFind = new_uf+                    , classes   = new_classes+                    , worklist  = new_worklist+                    , memo      = new_memo+                    }+                    & AM.modifyA new_eclass_id++        return ( new_eclass_id, egr1 )+{-# INLINABLE addM #-}++-- | Like 'merge', but for a monadic analysis+mergeM :: forall a l m. (AM.AnalysisM m a l, Language l) => ClassId -> ClassId -> EGraph a l -> m (ClassId, EGraph a l)+mergeM a b egr0 = do+  -- Canonical implementation is merge, this is just the monadic variant of it+  let+      a' = find a egr0+      b' = find b egr0+   in+   if a' == b'+     then return (a', egr0)+     else do+       let+           class_a = egr0 ^._class a'+           class_b = egr0 ^._class b'++           (leader, leader_class, sub, sub_class) =+               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++           (new_id, new_uf) = unionSets leader sub (unionFind egr0)+                                & first (\n -> assert (leader == n) n)++       new_data <- AM.joinA @m @a @l (leader_class^._data) (sub_class^._data)++       let+           updatedLeader = leader_class+                             & _parents %~ (sub_class^._parents <>)+                             & _nodes   %~ (sub_class^._nodes <>)+                             & _data    .~ new_data++           new_classes = (IM.insert leader updatedLeader . IM.delete sub) (classes egr0)++           -- 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 = 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 toListSL (leader_class^._parents)+                else mempty) <>+             analysisWorklist egr0++       -- ROMES:TODO: The code that makes the -1 * cos test pass when some other things are tweaked+       -- new_memo = foldr (`insertNM` leader) (memo egr0) (sub_class^._nodes)++       -- Build new e-graph+       egr1 <- egr0 { unionFind = new_uf+                    , classes   = new_classes+                    -- , memo      = new_memo+                    , worklist  = new_worklist+                    , analysisWorklist = new_analysis_worklist+                    }+                    & AM.modifyA new_id++       return (new_id, egr1)+{-# INLINEABLE mergeM #-}++-- | Like 'rebuild', but for a monadic analysis+rebuildM :: forall a l m. (AM.AnalysisM m a l, Language l) => EGraph a l -> m (EGraph a l)+rebuildM (EGraph uf cls mm wl awl) = do+  -- Canonical implementation is rebuild, this is just the monadic variant of it+  let+    emptiedEgr = EGraph uf cls mm mempty mempty++    wl' = nubOrd $ bimap (`find` emptiedEgr) (`canonicalize` emptiedEgr) <$> wl++  egr'  <- foldlM (flip repairM) emptiedEgr wl'++  let awl' = nubIntOn fst $ first (`find` egr') <$> awl++  egr'' <- foldlM (flip repairAnalM) egr' awl'++  if null (worklist egr'') && null (analysisWorklist egr'')+     then return egr''+     else rebuildM egr''+{-# INLINEABLE rebuildM #-}++-- | Like 'repair', but for a monadic analysis+repairM :: forall a l m. (AM.AnalysisM m a l, Language l) => (ClassId, ENode l) -> EGraph a l -> m (EGraph a l)+repairM (repair_id, node) egr =+  -- Canonical implementation is repair, this is just the monadic variant of it+   case insertLookupNM node repair_id (memo egr) of++      (Nothing, memo') -> return $ egr { memo = memo' }++      (Just existing_class, memo') -> snd <$> (mergeM existing_class repair_id egr{memo = memo'})+{-# INLINE repairM #-}++-- | Like 'repairAnal', but for a monadic analysis+repairAnalM :: forall a l m. (AM.AnalysisM m a l, Language l) => (ClassId, ENode l) -> EGraph a l -> m (EGraph a l)+repairAnalM (repair_id, node) egr = do+  -- Canonical implementation is repairAnal, this is just the monadic variant of it+    let c = egr^._class repair_id++    new_data <- AM.joinA @m @a @l (c^._data) =<< AM.makeA @m @a ((\i -> egr^._class i^._data @a) <$> unNode node)++    if c^._data /= new_data+       then+         egr { analysisWorklist = toListSL (c^._parents) <> analysisWorklist egr+             , classes = IM.adjust (_data .~ new_data) repair_id (classes egr)+             }+             & AM.modifyA repair_id+       else+        return egr+{-# INLINE repairAnalM #-}

@@ -23,22 +23,21 @@ import Data.GraphViz.Attributes (style, dotted, textLabel) import Data.GraphViz.Attributes.Complete -import Data.Equality.Saturation import Data.Equality.Graph-import Data.Equality.Matching-import Database+import Data.Equality.Graph.Internal +txt :: Show a => a -> Text txt = pack . show -writeDemo :: (Functor f, Foldable f, Show (ENode f)) => EGraph f -> IO ()+writeDemo :: (Language language, Show (ENode language)) => EGraph analysis language -> IO () writeDemo = writeDotFile "demo.gv" . toDotGraph -toDotGraph :: (Functor f, Foldable f, Show (ENode f)) => EGraph f -> DotGraph Text+toDotGraph :: (Language language, Show (ENode language)) => EGraph analysis language -> DotGraph Text toDotGraph eg = digraph (Str "egraph") $ do      graphAttrs [Compound True, ClusterRank Local] -    forM_ (IM.toList $ classes eg) $ \(class_id, EClass _ nodes parents) ->+    forM_ (IM.toList $ classes eg) $ \(class_id, EClass _ nodes _ _) ->          subgraph (Str ("cluster_" <> txt class_id)) $ do             graphAttrs [style dotted]@@ -46,17 +45,16 @@                 let n' = canonicalize n eg                 node (txt class_id <> "." <> txt (find i eg)) [textLabel (txt n')] -    forM_ (IM.toList $ classes eg) $ \(class_id, EClass _ nodes parents) -> do+    forM_ (IM.toList $ classes eg) $ \(class_id, EClass _ nodes _ _) -> do          forM_ (zip (S.toList nodes) [0..]) $ \(n, i_in_class) -> do              let n' = canonicalize n eg             let i_in_class' = find i_in_class eg -            forM_ (zip (children n') [0..]) $ \(child, arg_i) -> do+            forM_ (zip (children n') [(0 :: Integer)..]) $ \(child, arg_i) -> do                 -- TODO: On anchors and labels...?                 let child_leader = find child eg                 if child_leader == class_id-                   then edge (txt class_id <> "." <> txt i_in_class') (txt class_id <> "." <> txt i_in_class') [textLabel (txt arg_i)] -- LHead ("cluster_" <> txt class_id), +                   then edge (txt class_id <> "." <> txt i_in_class') (txt class_id <> "." <> txt i_in_class') [textLabel (txt arg_i)] -- LHead ("cluster_" <> txt class_id),                    else edge (txt class_id <> "." <> txt i_in_class') (txt child <> ".0") [LHead ("cluster_" <> txt child_leader), textLabel (txt arg_i)]-    

@@ -22,8 +22,11 @@ import Data.Equality.Graph.Classes import Data.Equality.Graph.ReprUnionFind --- | A 'Lens'' as defined in other lenses libraries+-- | A 'Lens'' as defined in lens type Lens' s a = forall f. Functor f => (a -> f a) -> (s -> f s)+-- | A 'Lens' as defined in lens+type Lens s t a b = forall f. Functor f => (a -> f b) -> (s -> f t)+-- | A 'Traversal' as defined in lens type Traversal s t a b = forall f. Applicative f => (a -> f b) -> (s -> f t)  -- outdated comment for "getClass":@@ -36,7 +39,7 @@ -- -- Invariant: The e-class exists. --- | Lens for the e-class with the given id in an e-graph+-- | Lens for the e-class at the representative of the given id in an e-graph -- -- Calls 'error' when the e-class doesn't exist _class :: ClassId -> Lens' (EGraph a l) (EClass a l)@@ -62,8 +65,8 @@ {-# INLINE _iclasses #-}  -- | Lens for the 'Domain' of an e-class-_data :: Lens' (EClass domain l) domain-_data afa EClass{..} = (\d1 -> EClass eClassId eClassNodes d1 eClassParents) <$> afa eClassData+_data :: Lens (EClass domain l) (EClass domain' l) domain domain'+_data afb EClass{..} = (\d1 -> EClass eClassId eClassNodes d1 eClassParents) <$> afb eClassData {-# INLINE _data #-}  -- | Lens for the parent e-classes of an e-class@@ -89,7 +92,7 @@ {-# INLINE (.~) #-}  -- | Synonym for @'over'@-(%~) :: Lens' s a -> (a -> a) -> (s -> s)+(%~) :: ASetter s t a b -> (a -> b) -> (s -> t) (%~) = over infixr 4 %~ {-# INLINE (%~) #-}@@ -105,13 +108,27 @@ {-# INLINE set #-}  -- | Applies a function to the target-over :: Lens' s a -> (a -> a) -> (s -> s)+over :: ASetter s t a b -> (a -> b) -> (s -> t) over ln f = runIdentity . ln (Identity . f) {-# INLINE over #-} +-- | Basically 'traverse' over a 'Traversal'+traverseOf :: Traversal s t a b -> forall f. Applicative f => (a -> f b) -> s -> f t +traverseOf t = t+{-# INLINE traverseOf #-}+ -- | Returns True if every target of a Traversable satisfies a predicate. allOf :: Traversal s t a b -> (a -> Bool) -> s -> Bool allOf trv f = getAll . getConst . trv (Const . All . f) {-# INLINE allOf #-} +-- * Utilities +-- We need to use 'ASetter' instead of 'Lens' in %~ to ensure type inference can+-- figure out the Functor or Applicative is 'Identity'. Otherwise, we won't be+-- able to use the 'Traversal' to modify something through a 'Lens'.++-- | Used instead of 'Lens' in 'over' and '%~' to ensure one can call those+-- combinators on 'Lens's and 'Traversal's. Essentially, it helps type+-- inference in such function applications+type ASetter s t a b = (a -> Identity b) -> s -> Identity t

@@ -4,7 +4,12 @@    Monadic interface to e-graph stateful computations  -} module Data.Equality.Graph.Monad-  ( egraph+  (+    -- * Threading e-graphs in a stateful computation+    --+    -- | These are the same operations over e-graphs as in 'Data.Equality.Graph',+    -- but defined in the context of a 'State' monad threading around the e-graph.+    egraph   , represent   , add   , merge@@ -13,9 +18,17 @@   , EG.find   , EG.emptyEGraph +    -- * E-graph transformations for monadic analysis+    --+    -- | The same e-graph operations in a stateful computation threading around+    -- the e-graph, but for 'Analysis' defined monadically ('AnalysisM').+  , representM, addM, mergeM, rebuildM+   -- * E-graph stateful computations   , EGraphM+  , EGraphMT   , runEGraphM+  , runEGraphMT    -- * E-graph definition re-export   , EG.EGraph@@ -30,11 +43,14 @@ import Data.Equality.Utils (Fix, cata)  import Data.Equality.Analysis+import qualified Data.Equality.Analysis.Monadic as AM import Data.Equality.Graph (EGraph, ClassId, Language, ENode(..)) import qualified Data.Equality.Graph as EG  -- | E-graph stateful computation type EGraphM a l = State (EGraph a l)+-- | E-graph stateful computation over an arbitrary monad+type EGraphMT a l = StateT (EGraph a l)  -- | Run EGraph computation on an empty e-graph --@@ -84,3 +100,30 @@ runEGraphM = flip runState {-# INLINE runEGraphM #-} +--------------------------------------------------------------------------------+-- Monadic Analysis interface++-- | Run 'EGraphM' computation on a given e-graph over a monadic analysis+runEGraphMT :: EGraph anl l -> EGraphMT anl l m a -> m (a, EGraph anl l)+runEGraphMT = flip runStateT+{-# INLINE runEGraphMT #-}++-- | Like 'represent', but for a monadic analysis+representM :: (AM.AnalysisM m anl l, Language l) => Fix l -> EGraphMT anl l m ClassId+representM = StateT . EG.representM+{-# INLINE representM #-}++-- | Like 'add', but for a monadic analysis+addM :: (AM.AnalysisM m anl l, Language l) => ENode l -> EGraphMT anl l m ClassId+addM = StateT . EG.addM+{-# INLINE addM #-}++-- | Like 'merge', but for a monadic analysis+mergeM :: (AM.AnalysisM m anl l, Language l) => ClassId -> ClassId -> EGraphMT anl l m ClassId+mergeM a b = StateT (EG.mergeM a b)+{-# INLINE mergeM #-}++-- | Like 'rebuild', but for a monadic analysis+rebuildM :: (AM.AnalysisM m anl l, Language l) => EGraphMT anl l m ()+rebuildM = StateT (fmap ((),) . EG.rebuildM)+{-# INLINE rebuildM #-}

@@ -4,6 +4,7 @@ {-# LANGUAGE OverloadedLists #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE LambdaCase #-}+{-# LANGUAGE CPP #-} {-|    Custom database implemented with trie-maps specialized to run conjunctive    queries using a (worst-case optimal) generic join algorithm.@@ -31,6 +32,9 @@ import Data.Maybe (mapMaybe) import Control.Monad +#if MIN_VERSION_base(4,20,0)+import Data.Foldable as F (toList, length)+#endif import Data.Foldable as F (toList, foldl', length) import qualified Data.Map.Strict    as M import qualified Data.IntMap.Strict as IM

@@ -123,7 +123,7 @@       -- 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 (!matches, newStats) = mconcat (fmap (\(rw_id,rw) -> first (map (rw,)) $ matchWithScheduler db i stats rw_id rw) (zip [1..] rewrites))        -- Write-only phase, temporarily break invariants       forM_ matches applyMatchesRhs@@ -141,10 +141,10 @@                 && IM.size afterClasses == IM.size beforeClasses)           (runEqualitySaturation' (i+1) newStats) -  matchWithScheduler :: Database l -> Int -> IM.IntMap (Stat schd) -> (Int, Rewrite a l) -> ([(Rewrite a 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+  matchWithScheduler :: Database l -> Int -> IM.IntMap (Stat schd) -> Int -> Rewrite a l -> ([Match], IM.IntMap (Stat schd))+  matchWithScheduler db i stats rw_id = \case+      rw  :| _ -> matchWithScheduler db i stats rw_id rw+      lhs := _ -> do           case IM.lookup rw_id stats of             -- If it's banned until some iteration, don't match this rule             -- against anything.@@ -159,7 +159,7 @@                 -- Backoff scheduler: update stats                 let newStats = updateStats schd i rw_id x stats matches' -                (map (lhs := rhs,) matches', newStats)+                (matches', newStats)    applyMatchesRhs :: (Rewrite a l, Match) -> EGraphM a l ()   applyMatchesRhs =

@@ -1,11 +1,14 @@-{-# LANGUAGE UnicodeSyntax, RankNTypes, QuantifiedConstraints, UndecidableInstances #-}+{-# LANGUAGE UnicodeSyntax, RankNTypes, QuantifiedConstraints, UndecidableInstances, CPP #-} {-|  Misc utilities used accross modules  -} module Data.Equality.Utils where  -- import GHC.Conc+#if MIN_VERSION_base(4,20,0)+#else import Data.Foldable+#endif import Data.Bits  -- import qualified Data.Set    as S

@@ -1,18 +1,43 @@ {-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE StandaloneDeriving, FlexibleInstances, DeriveAnyClass, RankNTypes, QuantifiedConstraints, UndecidableInstances, DeriveGeneric #-} import Test.Tasty.Bench +import GHC.Generics+import Control.DeepSeq+ import Data.Equality.Utils import Invariants import Sym-import Lambda-import SimpleSym +-- Instances for benchmarking. It's amazing this works!+deriving instance (forall a. Generic a => Generic (f a)) => Generic (Fix f)+deriving instance NFData UOp+deriving instance NFData BOp+deriving instance NFData a => NFData (Expr a)+deriving instance (forall a. NFData a => NFData (f a), forall a. Generic a => Generic (f a)) => NFData (Fix f)+ tests :: [Benchmark] tests = [ bgroup "Tests"-  [ symTests-  , lambdaTests-  , simpleSymTests-  , invariants+  [ bgroup "Symbolic bench"+    [ bench "i1" $+        nf rewrite (Fix $ BinOp Integral 1 "x")++    , bench "i2" $+        nf rewrite (Fix $ BinOp Integral (Fix $ UnOp Cos "x") "x")++    , bench "i3" $+        nf rewrite (Fix $ BinOp Integral (Fix $ BinOp Pow "x" 1) "x")++    , bench "i4" $+        nf rewrite (_i ((*) "x" (_cos "x")) "x")++    , bench "i5" $+        nf rewrite (_i ((*) (_cos "x") "x") "x")++    , bench "i6" $+        nf rewrite (_i (_ln "x") "x")+    ]+  -- , invariants   ] ]  main :: IO ()

@@ -68,7 +68,6 @@                            | _:xs <- queryHeadVars q ->                                L.sort xs == L.sort (vars p)                                  && length atoms == numNonVarPatterns p-                         _ -> error "impossible! testCompileToQuery"     where         numNonVarPatterns :: Foldable lang => Pattern lang -> Int         numNonVarPatterns (VariablePattern _) = 0

@@ -73,14 +73,16 @@ instance Analysis (Maybe (Lambda ())) Lambda where   makeA = evalL   joinA = (<|>)-  modifyA c = case c^._data of-                Nothing -> (c, [])-                Just v  -> (c, [f v])-                  where-                    f = \case-                      Bool b -> Fix $ Bool b-                      Num i  -> Fix $ Num i-                      _ -> error "impossible, lambda () can't construct this"+  modifyA c eg+    = case eg^._class c._data of+        Nothing -> eg+        Just v  -> let (c', eg') = represent (f v) eg+                    in snd $ merge c c' eg'+          where+            f = \case+              Bool b -> Fix $ Bool b+              Num i  -> Fix $ Num i+              _ -> error "impossible, lambda () can't construct this"                         -- Free variable analysis for lambda

@@ -13,7 +13,8 @@ import Data.Equality.Matching import Data.Equality.Saturation import Data.Equality.Analysis-import Data.Equality.Graph.Lens ((^.), _data)+import Data.Equality.Graph+import Data.Equality.Graph.Lens  data SymExpr a = Const Double                | Symbol String@@ -37,9 +38,12 @@   joinA Nothing Nothing  = Nothing   joinA (Just x) (Just y) = if x == y then Just x else error "ouch, that shouldn't have happened" -  modifyA c = case c^._data of-                Nothing -> (c, [])-                Just i  -> (c, [Fix (Const i)])+  modifyA c eg+    = case eg^._class c._data of+        Nothing -> eg+        Just i  ->+          let (c', eg') = represent (Fix (Const i)) eg+           in snd $ merge c c' eg'  cost :: CostFunction SymExpr Int cost = \case

@@ -6,8 +6,11 @@ {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE LambdaCase #-}+{-# LANGUAGE DeriveGeneric #-}+{-# LANGUAGE CPP #-} module Sym where +import GHC.Generics import Test.Tasty import Test.Tasty.HUnit @@ -18,9 +21,13 @@  import qualified Data.Foldable as F +#if MIN_VERSION_base(4,18,0)+#else import Control.Applicative (liftA2)+#endif import Control.Monad (unless) +import Data.Function ((&)) import Data.Equality.Graph.Lens import Data.Equality.Graph import Data.Equality.Extraction@@ -35,6 +42,7 @@             | BinOp !BOp !a !a             deriving ( Eq, Ord, Show                      , Functor, Foldable, Traversable+                     , Generic                      ) data BOp = Add          | Sub@@ -43,13 +51,13 @@          | Pow          | Diff          | Integral-        deriving (Eq, Ord, Show)+        deriving (Eq, Ord, Show, Generic)  data UOp = Sin          | Cos          | Sqrt          | Ln-         deriving (Eq, Ord, Show)+         deriving (Eq, Ord, Show, Generic)  instance IsString (Fix Expr) where     fromString = Fix . Sym@@ -112,17 +120,15 @@         !_ <- unless (a == b || (a == 0 && b == (-0)) || (a == (-0) && b == 0)) (error "Merged non-equal constants!")         return a -    modifyA cl = case cl^._data of-                 Nothing -> (cl, [])-                 Just d -> ((_nodes %~ S.filter (F.null .unNode)) cl, [Fix (Const d)])--    --         -- Add constant as e-node-    --         new_c <- represent (Fix $ Const d)-    --         _     <- GM.merge i new_c--    --         -- Prune all except leaf e-nodes-    --         modify (_class i._nodes %~ S.filter (F.null . unNode))-+    modifyA cl eg0+      = case eg0^._class cl._data of+          Nothing -> eg0+          Just d  ->+                -- Add constant as e-node+            let (new_c,eg1) = represent (Fix (Const d)) eg0+                (rep, eg2)  = merge cl new_c eg1+                -- Prune all except leaf e-nodes+             in eg2 & _class rep._nodes %~ S.filter (F.null .unNode)   evalConstant :: Expr (Maybe Double) -> Maybe Double

@@ -0,0 +1,47 @@+{-# LANGUAGE DeriveTraversable #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}+module T3 where++-- Some e-graph unit tests++import Prelude hiding (not)++-- import Test.Tasty.HUnit+import Data.Equality.Graph+import Data.Equality.Utils+import qualified Data.Equality.Graph.Monad as EGM++data Lang a = And a a+            | Or a a+            | Not a+            | ToElim a+            | Sym Int+            deriving (Functor, Foldable, Traversable, Eq, Ord, Show)++main :: IO ()+main = do+  let _ = EGM.egraph @Lang @() $ do+        id1 <- EGM.represent (Fix (Sym 1))+        id2 <- EGM.represent (Fix (Not (Fix (Not (Fix (Sym 1))))))+        id3 <- EGM.represent (Fix (Sym 2))+        a1  <- EGM.add (Node (And id3 id1))+        a2  <- EGM.add (Node (And id3 id2))+        _ <- EGM.merge a1 a2+        EGM.rebuild -- even rebuilding this fails...+        return (id1,id2)++  -- The children should now be in the same e-class?+  --+  -- Turns out, they don't. So this test should actually be the other way+  -- around (we do not learn s1 ~= s2 from merging a1 and a2)+  --+  -- A counter example:+  -- Consider `f` that returns its second argument (f _ x = x):+  -- so f(a,c) = f(b,c), but a != b+  --+  -- find s1 eg @=? find s2 eg++  return ()+

@@ -13,6 +13,7 @@  import qualified T1 import qualified T2+import qualified T3  tests :: TestTree tests = testGroup "Tests"@@ -22,6 +23,7 @@   , invariants   , testCase "T1" (T1.main `catch` (\(e :: SomeException) -> assertFailure (show e)))   , testCase "T2" (T2.main `catch` (\(e :: SomeException) -> assertFailure (show e)))+  , testCase "T3" (T3.main `catch` (\(e :: SomeException) -> assertFailure (show e)))   ]  main :: IO ()