hegg (empty) → 0.1.0.0
raw patch · 28 files changed
+3191/−0 lines, 28 filesdep +basedep +containersdep +deriving-compat
Dependencies added: base, containers, deriving-compat, graphviz, hegg, tasty, tasty-hunit, tasty-quickcheck, transformers
Files
- CHANGELOG.md +5/−0
- hegg.cabal +116/−0
- src/Data/Equality/Analysis.hs +61/−0
- src/Data/Equality/Extraction.hs +154/−0
- src/Data/Equality/Graph.hs +310/−0
- src/Data/Equality/Graph.hs-boot +22/−0
- src/Data/Equality/Graph/Classes.hs +35/−0
- src/Data/Equality/Graph/Classes.hs-boot +8/−0
- src/Data/Equality/Graph/Classes/Id.hs +17/−0
- src/Data/Equality/Graph/Dot.hs +62/−0
- src/Data/Equality/Graph/Lens.hs +101/−0
- src/Data/Equality/Graph/Monad.hs +83/−0
- src/Data/Equality/Graph/Nodes.hs +137/−0
- src/Data/Equality/Graph/ReprUnionFind.hs +133/−0
- src/Data/Equality/Language.hs +44/−0
- src/Data/Equality/Matching.hs +137/−0
- src/Data/Equality/Matching/Database.hs +323/−0
- src/Data/Equality/Matching/Pattern.hs +97/−0
- src/Data/Equality/Saturation.hs +187/−0
- src/Data/Equality/Saturation/Rewrites.hs +52/−0
- src/Data/Equality/Saturation/Scheduler.hs +89/−0
- src/Data/Equality/Utils.hs +58/−0
- src/Data/Equality/Utils/IntToIntMap.hs +132/−0
- test/Invariants.hs +218/−0
- test/Lambda.hs +146/−0
- test/SimpleSym.hs +65/−0
- test/Sym.hs +371/−0
- test/Test.hs +28/−0
+ CHANGELOG.md view
@@ -0,0 +1,5 @@+# Revision history for hsym++## 0.1.0.0 -- YYYY-mm-dd++* First version. Released on an unsuspecting world.
+ hegg.cabal view
@@ -0,0 +1,116 @@+cabal-version: 2.4+name: hegg+version: 0.1.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'),+ a data structure which efficiently represents a congruence+ relation over many expressions+ .+ Secondly, it provides functions for doing equality+ 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+ described in "Relational E-Matching"+ .+ For a walkthrough of writing a simple symbolic+ simplification program see the [hegg symbolic+ tutorial](https://github.com/alt-romes/hegg#equality-saturation-in-haskell).+ .+ Additional information can be found [in the README](https://github.com/alt-romes/hegg).++homepage: https://github.com/alt-romes/hegg+bug-reports: https://github.com/alt-romes/hegg/issues+license: BSD-3-Clause+author: Rodrigo Mesquita <romes>+maintainer: Rodrigo Mesquita <rodrigo.m.mesquita@gmail.com>+copyright: Copyright (C) 2022 Rodrigo Mesquita+category: Data+extra-source-files: CHANGELOG.md++source-repository head+ type: git+ location: https://github.com/alt-romes/hegg++library+ ghc-options: -Wall -Wcompat++ -- -fno-prof-auto++ -- -ddump-simpl+ -- -ddump-to-file+ -- -dsuppress-ticks+ -- -dsuppress-stg-exts+ -- -dsuppress-coercions+ -- -dsuppress-idinfo+ -- -dsuppress-unfoldings+ -- -dsuppress-module-prefixes+ -- -dsuppress-timestamps+ -- -dsuppress-uniques+ -- -dsuppress-var-kinds++ exposed-modules: Data.Equality.Graph,+ Data.Equality.Graph.ReprUnionFind,+ Data.Equality.Graph.Classes,+ Data.Equality.Graph.Classes.Id,+ Data.Equality.Graph.Nodes,+ Data.Equality.Graph.Lens,+ Data.Equality.Graph.Monad,+ Data.Equality.Matching,+ Data.Equality.Matching.Database,+ Data.Equality.Matching.Pattern,+ Data.Equality.Saturation,+ Data.Equality.Extraction,+ Data.Equality.Language,+ Data.Equality.Analysis,+ Data.Equality.Saturation.Scheduler,+ Data.Equality.Saturation.Rewrites,+ Data.Equality.Utils+ if impl(ghc >= 9.2)+ exposed-modules: Data.Equality.Utils.IntToIntMap++ if flag(vizdot)+ exposed-modules: Data.Equality.Graph.Dot++ -- Modules included in this library but not exported.+ -- other-modules:++ -- LANGUAGE extensions used by modules in this package.+ build-depends: base >= 4.4 && < 5,+ transformers >= 0.4 && < 0.7,+ containers >= 0.4 && < 0.7+ if flag(vizdot)+ build-depends: graphviz >= 2999.6 && < 2999.7+ hs-source-dirs: src+ default-language: Haskell2010++test-suite hegg-test+ ghc-options: -threaded -Wall+ -- -finfo-table-map -fdistinct-constructor-tables+ -- -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,+ deriving-compat >= 0.6 && < 0.7,+ tasty >= 1.4 && < 1.5,+ tasty-hunit >= 0.10 && < 0.11,+ tasty-quickcheck >= 0.10 && < 0.11++Flag vizdot+ Description: Compile 'Data.Equality.Graph.Dot' module to visualize e-graphs+ Manual: True+ Default: False
+ src/Data/Equality/Analysis.hs view
@@ -0,0 +1,61 @@+{-# LANGUAGE AllowAmbiguousTypes #-} -- joinA+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-|++E-class analysis, which allows the concise expression of a program analysis over+the e-graph.++An e-class analysis resembles abstract interpretation lifted to the e-graph+level, attaching analysis data from a semilattice to each e-class.++The e-graph maintains and propagates this data as e-classes get merged and new+e-nodes are added.++Analysis data can be used directly to modify the e-graph, to inform how or if+rewrites apply their right-hand sides, or to determine the cost of terms during+the extraction process.++References: https://arxiv.org/pdf/2004.03082.pdf++-}+module Data.Equality.Analysis where++import Data.Equality.Graph.Classes.Id+import Data.Equality.Graph.Nodes++import {-# SOURCE #-} Data.Equality.Graph (EGraph)++-- | The e-class analysis defined for a language @l@.+class Eq (Domain l) => Analysis l where++ -- | Domain of data stored in e-class according to e-class analysis+ type Domain l++ -- | 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, typically+ -- by accessing the associated data of n's children+ makeA :: ENode l -> EGraph l -> Domain l++ -- | 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 l -> Domain l -> Domain l++ -- | Optionaly 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)+ --+ -- === Example+ --+ -- Pruning an e-class with a constant value of all its nodes except for the+ -- leaf values+ --+ -- @+ -- -- Prune all except leaf e-nodes+ -- modify (_class i._nodes %~ S.filter (null . children))+ -- @+ modifyA :: ClassId -> EGraph l -> EGraph l+ modifyA _ = id+ {-# INLINE modifyA #-}
+ src/Data/Equality/Extraction.hs view
@@ -0,0 +1,154 @@+{-# LANGUAGE RecordWildCards #-}+{-# LANGUAGE GeneralizedNewtypeDeriving #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE DeriveFunctor #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE PatternSynonyms #-}+{-# LANGUAGE ViewPatterns #-}+{-|+ Given an e-graph representing expressions of our language, we might want to+ extract, out of all expressions represented by some equivalence class, /the best/+ expression (according to a 'CostFunction') represented by that class++ The function 'extractBest' allows us to do exactly that: get the best+ expression represented in an e-class of an e-graph given a 'CostFunction'+ -}+module Data.Equality.Extraction+ (+ -- * Extraction+ extractBest++ -- * Cost+ , CostFunction+ , Cost+ , depthCost+ ) where++import qualified Data.Set as S+import qualified Data.IntMap.Strict as IM++import Data.Equality.Utils+import Data.Equality.Graph++-- vvvv and necessarily all the best sub-expressions from children equilalence classes++-- | Extract the /best/ expression from an equivalence class according to a+-- 'CostFunction'+--+-- @+-- (i, egr) = ...+-- i <- represent expr+-- ...+--+-- bestExpr = extractBest egr 'depthCost' i+-- @+--+-- 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) = ++ -- 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++ in case findBest i allCosts of+ Just (CostWithExpr (_,n)) -> n+ Nothing -> error $ "Couldn't find a best node for e-class " <> show i++ 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 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 _ _) ->++ let+ currentCost = IM.lookup i' beingUpdated++ newCost = S.foldl' (\c n -> case (c, nodeTotalCost beingUpdated n) of+ (Nothing, Nothing) -> Nothing+ (Nothing, Just nc) -> Just nc+ (Just oc, Nothing) -> Just oc+ (Just oc, Just nc) -> Just (oc `min` nc)+ ) Nothing nodes+ -- Current cost + get lowest cost and corresponding node of an e-class if possible+ in case (currentCost, newCost) of++ (Nothing, Just new) -> (True, IM.insert i' new beingUpdated)++ (Just (CostWithExpr old), Just (CostWithExpr new))+ | fst new < fst old -> (True, IM.insert i' (CostWithExpr new) beingUpdated)++ _ -> acc++ -- If any class was modified, loop+ in if modified+ then findCosts eclasses updated+ else updated++ -- | Get the total cost of a node in an e-graph if possible at this stage of+ -- the extraction+ --+ -- 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 m (Node n) = do+ expr <- traverse ((`IM.lookup` m) . flip find g) n+ return $ CostWithExpr (cost ((fst . unCWE) <$> expr), (Fix $ (snd . unCWE) <$> expr))+ {-# INLINE nodeTotalCost #-}++{-# SCC extractBest #-}++-- | A cost function is used to attribute a cost to representations in the+-- e-graph and to extract the best one.+--+-- === Example+-- @+-- symCost :: Expr Cost -> Cost+-- symCost = \case+-- BinOp Integral e1 e2 -> e1 + e2 + 20000+-- BinOp Diff e1 e2 -> e1 + e2 + 500+-- BinOp x e1 e2 -> e1 + e2 + 3+-- UnOp x e1 -> e1 + 30+-- Sym _ -> 1+-- Const _ -> 1+-- @+type CostFunction l = l Cost -> Cost++-- | 'Cost' is simply an integer+type Cost = Int++-- | Simple cost function: the deeper the expression, the bigger the cost+depthCost :: Language l => CostFunction l+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 i = IM.lookup i+{-# INLINE findBest #-}++newtype CostWithExpr lang = CostWithExpr { unCWE :: (Cost, Fix lang) }++instance Eq (CostWithExpr lang) where+ (==) (CostWithExpr (a,_)) (CostWithExpr (b,_)) = a == b+ {-# INLINE (==) #-}++instance Ord (CostWithExpr lang) where+ compare (CostWithExpr (a,_)) (CostWithExpr (b,_)) = a `compare` b+ {-# INLINE compare #-}+
+ src/Data/Equality/Graph.hs view
@@ -0,0 +1,310 @@+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE TupleSections #-}+-- {-# LANGUAGE ApplicativeDo #-}+{-# LANGUAGE BlockArguments #-}+{-# LANGUAGE UndecidableInstances #-} -- tmp show+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeFamilies #-}+{-|+ An e-graph efficiently represents a congruence relation over many expressions.++ Based on \"egg: Fast and Extensible Equality Saturation\" https://arxiv.org/abs/2004.03082.+ -}+module Data.Equality.Graph+ (+ -- * Definition of e-graph+ EGraph(..)++ , Memo, Worklist++ -- * Functions on e-graphs+ , emptyEGraph++ -- ** Transformations+ , add, merge, rebuild+ -- , repair, repairAnal++ -- ** Querying+ , find, canonicalize++ -- * Re-exports+ , module Data.Equality.Graph.Classes+ , module Data.Equality.Graph.Nodes+ , module Data.Equality.Language+ ) where++-- import GHC.Conc++import Data.Function++import Data.Functor.Classes++import qualified Data.IntMap.Strict as IM+import qualified Data.Set as S++import Data.Equality.Graph.ReprUnionFind+import Data.Equality.Graph.Classes+import Data.Equality.Graph.Nodes+import Data.Equality.Analysis+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++ in case {-# SCC "-1" #-} lookupNM new_en (memo egr) of+ Just canon_enode_id -> {-# SCC "0" #-} (find canon_enode_id egr, egr)+ Nothing ->++ let++ -- Make new equivalence class with a new id in the union-find+ (new_eclass_id, new_uf) = makeNewSet (unionFind egr)++ -- New singleton e-class stores the e-node and its analysis data+ new_eclass = EClass new_eclass_id (S.singleton new_en) (makeA new_en egr) mempty++ -- TODO:Performance: All updates can be done to the map first? Parallelize?+ --+ -- Update e-classes by going through all e-node children and adding+ -- 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))+ (classes egr)+ (unNode new_en)++ -- TODO: From egg: Is this needed?+ -- This is required if we want math pruning to work. Unfortunately, it+ -- also makes the invariants tests x4 slower (because they aren't using+ -- analysis) I think there might be another way to ensure math analysis+ -- pruning to work without having this line here. Comment it out to+ -- check the result on the unit tests.+ -- + -- Update: I found a fix for that case: the modifyA function must add+ -- the parents of the pruned class to the worklist for them to be+ -- upward merged. I think it's a good compromise for requiring the user+ -- to do this. Adding the added node to the worklist everytime creates+ -- too much unnecessary work.+ --+ -- Actually I've found more bugs regarding this, and can't fix them+ -- there, so indeed this seems to be necessary for sanity with 'modifyA'+ --+ -- This way we also liberate the user from caring about the worklist+ --+ -- The hash cons invariants test suffer from this greatly but the+ -- saturation tests seem mostly fine?+ --+ -- And adding to the analysis worklist doesn't work, so maybe it's+ -- 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)++ -- 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)++ in ( new_eclass_id++ , egr { unionFind = new_uf+ , classes = new_classes+ , worklist = new_worklist+ , memo = new_memo+ }++ -- Modify created node according to analysis+ & {-# SCC "6" #-} modifyA new_eclass_id++ )+{-# SCC add #-}++-- | Merge 2 e-classes by id+merge :: forall l. Language l => ClassId -> ClassId -> EGraph l -> (ClassId, EGraph l)+merge a b egr0 =++ -- Use canonical ids+ let+ a' = find a egr0+ b' = find b egr0+ in+ if a' == b'+ then (a', egr0)+ else+ let+ -- Get classes being merged+ class_a = egr0 ^._class a'+ class_b = egr0 ^._class b'++ -- Leader is the class with more parents+ (leader, leader_class, sub, sub_class) =+ if (sizeNM (class_a^._parents)) < (sizeNM (class_b^._parents))+ then (b', class_b, a', class_a) -- b is leader+ else (a', class_a, b', class_b) -- a is leader++ -- Make leader the leader in the union find+ (new_id, new_uf) = unionSets leader sub (unionFind egr0)++ -- Update leader class with all e-nodes and parents from the+ -- subsumed class+ updatedLeader = leader_class & _parents %~ (<> sub_class^._parents)+ & _nodes %~ (<> sub_class^._nodes)+ & _data .~ new_data+ new_data = joinA @l (leader_class^._data) (sub_class^._data)++ -- Update leader in classes so that it has all nodes and parents+ -- from subsumed class, and delete the subsumed class+ 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 = 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 /= (leader_class^._data)+ then 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+ -- new_memo = foldr (`insertNM` leader) (memo egr0) (sub_class^._nodes)++ -- Build new e-graph+ new_egr = egr0+ { unionFind = new_uf+ , classes = new_classes+ -- , memo = new_memo+ , worklist = new_worklist+ , analysisWorklist = new_analysis_worklist+ }++ -- Modify according to analysis+ & modifyA new_id++ in (new_id, new_egr)+{-# SCC merge #-}+ ++-- | 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+-- larger algorithm like equality saturation.+rebuild :: Language l => EGraph l -> EGraph l+rebuild (EGraph uf cls mm wl awl) =+ -- empty worklists+ -- repair deduplicated e-classes+ let+ egr' = foldrWithKeyNM' repair (EGraph uf cls mm mempty mempty) wl+ egr'' = foldrWithKeyNM' repairAnal egr' awl+ in+ -- Loop until worklist is completely empty+ if null (worklist egr'') && null (analysisWorklist egr'')+ then egr''+ else rebuild egr''++{-# SCC 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 =++ 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)++ (Nothing, memo2) -> egr { memo = memo2 } -- Return new memo but delete uncanonicalized node++ (Just existing_class, memo2) -> snd (merge existing_class repair_id egr{memo = memo2})+{-# SCC repair #-}++-- | Repair a single analysis-worklist entry.+repairAnal :: forall l. Language l => ENode l -> ClassId -> EGraph l -> EGraph l+repairAnal node repair_id egr =+ let+ canon_id = find repair_id egr+ c = egr^._class canon_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+ }+ & _class canon_id._data .~ new_data+ & modifyA canon_id+ else egr+{-# SCC repairAnal #-}++-- | Canonicalize an e-node+--+-- Two e-nodes are equal when their canonical form is equal. Canonicalization+-- makes the list of e-class ids the e-node holds a list of canonical ids.+-- Meaning two seemingly different e-nodes might be equal when we figure out+-- that their e-class ids are represented by the same e-class canonical ids+--+-- canonicalize(𝑓(𝑎,𝑏,𝑐,...)) = 𝑓((find 𝑎), (find 𝑏), (find 𝑐),...)+canonicalize :: Functor l => ENode l -> EGraph l -> ENode l+canonicalize (Node enode) eg = Node $ fmap (`find` eg) enode+{-# SCC canonicalize #-}++-- | Find the canonical representation of an e-class id in the e-graph+-- Invariant: The e-class id always exists.+find :: ClassId -> EGraph l -> ClassId+find cid = findRepr cid . unionFind+{-# INLINE find #-}++-- | The empty e-graph. Nothing is represented in it yet.+emptyEGraph :: Language l => EGraph l+emptyEGraph = EGraph emptyUF mempty mempty mempty mempty+{-# INLINE emptyEGraph #-}
+ src/Data/Equality/Graph.hs-boot view
@@ -0,0 +1,22 @@+{-# 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
+ src/Data/Equality/Graph/Classes.hs view
@@ -0,0 +1,35 @@+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE UndecidableInstances #-}+{-|+ Module for the definition of 'EClass'.+-}+module Data.Equality.Graph.Classes+ ( module Data.Equality.Graph.Classes+ , module Data.Equality.Graph.Classes.Id+ ) where++import qualified Data.Set as S++import Data.Functor.Classes++import Data.Equality.Graph.Classes.Id+import Data.Equality.Graph.Nodes++import Data.Equality.Analysis++-- | An e-class (an equivalence class of terms) of a 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 EClass l = EClass+ { 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.+ }++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+
+ src/Data/Equality/Graph/Classes.hs-boot view
@@ -0,0 +1,8 @@+{-# LANGUAGE RoleAnnotations #-}+{-# LANGUAGE KindSignatures #-}+module Data.Equality.Graph.Classes where++import Data.Kind++type role EClass nominal+data EClass (l :: Type -> Type)
+ src/Data/Equality/Graph/Classes/Id.hs view
@@ -0,0 +1,17 @@+{-|++Type synonyms for e-class ids.++-}+module Data.Equality.Graph.Classes.Id+ ( ClassId+ , ClassIdMap+ ) where++import qualified Data.IntMap.Strict as IM++-- | Type synonym for e-class ids+type ClassId = Int++-- | Type synonym for a map from e-class ids to values+type ClassIdMap = IM.IntMap
+ src/Data/Equality/Graph/Dot.hs view
@@ -0,0 +1,62 @@+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE NoMonomorphismRestriction #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE TypeFamilies #-}++module Data.Equality.Graph.Dot+ ( module Data.Equality.Graph.Dot+ , writeDotFile+ )+ where++import Control.Monad++import Data.Text.Lazy (Text, pack)++import qualified Data.Set as S+import qualified Data.IntMap as IM++import Data.GraphViz.Commands.IO+import Data.GraphViz.Types.Generalised+import Data.GraphViz.Types.Monadic+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++txt = pack . show++writeDemo :: (Functor f, Foldable f, Show (ENode f)) => EGraph f -> IO ()+writeDemo = writeDotFile "demo.gv" . toDotGraph++toDotGraph :: (Functor f, Foldable f, Show (ENode f)) => EGraph f -> DotGraph Text+toDotGraph eg = digraph (Str "egraph") $ do++ graphAttrs [Compound True, ClusterRank Local]++ forM_ (IM.toList $ classes eg) $ \(class_id, EClass _ nodes parents) ->++ subgraph (Str ("cluster_" <> txt class_id)) $ do+ graphAttrs [style dotted]+ forM_ (zip (S.toList nodes) [0..]) $ \(n, i) -> do+ 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_ (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+ -- 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), + else edge (txt class_id <> "." <> txt i_in_class') (txt child <> ".0") [LHead ("cluster_" <> txt child_leader), textLabel (txt arg_i)]+
+ src/Data/Equality/Graph/Lens.hs view
@@ -0,0 +1,101 @@+{-# 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'.+ -}+module Data.Equality.Graph.Lens where++import qualified Data.IntMap.Strict as IM+import qualified Data.Set as S++import Data.Functor.Identity+import Data.Functor.Const++import Data.Equality.Graph.Classes.Id+import Data.Equality.Graph.Nodes+import Data.Equality.Graph.Classes+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)+++-- outdated comment for "getClass":+--+-- Get an e-class from an e-graph given its e-class id+--+-- Returns the canonical id of the class and the class itself+--+-- We'll find its canonical representation and then get it from the e-classes map+--+-- Invariant: The e-class exists.++-- | Lens for the e-class with the given id in an e-graph+--+-- Calls 'error' when the e-class doesn't exist+_class :: ClassId -> Lens' (EGraph l) (EClass l)+_class i afa s =+ let canon_id = find i 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)+_memo afa egr = (\m1 -> egr {memo = m1}) <$> afa (memo egr)+{-# INLINE _memo #-}++-- | Lens for the map of existing classes by id in an e-graph+_classes :: Lens' (EGraph l) (ClassIdMap (EClass l))+_classes afa egr = (\m1 -> egr {classes = m1}) <$> afa (classes egr)+{-# INLINE _classes #-}++-- | Lens for the 'Domain' of an e-class+_data :: Lens' (EClass l) (Domain l)+_data afa EClass{..} = (\d1 -> EClass eClassId eClassNodes d1 eClassParents) <$> afa eClassData+{-# 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+{-# INLINE _parents #-}++-- | Lens for the e-nodes in an e-class+_nodes :: Lens' (EClass l) (S.Set (ENode l))+_nodes afa EClass{..} = (\ns -> EClass eClassId ns eClassData eClassParents) <$> afa eClassNodes+{-# INLINE _nodes #-}++-- | Like @'view'@ but with the arguments flipped+(^.) :: s -> Lens' s a -> a+(^.) s ln = view ln s+infixl 8 ^.+{-# INLINE (^.) #-}++-- | Synonym for @'set'@+(.~) :: Lens' s a -> a -> (s -> s)+(.~) = set+infixr 4 .~+{-# INLINE (.~) #-}++-- | Synonym for @'over'@+(%~) :: Lens' s a -> (a -> a) -> (s -> s)+(%~) = over+infixr 4 %~+{-# INLINE (%~) #-}++-- | Applies a getter to a value+view :: Lens' s a -> (s -> a)+view ln = getConst . ln Const+{-# INLINE view #-}++-- | Applies a setter to a value+set :: Lens' s a -> a -> (s -> s)+set ln x = over ln (const x)+{-# INLINE set #-}++-- | Applies a function to the target+over :: Lens' s a -> (a -> a) -> (s -> s)+over ln f = runIdentity . ln (Identity . f)+{-# INLINE over #-}+
+ src/Data/Equality/Graph/Monad.hs view
@@ -0,0 +1,83 @@+{-# LANGUAGE TupleSections #-}+{-|+ Monadic interface to e-graph stateful computations+ -}+module Data.Equality.Graph.Monad+ ( egraph+ , represent+ , add+ , merge+ , rebuild+ , EG.canonicalize+ , EG.find+ , EG.emptyEGraph++ -- * E-graph stateful computations+ , EGraphM+ , runEGraphM++ -- * E-graph definition re-export+ , EG.EGraph++ -- * 'State' monad re-exports+ , modify, get, gets+ ) where++import Control.Monad ((>=>))+import Control.Monad.Trans.State.Strict++import Data.Equality.Utils (Fix, cata)++import Data.Equality.Graph (EGraph, ClassId, Language, ENode(..))+import qualified Data.Equality.Graph as EG++-- | E-graph stateful computation+type EGraphM s = State (EGraph s)++-- | Run EGraph computation on an empty e-graph+--+-- === Example+-- @+-- egraph $ do+-- id1 <- represent t1+-- id2 <- represent t2+-- merge id1 id2+-- @+egraph :: Language l => EGraphM l a -> (a, EGraph l)+egraph = runEGraphM EG.emptyEGraph+{-# INLINE egraph #-}++-- | Represent an expression (@Fix l@) in an e-graph by recursively+-- representing sub expressions+represent :: Language l => Fix l -> EGraphM l ClassId+represent = cata $ sequence >=> add . Node+{-# INLINE represent #-}++-- | Add an e-node to the e-graph+add :: Language l => ENode l -> EGraphM l ClassId+add = StateT . fmap pure . EG.add+{-# INLINE add #-}++-- | Merge two e-classes by id+--+-- E-graph invariants may be broken by merging, and 'rebuild' should be used+-- /eventually/ to restore them+merge :: Language l => ClassId -> ClassId -> EGraphM l ClassId+merge a b = StateT (pure <$> EG.merge a b)+{-# INLINE merge #-}++-- | Rebuild: Restore e-graph invariants+--+-- E-graph invariants are traditionally maintained after every merge, but we+-- allow operations to temporarilly break the invariants (specifically, until we call+-- 'rebuild')+--+-- The paper describing rebuilding in detail is https://arxiv.org/abs/2004.03082+rebuild :: Language l => EGraphM l ()+rebuild = StateT (pure . ((),). EG.rebuild)+{-# INLINE rebuild #-}++-- | Run 'EGraphM' computation on a given e-graph+runEGraphM :: EGraph l -> EGraphM l a -> (a, EGraph l)+runEGraphM = flip runState+{-# INLINE runEGraphM #-}
+ src/Data/Equality/Graph/Nodes.hs view
@@ -0,0 +1,137 @@+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE TupleSections #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE DeriveGeneric #-}+{-# LANGUAGE ViewPatterns #-}+{-# LANGUAGE GeneralizedNewtypeDeriving #-}+{-# LANGUAGE DeriveTraversable #-}+{-|++Module defining e-nodes ('ENode'), the e-node function symbol ('Operator'), and+mappings from e-nodes ('NodeMap').++-}+module Data.Equality.Graph.Nodes where++import Data.Functor.Classes+import Data.Foldable+import Data.Bifunctor++import Data.Kind++import Control.Monad (void)++import qualified Data.Map.Strict as M++import Data.Equality.Graph.Classes.Id+++-- * E-node++-- | An e-node is a function symbol paired with a list of children e-classes.+-- +-- We define an e-node to be the base functor of some recursive data type+-- parametrized over 'ClassId', i.e. all recursive fields are rather e-class ids.+newtype ENode l = Node { unNode :: l ClassId }++-- | Get the children e-class ids of an e-node+children :: Traversable l => ENode l -> [ClassId]+children = toList . unNode+{-# SCC children #-}++-- * Operator++-- | An operator is solely the function symbol part of the e-node. Basically,+-- this means children e-classes are ignored.+newtype Operator l = Operator { unOperator :: l () }++-- | Get the operator (function symbol) of an e-node+operator :: Traversable l => ENode l -> Operator l+operator = Operator . void . unNode+{-# SCC operator #-}++instance Eq1 l => (Eq (ENode l)) where+ (==) (Node a) (Node b) = liftEq (==) a b+ {-# INLINE (==) #-}++instance Ord1 l => (Ord (ENode l)) where+ compare (Node a) (Node b) = liftCompare compare a b+ {-# INLINE compare #-}++instance Show1 l => (Show (ENode l)) where+ showsPrec p (Node l) = liftShowsPrec showsPrec showList p l++instance Eq1 l => (Eq (Operator l)) where+ (==) (Operator a) (Operator b) = liftEq (\_ _ -> True) a b+ {-# INLINE (==) #-}++instance Ord1 l => (Ord (Operator l)) where+ compare (Operator a) (Operator b) = liftCompare (\_ _ -> EQ) a b+ {-# INLINE compare #-}++instance Show1 l => (Show (Operator l)) where+ showsPrec p (Operator l) = liftShowsPrec (const . const $ showString "") (const $ showString "") p l++-- * 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 }+-- 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++-- | 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)+{-# INLINE insertNM #-}++-- | Lookup an e-node in a 'NodeMap'+lookupNM :: Ord1 l => ENode l -> NodeMap l a -> Maybe a+lookupNM e = M.lookup e . unNodeMap+{-# INLINE lookupNM #-}++-- | 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)+{-# 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+{-# INLINE insertLookupNM #-}++-- | As 'Data.Map.foldlWithKeyNM'' but in a 'NodeMap'+foldlWithKeyNM' :: Ord1 l => (b -> ENode l -> a -> b) -> b -> NodeMap l a -> b +foldlWithKeyNM' f b = M.foldlWithKey' f b . unNodeMap+{-# INLINE foldlWithKeyNM' #-}++-- | As 'Data.Map.foldrWithKeyNM'' but in a 'NodeMap'+foldrWithKeyNM' :: Ord1 l => (ENode l -> a -> b -> b) -> b -> NodeMap l a -> b +foldrWithKeyNM' f b = M.foldrWithKey' f b . unNodeMap+{-# INLINE foldrWithKeyNM' #-}++-- | Get the number of entries in a 'NodeMap'.+--+-- This operation takes constant time (__O(1)__)+sizeNM :: NodeMap l a -> Int+sizeNM = sizeNodeMap+{-# 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+{-# INLINE traverseWithKeyNM #-}++-- Node Set++-- newtype NodeSet l a = NodeSet { unNodeSet :: IM.IntMap (a, ENode l) }+-- deriving (Semigroup, Monoid)++-- insertNS :: Hashable1 l => ENode l -> NodeSet l -> NodeSet l+-- insertNS v = NodeSet . IM.insert (hashNode v) v . unNodeSet
+ src/Data/Equality/Graph/ReprUnionFind.hs view
@@ -0,0 +1,133 @@+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE CPP #-}+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE BangPatterns #-}+{-|++Union-find-like data structure that defines equivalence classes of e-class ids.++-}+module Data.Equality.Graph.ReprUnionFind+ ( ReprUnionFind+ , emptyUF+ , makeNewSet+ , unionSets+ , findRepr+ ) where++import Data.Equality.Graph.Classes.Id++#if __GLASGOW_HASKELL__ >= 902++import qualified Data.Equality.Utils.IntToIntMap as IIM+import GHC.Exts ((+#), Int(..), Int#)++type RUFSize = Int#++-- | A union find for equivalence classes of e-class ids.+data ReprUnionFind = RUF IIM.IntToIntMap -- ^ Map every id to either 0# (meaning its the representative) or to another int# (meaning its represented by some other id)+ RUFSize -- ^ Counter for new ids++ -- !(IM.IntMap [ClassId]) -- ^ Mapping from an id to all its children: This is used for "rebuilding" (compress all paths) when merging. Its a hashcons?+ -- [ClassId] -- ^ Ids that can be safely deleted after the e-graph is rebuilt+#else++import qualified Data.IntMap.Internal as IIM (IntMap(..))+import qualified Data.IntMap.Strict as IIM++-- | A union find for equivalence classes of e-class ids.+data ReprUnionFind = RUF (IIM.IntMap Int) -- ^ Map every id to either 0# (meaning its the representative) or to another int# (meaning its represented by some other id)+ {-# UNPACK #-} !Int -- ^ Counter for new ids++#endif++-- Note that there's no value associated with identifier, so this union find+-- serves only to find the representative of an e-class id++instance Show ReprUnionFind where+ show (RUF _ _) = "Warning: Incomplete show: ReprUnionFind"++-- | An @id@ can be represented by another @id@ or be canonical, meaning it+-- represents itself.+--+-- @(x, Represented y)@ would mean x is represented by y+-- @(x, Canonical)@ would mean x is canonical -- represents itself+newtype Repr+ = Represented { unRepr :: ClassId } -- ^ @Represented x@ is represented by @x@+-- | Canonical -- ^ @Canonical x@ is the canonical representation, meaning @find(x) == x@+ deriving Show++-- | The empty 'ReprUnionFind'.+emptyUF :: ReprUnionFind+-- TODO: If I can make an instance of 'ReprUnionFind' for Monoid(?), this is 'mempty'+emptyUF = RUF IIM.Nil+#if __GLASGOW_HASKELL__ >= 902+ 1# -- Must start with 1# since 0# means "Canonical"+#else+ 1+#endif++-- | Create a new e-class id in the given 'ReprUnionFind'.+makeNewSet :: ReprUnionFind+ -> (ClassId, ReprUnionFind) -- ^ Newly created e-class id and updated 'ReprUnionFind'+#if __GLASGOW_HASKELL__ >= 902+makeNewSet (RUF im si) = ((I# si), RUF (IIM.insert si 0# im) ((si +# 1#)))+#else+makeNewSet (RUF im si) = (si, RUF (IIM.insert si 0 im) (si + 1))+#endif+{-# SCC makeNewSet #-}++-- | Union operation of the union find.+--+-- Given two leader ids, unions the two eclasses making @a@ the leader, that+-- is, @b@ is now represented by @a@+unionSets :: ClassId -- ^ E-class id @a@+ -> ClassId -- ^ E-class id @b@+ -> ReprUnionFind -- ^ Union-find containing @a@ and @b@+ -> (ClassId, ReprUnionFind) -- ^ The new leader (always @a@) and the updated union-find+#if __GLASGOW_HASKELL__ >= 902+unionSets a@(I# a#) (I# b#) (RUF im si) = (a, RUF (IIM.insert b# a# im) si)+#else+unionSets a b (RUF im si) = (a, RUF (IIM.insert b a im) si)+#endif+ -- where+ -- 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 #-}++-- | 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+ 0# -> v+ x -> findRepr (I# x) (RUF m s)+#else+findRepr v (RUF m s) =+ case {-# SCC "findRepr_TAKE" #-} m IIM.! v of+ 0 -> v+ x -> findRepr x (RUF m s)+#endif++-- ROMES:TODO: Path compression in immutable data structure? Is it worth+-- the copy + threading?+--+-- ANSWER: According to my tests, findRepr is always quite shallow, going only+-- (from what I saw) until, at max, depth 3!+--+-- When using the ad-hoc path compression in `unionSets`, the depth of+-- recursion never even goes above 1!+{-# SCC findRepr #-}+++-- {-# RULES+-- "union/find" forall a b x im. findRepr (I# b) (RUF (IIM.insert b a im) x) = I# a+-- #-}++-- -- | Delete nodes that have been merged after e-graph has been rebuilt+-- rebuildUF :: ReprUnionFind -> ReprUnionFind+-- rebuildUF (RUF m' a b dl) = RUF (IIM.unliftedFoldr (\(I# x) -> IIM.delete x) m' dl) a b mempty
+ src/Data/Equality/Language.hs view
@@ -0,0 +1,44 @@+{-# LANGUAGE FlexibleContexts #-}+{-|++Defines 'Language', which is the required constraint on /expressions/ that are+to be represented in e-graph and on which equality saturation can be run.++=== Example+@+data Expr a = Sym String+ | Const Double+ | UnOp UOp a+ | BinOp BOp a a+ deriving ( Eq, Ord, Functor+ , Foldable, Traversable)++instance Eq1 Expr where+ ...+instance Ord1 Expr where+ ...++instance Analysis Expr where+ ...++-- meaning we satisfy all other constraints and Expr is! a language+instance Language Expr++@+-}+module Data.Equality.Language where++import Data.Functor.Classes++import Data.Equality.Analysis++-- | A 'Language' is the required constraint on /expressions/ that are to be+-- represented in an e-graph.+--+-- Recursive data types must be expressed in its functor form to instance+-- 'Language'. Additionally, for a datatype to be a 'Language' (used in+-- e-graphs), note that it must satisfy the other class constraints. In+-- particular an 'Data.Equality.Analysis.Analysis' must be defined for the+-- language.+class (Analysis l, Traversable l, Ord1 l) => Language l where+
+ src/Data/Equality/Matching.hs view
@@ -0,0 +1,137 @@+{-# LANGUAGE RecordWildCards #-}+{-# LANGUAGE ViewPatterns #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-|+ Equality-matching, implemented using a relational database+ (defined in 'Data.Equality.Matching.Database') according to the paper+ \"Relational E-Matching\" https://arxiv.org/abs/2108.02290.+ -}+module Data.Equality.Matching+ ( ematch+ , eGraphToDatabase+ , Match(..)+ , compileToQuery++ , module Data.Equality.Matching.Pattern+ )+ where++import Data.Maybe (mapMaybe)+import Data.Foldable (toList)+import Data.Containers.ListUtils++import Control.Monad+import Control.Monad.Trans.State.Strict++import qualified Data.Map.Strict as M+import qualified Data.IntMap.Strict as IM+import qualified Data.IntSet as IS++import Data.Equality.Graph+import Data.Equality.Matching.Database+import Data.Equality.Matching.Pattern++-- | Matching a pattern on an e-graph returns the e-class in which the pattern+-- was matched and an e-class substitution for every 'VariablePattern' in the pattern.+data Match = Match+ { matchSubst :: !Subst+ , matchClassId :: {-# UNPACK #-} !ClassId+ }++-- TODO: Perhaps e-graph could carry database and rebuild it on rebuild++-- | Match a pattern against a 'Database', which can be gotten from an 'EGraph' with 'eGraphToDatabase'+--+-- Returns a list of matches, one 'Match' for each set of valid substitutions+-- for all variables and the equivalence class in which the pattern was matched.+--+-- 'ematch' takes a 'Database' instead of an 'EGraph' because the 'Database'+-- could be constructed only once and shared accross matching.+ematch :: Language l+ => Database l+ -> Pattern l+ -> [Match]+ematch db patr =+ let+ (q, root) = compileToQuery patr++ -- | Convert each substitution into a match by getting the class-id+ -- where we matched from the subst+ --+ -- If the substitution is empty there is no match+ f :: Subst -> Maybe Match+ f s = if IM.null s then Nothing+ else case IM.lookup root s of+ Nothing -> error "how is root not in map?"+ Just found -> pure (Match s found)++ in mapMaybe f (genericJoin db q)++-- | Convert an e-graph into a database+eGraphToDatabase :: Language l => EGraph l -> Database l+eGraphToDatabase EGraph{..} = foldrWithKeyNM' addENodeToDB (DB mempty) memo+ where++ -- Add an enode in an e-graph, given its class, to a database+ addENodeToDB :: Language l => ENode l -> ClassId -> Database l -> Database l+ addENodeToDB enode classid (DB m) =+ -- 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+ populate :: [ClassId] -> Maybe IntTrie -> IntTrie+ -- If trie map entry doesn't exist yet, populate an empty map with the remaining ids+ populate [] Nothing = MkIntTrie mempty mempty+ populate (x:xs) Nothing = MkIntTrie (IS.singleton x) $ IM.singleton x (populate xs Nothing)+ -- 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 #-}+++-- * Database related internals++-- | Auxiliary result in 'compileToQuery' algorithm+data AuxResult lang = {-# UNPACK #-} !Var :~ [Atom lang]++-- | Compiles a 'Pattern' to a 'Query' and returns the query root variable with+-- it.+-- The root variable's substitutions are the e-classes where the pattern+-- matched+compileToQuery :: (Traversable lang) => Pattern lang -> (Query lang, Var)+compileToQuery (VariablePattern x) = (SelectAllQuery x, x)+compileToQuery pa@(NonVariablePattern _) =++ let root :~ atoms = evalState (aux pa) 0+ in (Query (nubInt $ root:vars pa) atoms, root)++ where++ aux :: (Traversable lang) => Pattern lang -> State Int (AuxResult lang)+ aux (VariablePattern x) = return (x :~ []) -- from definition in relational e-matching paper (needed for as base case for recursion)+ aux (NonVariablePattern p) = do+ v <- get+ modify' (+1)+ (toList -> auxs) <- traverse aux p+ let boundVars = map (\(b :~ _) -> b) auxs+ atoms = join $ map (\(_ :~ a) -> a) auxs+ -- Number of bound vars should match number of children of this+ -- lang. We can traverse the pattern and replace sub-patterns with+ -- their corresponding bound variable+ p' = evalState (subPatsToVars p boundVars) 0+ return (v :~ (Atom (CVar v) (fmap CVar p'):atoms))+ where+ -- State keeps track of the index of the variable we're+ -- taking from the bound vars array+ subPatsToVars :: Traversable lang => lang (Pattern lang) -> [Var] -> State Int (lang Var)+ subPatsToVars p' boundVars = traverse (const $ (boundVars !!) <$> (get >>= \i -> modify' (+1) >> return i)) p'++ -- | Return distinct variables in a pattern+ vars :: Foldable lang => Pattern lang -> [Var]+ vars (VariablePattern x) = [x]+ vars (NonVariablePattern p) = nubInt $ join $ map vars $ toList p+{-# SCC compileToQuery #-}
+ src/Data/Equality/Matching/Database.hs view
@@ -0,0 +1,323 @@+{-# LANGUAGE ViewPatterns #-}+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE OverloadedLists #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE TupleSections #-}+{-|+ Custom database implemented with trie-maps specialized to run conjunctive+ queries using a (worst-case optimal) generic join algorithm.++ Used in e-matching ('Data.Equality.Matching') as described by \"Relational+ E-Matching\" https://arxiv.org/abs/2108.02290.++ You probably don't need this module.+ -}+module Data.Equality.Matching.Database+ (+ genericJoin++ , Database(..)+ , Query(..)+ , IntTrie(..)+ , Subst+ , Var+ , Atom(..)+ , ClassIdOrVar(..)+ ) where++import Data.List (sortBy)+import Data.Function (on)+import Data.Maybe (mapMaybe)+import Control.Monad++import Data.Foldable as F (toList, foldl', length)+import qualified Data.Map.Strict as M+import qualified Data.IntMap.Strict as IM+import qualified Data.IntSet as IS++import Data.Equality.Graph.Classes.Id+import Data.Equality.Graph.Nodes+import Data.Equality.Language++-- | A variable in a query is identified by an 'Int'.+-- This is much more efficient than using e.g. a 'String'.+--+-- As a consequence, patterns also use 'Int' to represent a variable, but we+-- can still have an 'Data.String.IsString' instance for variable patterns by hashing the+-- string into a unique number.+type Var = Int++-- | Mapping from 'Var' to 'ClassId'. In a 'Subst' there is only one+-- substitution for each variable+type Subst = IM.IntMap ClassId++-- | A value which is either a 'ClassId' or a 'Var'+data ClassIdOrVar = CClassId {-# UNPACK #-} !ClassId+ | CVar {-# UNPACK #-} !Var+ deriving (Show, Eq, Ord)++-- | An 'Atom' 𝑅ᵢ(𝑣, 𝑣1, ..., 𝑣𝑘) is defined by the relation 𝑅ᵢ and by the+-- class-ids or variables 𝑣, 𝑣1, ..., 𝑣𝑘. It represents one conjunctive query's body atom.+data Atom lang+ = Atom+ !ClassIdOrVar -- ^ Represents 𝑣+ !(lang ClassIdOrVar) -- ^ Represents 𝑅ᵢ(𝑣1, ..., 𝑣𝑘). Note how 𝑣 isn't included since the arity of the constructor is 𝑘 instead of 𝑘+1.++-- | A conjunctive query to be run on the database+data Query lang+ = Query ![Var] ![Atom lang]+ | SelectAllQuery {-# UNPACK #-} !Var++-- | The relational representation of an e-graph, as described in section 3.1+-- of \"Relational E-Matching\".+--+-- Every e-node with symbol 𝑓 in the e-graph corresponds to a tuple in the relation 𝑅𝑓 in the database.+-- If 𝑓 has arity 𝑘, then 𝑅𝑓 will have arity 𝑘 + 1; its first attribute is the e-class id that contains the+-- corresponding e-node , and the remaining attributes are the 𝑘 children of the 𝑓 e-node+--+-- For every existing symbol in the e-graph the 'Database' has a table.+--+-- In concrete, we map 'Operator's to 'IntTrie's -- each operator has one table+-- represented by an 'IntTrie'+newtype Database lang+ = DB (M.Map (Operator lang) IntTrie)++-- | An integer triemap that keeps a cache of all keys in at each level.+--+-- As described in the paper:+-- Generic join requires two important performance bounds to be met in order for its own run time+-- to meet the AGM bound. First, the intersection [...] must run in 𝑂 (min(|𝑅𝑗 .𝑥 |)) time. Second,+-- the residual relations should be computed in constant time, i.e., computing from the relation 𝑅(𝑥, 𝑦)+-- the relation 𝑅(𝑣𝑥 , 𝑦) for some 𝑣𝑥 ∈ 𝑅(𝑥, 𝑦).𝑥 must take constant time. Both of these can be solved by+-- using tries (sometimes called prefix or suffix trees) as an indexing data structure.+data IntTrie = MkIntTrie+ { tkeys :: !IS.IntSet+ , trie :: !(IM.IntMap IntTrie)+ }+++-- TODO use this somehow?+-- queryHeadVars :: Foldable lang => Query lang -> [Var]+-- queryHeadVars (SelectAllQuery x) = [x]+-- queryHeadVars (Query qv _) = qv+-- {-# INLINE queryHeadVars #-}++-- | Run a conjunctive 'Query' on a 'Database'+--+-- Produce the list of valid substitutions from query variables to the+-- query-matching class ids.+genericJoin :: forall l. Language l => Database l -> Query l -> [Subst]+-- ROMES:TODO a less ad-hoc/specialized implementation of generic join...+-- ROMES:TODO query ordering is very important!++-- We want to match against ANYTHING, so we return a valid substitution for+-- all existing e-class: get all relations and make a substition for each class in that relation, then join all substitutions across all classes+genericJoin (DB m) (SelectAllQuery x) = concatMap (map (IM.singleton x) . IS.toList . tkeys) (M.elems m)++-- This is the last variable, so we return a valid substitution for every+-- possible value for the variable (hence, we prepend @x@ to each and make it+-- its own substitution)+-- ROMES:TODO: Start here. Map vars to indexs in an array and substitute in the resulting subst+genericJoin d q@(Query _ atoms) = genericJoin' atoms (orderedVarsInQuery q)++ where+ genericJoin' :: [Atom l] -> [Var] -> [Subst]+ genericJoin' !atoms' = \case++ [] -> map 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' #-}++ atomsWithX :: Var -> [Atom l] -> [Atom l]+ atomsWithX x = filter (x `elemOfAtom`)+ {-# INLINE atomsWithX #-}++ domainX :: Var -> [Atom l] -> IS.IntSet+ domainX x = intersectAtoms x d . atomsWithX x+ {-# INLINE domainX #-}++{-# INLINABLE genericJoin #-}+{-# SCC genericJoin #-}+++-- ROMES:TODO: Batching? How? https://arxiv.org/pdf/2108.02290.pdf++-- | Extract a list of unique variables from a 'Query', ordered by prioritizing+-- variables that occur in many relations, and secondly by prioritizing+-- variables that occur in small relations.+--+-- We use these heuristics because the variables' ordering is significant in+-- the query run-time performance.+--+-- This extraction could still be improved as some other strategies are+-- described in the paper (such as batching)+orderedVarsInQuery :: (Functor lang, Foldable lang) => Query lang -> [Var]+orderedVarsInQuery (SelectAllQuery x) = [x]+orderedVarsInQuery (Query _ atoms) = IS.toList . IS.fromAscList $ sortBy (compare `on` varCost) $ mapMaybe toVar $ foldl' f mempty atoms+ where++ f :: Foldable lang => [ClassIdOrVar] -> Atom lang -> [ClassIdOrVar]+ f s (Atom v (toList -> l)) = v:(l <> s)+ {-# INLINE f #-}++ -- First, prioritize variables that occur in many relations; second,+ -- prioritize variables that occur in small relations+ varCost :: Var -> Int+ varCost v = foldl' (\acc a -> if v `elemOfAtom` a then acc - 100 + atomLength a else acc) 0 atoms+ {-# INLINE varCost #-}++ -- | Get the size of an atom+ atomLength :: Foldable lang => Atom lang -> Int+ atomLength (Atom _ l) = 1 + F.length l+ {-# SCC atomLength #-}++ -- | Extract 'Var' from 'ClassIdOrVar'+ toVar :: ClassIdOrVar -> Maybe Var+ toVar (CVar v) = Just v+ 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+-- @D_x@, the domain of variable x, that is, the values x can take+--+-- Returns the class id set of classes forming the domain of var @x@+intersectAtoms :: forall l. Language l => Var -> Database l -> [Atom l] -> IS.IntSet+intersectAtoms !var (DB db) (a:atoms) = foldr (\x xs -> (f x) `IS.intersection` xs) (f a) atoms+ where+ -- Get the matching ids for an atom+ f :: Atom l -> IS.IntSet+ f (Atom v l) = case M.lookup (Operator $ void l) db of++ -- If needed relation doesn't exist altogether, return the matching+ -- class ids (none). When intersecting, nothing will be available -- as expected+ Nothing -> mempty++ -- If needed relation does exist, find intersection in it+ -- Add list of found intersections to existing+ Just r -> case intersectInTrie var mempty r (v:toList l) of+ Nothing -> error "intersectInTrie should return valid substitution for variable query"+ 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+--+-- Invalid substitutions are represented as Nothing+--+-- The intersection might be invalid while assuming values for variables. If+-- we're looking for the domain of some variables we should never get an+-- invalid substitution, but rather an empty list saying that the query+-- intersection is valid but empty.+--+--+-- If R_f(1,y,z), this function receives [1,y,z] :: [ClassIdOrVar] and+-- intersects the trie map of R_f with this prefix+--+-- TODO: write a note for this...+--+--+-- TODO: Really, a valid substitution is one which isn't empty...+intersectInTrie :: Var -- ^ The variable whose domain we are looking for+ -> IM.IntMap ClassId -- ^ A mapping from variables that have been substituted+ -> IntTrie -- ^ The trie+ -> [ClassIdOrVar] -- ^ The "query"+ -> Maybe IS.IntSet -- ^ The resulting domain for a valid substitution+intersectInTrie !var !substs (MkIntTrie trieKeys m) = \case++ [] -> pure []++ -- Looking for a class-id, so if it exists in map the intersection is+ -- valid and we simply continue the search for the domain+ CClassId x:xs ->+ IM.lookup x m >>= \next -> intersectInTrie var substs next xs++ -- Looking for a var. It might be one of the following:+ --+ -- (1) The variable whose domain we're looking for, and this is the+ -- first time we found it. In this case we'll assume all substitutions+ -- are valid, and try to get a valid substitution with that+ -- assumption. If the substitution is valid, the substitution is an+ -- element of the domain.+ --+ -- (2) The variable whose domain we're looking for, but we've already+ -- assumed a value for it in this branch, so we continue the recursion+ -- guaranteeing the assumption results in a valid substitution+ --+ -- (3) A bound variable, and this is the first time we find it. We+ -- assume its value for all branches and concatenate the result of all+ -- valid domain elements for each branch that resulted in a valid+ -- substitution+ --+ -- (4) A bound variable, but we've assumed a value for it, so we+ -- continue the recursion again to validate the assumption and+ -- possibly find the domain of the variable we're looking for ahead+ --+ CVar x:xs -> case IM.lookup x substs of+ -- (2) or (4), we simply continue+ Just varVal -> IM.lookup varVal m >>= \next -> intersectInTrie var substs next xs+ -- (1) or (3)+ Nothing -> pure $ if x == var+ -- (1)+ then+ -- If this is the var we're looking for, and the remaining @xs@+ -- suffix only consists of variables modulo the var we're looking+ -- for, we can simply return all possible keys for this since it is+ -- the correct variable. This is quite important for performance!+ if all (isVarDifferentFrom x) xs+ then trieKeys+ else IM.foldrWithKey (\k ls (!acc) ->+ case intersectInTrie var (IM.insert x k substs) ls xs of+ Nothing -> acc+ 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 IM.foldrWithKey (\k ls (!acc) ->+ case intersectInTrie var (IM.insert x k substs) ls xs of+ Nothing -> acc+ Just rs -> rs <> acc) mempty m+ where++ -- | Returns True if given 'ClassIdOrVar' holds a 'Var' and is different from given 'Var'.+ isVarDifferentFrom :: Var -> ClassIdOrVar -> Bool+ isVarDifferentFrom _ (CClassId _) = False+ isVarDifferentFrom x (CVar y) = x /= y+ {-# INLINE isVarDifferentFrom #-}++{-# INLINABLE intersectInTrie #-}+{-# SCC intersectInTrie #-}
+ src/Data/Equality/Matching/Pattern.hs view
@@ -0,0 +1,97 @@+{-|+ Definition of 'Pattern' for use in equality matching+ ('Data.Equality.Matching'), where patterns are matched against the e-graph+ -}+module Data.Equality.Matching.Pattern where++import Data.Functor.Classes+import Data.String++import Data.Equality.Utils+import Data.Equality.Matching.Database++-- | A pattern can be either a variable or an non-variable expression of+-- patterns.+--+-- A 'NonVariablePattern' will only match an expression if the @lang@ constructor matches an expression and all child patterns match the expression children.+-- A 'VariablePattern' matches any expression.+--+-- === Example+--+-- The expression+--+-- @+-- expr :: Fix Sym+-- expr = BinOp Add (Sym "x") (Const 2.0) -- i.e. x + 2+-- @+--+-- Would be matched against the following patterns+--+-- @+-- pat1 :: Pattern Sym+-- pat1 = VariablePattern 1+--+-- pat2 :: Pattern Sym+-- pat2 = NonVariablePattern (BinOp Add (VariablePattern 1) (VariablePattern 2))+--+-- pat3 :: Pattern Sym+-- pat3 = NonVariablePattern (BinOp Add (VariablePattern 1) (NonVariablePattern (Const 2)))+-- @+--+-- But would not be matched against the following patterns+-- +-- @+-- pat4 :: Pattern Sym+-- pat4 = NonVariablePattern (Const 5)+--+-- pat5 :: Pattern Sym+-- pat5 = NonVariablePattern (BinOp Add (NonVariablePattern (Sym "y")) (NonVariablePattern (Const 2)))+--+-- pat6 :: Pattern Sym+-- pat6 = NonVariablePattern (BinOp Add (NonVariablePattern (Sym "x")) (NonVariablePattern (Const 3)))+-- @+--+-- === IsString+-- 'Pattern' instances 'IsString', which means one can write a variable pattern simply as a string.+--+-- It works by using 'Data.Equality.Utils.hashString' to create a unique integer for a 'VariablePattern'+--+-- For example, we could write the following pattern that would match @a+a@ and @b+b@ but not @a+b@+--+-- @+-- pat7 :: Pattern Sym+-- pat7 = 'pat' (BinOp Add "x" "x")+-- @+data Pattern lang+ = NonVariablePattern (lang (Pattern lang))+ | VariablePattern Var -- ^ Should be a >0 positive number++-- | Synonym for 'NonVariablePattern'.+--+-- Example+--+-- @+-- pat8 :: Pattern Sym+-- pat8 = pat (BinOp Mul "y" (pat (Const 2))) -- matches any product of an expression by 2+-- @+pat :: lang (Pattern lang) -> Pattern lang+pat = NonVariablePattern++instance Eq1 l => (Eq (Pattern l)) where+ (==) (NonVariablePattern a) (NonVariablePattern b) = liftEq (==) a b+ (==) (VariablePattern a) (VariablePattern b) = a == b + (==) _ _ = False++instance Ord1 l => (Ord (Pattern l)) where+ compare (VariablePattern _) (NonVariablePattern _) = LT+ compare (NonVariablePattern _) (VariablePattern _) = GT+ compare (VariablePattern a) (VariablePattern b) = compare a b+ compare (NonVariablePattern a) (NonVariablePattern b) = liftCompare compare a b++instance Show1 lang => Show (Pattern lang) where+ showsPrec _ (VariablePattern s) = showString (show s) -- ROMES:TODO don't ignore prec?+ showsPrec d (NonVariablePattern x) = liftShowsPrec showsPrec showList d x++instance IsString (Pattern lang) where+ fromString = VariablePattern . hashString+
+ src/Data/Equality/Saturation.hs view
@@ -0,0 +1,187 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE TupleSections #-}+{-# LANGUAGE BlockArguments #-}+{-|+ Given an input program 𝑝, equality saturation constructs an e-graph 𝐸 that+ represents a large set of programs equivalent to 𝑝, and then extracts the+ “best” program from 𝐸.++ The e-graph is grown by repeatedly applying pattern-based rewrites.+ Critically, these rewrites only add information to the e-graph, eliminating+ the need for careful ordering.++ Upon reaching a fixed point (saturation), 𝐸 will represent all equivalent+ ways to express 𝑝 with respect to the given rewrites.++ After saturation (or timeout), a final extraction procedure analyzes 𝐸 and+ selects the optimal program according to a user-provided cost function.+ -}+module Data.Equality.Saturation+ (+ -- * Equality saturation+ equalitySaturation, equalitySaturation'++ -- * Re-exports for equality saturation++ -- ** Writing rewrite rules+ , Rewrite(..), RewriteCondition++ -- ** Writing cost functions+ --+ -- | 'CostFunction' re-exported from 'Data.Equality.Extraction' since they are required to do equality saturation+ , CostFunction --, Cost, depthCost++ -- ** Writing expressions+ -- + -- | Expressions must be written in their fixed-point form, since the+ -- 'Language' must be given in its base functor form+ , Fix(..), cata++ ) where++import qualified Data.IntMap.Strict as IM++import Data.Bifunctor+import Control.Monad++import Data.Proxy++import Data.Equality.Utils+import qualified Data.Equality.Graph as G+import Data.Equality.Graph.Monad+import Data.Equality.Language+import Data.Equality.Graph.Classes+import Data.Equality.Matching+import Data.Equality.Matching.Database+import Data.Equality.Extraction++import Data.Equality.Saturation.Rewrites+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 = equalitySaturation' (Proxy @BackoffScheduler)+++-- | Run equality saturation on an expression given a list of rewrites, and+-- 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++ -- 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++ -- Extract best solution from the e-class of the original expression+ gets $ \g -> extractBest g cost origClass++ 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++ egr@G.EGraph{ G.memo = beforeMemo, G.classes = beforeClasses } <- get++ let db = eGraphToDatabase egr++ -- 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))++ -- Write-only phase, temporarily break invariants+ forM_ matches applyMatchesRhs++ -- Restore the invariants once per iteration+ rebuild+ + G.EGraph { G.memo = afterMemo, G.classes = afterClasses } <- get++ -- 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)++ 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)++ -- Otherwise, match and update stats+ x -> do++ -- Match pattern+ let matches' = ematch db lhs -- Add rewrite to the e-match substitutions++ -- Backoff scheduler: update stats+ let newStats = updateStats @schd i rw_id x stats matches'++ (map (lhs := rhs,) matches', newStats)++ 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)++ (_ := 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 ()++ (_ := 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 . 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' #-}+
+ src/Data/Equality/Saturation/Rewrites.hs view
@@ -0,0 +1,52 @@+{-|++Definition of 'Rewrite' and 'RewriteCondition' used to define rewrite rules.++Rewrite rules are applied to all represented expressions in an e-graph every+iteration of equality saturation.++-}+module Data.Equality.Saturation.Rewrites where++import Data.Equality.Graph+import Data.Equality.Matching+import Data.Equality.Matching.Database++-- | A rewrite rule that might have conditions for being applied+--+-- === __Example__+-- @+-- rewrites :: [Rewrite Expr] -- from Sym.hs+-- rewrites =+-- [ "x"+"y" := "y"+"x"+-- , "x"*("y"*"z") := ("x"*"y")*"z"+--+-- , "x"*0 := 0+-- , "x"*1 := "x"+--+-- , "a"-"a" := 1 -- cancel sub+-- , "a"/"a" := 1 :| is_not_zero "a"+-- ]+-- @+--+-- See the definition of @is_not_zero@ in the documentation for+-- 'RewriteCondition'+data Rewrite lang = !(Pattern lang) := !(Pattern lang) -- ^ Trivial Rewrite+ | !(Rewrite lang) :| !(RewriteCondition lang) -- ^ Conditional Rewrite+infix 3 :=+infixl 2 :|++-- | A rewrite condition. With a substitution from bound variables in the+-- pattern to e-classes and with the e-graph, return 'True' if the condition is+-- satisfied+--+-- === Example+-- @+-- is_not_zero :: String -> RewriteCondition Expr+-- is_not_zero v subst egr =+-- case lookup v subst of+-- Just class_id ->+-- egr^._class class_id._data /= Just 0+-- @+type RewriteCondition lang = Subst -> EGraph lang -> Bool+
+ src/Data/Equality/Saturation/Scheduler.hs view
@@ -0,0 +1,89 @@+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE ViewPatterns #-}+{-# LANGUAGE AllowAmbiguousTypes #-} -- Scheduler+{-# LANGUAGE TypeFamilies #-}+{-|++Definition of 'Scheduler' as a way to control application of rewrite rules.++The 'BackoffScheduler' is a scheduler which implements exponential rule backoff+and is used by default in 'Data.Equality.Saturation.equalitySaturation'++-}+module Data.Equality.Saturation.Scheduler+ ( Scheduler(..), BackoffScheduler+ ) where++import qualified Data.IntMap.Strict as IM+import Data.Equality.Matching++-- | A 'Scheduler' determines whether a certain rewrite rule is banned from+-- being used based on statistics it defines and collects on applied rewrite+-- rules.+class Scheduler s where+ type Stat s++ -- | Scheduler: update stats+ updateStats :: Int -- ^ Iteration we're in+ -> Int -- ^ Index of rewrite rule we're updating+ -> Maybe (Stat s) -- ^ Current stat for this rewrite rule (we already got it so no point in doing a lookup again)+ -> IM.IntMap (Stat s) -- ^ The current stats map+ -> [Match] -- ^ The list of matches resulting from matching this rewrite rule+ -> IM.IntMap (Stat s) -- ^ The updated map with new stats++ -- Decide whether to apply a matched rule based on its stats and current iteration+ isBanned :: Int -- ^ Iteration we're in+ -> Stat s -- ^ Stats for the rewrite rule+ -> Bool -- ^ Whether the rule should be applied or not++-- | A 'Scheduler' that implements exponentional rule backoff.+--+-- For each rewrite, there exists a configurable initial match limit. If a rewrite+-- search yield more than this limit, then we ban this rule for number of+-- iterations, double its limit, and double the time it will be banned next time.+--+-- This seems effective at preventing explosive rules like associativity from+-- taking an unfair amount of resources.+--+-- Originaly in [egg](https://docs.rs/egg/0.6.0/egg/struct.BackoffScheduler.html)+data BackoffScheduler+instance Scheduler BackoffScheduler where+ type Stat BackoffScheduler = BoSchStat++ updateStats i rw currentStat stats matches =++ if total_len > threshold++ then+ IM.alter updateBans rw stats++ else+ stats++ where++ -- TODO: Overall difficult, and buggy at the moment.+ total_len = sum (map (length . matchSubst) matches)++ defaultMatchLimit = 1000+ defaultBanLength = 10++ bannedN = case currentStat of+ Nothing -> 0;+ Just (timesBanned -> n) -> n++ threshold = defaultMatchLimit * (2^bannedN)++ ban_length = defaultBanLength * (2^bannedN)++ 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+++data BoSchStat = BSS { bannedUntil :: {-# UNPACK #-} !Int+ , timesBanned :: {-# UNPACK #-} !Int+ } deriving Show
+ src/Data/Equality/Utils.hs view
@@ -0,0 +1,58 @@+{-# LANGUAGE StandaloneDeriving #-}+{-|+ Misc utilities used accross modules+ -}+module Data.Equality.Utils where++-- import GHC.Conc+import Data.Foldable+import Data.Bits++-- import qualified Data.Set as S+-- import qualified Data.IntSet as IS+import Data.Functor.Classes++-- | Fixed point newtype.+--+-- Ideally we should use the data-fix package, but right now we're rolling our+-- own due to an initial idea to avoid dependencies to be easier to upstream+-- into GHC (for improvements to the pattern match checker involving equality+-- graphs). I no longer think we can do that without vendoring in some part of+-- just e-graphs, but until I revert the decision we use this type.+newtype Fix f = Fix { unFix :: f (Fix f) }++instance Eq1 f => Eq (Fix f) where+ (==) (Fix a) (Fix b) = liftEq (==) a b+ {-# INLINE (==) #-}++instance Show1 f => Show (Fix f) where+ showsPrec d (Fix f) = liftShowsPrec showsPrec showList d f+ {-# INLINE showsPrec #-}++-- | Catamorphism+cata :: Functor f => (f a -> a) -> (Fix f -> a)+cata f = f . fmap (cata f) . unFix+{-# INLINE cata #-}++-- | Get the hash of a string.+--+-- This util is currently used to generate an 'Int' used for the internal+-- pattern variable representation from the external pattern variable+-- representation ('String')+hashString :: String -> Int+hashString = foldl' (\h c -> 33*h `xor` fromEnum c) 5381+{-# INLINE hashString #-}++-- -- | We don't have the parallel package, so roll our own simple parMap+-- parMap :: (a -> b) -> [a] -> [b]+-- parMap _ [] = []+-- parMap f (x:xs) = fx `par` (fxs `pseq` (fx : fxs))+-- where fx = f x; fxs = parMap f xs++-- toSet :: (Ord a, Foldable f) => f a -> S.Set a+-- toSet = foldl' (flip S.insert) mempty+-- {-# INLINE toSet #-}++-- toIntSet :: (Foldable f) => f Int -> IS.IntSet+-- toIntSet = foldl' (flip IS.insert) mempty+-- {-# INLINE toIntSet #-}
+ src/Data/Equality/Utils/IntToIntMap.hs view
@@ -0,0 +1,132 @@+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE ExplicitForAll #-}+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE PolyKinds #-}+{-# LANGUAGE UnliftedDatatypes #-}+{-# LANGUAGE ViewPatterns #-}+{-# LANGUAGE StandaloneKindSignatures #-}+{-|+ This module defines 'IntToIntMap', a variant of 'Data.IntMap' in which the+ values are fixed to 'Int'.++ We make use of this structure in 'Data.Equality.Graph.ReprUnionFind' to+ improve performance by a constant factor+ -}+module Data.Equality.Utils.IntToIntMap+ ( IntToIntMap(Nil)+ , Key, Val+ , find, insert, (!)+ , unliftedFoldr+ ) where++import GHC.Exts+import Data.Bits++-- | A map of integers to integers+type IntToIntMap :: TYPE ('BoxedRep 'Unlifted)+data IntToIntMap = Bin Prefix Mask IntToIntMap IntToIntMap+ | Tip InternalKey Val+ | Nil -- ^ An empty 'IntToIntMap'. Ideally this would be defined as a function instead of an exported constructor, but it's currently not possible to have top-level bindings for unlifted datatypes++type Prefix = Word#+type Mask = Word#+type InternalKey = Word#++-- | Key type synonym in an 'IntToIntMap'+type Key = Int#+-- | Value type synonym in an 'IntToIntMap'+type Val = Int#++-- | \(O(\min(n,W))\). Find the value at a key.+-- Calls 'error' when the element can not be found.+(!) :: IntToIntMap -> Key -> Val+(!) m k = find k m+{-# INLINE (!) #-}++-- | Find the 'Val' for a 'Key' in an 'IntToIntMap'+find :: Key -> IntToIntMap -> Val+find (int2Word# -> k) = find' k+{-# INLINE find #-}++-- | Insert a 'Val' at a 'Key' in an 'IntToIntMap'+insert :: Key -> Val -> IntToIntMap -> IntToIntMap+insert k = insert' (int2Word# k)+{-# INLINE insert #-}++insert' :: InternalKey -> Val -> IntToIntMap -> IntToIntMap+insert' k x t@(Bin p m l r)+ | nomatch k p m = link k (Tip k x) p t+ | zero k m = Bin p m (insert' k x l) r+ | otherwise = Bin p m l (insert' k x r)+insert' k x t@(Tip ky _)+ | isTrue# (k `eqWord#` ky) = Tip ky x+ | otherwise = link k (Tip k x) ky t+insert' k x Nil = Tip k x++-- DANGEROUS NOTE:+-- Since this is the function that currently takes 10% of runtime, we want to+-- improve constant factors: we'll remove the comparison that checks that the+-- tip we found is the tip we are looking for. This is a very custom map,+-- we will assume the tip we find is ALWAYS the one we are looking for. This,+-- of course, will return wrong results instead of blow up if we use it+-- unexpectedly. Hopefully the testsuite will serve to warn us of this+--+-- Update: The speedup is not noticeable, so we don't do it, but I'll leave the comment here for now+find' :: InternalKey -> IntToIntMap -> Val+find' k (Bin _p m l r)+ | zero k m = find' k l+ | 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++link :: Prefix -> IntToIntMap -> Prefix -> IntToIntMap -> IntToIntMap+link p1 t1 p2 t2 = linkWithMask (highestBitMask (p1 `xor#` p2)) p1 t1 {-p2-} t2+{-# INLINE link #-}++-- `linkWithMask` is useful when the `branchMask` has already been computed+linkWithMask :: Mask -> Prefix -> IntToIntMap -> IntToIntMap -> IntToIntMap+linkWithMask m p1 t1 t2+ | zero p1 m = Bin p m t1 t2+ | otherwise = Bin p m t2 t1+ where+ p = maskW p1 m+{-# INLINE linkWithMask #-}+++-- The highestBitMask implementation is based on+-- http://graphics.stanford.edu/~seander/bithacks.html#RoundUpPowerOf2+-- which has been put in the public domain.++-- | Return a word where only the highest bit is set.+highestBitMask :: Word# -> Word#+highestBitMask w =+ case finiteBitSize (0 :: Word) of+ I# wordSize -> shiftL# (int2Word# 1#) (wordSize -# 1# -# (word2Int# (clz# w)))+{-# INLINE highestBitMask #-}++nomatch :: InternalKey -> Prefix -> Mask -> Bool+nomatch i p m+ = isTrue# ((maskW i m) `neWord#` p)+{-# INLINE nomatch #-}++-- | The prefix of key @i@ up to (but not including) the switching+-- bit @m@.+maskW :: Word# -> Word# -> Prefix+maskW i m+ = (i `and#` ((int2Word# (negateInt# (word2Int# m))) `xor#` m))+{-# INLINE maskW #-}++zero :: InternalKey -> Mask -> Bool+zero i m+ = isTrue# ((i `and#` m) `eqWord#` (int2Word# 0#))+{-# INLINE zero #-}++-- | A 'foldr' in which the accumulator is unlifted+unliftedFoldr :: forall a {b :: TYPE ('BoxedRep 'Unlifted)} . (a -> b -> b) -> b -> [a] -> b +unliftedFoldr k z = go+ where+ go [] = z+ go (y:ys) = y `k` go ys
+ test/Invariants.hs view
@@ -0,0 +1,218 @@+{-# OPTIONS_GHC -Wno-orphans #-} -- Arbitrary+{-# LANGUAGE RoleAnnotations #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE GeneralizedNewtypeDeriving #-}+{-# LANGUAGE DeriveTraversable #-}+{-# LANGUAGE DerivingStrategies #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE RecordWildCards #-}+module Invariants where++import Test.Tasty+import Test.Tasty.QuickCheck as QC hiding (classes)++import Data.Functor.Classes+import Control.Monad++import qualified Data.Containers.ListUtils as LU+import qualified Data.Foldable as F+import qualified Data.List as L+import qualified Data.Set as S+import qualified Data.IntMap.Strict as IM++import Data.Equality.Graph.Monad as GM+import Data.Equality.Graph+import Data.Equality.Analysis+import Data.Equality.Extraction+import Data.Equality.Saturation+import Data.Equality.Matching+import Data.Equality.Matching.Database+import Sym++-- | Newtype deriving via Expr to be able to define a different analysis+-- TODO: Use type level symbol to define the analysis+type role SimpleExpr nominal+newtype SimpleExpr l = SE (Expr l)+ deriving (Functor, Foldable, Traversable, Show1, Eq1, Ord1, Language)++instance Analysis SimpleExpr where+ type Domain SimpleExpr = ()+ makeA _ _ = ()+ joinA = (<>)+ modifyA _ = id++-- | When a rewrite of type "x":=c where x is a pattern variable and c is a+-- constant is used in equality saturation of any expression, all e-classes+-- should be merged into a single one, since all classes are equal to c and+-- therefore equivalent to themselves+patFoldAllClasses :: forall l. (Language l, Num (Pattern l))+ => Fix l -> Integer -> Bool+patFoldAllClasses expr i =+ case IM.toList $ classes eg of+ [_] -> True+ _ -> False+ where+ eg :: EGraph l+ eg = snd $ equalitySaturation expr [VariablePattern 1:=fromInteger i] (error "Cost function shouldn't be used")++-- | Test 'compileToQuery'.+--+-- Every pattern compiled to a query should have the same number of free variables (except for the root variable)+-- as the pattern+--+-- The number of atoms should also match the number of non variable patterns+-- since we should create an additional atom (with a new bound variable) for each. +testCompileToQuery :: Traversable lang => Pattern lang -> Bool+testCompileToQuery p = case fst $ compileToQuery p of+ -- Handle special case for selectAll queries...+ SelectAllQuery x -> [x] == vars p && numNonVarPatterns p == 0+ q@(Query _ atoms)+ | [] <- queryHeadVars q -> False+ | _: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+ numNonVarPatterns (NonVariablePattern l) = F.foldl' (flip $ (+) . numNonVarPatterns) 1 l++ queryHeadVars :: Foldable lang => Query lang -> [Var]+ queryHeadVars (SelectAllQuery x) = [x]+ queryHeadVars (Query qv _) = qv++ -- | Return distinct variables in a pattern+ vars :: Foldable lang => Pattern lang -> [Var]+ vars (VariablePattern x) = [x]+ vars (NonVariablePattern p') = LU.nubInt $ join $ map vars $ F.toList p'++-- | If we match a singleton variable pattern against an e-graph, we should get+-- a match on all e-classes in the e-graph+ematchSingletonVar :: Language lang => Var -> EGraph lang -> Bool+ematchSingletonVar v eg =+ let+ db = eGraphToDatabase eg+ matches = S.fromList $ map matchClassId $ ematch db (VariablePattern v)+ eclasses = S.fromList $ map fst $ IM.toList $ classes eg+ in+ matches == eclasses +++-- | Property test for 'genericJoin'.+--+-- If we search a database with an expression in which all patterns are+-- variables (the only non-variable pattern is the top one), then, altogether,+-- we should get a list of all e-classes +-- genericJoinAll :: Database lang -> +++-- The equivalence relation over e-nodes must be closed over congruence after rebuilding+-- congruenceInvariant :: Testable m (EGraph lang) => Property m+++-- The hashcons 𝐻 must map all canonical e-nodes to their e-class ids+--+-- Note: the e-graph argument must have been rebuilt -- checking the property+-- when invariants are broken for sure doesn't make much sense+--+-- 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)+ where+ -- e-node 𝑛 ∈ 𝑀 [𝑎] ⇐⇒ 𝐻 [canonicalize(𝑛)] = find(𝑎)+ f (i, EClass _ nodes _ _) = all g nodes+ where+ g en = case lookupNM (canonicalize en 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+benchSaturate rws cost expr =+ equalitySaturation expr rws cost `seq` True+++-- ROMES:TODO: Property: Extract expression after equality saturation is always better or equal to the original expression++-- ROMES:TODO: Use action trick https://jaspervdj.be/posts/2015-03-13-practical-testing-in-haskell.html+instance Arbitrary (EGraph SimpleExpr) where+ arbitrary = sized $ \n -> do+ exps <- forM [0..n] $ const arbitrary+ -- rws :: [Rewrite Expr] <- forM [0..n] $ const arbitrary+ (ids, eg) <- return $ egraph $+ mapM represent exps+ ids1 <- sublistOf ids+ ids2 <- sublistOf ids+ return $ snd $ runEGraphM eg $ do+ forM_ (zip ids1 ids2) $ \(a,b) -> do+ GM.merge a b+ GM.rebuild++instance Arbitrary BOp where+ arbitrary = oneof [ return Add+ , return Sub+ , return Mul+ , return Div ]++instance Arbitrary UOp where+ arbitrary = oneof [ return Sin+ , return Cos+ ]++instance Arbitrary a => Arbitrary (SimpleExpr a) where+ arbitrary = SE <$> arbitrary++instance Arbitrary a => Arbitrary (Expr a) where+ arbitrary = sized expr'+ where+ expr' :: Int -> Gen (Expr a)+ expr' 0 = oneof [ Sym . un <$> arbitrary+ , Const . fromInteger <$> arbitrary+ ]+ expr' n+ | n > 0 = oneof [ BinOp <$> arbitrary <*> resize (n `div` 2) arbitrary <*> resize (n `div` 2) arbitrary+ , UnOp <$> arbitrary <*> resize (n - 1) arbitrary ]+ expr' _ = error "size is negative?"++instance Arbitrary (Fix SimpleExpr) where+ arbitrary = Fix <$> arbitrary++instance Arbitrary (Fix Expr) where+ arbitrary = Fix <$> arbitrary++instance Arbitrary (Pattern SimpleExpr) where+ arbitrary = sized p'+ where+ p' 0 = VariablePattern <$> oneof (return <$> [1..16])+ p' n = NonVariablePattern <$> resize (n `div` 2) arbitrary++newtype Name = Name { un :: String }++instance Arbitrary Name where+ arbitrary = oneof (return . Name . (:[]) <$> ['a'..'l'])++instance Num (Pattern SimpleExpr) where+ fromInteger = NonVariablePattern . SE . Const . fromInteger+ (+) = error "Should use @Expr or have other way to switch analysis"+ (*) = error "Should use @Expr or have other way to switch analysis"+ (-) = error "Should use @Expr or have other way to switch analysis"+ abs = error "Should use @Expr or have other way to switch analysis"+ signum = error "Should use @Expr or have other way to switch analysis"++invariants :: TestTree+invariants = testGroup "Invariants"+ [ QC.testProperty "Compile to query" (testCompileToQuery @SimpleExpr)+ -- TODO: This bench is still failing because of the bad rewrite scheduler+ -- TODO: Much infinite looping ...+ -- , QC.testProperty "Bench saturation @Expr" (withMaxSuccess 10 (benchSaturate @Expr rewrites symCost))+ , QC.testProperty "Singleton variable matches all" (ematchSingletonVar @SimpleExpr)+ , QC.testProperty "Hash Cons Invariant" (hashConsInvariant @SimpleExpr)+ , QC.testProperty "Fold all classes with x:=c" (patFoldAllClasses @SimpleExpr)+ ]+
+ test/Lambda.hs view
@@ -0,0 +1,146 @@+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE DeriveGeneric #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE ViewPatterns #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE TemplateHaskell #-}+{-# LANGUAGE DeriveTraversable #-}+module Lambda where++import Test.Tasty+import Test.Tasty.HUnit++import qualified Data.Set as S++import Control.Applicative ((<|>))++import Data.Eq.Deriving+import Data.Ord.Deriving+import Text.Show.Deriving++import Data.Equality.Graph.Lens+import Data.Equality.Graph.Monad as GM+import Data.Equality.Graph+import Data.Equality.Extraction+import Data.Equality.Analysis+import Data.Equality.Saturation+import Data.Equality.Matching++data Lambda a+ = Bool Bool+ | Num Int+ | Var a+ | Add a a+ | Eq a a+ | App a a+ | Lam a a+ | Let a a a+ | LFix a a+ | If a a a+ | Symbol String+ deriving ( Eq, Ord, Functor+ , Foldable, Traversable+ )++deriveEq1 ''Lambda+deriveOrd1 ''Lambda+deriveShow1 ''Lambda++data Data = Data { free :: S.Set ClassId+ , constant :: Maybe (Fix Lambda)+ } deriving Eq++evalL :: EGraph Lambda -> Lambda ClassId -> Maybe (Fix Lambda)+evalL egr = \case+ Bool n -> Just (Fix $ Bool n)+ Num n -> Just (Fix $ Num n)+ Add a b -> do+ a' <- constant (egr^._class a._data) >>= num+ b' <- constant (egr^._class b._data) >>= num+ return (Fix $ Num $ a' + b')+ Eq a b -> do+ a' <- constant (egr^._class a._data)+ b' <- constant (egr^._class b._data)+ return (Fix $ Bool $ a' == b')+ _ -> Nothing+ where+ num :: Fix Lambda -> Maybe Int+ num = \case+ Fix (Num i) -> Just i+ _ -> Nothing++instance Analysis Lambda where+ type Domain Lambda = Data++ makeA n egr =+ let+ freeVs = case unNode n of+ Var x -> S.singleton x+ Let v a b ->+ free (egr^._class a._data) <> S.delete v (free (egr^._class b._data))+ Lam v a -> S.delete v (free (egr^._class a._data))+ LFix v a -> S.delete v (free (egr^._class a._data))+ _ -> mconcat (map (\i -> free $ egr^._class i._data) (children n))++ cnst = evalL egr (unNode n)+ in+ Data freeVs cnst++ joinA (Data fv1 c1) (Data fv2 c2) =+ Data (fv1 `S.intersection` fv2) (c1 <|> c2)++ -- modifyA :: ClassId -> EGraph l -> EGraph l+ modifyA i egr = + case constant (egr^._class i._data) of+ Nothing -> egr+ Just c -> snd $ runEGraphM egr $ do+ new_c <- represent c+ GM.merge i new_c++instance Language Lambda++instance Num (Fix Lambda) where+ fromInteger = Fix . Num . fromInteger+ (+) = error "todo..."+ (-) = error "todo..."+ (*) = error "todo..."+ abs = error "todo..."+ signum = error "todo..."++rules :: [Rewrite Lambda]+rules =+ [ ifP trP "x" "y" := "x"+ , ifP flP "x" "y" := "y"+ -- , ifP (pat $ eq (varP "x") "e" "then" "else") := "else" :| if ...+ ]++rewrite :: Fix Lambda -> Fix Lambda+rewrite e = fst $ equalitySaturation e rules depthCost++lambdaTests :: TestTree+lambdaTests = testGroup "Lambda"+ [ testCase "if tr" $+ rewrite (ifL tr 1 2) @?= 1++ , testCase "if fl" $+ rewrite (ifL fl 1 2) @?= 2+ ]+++++ifP :: Pattern Lambda -> Pattern Lambda -> Pattern Lambda -> Pattern Lambda+ifP a b c = pat (If a b c)+trP, flP :: Pattern Lambda+trP = pat (Bool True)+flP = pat (Bool False)+varP :: Pattern Lambda -> Pattern Lambda+varP x = pat (Var x)++-- TODO: recursion-schemes extension in separate package+ifL :: Fix Lambda -> Fix Lambda -> Fix Lambda -> Fix Lambda+ifL a b c = Fix (If a b c)+tr, fl :: Fix Lambda+tr = Fix $ Bool True+fl = Fix $ Bool False
+ test/SimpleSym.hs view
@@ -0,0 +1,65 @@+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TemplateHaskell #-}+{-# LANGUAGE DeriveTraversable #-}+module SimpleSym where++import Test.Tasty+import Test.Tasty.HUnit++import Data.Eq.Deriving+import Data.Ord.Deriving+import Text.Show.Deriving++import Data.Equality.Utils+import Data.Equality.Matching+import Data.Equality.Saturation+import Data.Equality.Language+import Data.Equality.Analysis++data SymExpr a = Const Double+ | Symbol String+ | a :+: a+ | a :*: a+ | a :/: a+ deriving (Functor, Foldable, Traversable)+infix 6 :+:+infix 7 :*:, :/:++deriveEq1 ''SymExpr+deriveOrd1 ''SymExpr+deriveShow1 ''SymExpr++instance Analysis SymExpr where+ type Domain SymExpr = ()+ makeA _ _ = ()+ joinA _ _ = ()++instance Language SymExpr++cost :: CostFunction SymExpr+cost = \case+ Const _ -> 1+ Symbol _ -> 1+ c1 :+: c2 -> c1 + c2 + 2+ c1 :*: c2 -> c1 + c2 + 3+ c1 :/: c2 -> c1 + c2 + 4++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"+ ]++rewrite :: Fix SymExpr -> Fix SymExpr+rewrite e = fst (equalitySaturation e rewrites cost)++e1 :: Fix SymExpr+e1 = Fix (Fix (Fix (Symbol "x") :*: Fix (Const 2)) :/: (Fix (Const 2))) -- (x*2)/2++simpleSymTests :: TestTree+simpleSymTests = testGroup "Simple Sym"+ [ testCase "(a*2)/2 = a" $ rewrite e1 @?= Fix (Symbol "x")+ ]
+ test/Sym.hs view
@@ -0,0 +1,371 @@+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE TemplateHaskell #-}+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE DeriveTraversable #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE LambdaCase #-}+module Sym where++import Test.Tasty+import Test.Tasty.HUnit++import qualified Data.IntMap.Strict as IM+import qualified Data.Set as S+import Data.String+import Data.Maybe (isJust)++import Data.Eq.Deriving+import Data.Ord.Deriving+import Text.Show.Deriving++import Control.Applicative (liftA2)+import Control.Monad (unless)++import Data.Equality.Graph.Monad as GM+import Data.Equality.Graph.Lens+import Data.Equality.Graph+import Data.Equality.Extraction+import Data.Equality.Analysis+import Data.Equality.Matching+import Data.Equality.Matching.Database+import Data.Equality.Saturation++data Expr a = Sym !String+ | Const !Double+ | UnOp !UOp !a+ | BinOp !BOp !a !a+ deriving ( Eq, Ord, Functor+ , Foldable, Traversable+ )+data BOp = Add+ | Sub+ | Mul+ | Div+ | Pow+ | Diff+ | Integral+ deriving (Eq, Ord, Show)++data UOp = Sin+ | Cos+ | Sqrt+ | Ln+ deriving (Eq, Ord, Show)++deriveEq1 ''Expr+deriveOrd1 ''Expr+deriveShow1 ''Expr++instance Language Expr++instance IsString (Fix Expr) where+ fromString = Fix . Sym++instance Num (Fix Expr) where+ (+) a b = Fix (BinOp Add a b)+ (-) a b = Fix (BinOp Sub a b)+ (*) a b = Fix (BinOp Mul a b)+ fromInteger = Fix . Const . fromInteger+ negate = error "DONT USE"+ abs = error "abs"+ signum = error "signum"++instance Fractional (Fix Expr) where+ (/) a b = Fix (BinOp Div a b)+ fromRational = Fix . Const . fromRational++symCost :: Expr Cost -> Cost+symCost = \case+ BinOp Pow e1 e2 -> e1 + e2 + 6+ BinOp Div e1 e2 -> e1 + e2 + 5+ BinOp Sub e1 e2 -> e1 + e2 + 4+ BinOp Mul e1 e2 -> e1 + e2 + 4+ BinOp Add e1 e2 -> e1 + e2 + 2+ BinOp Diff e1 e2 -> e1 + e2 + 500+ BinOp Integral e1 e2 -> e1 + e2 + 20000+ UnOp Sin e1 -> e1 + 20+ UnOp Cos e1 -> e1 + 20+ UnOp Sqrt e1 -> e1 + 30+ UnOp Ln e1 -> e1 + 30+ Sym _ -> 1+ Const _ -> 1++instance Num (Pattern Expr) where+ (+) a b = NonVariablePattern $ BinOp Add a b+ (-) a b = NonVariablePattern $ BinOp Sub a b+ (*) a b = NonVariablePattern $ BinOp Mul a b+ fromInteger = NonVariablePattern . Const . fromInteger+ negate = error "DONT USE" -- NonVariablePattern. BinOp Mul (fromInteger $ -1)+ abs = error "abs"+ signum = error "signum"++instance Fractional (Pattern Expr) where+ (/) a b = NonVariablePattern $ BinOp Div a b+ fromRational = NonVariablePattern . Const . fromRational++-- | Define analysis for the @Expr@ language over domain @Maybe Double@ for+-- constant folding+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+ -- this assertion only seemed to be triggering when using bogus+ -- constant assignments for "Fold all classes with x:=c"+ -- 0 bug found by property checking+ !_ <- 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+ Just d -> snd $ runEGraphM egr $ do++ -- 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 (null . children))++++evalConstant :: Expr (Maybe Double) -> Maybe Double+evalConstant = \case+ -- Exception: Negative exponent: BinOp Pow e1 e2 -> liftA2 (^) e1 (round <$> e2 :: Maybe Integer)+ BinOp Div e1 e2 -> liftA2 (/) e1 e2+ BinOp Sub e1 e2 -> liftA2 (-) e1 e2+ BinOp Mul e1 e2 -> liftA2 (*) e1 e2+ BinOp Add e1 e2 -> liftA2 (+) e1 e2+ BinOp Pow _ _ -> Nothing+ BinOp Diff _ _ -> Nothing+ BinOp Integral _ _ -> Nothing+ UnOp Sin e1 -> sin <$> e1+ UnOp Cos e1 -> cos <$> e1+ UnOp Sqrt e1 -> sqrt <$> e1+ UnOp Ln _ -> Nothing+ Sym _ -> Nothing+ Const x -> Just x+ +unsafeGetSubst :: Pattern Expr -> Subst -> ClassId+unsafeGetSubst (NonVariablePattern _) _ = error "unsafeGetSubst: NonVariablePattern; expecting VariablePattern"+unsafeGetSubst (VariablePattern v) subst = case IM.lookup v subst of+ Nothing -> error "Searching for non existent bound var in conditional"+ Just class_id -> class_id++is_not_zero :: Pattern Expr -> RewriteCondition Expr+is_not_zero v subst egr =+ egr^._class (unsafeGetSubst v subst)._data /= Just 0++is_sym :: Pattern Expr -> RewriteCondition Expr+is_sym v subst egr =+ any ((\case (Sym _) -> True; _ -> False) . unNode) (egr^._class (unsafeGetSubst v subst)._nodes)++is_const :: Pattern Expr -> RewriteCondition Expr+is_const v subst egr =+ isJust (egr^._class (unsafeGetSubst v subst)._data)++is_const_or_distinct_var :: Pattern Expr -> Pattern Expr -> RewriteCondition Expr+is_const_or_distinct_var v w subst egr =+ let v' = unsafeGetSubst v subst+ w' = unsafeGetSubst w subst+ in (eClassId (egr^._class v') /= eClassId (egr^._class w'))+ && (isJust (egr^._class v'._data)+ || any ((\case (Sym _) -> True; _ -> False) . unNode) (egr^._class v'._nodes))++rewrites :: [Rewrite Expr]+rewrites =+ [ "a"+"b" := "b"+"a" -- comm add+ , "a"*"b" := "b"*"a" -- comm mul+ , "a"+("b"+"c") := ("a"+"b")+"c" -- assoc add+ , "a"*("b"*"c") := ("a"*"b")*"c" -- assoc mul++ , "a"-"b" := "a"+(fromInteger (-1) * "b") -- sub cannon+ , "a"/"b" := "a"*powP "b" (fromInteger $ -1) :| is_not_zero "b" -- div cannon++ -- identities+ , "a"+0 := "a"+ , "a"*0 := 0+ , "a"*1 := "a"++ -- TODO This causes many problems+ -- , "a" := "a"+0++ -- This already works+ , "a" := "a"*1++ , "a"-"a" := 0 -- cancel sub+ , "a"/"a" := 1 :| is_not_zero "a" -- cancel div++ , "a"*("b"+"c") := ("a"*"b")+("a"*"c") -- distribute+ , ("a"*"b")+("a"*"c") := "a"*("b"+"c") -- factor++ , powP "a" "b"*powP "a" "c" := powP "a" ("b" + "c") -- pow mul+ , powP "a" 0 := 1 :| is_not_zero "a"+ , powP "a" 1 := "a"+ , powP "a" 2 := "a"*"a"+ , powP "a" (fromInteger $ -1) := 1/"a" :| is_not_zero "a"++ , "x"*(1/"x") := 1 :| is_not_zero "x"++ , diffP "x" "x" := 1 :| is_sym "x"+ , diffP "x" "c" := 0 :| is_sym "x" :| is_const_or_distinct_var "c" "x"++ , diffP "x" ("a" + "b") := diffP "x" "a" + diffP "x" "b"+ , diffP "x" ("a" * "b") := ("a"*diffP "x" "b") + ("b"*diffP "x" "a")++ , diffP "x" (sinP "x") := cosP "x"+ , diffP "x" (cosP "x") := fromInteger (-1) * sinP "x"++ , diffP "x" (lnP "x") := 1/"x" :| is_not_zero "x"++ -- diff-power+ , diffP "x" (powP "f" "g") := powP "f" "g" * ((diffP "x" "f" * ("g" / "f")) ++ (diffP "x" "g" * lnP "f")) :| is_not_zero "f" :| is_not_zero "g"++ -- i-one+ , intP 1 "x" := "x"++ -- i power const+ , intP (powP "x" "c") "x" := (/) (powP "x" ((+) "c" 1)) ((+) "c" 1) :| is_const "c"++ , intP (cosP "x") "x" := sinP "x"+ , intP (sinP "x") "x" := fromInteger (-1)*cosP "x"++ , intP ("f" + "g") "x" := intP "f" "x" + intP "g" "x"++ , intP ("f" - "g") "x" := intP "f" "x" - intP "g" "x"++ , intP ("a" * "b") "x" := (-) ((*) "a" (intP "b" "x")) (intP ((*) (diffP "x" "a") (intP "b" "x")) "x")++ -- Additional ad-hoc: because of negate representations?+ , "a"-(fromInteger (-1)*"b") := "a"+"b"++ ]++rewrite :: Fix Expr -> Fix Expr+rewrite e = fst $ equalitySaturation e rewrites symCost++symTests :: TestTree+symTests = testGroup "Symbolic"+ [ testCase "(a*2)/2 = a (custom rules)" $+ fst (equalitySaturation (("a"*2)/2) [ ("x"*"y")/"z" := "x"*("y"/"z")+ , "y"/"y" := 1+ , "x"*1 := "x"] symCost) @?= "a"++ , testCase "(a/2)*2 = a (all rules)" $+ rewrite (("a"/2)*2) @?= "a"++ , testCase "(a+a)/2 = a (extra rules)" $+ rewrite (("a"+"a")/2) @?= "a"++ , testCase "x/y (custom rules)" $+ -- without backoff scheduler this will loop forever+ fst (equalitySaturation+ ("x"/"y")++ [ "x"/"y" := "x"*(1/"y")+ , "x"*("y"*"z") := ("x"*"y")*"z"+ ]++ symCost) @?= ("x"/"y")++ , testCase "0+1 = 1 (all rules)" $+ fst (equalitySaturation (0+1) rewrites symCost) @?= 1++ , testCase "b*(1/b) = 1 (custom rules)" $+ fst (equalitySaturation ("b"*(1/"b")) [ "a"*(1/"a") := 1 ] symCost) @?= 1++ , testCase "1+1=2 (constant folding)" $+ fst (equalitySaturation (1+1) [] symCost) @?= 2++ , testCase "a*(2-1) (1 rule + constant folding)" $+ fst (equalitySaturation ("a" * (2-1)) ["x"*1:="x"] symCost) @?= "a"++ , testCase "1+a*(2-1) = 1+a (all + constant folding)" $+ rewrite (1+("a"*(2-1))) @?= (1+"a")++ , testCase "1+a*(2-1) = 1+a (all + constant f.)" $+ rewrite (fromInteger(-3)+fromInteger(-3)-6) @?= Fix (Const $ -12)++ , testCase "1+a-a*(2-1) = 1 (all + constant f.)" $+ rewrite (1 + "a" - "a"*(2-1)) @?= 1++ , testCase "1+(a-a*(2-1)) = 1 (all + constant f.)" $+ rewrite ("a" - "a"*(4-1)) @?= "a"*(Fix . Const $ -2)++ , testCase "x + x + x + x = 4*x" $+ rewrite ("a"+"a"+"a"+"a") @?= "a"*4++ , testCase "math powers" $+ rewrite (Fix (BinOp Pow 2 "x")*Fix (BinOp Pow 2 "y")) @?= Fix (BinOp Pow 2 ("x" + "y"))++ , testCase "d1" $+ rewrite (Fix $ BinOp Diff "a" "a") @?= 1++ , testCase "d2" $+ rewrite (Fix $ BinOp Diff "a" "b") @?= 0++ , testCase "d3" $+ rewrite (Fix $ BinOp Diff "x" (1 + 2*"x")) @?= 2++ , testCase "d4" $+ rewrite (Fix $ BinOp Diff "x" (1 + "y"*"x")) @?= "y"++ , testCase "d5" $+ rewrite (Fix $ BinOp Diff "x" (Fix $ UnOp Ln "x")) @?= 1/"x"++ , testCase "i1" $+ rewrite (Fix $ BinOp Integral 1 "x") @?= "x"++ , testCase "i2" $+ rewrite (Fix $ BinOp Integral (Fix $ UnOp Cos "x") "x") @?= Fix (UnOp Sin "x")++ , testCase "i3" $+ rewrite (Fix $ BinOp Integral (Fix $ BinOp Pow "x" 1) "x") @?= "x"*("x"*0.5)++ , testCase "i4" $+ rewrite (_i ((*) "x" (_cos "x")) "x") @?= (+) (_cos "x") ((*) "x" (_sin "x"))++ , testCase "i5" $+ rewrite (_i ((*) (_cos "x") "x") "x") @?= (+) (_cos "x") ((*) "x" (_sin "x"))++ -- TODO: How does this even work ?+ , testCase "i6" $+ rewrite (_i (_ln "x") "x") @?= "x"*(_ln "x" + fromInteger(-1))++ ]++_i :: Fix Expr -> Fix Expr -> Fix Expr+_i a b = Fix (BinOp Integral a b)+_ln, _cos, _sin :: Fix Expr -> Fix Expr+_ln a = Fix (UnOp Ln a)+_cos a = Fix (UnOp Cos a)+_sin a = Fix (UnOp Sin a)++powP :: Pattern Expr -> Pattern Expr -> Pattern Expr+powP a b = NonVariablePattern (BinOp Pow a b)++diffP :: Pattern Expr -> Pattern Expr -> Pattern Expr+diffP a b = NonVariablePattern (BinOp Diff a b)++intP :: Pattern Expr -> Pattern Expr -> Pattern Expr+intP a b = NonVariablePattern (BinOp Integral a b)++cosP :: Pattern Expr -> Pattern Expr+cosP a = NonVariablePattern (UnOp Cos a)++sinP :: Pattern Expr -> Pattern Expr+sinP a = NonVariablePattern (UnOp Sin a)++lnP :: Pattern Expr -> Pattern Expr+lnP a = NonVariablePattern (UnOp Ln a)
+ test/Test.hs view
@@ -0,0 +1,28 @@+{-# LANGUAGE OverloadedStrings #-}+import Test.Tasty++-- import Data.Equality.Utils+import Invariants+import Sym+import Lambda+import SimpleSym++tests :: TestTree+tests = testGroup "Tests"+ [ symTests+ , lambdaTests+ , simpleSymTests+ , invariants+ ]++main :: IO ()+main = defaultMain tests++-- main :: IO ()+-- main = do+-- print $ Sym.rewrite (Fix $ BinOp Integral (Fix $ BinOp Pow "x" 1) "x")++-- main :: IO ()+-- main = do+-- print $ Sym.rewrite (_i (_ln "x") "x")+-- putStrLn "Expecting: x*ln(x) + (-1)"