catlog/
egglog_util.rs

1//! Utilities for working with e-graphs as implemented in egglog.
2
3use std::fmt::Display;
4
5use egglog::{EGraph, Error, ast::*, span};
6use ref_cast::RefCast;
7
8/** An egglog program.
9
10This struct is just a newtype wrapper around a vector of [`Command`]s with a
11builder interface.
12 */
13#[derive(Clone, Debug, RefCast)]
14#[repr(transparent)]
15pub struct Program(pub Vec<Command>);
16
17/// Displays the program as a newline-separated sequence of S-expressions.
18impl Display for Program {
19    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
20        for command in self.0.iter() {
21            writeln!(f, "{}", command)?;
22        }
23        Ok(())
24    }
25}
26
27impl Program {
28    /// Checks equality of two expressions, possibly after running a schedule.
29    pub fn check_equal(&mut self, lhs: Expr, rhs: Expr, schedule: Option<Schedule>) {
30        let prog = &mut self.0;
31        if let Some(schedule) = schedule {
32            // Should we assign to gensym-ed bindings rather than cloning?
33            prog.push(Command::Action(Action::Expr(span!(), lhs.clone())));
34            prog.push(Command::Action(Action::Expr(span!(), rhs.clone())));
35            prog.push(Command::RunSchedule(schedule));
36        }
37        prog.push(Command::Check(span!(), vec![Fact::Eq(span!(), lhs, rhs)]));
38    }
39
40    /// Unions the two expressions.
41    pub fn union(&mut self, lhs: Expr, rhs: Expr) {
42        self.0.push(Command::Action(Action::Union(span!(), lhs, rhs)));
43    }
44
45    /// Runs the program in the given e-graph, consuming the program.
46    pub fn run_in(self, egraph: &mut EGraph) -> Result<Vec<String>, Error> {
47        egraph.run_program(self.0)
48    }
49
50    /** Runs the program and returns whether all checks were successful.
51
52    Any errors besides check failures are handled normally.
53     */
54    pub fn check_in(self, egraph: &mut EGraph) -> Result<bool, Error> {
55        match self.run_in(egraph) {
56            Ok(_) => Ok(true),
57            Err(Error::CheckError(_, _)) => Ok(false),
58            Err(error) => Err(error),
59        }
60    }
61}
62
63/// Simplified egglog AST node for a rewrite.
64pub struct CommandRewrite {
65    /// Rule set to which the rewrite belongs.
66    pub ruleset: Symbol,
67    /// Left-hand side of rewrite.
68    pub lhs: Expr,
69    /// Right-hand side of rewrite.
70    pub rhs: Expr,
71}
72
73impl From<CommandRewrite> for Command {
74    fn from(rule: CommandRewrite) -> Self {
75        Command::Rewrite(
76            rule.ruleset,
77            Rewrite {
78                span: span!(),
79                lhs: rule.lhs,
80                rhs: rule.rhs,
81                conditions: vec![],
82            },
83            false,
84        )
85    }
86}
87
88/// Simplified egglog AST node for a rule.
89pub struct CommandRule {
90    /// Rule set to which the rule belongs.
91    pub ruleset: Symbol,
92    /// Head of rule.
93    pub head: Vec<Action>,
94    /// Body of rule.
95    pub body: Vec<Fact>,
96}
97
98impl From<CommandRule> for Command {
99    fn from(rule: CommandRule) -> Self {
100        Command::Rule {
101            name: "".into(),
102            ruleset: rule.ruleset,
103            rule: Rule {
104                span: span!(),
105                head: Actions::new(rule.head),
106                body: rule.body,
107            },
108        }
109    }
110}