egraphs: Undo changes to union find and gvn map structures when backtracking

cranelift/codegen/src/ctxhash.rs

@@ -51,13 +51,15 @@ impl<V: Eq + Hash> CtxHash<V> for NullCtx {
 /// the hashcode, for memory efficiency: in common use, `K` and `V`
 /// are often 32 bits also, and a 12-byte bucket is measurably better
 /// than a 16-byte bucket.
+#[derive(Clone)]
 struct BucketData<K, V> {
     hash: u32,
     k: K,
     v: V,
 }
 
 /// A HashMap that takes external context for all operations.
+#[derive(Clone)]
 pub struct CtxHashMap<K, V> {
     raw: RawTable<BucketData<K, V>>,
 }

cranelift/codegen/src/egraph.rs

@@ -497,15 +497,15 @@ impl<'a> EgraphPass<'a> {
         let root = self.domtree_children.root();
         enum StackEntry {
             Visit(Block),
-            Pop,
+            Pop(UnionFind<Value>, CtxHashMap<(Type, InstructionData), Value>),
         }
         let mut block_stack = vec![StackEntry::Visit(root)];
         while let Some(entry) = block_stack.pop() {
             match entry {
                 StackEntry::Visit(block) => {
                     // We popped this block; push children
                     // immediately, then process this block.
-                    block_stack.push(StackEntry::Pop);
+                    block_stack.push(StackEntry::Pop(self.eclasses.clone(), gvn_map.clone()));
                     block_stack
                         .extend(self.domtree_children.children(block).map(StackEntry::Visit));
                     effectful_gvn_map.increment_depth();
@@ -579,8 +579,13 @@ impl<'a> EgraphPass<'a> {
                         }
                     }
                 }
-                StackEntry::Pop => {
+                StackEntry::Pop(eclasses, saved_gvn_map) => {
                     effectful_gvn_map.decrement_depth();
+                    // TODO: do we need to save both eclasses and gvn_map? All the tests pass
+                    // without restoring eclasses, but perhaps that's a limitation of our test
+                    // suite.
+                    self.eclasses = eclasses;
+                    gvn_map = saved_gvn_map;
                 }
             }
         }
@@ -701,5 +706,4 @@ pub(crate) struct Stats {
     pub(crate) elaborate_func: u64,
     pub(crate) elaborate_func_pre_insts: u64,
     pub(crate) elaborate_func_post_insts: u64,
-    pub(crate) elaborate_best_cost_fixpoint_iters: u64,
 }

cranelift/codegen/src/egraph/cost.rs

@@ -74,7 +74,7 @@ impl Cost {
     const DEPTH_BITS: u8 = 8;
     const DEPTH_MASK: u32 = (1 << Self::DEPTH_BITS) - 1;
     const OP_COST_MASK: u32 = !Self::DEPTH_MASK;
-    const MAX_OP_COST: u32 = Self::OP_COST_MASK >> Self::DEPTH_BITS;
+    const MAX_OP_COST: u32 = (Self::OP_COST_MASK >> Self::DEPTH_BITS) - 1;
 
     pub(crate) fn infinity() -> Cost {
         // 2^32 - 1 is, uh, pretty close to infinite... (we use `Cost`
@@ -86,16 +86,14 @@ impl Cost {
         Cost(0)
     }
 
-    /// Construct a new `Cost` from the given parts.
+    /// Construct a new finite cost from the given parts.
     ///
-    /// If the opcode cost is greater than or equal to the maximum representable
-    /// opcode cost, then the resulting `Cost` saturates to infinity.
-    fn new(opcode_cost: u32, depth: u8) -> Cost {
-        if opcode_cost >= Self::MAX_OP_COST {
-            Self::infinity()
-        } else {
-            Cost(opcode_cost << Self::DEPTH_BITS | u32::from(depth))
-        }
+    /// The opcode cost is clamped to the maximum value representable.
+    fn new_finite(opcode_cost: u32, depth: u8) -> Cost {
+        let opcode_cost = std::cmp::min(opcode_cost, Self::MAX_OP_COST);
+        let cost = Cost((opcode_cost << Self::DEPTH_BITS) | u32::from(depth));
+        debug_assert_ne!(cost, Cost::infinity());
+        cost
     }
 
     fn depth(&self) -> u8 {
@@ -113,7 +111,7 @@ impl Cost {
     /// that satisfies `inst_predicates::is_pure_for_egraph()`.
     pub(crate) fn of_pure_op(op: Opcode, operand_costs: impl IntoIterator<Item = Self>) -> Self {
         let c = pure_op_cost(op) + operand_costs.into_iter().sum();
-        Cost::new(c.op_cost(), c.depth().saturating_add(1))
+        Cost::new_finite(c.op_cost(), c.depth().saturating_add(1))
     }
 }
 
@@ -133,9 +131,12 @@ impl std::ops::Add<Cost> for Cost {
     type Output = Cost;
 
     fn add(self, other: Cost) -> Cost {
-        let op_cost = self.op_cost().saturating_add(other.op_cost());
+        let op_cost = std::cmp::min(
+            self.op_cost().saturating_add(other.op_cost()),
+            Self::MAX_OP_COST,
+        );
         let depth = std::cmp::max(self.depth(), other.depth());
-        Cost::new(op_cost, depth)
+        Cost::new_finite(op_cost, depth)
     }
 }
 
@@ -146,11 +147,11 @@ impl std::ops::Add<Cost> for Cost {
 fn pure_op_cost(op: Opcode) -> Cost {
     match op {
         // Constants.
-        Opcode::Iconst | Opcode::F32const | Opcode::F64const => Cost::new(1, 0),
+        Opcode::Iconst | Opcode::F32const | Opcode::F64const => Cost::new_finite(1, 0),
 
         // Extends/reduces.
         Opcode::Uextend | Opcode::Sextend | Opcode::Ireduce | Opcode::Iconcat | Opcode::Isplit => {
-            Cost::new(2, 0)
+            Cost::new_finite(2, 0)
         }
 
         // "Simple" arithmetic.
@@ -162,52 +163,9 @@ fn pure_op_cost(op: Opcode) -> Cost {
         | Opcode::Bnot
         | Opcode::Ishl
         | Opcode::Ushr
-        | Opcode::Sshr => Cost::new(3, 0),
+        | Opcode::Sshr => Cost::new_finite(3, 0),
 
         // Everything else (pure.)
-        _ => Cost::new(4, 0),
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn add_cost() {
-        let a = Cost::new(5, 2);
-        let b = Cost::new(37, 3);
-        assert_eq!(a + b, Cost::new(42, 3));
-        assert_eq!(b + a, Cost::new(42, 3));
-    }
-
-    #[test]
-    fn add_infinity() {
-        let a = Cost::new(5, 2);
-        let b = Cost::infinity();
-        assert_eq!(a + b, Cost::infinity());
-        assert_eq!(b + a, Cost::infinity());
-    }
-
-    #[test]
-    fn op_cost_saturates_to_infinity() {
-        let a = Cost::new(Cost::MAX_OP_COST - 10, 2);
-        let b = Cost::new(11, 2);
-        assert_eq!(a + b, Cost::infinity());
-        assert_eq!(b + a, Cost::infinity());
-    }
-
-    #[test]
-    fn depth_saturates_to_max_depth() {
-        let a = Cost::new(10, u8::MAX);
-        let b = Cost::new(10, 1);
-        assert_eq!(
-            Cost::of_pure_op(Opcode::Iconst, [a, b]),
-            Cost::new(21, u8::MAX)
-        );
-        assert_eq!(
-            Cost::of_pure_op(Opcode::Iconst, [b, a]),
-            Cost::new(21, u8::MAX)
-        );
+        _ => Cost::new_finite(4, 0),
     }
 }