E-Graphs in Rust

E-Graphs (aka equality graphs), are a powerful data structure used in program optimization and equivalence checking. They represent a set of expressions and their equivalences compactly, allowing for efficient manipulation and transformation. Let's discuss their implementation in Rust using the egg crate, and provide practical examples to illustrate their use. E-Graphs are designed to handle the problem of reasoning about large sets of equivalent expressions. They are particularly useful in compiler optimizations, where they can be used to detect and transform equivalent expressions to more optimized forms.

Notably they consist of:

• E-Nodes: Representing individual expressions.
• E-Classes: Representing sets of equivalent expressions.
• E-Class IDs: Identifying e-classes uniquely.

The key feature of E-Graphs is their ability to compactly represent and efficiently manipulate large equivalence sets, making them a valuable tool in various optimization-related tasks. To get started with E-Graphs in Rust, we will use the egg library. First, add the egg crate to your Cargo.toml:

[dependencies]
egg = "0.6"

We will start by defining a simple expression language and using egg to build and manipulate E-Graphs. Let's define an expression language consisting of variables and basic arithmetic operations:

use egg::{define_language, EGraph, Rewrite, Runner};

define_language! {
    enum Expr {
        "+" = Add([Id; 2]),
        "*" = Mul([Id; 2]),
        Num(i32),
        Var(String),
    }
}

fn main() {
    let mut egraph = EGraph::default();

    // Add some initial expressions to the e-graph
    let a = egraph.add(Expr::Var("a".to_string()));
    let b = egraph.add(Expr::Var("b".to_string()));
    let expr1 = egraph.add(Expr::Add([a, b])); // a + b
    let expr2 = egraph.add(Expr::Mul([a, b])); // a * b

    println!("Initial e-graph: {:#?}", egraph);

    // Add some rewrite rules (e.g., distributive property)
    let rewrite1 = Rewrite::new(
        "mul_to_add",
        expr!(a * (b + c)),
        expr!((a * b) + (a * c)),
    )
    .unwrap();

    let runner = Runner::default()
        .with_expr(&expr)
        .run(&rules);

    let egraph = runner.egraph;

    // Print the transformed e-graph
    println!("Transformed e-graph: {:#?}", egraph);
}

Let's take a more detailed example where we use E-Graphs to perform optimizations on arithmetic expressions. Suppose we want to simplify expressions using standard arithmetic rules like associativity and distributivity. First, we define our rewrite rules:

# [dependencies]
# egg = "0.6"
use egg::{rewrite as rw, *};

rewrite! {
    comm_assoc_mul;
    "(+ ?x ?y)" => "(+ ?y ?x)"
}

rewrite! {
    comm_assoc_add;
    "(+ (+ ?x ?y) ?z)" => "(+ ?x (+ ?y ?z))"
}

rewrite! {
    distributive;
    "(* ?x (+ ?y ?z))" => "(+ (* ?x ?y) (* ?x ?z))"
}

Next, we create an E-Graph, add our expressions, and apply the rewrite rules:

fn main() {
    // Define the initial expression: a * (b + c)
    let a = Expr::Var("a".to_string());
    let b = Expr::Var("b".to_string());
    let c = Expr::Var("c".to_string());
    let expr = Expr::Mul(Box::new(a), Box::new(Expr::Add(Box::new(b), Box::new(c))));

    // Create an E-Graph and import the expression
    let mut egraph = EGraph::default();
    let _id = egraph.add(expr);

    // Define rewrite rules
    let rules: &[&Rewrite<Expr, ()>] = &[
        &comm_assoc_mul(),
        &comm_assoc_add(),
        &distributive(),
    ];

    // Run the rewriter
    let runner = Runner::default().with_expr(&expr).run(rules);

    // Extract the best expression using cost function
    let root = runner.roots[0];
    let extractor = Extractor::new(&runner.egraph, AstSize);
    let (best_cost, best_expr) = extractor.find_best(root);

    // Print the optimized expression
    println!("Best expression: {}", best_expr);
}

In this example, we start with the expression a * (b + c) and use rewrite rules to transform it using the distributive property, resulting in the optimized form (a * b) + (a * c).

If you're interested in learning more, the egg library's documentation provides further insights and examples to help you get started with E-Graphs in Rust. Happy coding!