Data.Equality.Saturation

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.

Equality saturation

equalitySaturation

Arguments

:: forall l cost. (Language l, Ord cost)

=> Fix l

Expression to run equality saturation on

-> [Rewrite l]

List of rewrite rules

-> CostFunction l cost

Cost function to extract the best equivalent representation

-> (Fix l, EGraph l)

Best equivalent expression and resulting e-graph

Equality saturation with defaults

equalitySaturation'

Arguments

:: forall l schd cost. (Language l, Scheduler schd, Ord cost)

=> Proxy schd

Proxy for the scheduler to use

-> Fix l

Expression to run equality saturation on

-> [Rewrite l]

List of rewrite rules

-> CostFunction l cost

Cost function to extract the best equivalent representation

-> (Fix l, EGraph l)

Best equivalent expression and resulting e-graph

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

runEqualitySaturation

Arguments

:: forall l schd. (Language l, Scheduler schd)

=> Proxy schd

Proxy for the scheduler to use

-> [Rewrite l]

List of rewrite rules

-> EGraphM l ()

Run equality saturation on an e-graph by non-destructively applying all given rewrite rules until saturation (using the given Scheduler)

Re-exports for equality saturation

Writing rewrite rules

data Rewrite lang

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

Constructors

!(Pattern lang) := !(Pattern lang) infix 3

Trivial Rewrite

!(Rewrite lang) :| !(RewriteCondition lang) infixl 2

Conditional Rewrite

type RewriteCondition lang = Subst -> EGraph lang -> Bool

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

Writing cost functions

CostFunction re-exported from Extraction since they are required to do equality saturation

type CostFunction l cost = l cost -> cost

A cost function is used to attribute a cost to representations in the e-graph and to extract the best one.

The cost function is polymorphic over the type used for the cost, however cost must instance Ord in order for the defined CostFunction to fulfill its purpose. That's why we have an Ord cost constraint in equalitySaturation and extractBest

Example

symCost :: Expr Int -> Int
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

Writing expressions

Expressions must be written in their fixed-point form, since the Language must be given in its base functor form

newtype Fix f

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.

Constructors

Fix

Fields

• unFix :: f (Fix f)

Instances

Show1 f => Show (Fix f)

Defined in Data.Equality.Utils

Methods

showsPrec :: Int -> Fix f -> ShowS

show :: Fix f -> String

showList :: [Fix f] -> ShowS

Eq1 f => Eq (Fix f)

Defined in Data.Equality.Utils

Methods

(==) :: Fix f -> Fix f -> Bool

(/=) :: Fix f -> Fix f -> Bool

cata :: Functor f => (f a -> a) -> Fix f -> a

Catamorphism