Add memoization when computing dag and tree cost

Some of the tensat examples have a lot more paths than nodes, so this helps them
finish faster.

src/extract/mod.rs

@@ -1,3 +1,5 @@
+use std::collections::HashMap;
+
 pub use crate::*;
 
 pub mod bottom_up;
@@ -19,34 +21,41 @@ pub trait Extractor: Sync {
 #[derive(Clone)]
 pub struct ExtractionResult {
     pub choices: Vec<Id>,
+    tree_memo: HashMap<Id, Cost>,
 }
 
 impl ExtractionResult {
     pub fn new(n_classes: usize) -> Self {
         ExtractionResult {
             choices: vec![0; n_classes],
+            tree_memo: HashMap::default(),
         }
     }
 
-    pub fn tree_cost(&self, egraph: &SimpleEGraph, root: Id) -> Cost {
+    pub fn tree_cost(&mut self, egraph: &SimpleEGraph, root: Id) -> Cost {
+        if let Some(&cost) = self.tree_memo.get(&root) {
+            return cost;
+        }
         let node = &egraph[root].nodes[self.choices[root]];
         let mut cost = node.cost;
         for &child in &node.children {
             cost += self.tree_cost(egraph, child);
         }
+        self.tree_memo.insert(root, cost);
         cost
     }
 
     // this will loop if there are cycles
     pub fn dag_cost(&self, egraph: &SimpleEGraph, root: Id) -> Cost {
         let mut costs = vec![INFINITY; egraph.classes.len()];
         let mut todo = vec![root];
-        while !todo.is_empty() {
-            let i = todo.pop().unwrap();
+        while let Some(i) = todo.pop() {
             let node = &egraph[i].nodes[self.choices[i]];
             costs[i] = node.cost;
             for &child in &node.children {
-                todo.push(child);
+                if costs[child] == INFINITY {
+                    todo.push(child);
+                }
             }
         }
         costs.iter().filter(|c| **c != INFINITY).sum()

src/main.rs

@@ -55,7 +55,7 @@ fn main() {
 
             for (ext_name, extractor) in &extractors {
                 let start_time = std::time::Instant::now();
-                let result = extractor.extract(&egraph, &egraph.roots);
+                let mut result = extractor.extract(&egraph, &egraph.roots);
                 let elapsed = start_time.elapsed();
                 for &root in &egraph.roots {
                     let msg = format!(