A better colored union find for parallel runs

src/colors.rs

@@ -7,6 +7,7 @@ use itertools::Itertools;
 use std::fmt::Formatter;
 use indexmap::{IndexMap, IndexSet};
 use crate::colored_union_find::ColoredUnionFind;
+use crate::unionfind::{SimpleUnionFind, UnionFindWrapper};
 
 global_counter!(COLOR_IDS, usize, usize::default());
 
@@ -19,7 +20,7 @@ pub struct Color<L: Language, N: Analysis<L>> {
     pub(crate) equality_classes: IndexMap<Id, IndexSet<Id>>,
     /// Used to implement a union find. Opposite function of `equality_classes`.
     /// Supports removal of elements when they are not needed.
-    union_find: ColoredUnionFind,
+    union_find: UnionFindWrapper<(), Id>,
     /// Used to determine for each a colored equality class what is the black colored class.
     /// Relevant when a colored edge was added.
     pub(crate) black_colored_classes: IndexMap<Id, Id>,
@@ -74,7 +75,7 @@ impl<L: Language, N: Analysis<L>> Color<L, N> {
     pub(crate) fn verify_uf_minimal(&self, egraph: &EGraph<L, N>) {
         let mut parents: IndexMap<Id, usize> = IndexMap::default();
         for (k, _v) in self.union_find.iter() {
-            let v = self.find(egraph, k);
+            let v = self.find(egraph, *k);
             *parents.entry(v).or_default() += 1;
         }
         for (k, v) in parents {
@@ -153,8 +154,8 @@ impl<L: Language, N: Analysis<L>> Color<L, N> {
             // This part only needs to happen if one of the two is in the union find.
             let orig_to = orig_to.unwrap_or(base_to);
             let orig_from = orig_from.unwrap_or(base_from);
-            self.union_find.insert(orig_to);
-            self.union_find.insert(orig_from);
+            self.union_find.insert(orig_to, ());
+            self.union_find.insert(orig_from, ());
             self.union_find.union(&orig_to, &orig_from).unwrap()
         } else {
             (base_to, base_from)
@@ -195,8 +196,8 @@ impl<L: Language, N: Analysis<L>> Color<L, N> {
     // Assumed id1 and id2 are parent canonized
     pub(crate) fn inner_colored_union(&mut self, id1: Id, id2: Id) -> (Id, Id, bool, Vec<(Id, Id)>) {
         // Parent classes will be updated in black union to come.
-        self.union_find.insert(id1);
-        self.union_find.insert(id2);
+        self.union_find.insert(id1, ());
+        self.union_find.insert(id2, ());
         let (to, from) = self.union_find.union(&id1, &id2).unwrap();
         let changed = to != from;
         let g_todo = self.update_black_classes(to, from).into_iter().collect_vec();

src/egraph.rs

@@ -19,13 +19,14 @@ use crate::Pattern;
 use crate::RecExpr;
 use crate::Searcher;
 use crate::Subst;
-use crate::UnionFind;
+use crate::SimpleUnionFind;
 use crate::{OpId, SymbolLang};
 
 pub use crate::colors::{Color, ColorId};
 use itertools::Itertools;
 use multimap::MultiMap;
 use serde::{Deserialize, Serialize};
+use crate::unionfind::UnionFind;
 use crate::util::UniqueQueue;
 
 /** A data structure to keep track of equalities between expressions.
@@ -195,7 +196,7 @@ pub struct EGraph<L: Language, N: Analysis<L>> {
     /// The `Analysis` given when creating this `EGraph`.
     pub analysis: N,
     pub(crate) memo: IndexMap<L, Id>,
-    unionfind: UnionFind,
+    unionfind: SimpleUnionFind,
     classes: SparseVec<EClass<L, N::Data>>,
     /// Nodes which need to be processed for rebuilding. The `Id` is the `Id` of the enode,
     /// not the canonical id of the eclass.
@@ -248,7 +249,7 @@ impl<L: Language, N: Analysis<L>> EGraph<L, N> {
 
 impl<L: Language, N: Analysis<L>> EGraph<L, N> {
     pub(crate) fn inner_new(
-        uf: UnionFind,
+        uf: SimpleUnionFind,
         classes: Vec<Option<Box<EClass<SymbolLang, ()>>>>,
         memo: IndexMap<SymbolLang, Id>,
     ) -> EGraph<SymbolLang, ()> {

src/lib.rs

@@ -108,7 +108,7 @@ impl std::fmt::Display for ColorId {
     }
 }
 
-pub(crate) use unionfind::UnionFind;
+pub(crate) use unionfind::SimpleUnionFind;
 
 pub use {
     dot::Dot,

src/ser.rs

@@ -2,7 +2,7 @@ use std::io::{Read, Write, BufReader, Result, BufRead};
 use indexmap::IndexMap;
 use itertools::Itertools;
 use crate::{EGraph, Analysis, Language, Id, EClass, SymbolLang, ColorId};
-use crate::unionfind::UnionFind;
+use crate::unionfind::SimpleUnionFind;
 
 /// A trait for EGraphs that can be serialized.
 pub trait Serialization {
@@ -149,7 +149,7 @@ impl Deserialization for EGraph<SymbolLang, ()> {
 
 #[derive(Debug, Clone, Default)]
 struct EGraphBuilder {
-    unionfind: UnionFind,
+    unionfind: SimpleUnionFind,
     classes: Vec<Option<Box<EClass<SymbolLang, ()>>>>,
     memo: IndexMap<SymbolLang, Id>,
     _palette: ColorPalette,