Disallow values from jumping between blocks inappropriately

Undo changes to the union find structure and gvn_map when popping
the node stack in `remove_pure_and_optimize`, ensuring that we don't
propagate values in a way that would invalidate the invariant that the
`ValueData::Union` structure provides.

Additionally, remove all uses of the union find structure from
elaboration, as we have already computed the best values at that point.

Co-authored-by: Jamey Sharp <jsharp@fastly.com>
Co-authored-by: L Pereira <l.pereira@fastly.com>
Co-authored-by: Nick Fitzgerald <fitzgen@gmail.com>

cranelift/codegen/src/ctxhash.rs

@@ -51,6 +51,7 @@ 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,
@@ -58,6 +59,7 @@ struct BucketData<K, 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,7 +497,7 @@ 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() {
@@ -505,7 +505,7 @@ impl<'a> EgraphPass<'a> {
                 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;
                 }
             }
         }

cranelift/codegen/src/egraph/elaborate.rs

@@ -21,7 +21,6 @@ pub(crate) struct Elaborator<'a> {
     domtree: &'a DominatorTree,
     domtree_children: &'a DomTreeWithChildren,
     loop_analysis: &'a LoopAnalysis,
-    eclasses: &'a mut UnionFind<Value>,
     /// Map from Value that is produced by a pure Inst (and was thus
     /// not in the side-effecting skeleton) to the value produced by
     /// an elaborated inst (placed in the layout) to whose results we
@@ -137,7 +136,7 @@ impl<'a> Elaborator<'a> {
         domtree_children: &'a DomTreeWithChildren,
         loop_analysis: &'a LoopAnalysis,
         remat_values: &'a FxHashSet<Value>,
-        eclasses: &'a mut UnionFind<Value>,
+        _eclasses: &'a mut UnionFind<Value>,
         stats: &'a mut Stats,
     ) -> Self {
         let num_values = func.dfg.num_values();
@@ -149,7 +148,6 @@ impl<'a> Elaborator<'a> {
             domtree,
             domtree_children,
             loop_analysis,
-            eclasses,
             value_to_elaborated_value: ScopedHashMap::with_capacity(num_values),
             value_to_best_value,
             loop_stack: smallvec![],
@@ -325,14 +323,14 @@ impl<'a> Elaborator<'a> {
                     let value = self.func.dfg.resolve_aliases(value);
 
                     self.stats.elaborate_visit_node += 1;
-                    let canonical_value = self.eclasses.find_and_update(value);
-                    debug_assert_ne!(canonical_value, Value::reserved_value());
-                    trace!(
-                        "elaborate: value {} canonical {} before {}",
-                        value,
-                        canonical_value,
-                        before
-                    );
+                    // let canonical_value = self.eclasses.find_and_update(value);
+                    // debug_assert_ne!(canonical_value, Value::reserved_value());
+                    // trace!(
+                    //     "elaborate: value {} canonical {} before {}",
+                    //     value,
+                    //     canonical_value,
+                    //     before
+                    // );
 
                     // Get the best option; we use `value` (latest
                     // value) here so we have a full view of the
@@ -342,7 +340,8 @@ impl<'a> Elaborator<'a> {
                     trace!("elaborate: value {} -> best {}", value, best_value);
                     debug_assert_ne!(best_value, Value::reserved_value());
 
-                    if let Some(elab_val) = self.value_to_elaborated_value.get(&canonical_value) {
+                    // if let Some(elab_val) = self.value_to_elaborated_value.get(&canonical_value) {
+                    if let Some(elab_val) = self.value_to_elaborated_value.get(&best_value) {
                         // Value is available; use it.
                         trace!("elaborate: value {} -> {:?}", value, elab_val);
                         self.stats.elaborate_memoize_hit += 1;
@@ -568,15 +567,16 @@ impl<'a> Elaborator<'a> {
                                 value: new_result,
                                 in_block: insert_block,
                             };
-                            let canonical_result = self.eclasses.find_and_update(result);
+                            // let canonical_result = self.eclasses.find_and_update(result);
+                            let BestEntry(_, canonical_result) = self.value_to_best_value[result];
                             self.value_to_elaborated_value.insert_if_absent_with_depth(
                                 canonical_result,
                                 elab_value,
                                 scope_depth,
                             );
 
-                            self.eclasses.add(new_result);
-                            self.eclasses.union(result, new_result);
+                            // self.eclasses.add(new_result);
+                            // self.eclasses.union(result, new_result);
                             self.value_to_best_value[new_result] = self.value_to_best_value[result];
 
                             trace!(
@@ -596,7 +596,8 @@ impl<'a> Elaborator<'a> {
                                 value: result,
                                 in_block: insert_block,
                             };
-                            let canonical_result = self.eclasses.find_and_update(result);
+                            // let canonical_result = self.eclasses.find_and_update(result);
+                            let BestEntry(_, canonical_result) = self.value_to_best_value[result];
                             self.value_to_elaborated_value.insert_if_absent_with_depth(
                                 canonical_result,
                                 elab_value,
@@ -684,7 +685,8 @@ impl<'a> Elaborator<'a> {
             // map now.
             for &result in self.func.dfg.inst_results(inst) {
                 trace!(" -> result {}", result);
-                let canonical_result = self.eclasses.find_and_update(result);
+                // let canonical_result = self.eclasses.find_and_update(result);
+                let BestEntry(_, canonical_result) = self.value_to_best_value[result];
                 self.value_to_elaborated_value.insert_if_absent(
                     canonical_result,
                     ElaboratedValue {