eggplant

eggplant is the High-Level Rust API repo for the egglog tool accompanying the paper "Better Together: Unifying Datalog and Equality Saturation" (ACM DL, arXiv).

It is hard to do Graph operations directly on EGraph because EGraph is a highly compressed data structure.

Based on that fact, eggplant provides out-of-box Graph API that allows you to do revisions on EGraph intuitively.

There is also a Proc-Macro library for users to quickly define a suite of DSL.

Transpiler Macros

Eggplant provides transpiler macros that convert egglog DSL syntax directly to Rust code:

datatype! Macro

Converts egglog datatype definitions to Rust code:

datatype! {
    (datatype Math (MNum i64:args_name "num") (MAdd Math Math:args_name "l,r"))
}

This generates the corresponding Rust enum and associated types.

rule! Macro

Converts egglog rewrite rules to Rust code:

rule! {
    (datatype Math (MNum i64:args_name "num") (MAdd Math Math:args_name "l,r"))
    (rewrite (MAdd x y) (MAdd y x))
}

This generates the corresponding pattern matching and rewrite rule code. This will be marked error by compiler, and you should inline it to get generated code.

Motivation

Recently, I've been researching the implementation of egglog. Since the egglog ecosystem is not yet mature, I spent some time providing a better API implementation with procedural macros and compile-time type checking for rust-native. This includes:

1. DSL Define API: Quickly define a set of DSL
2. Pattern Define API: Pattern definition for pattern matching, using declarative definition approach
3. Commit API: Since egraph is a highly compressed structure that compresses node relationship information and is naturally averse to deletion operations, we need a bridge between normal graph structures and egraph. I borrowed from git's commit API for batch commits and version control.
4. Insert API: Insert API with compile-time type checking added on top of the original, so you no longer need to nervously add nodes like with the native API
5. Run Rule API: Run Pattern Rewrite

Finally, this forms a framework that allows users to implement a pattern rewriting framework supporting integer addition, subtraction, multiplication, and division with constant propagation in just fifty lines of code.

DSL Define API

#[eggplant::dsl]
pub enum Expr {
    Const { num: i64 },
    Mul { l: Expr, r: Expr },
    Sub { l: Expr, r: Expr },
    Add { l: Expr, r: Expr },
    Div { l: Expr, r: Expr },
}

Yes, you read that correctly. Our addition, subtraction, multiplication, and division constant DSL is defined just like that. Using Rust's Sum Type can represent a set of DSL very well, which is very intuitive.

You can also specify cost attributes for DSL variants to guide extraction of optimal values:

#[eggplant::dsl]
pub enum Expr {
    #[cost(0)]
    Const { num: i64 },
    Mul { l: Expr, r: Expr },
}

The #[cost(0)] attribute specifies that Const nodes have a cost of 0, which helps the system extract the most optimal expressions during pattern rewriting.

Pattern Define API

Here's an example of addition pattern definition API:

#[eggplant::pat_vars]
struct AddPat {
    l: Const,
    r: Const,
    p: Add,
}

let pat = || {
    // mention here we use query() since it's a variant type, if we want to query an expression
    // we should use query_leaf()
    let l = Const::query(); 
    let r = Const::query();
    let p = Add::query(&l, &r);
    AddPat::new(l, r, p)
}

Note that currently we must add a generic parameter after AddPat. I think we can use procedural macro tricks to omit this generic parameter later :) (Now supports omitting generics).

Defining a Pattern is actually defining two things: what to extract from the Pattern for computation and what this Pattern looks like. From the user's perspective, using a declarative approach to define is definitely more convenient. Here we use a closure to define a Pattern, but as the generic parameter exposed by the struct above shows, there's a global singleton managing all Patterns and submitting them to the database at the appropriate time.

Addition and Pattern Rewrite Definition!

let action = |ctx, values| {
    let cal = ctx.devalue(values.l.num) + ctx.devalue(values.r.num); // Used as struct rather than enum
    let add_value = ctx.insert_const(cal);
    ctx.union(values.p, add_value);
},

You can see that we still need to use ctx.devalue to convert from database ID to real data. Through compile-time type information deduction, we can save the type annotation for devalue. Also, the procedural macro generates APIs corresponding to the DSL for ctx, maximizing the power of intellisense and compile-time type information.

Running

tx_rx_vt_pr!(MyTx, MyPatRec);  // Global singleton definition for storing graph and patterns
fn main() {
    let expr: Expr<MyTx, AddTy> =
        Add::new(&Mul::new(&Const::new(3), &Const::new(2)), &Const::new(4));
    expr.commit();

    let ruleset = MyTx::new_ruleset("constant_prop");
    prop!(Add,+,AddPat,ruleset);
    prop!(Sub,-,SubPat,ruleset);
    prop!(Mul,*,MulPat,ruleset);
    prop!(Div,/,DivPat,ruleset);
    MyTx::run_ruleset(ruleset, RunConfig::Sat);
    MyTx::egraph_to_dot("egraph.dot".into());
}

Finally, the following EGraph is generated, and you can see that the root node value is directly derived.

Note that the execution count of run_ruleset is not the number of matches, but should be less than the tree depth.

Dependencies

Using eggplant requires

cargo add eggplant

and then add this to toml

derive_more = { version = "2.0.1", features = [
    "deref_mut",
    "deref",
    "into_iterator",
    "debug",
] }
strum = { version = "0.27.2", features = ["strum_macros"] }
strum_macros = "0.27.2"

Future Work

1. Support slotted egraph
2. Proof & Viewer

Here's the complete code for implementing addition, subtraction, multiplication, and division constant propagation:

use eggplant::{prelude::*, tx_rx_vt_pr};

#[eggplant::ty]
pub enum Expr {
    Const { num: i64 },
    Mul { l: Expr, r: Expr },
    Sub { l: Expr, r: Expr },
    Add { l: Expr, r: Expr },
    Div { l: Expr, r: Expr },
}

tx_rx_vt_pr!(MyTx, MyPatRec);

macro_rules! prop {
    ($ty:ident,$op:tt,$pat_name:ident,$ruleset:ident) => {
        #[eggplant::pat_vars]
        struct $pat_name {
            l: Const,
            r: Const,
            p: $ty,
        }
        MyTx::add_rule(
            stringify!($pat_name),
            $ruleset,
            || {
                let l = Const::query();
                let r = Const::query();
                let p = $ty::query(&l, &r);
                $pat_name::new(l, r, p)
            },
            |ctx, values| {
                let cal = ctx.devalue(values.l.num) $op ctx.devalue(values.r.num);
                let op_value = ctx.insert_const(cal);
                ctx.union(values.p, op_value);
            },
        );
    };
}

fn main() {
    let expr: Expr<MyTx, AddTy> =
        Add::new(&Mul::new(&Const::new(3), &Const::new(2)), &Const::new(4));
    expr.commit();

    let ruleset = MyTx::new_ruleset("constant_prop");
    prop!(Add,+,AddPat,ruleset);
    prop!(Sub,-,SubPat,ruleset);
    prop!(Mul,*,MulPat,ruleset);
    prop!(Div,/,DivPat,ruleset);
    MyTx::run_ruleset(ruleset, RunConfig::Sat);
    MyTx::sgl().egraph_to_dot("egraph.dot".into());
}

Examples

The project includes several example files demonstrating different features of eggplant:

Basic Examples

• examples/constrain_basic.rs: Demonstrates basic constant propagation with pattern matching. Shows how to define addition patterns and apply constant folding rules.

• examples/constrain_complex.rs: Shows more complex pattern matching with equality constraints. Includes addition and multiplication patterns with handle equality checks.

Constraint Examples

• examples/constrain_lt100_test.rs: Demonstrates pattern matching with numerical constraints. Shows how to apply rules only when operands satisfy specific conditions (e.g., both operands < 100).

Container Examples

• examples/container.rs: Shows how to use container types (arrays and sets) in DSL definitions. Demonstrates vector summation and set operations with custom container types.

Base Type Examples

• examples/base_ty_def.rs: Demonstrates the use of base types in DSL definitions. Shows how to define an operation type enum and use it in binary expressions with pattern matching.

Documentation

To view documentation, run cargo doc --open.

Contributing

Welcome to submit issues! Hope you have fun!