| //! Simple union-find data structure. |
| |
| use crate::trace; |
| use cranelift_entity::{packed_option::ReservedValue, EntityRef, SecondaryMap}; |
| use std::hash::Hash; |
| |
| /// A union-find data structure. The data structure can allocate |
| /// `Id`s, indicating eclasses, and can merge eclasses together. |
| #[derive(Clone, Debug, PartialEq)] |
| pub struct UnionFind<Idx: EntityRef> { |
| parent: SecondaryMap<Idx, Val<Idx>>, |
| } |
| |
| #[derive(Clone, Debug, PartialEq)] |
| struct Val<Idx>(Idx); |
| impl<Idx: EntityRef + ReservedValue> Default for Val<Idx> { |
| fn default() -> Self { |
| Self(Idx::reserved_value()) |
| } |
| } |
| |
| impl<Idx: EntityRef + Hash + std::fmt::Display + Ord + ReservedValue> UnionFind<Idx> { |
| /// Create a new `UnionFind` with the given capacity. |
| pub fn with_capacity(cap: usize) -> Self { |
| UnionFind { |
| parent: SecondaryMap::with_capacity(cap), |
| } |
| } |
| |
| /// Add an `Idx` to the `UnionFind`, with its own equivalence class |
| /// initially. All `Idx`s must be added before being queried or |
| /// unioned. |
| pub fn add(&mut self, id: Idx) { |
| debug_assert!(id != Idx::reserved_value()); |
| self.parent[id] = Val(id); |
| } |
| |
| /// Find the canonical `Idx` of a given `Idx`. |
| pub fn find(&self, mut node: Idx) -> Idx { |
| while node != self.parent[node].0 { |
| node = self.parent[node].0; |
| } |
| node |
| } |
| |
| /// Find the canonical `Idx` of a given `Idx`, updating the data |
| /// structure in the process so that future queries for this `Idx` |
| /// (and others in its chain up to the root of the equivalence |
| /// class) will be faster. |
| pub fn find_and_update(&mut self, mut node: Idx) -> Idx { |
| // "Path halving" mutating find (Tarjan and Van Leeuwen). |
| debug_assert!(node != Idx::reserved_value()); |
| while node != self.parent[node].0 { |
| let next = self.parent[self.parent[node].0].0; |
| debug_assert!(next != Idx::reserved_value()); |
| self.parent[node] = Val(next); |
| node = next; |
| } |
| debug_assert!(node != Idx::reserved_value()); |
| node |
| } |
| |
| /// Merge the equivalence classes of the two `Idx`s. |
| pub fn union(&mut self, a: Idx, b: Idx) { |
| let a = self.find_and_update(a); |
| let b = self.find_and_update(b); |
| let (a, b) = (std::cmp::min(a, b), std::cmp::max(a, b)); |
| if a != b { |
| // Always canonicalize toward lower IDs. |
| self.parent[b] = Val(a); |
| trace!("union: {}, {}", a, b); |
| } |
| } |
| } |