coverage: Replace `FrozenUnionFind` with a plain IndexVec

This dedicated type seemed like a good idea at the time, but if we want to
store this information in a query result then a plainer data type is more
convenient.

2 files changed, 11 insertions, 31 deletions

diff --git a/compiler/rustc_mir_transform/src/coverage/counters/node_flow.rs b/compiler/rustc_mir_transform/src/coverage/counters/node_flow.rs
index a022ae67c0b..7f767810f31 100644
--- a/compiler/rustc_mir_transform/src/coverage/counters/node_flow.rs
+++ b/compiler/rustc_mir_transform/src/coverage/counters/node_flow.rs
@@ -12,7 +12,7 @@ use rustc_index::{Idx, IndexVec};
 use rustc_middle::mir::coverage::Op;
 
 use crate::coverage::counters::iter_nodes::IterNodes;
-use crate::coverage::counters::union_find::{FrozenUnionFind, UnionFind};
+use crate::coverage::counters::union_find::UnionFind;
 
 #[cfg(test)]
 mod tests;
@@ -32,7 +32,7 @@ mod tests;
 pub(crate) struct MergedNodeFlowGraph<Node: Idx> {
     /// Maps each node to the supernode that contains it, indicated by some
     /// arbitrary "root" node that is part of that supernode.
-    supernodes: FrozenUnionFind<Node>,
+    supernodes: IndexVec<Node, Node>,
     /// For each node, stores the single supernode that all of its successors
     /// have been merged into.
     ///
@@ -66,11 +66,11 @@ impl<Node: Idx> MergedNodeFlowGraph<Node> {
             })
             .collect::<IndexVec<G::Node, G::Node>>();
 
-        // Now that unification is complete, freeze the supernode forest,
+        // Now that unification is complete, take a snapshot of the supernode forest,
         // and resolve each arbitrarily-chosen successor to its canonical root.
         // (This avoids having to explicitly resolve them later.)
-        let supernodes = supernodes.freeze();
-        let succ_supernodes = successors.into_iter().map(|succ| supernodes.find(succ)).collect();
+        let supernodes = supernodes.snapshot();
+        let succ_supernodes = successors.into_iter().map(|succ| supernodes[succ]).collect();
 
         Self { supernodes, succ_supernodes }
     }
@@ -80,7 +80,7 @@ impl<Node: Idx> MergedNodeFlowGraph<Node> {
     }
 
     fn is_supernode(&self, node: Node) -> bool {
-        self.supernodes.find(node) == node
+        self.supernodes[node] == node
     }
 
     /// Using the information in this merged graph, together with a given
@@ -225,7 +225,7 @@ impl<'a, Node: Idx> SpantreeBuilder<'a, Node> {
 
         // Get the supernode containing `this`, and make it the root of its
         // component of the spantree.
-        let this_supernode = self.graph.supernodes.find(this);
+        let this_supernode = self.graph.supernodes[this];
         self.yank_to_spantree_root(this_supernode);
 
         // Get the supernode containing all of this's successors.
diff --git a/compiler/rustc_mir_transform/src/coverage/counters/union_find.rs b/compiler/rustc_mir_transform/src/coverage/counters/union_find.rs
index 2da4f5f5fce..a826a953fa6 100644
--- a/compiler/rustc_mir_transform/src/coverage/counters/union_find.rs
+++ b/compiler/rustc_mir_transform/src/coverage/counters/union_find.rs
@@ -88,29 +88,9 @@ impl<Key: Idx> UnionFind<Key> {
         a
     }
 
-    /// Creates a snapshot of this disjoint-set forest that can no longer be
-    /// mutated, but can be queried without mutation.
-    pub(crate) fn freeze(&mut self) -> FrozenUnionFind<Key> {
-        // Just resolve each key to its actual root.
-        let roots = self.table.indices().map(|key| self.find(key)).collect();
-        FrozenUnionFind { roots }
-    }
-}
-
-/// Snapshot of a disjoint-set forest that can no longer be mutated, but can be
-/// queried in O(1) time without mutation.
-///
-/// This is really just a wrapper around a direct mapping from keys to roots,
-/// but with a [`Self::find`] method that resembles [`UnionFind::find`].
-#[derive(Debug)]
-pub(crate) struct FrozenUnionFind<Key: Idx> {
-    roots: IndexVec<Key, Key>,
-}
-
-impl<Key: Idx> FrozenUnionFind<Key> {
-    /// Returns the "root" key of the disjoint-set containing the given key.
-    /// If two keys have the same root, they belong to the same set.
-    pub(crate) fn find(&self, key: Key) -> Key {
-        self.roots[key]
+    /// Takes a "snapshot" of the current state of this disjoint-set forest, in
+    /// the form of a vector that directly maps each key to its current root.
+    pub(crate) fn snapshot(&mut self) -> IndexVec<Key, Key> {
+        self.table.indices().map(|key| self.find(key)).collect()
     }
 }