Sign in

    egg

    E-graphs and equality saturation in MoonBit, ported from egg

    Download zip
    Author
    Version
    0.1.0
    License
    MIT
    Last updated
    11 hours ago
    Downloads
    2

    #egg.mbt

    E-graphs and equality saturation in pure MoonBit, ported from the core algorithms of Rust egg. The module is bobzhang/egg. No external MoonBit dependencies; runs on JavaScript, native, Wasm, and Wasm GC.

    An e-graph stores many equivalent expressions together. Rewrites add alternatives without discarding earlier choices; extraction picks a low-cost expression.

    ///|
    test "simplify with equality saturation" {
    let graph = @egg.EGraph()
    let root = graph.add_expr(@egg.Expr::parse("(+ 0 (* 1 x))"))
    let rules = [
    @egg.Rewrite::parse("commute", "(+ ?a ?b)", "(+ ?b ?a)"),
    @egg.Rewrite::parse("add-zero", "(+ ?a 0)", "?a"),
    @egg.Rewrite::parse("multiply-one", "(* 1 ?a)", "?a"),
    ]
    let report = graph.run(rules)
    assert_eq(report.reason, Saturated)
    let best = graph.extract(root).unwrap()
    assert_eq(best.expression.to_string(), "x")
    assert_eq(best.cost, 1L)
    }

    #Installation

    moon add bobzhang/egg@0.1.0

    #Core API

    • EGraph() creates a graph. Node(op, children=[...]) adds an operator with ordered child IDs. Operators can be any immutable Eq + Hash type, including your own enums. String operators provide the analogue of egg's SymbolLang.
    • Expr() builds expression DAGs in child-before-parent order. Expr::parse and to_string support S-expressions with string operators.
    • add, add_expr, and union change a graph. Call rebuild() before queries. Rebuilding propagates congruence through parent links and restores hash-consing.
    • find, lookup, lookup_expr, class_ids, and class_nodes inspect it. Returned child arrays are copies. Internal storage is private.
    • Pattern::parse("(+ ?x ?x)") matches equal children. search returns substitutions and a truncated flag; instantiate adds a substituted pattern.
    • Rewrite::parse(name, lhs, rhs) rejects unbound RHS variables. Use Rewrite(name, lhs, rhs, condition=...) for a pure conditional rule.
    • run searches all rules before applying matches, rebuilds, and repeats. Its report distinguishes saturation, iteration/node/matching limits, and caller cancellation. Per-iteration counts record successful rule unions.
    • extract(root) minimizes AST size. cost=@egg.ast_depth minimizes depth; custom callbacks receive a node and child costs. The returned DAG shares repeated subexpressions even though AST cost counts each occurrence.
    • analyze(make, merge) computes a bounded e-class analysis fixed point. A run hook can use this to introduce derived constants or equalities. See analysis_test.mbt for constant folding.

    #Scope and contracts

    This is a usable core port, not complete upstream API or performance parity. See PORTING.md for the feature matrix and deliberate differences.

    IDs belong to one graph. Operators and analysis data must be immutable. Rewrite conditions and cost/analysis callbacks must be deterministic and obey their method contracts. Mathematical rewrites need valid domain assumptions: for example, x / x = 1 is not valid when x can be zero.

    Search limits are conservative and may stop at exactly the configured threshold. Node limits count allocated IDs, including merged nodes, and are checked between applications; one RHS may overshoot. A caller hook owns its own resource limits. Cancellation is checked between iterations/applications, not within callbacks. There is no wall-clock timeout or automatic eviction. Use bounded graphs and replace them when their lifetime ends.

    Parsed patterns have a 256-level nesting limit. Programmatically constructed patterns and large-arity matching are limited by backend call stacks. Expression DAG parsing and extraction reconstruction are iterative. Printing expands shared subexpressions and can produce exponentially large text; inspect DAG nodes for large results instead. AST costs use Int64 and abort on overflow. Custom costs must be nonnegative, monotone, and strictly greater than each child's cost.

    #Build and validate

    moon check --target all --warn-list +unnecessary_annotation --deny-warn moon test --target all moon run cmd/demo moon fmt --check moon info

    The README example is a test. Validation includes upstream arithmetic examples, cyclic classes, repeated variables, conditional rules, analysis, parser failures, explicit limits, and randomized congruence histories checked against an independent Boolean-matrix oracle.

    For cross-language validation, clone the reference as described in UPSTREAM.md, then run:

    python3 validation/compare.py

    This builds the unmodified pinned Rust library and the MoonBit JS export, then checks 512 seeded expressions against both implementations and an independent arithmetic normalizer. It compares extracted costs and evaluates the resulting expressions under three variable assignments. Equal-cost expression spelling may differ. Recorded results are correctness results, not performance measurements.

    Tested with MoonBit compiler 0.10.15+943464bfa. CI checks pinned/latest MoonBit on Linux and macOS, plus the pinned Rust differential harness.

    MIT licensed; upstream attribution is retained in LICENSE.

    ParseError

    pub suberror ParseError {
    ParseError(String)
    } derive(
    Debug
    )

    AnalysisResult

    pub(all) struct AnalysisResult[Data] {
    data : Map[Int, Data]
    converged : Bool
    iterations : Int
    }

    EGraph

    pub struct EGraph[Op] {
    // private fields
    }

    E-graph with hash-consing, path compression, and deferred congruence repair. IDs are local to this graph. Call rebuild after mutations before querying.

    EGraph::EGraph

    fn[Op : Eq + Hash] EGraph::EGraph() -> EGraph[Op]

    EGraph::add

    fn[Op : Eq + Hash] EGraph::add(self : EGraph[Op], node : Node[Op]) -> Int

    EGraph::add_expr

    fn[Op : Eq + Hash] EGraph::add_expr(self : EGraph[Op], expr : Expr[Op]) -> Int

    EGraph::allocated_size

    fn[Op] EGraph::allocated_size(self : EGraph[Op]) -> Int

    Number of allocated node IDs, including nodes merged into other classes.

    EGraph::analyze

    fn[Op : Eq + Hash, Data : Eq] EGraph::analyze(self : EGraph[Op], make : (Node[Op], Array[Data?]) -> Data?, merge : (Data, Data) -> Data, iteration_limit? : Int) -> AnalysisResult[Data]

    Computes e-class analysis to a bounded fixed point. make observes the current optional child data, and merge joins contributions. Both callbacks must be pure; merge must be associative, commutative, idempotent, and monotone. Data values must be immutable. A nonconverged result must not be treated as complete. Run this from the saturation hook to add derived constants/unions.

    EGraph::class_ids

    fn[Op] EGraph::class_ids(self : EGraph[Op]) -> Array[Int]

    EGraph::class_nodes

    fn[Op : Eq + Hash] EGraph::class_nodes(self : EGraph[Op], id : Int) -> Array[Node[Op]]

    Returns detached, canonicalized nodes, deduplicated within the class.

    EGraph::extract

    fn[Op : Eq + Hash] EGraph::extract(self : EGraph[Op], root : Int, cost? : (Node[Op], Array[Int64]) -> Int64) -> Extraction[Op]?

    Finds a cheapest finite expression. Costs must be deterministic, monotone in child costs, nonnegative, and strictly greater than each child cost. This excludes cyclic solutions. None means no finite representative was found.

    EGraph::find

    fn[Op] EGraph::find(self : EGraph[Op], id : Int) -> Int

    Finds a representative, compressing the path. Existing IDs remain valid.

    EGraph::instantiate

    fn[Op : Eq + Hash] EGraph::instantiate(self : EGraph[Op], pattern : Pattern[Op], subst : Map[String, Int]) -> Int

    Adds a substituted pattern. Rejects missing bindings before mutating the graph.

    EGraph::is_clean

    fn[Op] EGraph::is_clean(self : EGraph[Op]) -> Bool

    EGraph::lookup

    fn[Op : Eq + Hash] EGraph::lookup(self : EGraph[Op], node : Node[Op]) -> Int?

    EGraph::lookup_expr

    fn[Op : Eq + Hash] EGraph::lookup_expr(self : EGraph[Op], expr : Expr[Op]) -> Int?

    EGraph::number_of_classes

    fn[Op] EGraph::number_of_classes(self : EGraph[Op]) -> Int

    EGraph::rebuild

    fn[Op : Eq + Hash] EGraph::rebuild(self : EGraph[Op]) -> Int

    Restores congruence and hash-cons uniqueness. Returns extra congruence unions.

    EGraph::run

    fn[Op : Eq + Hash] EGraph::run(self : EGraph[Op], rules : Array[Rewrite[Op]], iteration_limit? : Int, node_limit? : Int, match_limit? : Int, search_steps? : Int, hook? : (EGraph[Op]) -> Unit, cancelled? : () -> Bool) -> RunReport

    Equality saturation: search every rule before applying any matches, then rebuild and invoke the optional analysis hook. Limits are checked at rule applications; node count can overshoot by at most one instantiated RHS. Match truncation stops explicitly rather than claiming saturation.

    EGraph::total_size

    fn[Op] EGraph::total_size(self : EGraph[Op]) -> Int

    Number of distinct canonical e-nodes after rebuild.

    EGraph::union

    fn[Op] EGraph::union(self : EGraph[Op], left : Int, right : Int) -> Bool

    Union classes, deferring parent congruences until rebuild. Returns false when both IDs already denote the same class.

    Expr

    pub struct Expr[Op] {
    // private fields
    }

    A topologically ordered expression DAG; the last node is its root.

    Expr::Expr

    fn[Op] Expr::Expr() -> Expr[Op]

    Expr::add

    fn[Op] Expr::add(self : Expr[Op], node : Node[Op]) -> Int

    Appends a node. Every child must refer to an earlier node.

    Expr::length

    fn[Op] Expr::length(self : Expr[Op]) -> Int

    Expr::node

    fn[Op] Expr::node(self : Expr[Op], id : Int) -> Node[Op]

    Returns a detached copy of the node's child array.

    Expr::parse

    fn Expr::parse(text : String) -> Expr[String] raise ParseError

    Parses one S-expression, using an iterative stack. Quoted atoms support quote, backslash, newline, carriage return, and tab escapes.

    Expr::root

    fn[Op] Expr::root(self : Expr[Op]) -> Int

    Expr::to_string

    fn Expr::to_string(self : Expr[String]) -> String

    Prints the expression rooted at its last node. Shared DAG nodes are expanded.

    Extraction

    pub(all) struct Extraction[Op] {
    cost : Int64
    expression : Expr[Op]
    }

    Iteration

    pub(all) struct Iteration {
    index : Int
    classes : Int
    nodes : Int
    applied : Map[String, Int]
    }

    Match

    pub(all) struct Match {
    eclass : Int
    subst : Map[String, Int]
    }

    Node

    pub(all) struct Node[Op] {
    op : Op
    children : Array[Int]
    } derive(Eq, Hash)

    An operator with ordered child IDs. Operators must be immutable map keys.

    Node::Node

    fn[Op] Node::Node(op : Op, children? : Array[Int]) -> Node[Op]

    Node::equal

    fn[Op : Eq] Node::equal(Node[Op], Node[Op]) -> Bool

    Node::hash

    fn[Op : Hash] Node::hash(self : Node[Op]) -> Int

    Node::hash_combine

    fn[Op : Hash] Node::hash_combine(Node[Op], Hasher) -> Unit

    Node::not_equal

    fn[Op : Eq] Node::not_equal(x : Node[Op], y : Node[Op]) -> Bool

    Pattern

    pub(all) enum Pattern[Op] {
    Variable(String)
    Apply(Op, Array[Pattern[Op]])
    }

    Immutable pattern trees. Variable names conventionally start with '?'.

    Pattern::parse

    fn Pattern::parse(text : String) -> Pattern[String] raise ParseError

    Parses patterns up to 256 levels deep. A '?name' leaf is a variable.

    Pattern::search

    fn[Op : Eq + Hash] Pattern::search(self : Pattern[Op], graph : EGraph[Op], limit? : Int, steps? : Int, eclass? : Int) -> SearchResult

    Searches a clean graph. Limits are conservative: reaching either limit marks the result truncated, so a runner cannot falsely report saturation.

    Pattern::variables

    fn[Op] Pattern::variables(self : Pattern[Op]) -> Array[String]

    Rewrite

    pub struct Rewrite[Op] {
    // private fields
    }

    Rewrite::Rewrite

    fn[Op] Rewrite::Rewrite(name : String, left : Pattern[Op], right : Pattern[Op], condition? : (EGraph[Op], Map[String, Int]) -> Bool) -> Rewrite[Op] raise ParseError

    Creates a rewrite, rejecting RHS variables not bound on the LHS. Conditions must not mutate the graph. Operators must remain immutable.

    Rewrite::name

    fn[Op] Rewrite::name(self : Rewrite[Op]) -> String

    Rewrite::parse

    fn Rewrite::parse(name : String, left : String, right : String) -> Rewrite[String] raise ParseError

    RunReport

    pub(all) struct RunReport {
    reason : StopReason
    iterations : Array[Iteration]
    }

    SearchResult

    pub(all) struct SearchResult {
    matches : Array[Match]
    truncated : Bool
    }

    StopReason

    pub(all) enum StopReason {
    Saturated
    IterationLimit
    NodeLimit
    MatchLimit
    Cancelled
    } derive(Eq,
    Debug
    )

    StopReason::equal

    fn StopReason::equal(StopReason, StopReason) -> Bool

    StopReason::not_equal

    fn StopReason::not_equal(x : StopReason, y : StopReason) -> Bool

    ast_depth

    fn[Op] ast_depth(_node : Node[Op], children : Array[Int64]) -> Int64

    ast_size

    fn[Op] ast_size(_node : Node[Op], children : Array[Int64]) -> Int64

    AST size counts repeated children separately, even when the output DAG shares their storage. Costs overflow explicitly rather than silently wrapping.