Skip eclasses that do not contain eclass with operator

Use an egraph-wide map to remember which eclasses contain a particular operator.

PR #21

Author: mwillsey

src/egraph.rs

@@ -127,6 +127,7 @@ pub struct EGraph<L, M> {
     memo: IndexMap<ENode<L>, Id>,
     classes: UnionFind<Id, EClass<L, M>>,
     unions_since_rebuild: usize,
+    pub(crate) classes_by_op: IndexMap<(L, usize), Vec<Id>>,
 }
 
 // manual debug impl to avoid L: Language bound on EGraph defn
@@ -147,6 +148,7 @@ impl<L, M> Default for EGraph<L, M> {
             memo: IndexMap::default(),
             classes: UnionFind::default(),
             unions_since_rebuild: 0,
+            classes_by_op: IndexMap::default(),
         }
     }
 }
@@ -327,6 +329,7 @@ impl<L: Language, M: Metadata<L>> EGraph<L, M> {
             #[cfg(feature = "parent-pointers")]
             parents: IndexSet::new(),
         };
+
         M::modify(&mut class);
         let next_id = self.classes.make_set(class);
         trace!("Added  {:4}: {:?}", next_id, enode);
@@ -406,9 +409,11 @@ impl<L: Language, M: Metadata<L>> EGraph<L, M> {
     }
 
     fn rebuild_classes(&mut self) -> usize {
+        let (find, mut_values) = self.classes.split();
+
+        self.classes_by_op.clear();
         let mut trimmed = 0;
 
-        let (find, mut_values) = self.classes.split();
         for class in mut_values {
             let old_len = class.len();
 
@@ -422,6 +427,19 @@ impl<L: Language, M: Metadata<L>> EGraph<L, M> {
 
             class.nodes.clear();
             class.nodes.extend(unique);
+
+            let unique_op: IndexSet<(&L, usize)> = class
+                .nodes
+                .iter()
+                .map(|node| (&node.op, node.children.len()))
+                .collect();
+
+            for op in unique_op {
+                self.classes_by_op
+                    .entry((op.0.clone(), op.1))
+                    .and_modify(|ids| ids.push(class.id))
+                    .or_insert(vec![class.id]);
+            }
         }
 
         trimmed
@@ -467,11 +485,6 @@ impl<L: Language, M: Metadata<L>> EGraph<L, M> {
     /// ```
     #[cfg(not(feature = "parent-pointers"))]
     pub fn rebuild(&mut self) -> usize {
-        if self.unions_since_rebuild == 0 {
-            info!("Skipping rebuild!");
-            return 0;
-        }
-
         self.unions_since_rebuild = 0;
 
         let old_hc_size = self.memo.len();

src/pattern.rs

@@ -49,6 +49,9 @@ use crate::{
 /// // variable in the Pattern
 /// let same_add: Pattern<Math> = "(+ ?a ?a)".parse().unwrap();
 ///
+/// // Rebuild before searching
+/// egraph.rebuild();
+///
 /// // This is the search method from the Searcher trait
 /// let matches = same_add.search(&egraph);
 /// let matched_eclasses: Vec<Id> = matches.iter().map(|m| m.eclass).collect();
@@ -144,10 +147,19 @@ where
     M: Metadata<L>,
 {
     fn search(&self, egraph: &EGraph<L, M>) -> Vec<SearchMatches> {
-        egraph
-            .classes()
-            .filter_map(|e| self.search_eclass(egraph, e.id))
-            .collect()
+        match &self.ast {
+            PatternAst::ENode(e) => {
+                let key = (e.op.clone(), e.children.len());
+                let ids: &[Id] = egraph.classes_by_op.get(&key).map_or(&[], Vec::as_slice);
+                ids.iter()
+                    .filter_map(|&id| self.search_eclass(egraph, id))
+                    .collect()
+            }
+            PatternAst::Var(_) => egraph
+                .classes()
+                .filter_map(|e| self.search_eclass(egraph, e.id))
+                .collect(),
+        }
     }
 
     fn search_eclass(&self, egraph: &EGraph<L, M>, eclass: Id) -> Option<SearchMatches> {

src/rewrite.rs

@@ -472,7 +472,9 @@ mod tests {
             "fold_add"; "(+ ?a ?b)" => { Appender }
         );
 
+        egraph.rebuild();
         fold_add.run(&mut egraph);
+        egraph.rebuild();
         assert_eq!(egraph.equivs(&start, &goal), vec![egraph.find(root)]);
     }
 }

src/run.rs

@@ -305,6 +305,7 @@ where
     /// [`stop_reason`](#structfield.stop_reason) is guaranteeed to be
     /// set.
     pub fn run(mut self, rules: &[Rewrite<L, M>]) -> Self {
+        self.egraph.rebuild();
         // TODO check that we haven't
         loop {
             if let Err(stop_reason) = self.run_one(rules) {

tests/prop.rs

@@ -149,6 +149,7 @@ fn const_fold() {
     let start_expr = start.parse().unwrap();
     let end = "false";
     let end_expr = end.parse().unwrap();
-    let (eg, _) = EGraph::from_expr(&start_expr);
+    let (mut eg, _) = EGraph::from_expr(&start_expr);
+    eg.rebuild();
     assert!(!eg.equivs(&start_expr, &end_expr).is_empty());
 }