Move union find to utils

hercules_opt/src/array_to_prod.rs

@@ -6,7 +6,6 @@ use bitvec::prelude::*;
 use crate::*;
 
 use std::collections::{HashMap, HashSet};
-use std::marker::PhantomData;
 
 /*
  * Top level function for array to product which will convert constant
@@ -200,7 +199,7 @@ fn array_usage_analysis(
     // that returns a collection (that in reference semantics returns the same reference as some of
     // its inputs) union with all of its users. The nodes that matter in this are arguments,
     // constants, writes, phis, selects, and reduces with array types.
-    let mut analysis = UnionFind::new();
+    let mut analysis = UnionFindVec::new();
     for node_idx in 0..num_nodes {
         let node_id = NodeID::new(node_idx);
         if editor.get_type(types[node_idx]).is_array() {
@@ -215,7 +214,7 @@ fn array_usage_analysis(
                 }
                 | Node::Write { .. } => {
                     for user in editor.get_users(node_id) {
-                        analysis.union(node_id, user);
+                        analysis.union(&node_id, &user);
                     }
                 }
                 _ => {}
@@ -266,87 +265,3 @@ fn can_replace(editor: &FunctionEditor, node: NodeID) -> bool {
         _ => false,
     }
 }
-
-define_id_type!(SetID);
-
-#[derive(Clone, Debug)]
-struct UnionFindNode {
-    parent: SetID,
-    rank: usize,
-}
-
-#[derive(Clone, Debug)]
-struct UnionFind<T> {
-    sets: Vec<UnionFindNode>,
-    _phantom: PhantomData<T>,
-}
-
-impl<T: ID> UnionFind<T> {
-    pub fn new() -> Self {
-        UnionFind {
-            sets: vec![],
-            _phantom: PhantomData,
-        }
-    }
-
-    fn extend_past(&mut self, size: usize) {
-        for i in self.sets.len()..=size {
-            // The new nodes we add are in their own sets and have rank 0
-            self.sets.push(UnionFindNode {
-                parent: SetID::new(i),
-                rank: 0,
-            });
-        }
-    }
-
-    pub fn find(&mut self, x: T) -> SetID {
-        self.extend_past(x.idx());
-        self.find_set(x.idx())
-    }
-
-    fn find_set(&mut self, x: usize) -> SetID {
-        let mut parent = self.sets[x].parent;
-        if parent.idx() != x {
-            parent = self.find_set(parent.idx());
-            self.sets[x].parent = parent;
-        }
-        parent
-    }
-
-    pub fn union(&mut self, x: T, y: T) {
-        let x = self.find(x);
-        let y = self.find(y);
-        self.link(x, y);
-    }
-
-    fn link(&mut self, x: SetID, y: SetID) {
-        if self.sets[x.idx()].rank > self.sets[y.idx()].rank {
-            self.sets[y.idx()].parent = x;
-        } else {
-            self.sets[x.idx()].parent = y;
-            if self.sets[x.idx()].rank == self.sets[y.idx()].rank {
-                self.sets[y.idx()].rank += 1;
-            }
-        }
-    }
-
-    pub fn sets(&mut self, keys: &[T]) -> Vec<Vec<T>> {
-        let key_index = keys
-            .iter()
-            .enumerate()
-            .map(|(i, k)| (self.find(*k), i))
-            .collect::<HashMap<SetID, usize>>();
-        let mut result = vec![vec![]; keys.len()];
-
-        let num_elements = self.sets.len();
-        for i in 0..num_elements {
-            let set = self.find_set(i);
-            let Some(idx) = key_index.get(&set) else {
-                continue;
-            };
-            result[*idx].push(T::new(i));
-        }
-
-        result
-    }
-}

hercules_opt/src/utils.rs

 use std::collections::{HashMap, HashSet};
+use std::hash::Hash;
+use std::marker::PhantomData;
 
 use nestify::nest;
 
@@ -534,3 +536,161 @@ where
         nodes[fork.idx()].try_fork().unwrap().1.len() == rep_dims.len()
     })
 }
+
+/*
+ * A generic union-find data structure that allows the use of different
+ * underlying data structures used to store the lookup from keys to sets
+ */
+#[derive(Clone, Debug)]
+pub struct UnionFindNode<S> {
+    parent: S,
+    rank: usize,
+}
+
+impl<S> UnionFindNode<S> {
+    pub fn new(set: S) -> Self {
+        UnionFindNode {
+            parent: set,
+            rank: 0,
+        }
+    }
+}
+
+pub trait UnionFind<E, S>
+where
+    S: Eq + Hash + Copy,
+{
+    // Converts the set element to its set identifier
+    fn element_set(&mut self, e: &E) -> S;
+    // Given a set number retrieves the UnionFindNode for it
+    fn get_set(&mut self, s: S) -> &mut UnionFindNode<S>;
+    // Get a list of all the elements that have sets defined
+    fn get_elements(&self) -> Vec<E>;
+
+    fn find(&mut self, e: &E) -> S {
+        let s = self.element_set(e);
+        self.find_set(s)
+    }
+
+    fn find_set(&mut self, s: S) -> S {
+        let mut parent = self.get_set(s).parent;
+        if parent != s {
+            parent = self.find_set(parent);
+            self.get_set(s).parent = parent;
+        }
+        parent
+    }
+
+    fn union(&mut self, e: &E, f: &E) {
+        let e = self.element_set(e);
+        let f = self.element_set(f);
+        self.link(e, f);
+    }
+
+    fn link(&mut self, x: S, y: S) {
+        if x != y {
+            let rank_x = self.get_set(x).rank;
+            let rank_y = self.get_set(y).rank;
+            if rank_x > rank_y {
+                self.get_set(y).parent = x;
+            } else {
+                self.get_set(x).parent = y;
+                if rank_x == rank_y {
+                    self.get_set(y).rank += 1;
+                }
+            }
+        }
+    }
+
+    fn sets(&mut self, elems: &[E]) -> Vec<Vec<E>> {
+        let elem_index = elems
+            .iter()
+            .enumerate()
+            .map(|(i, e)| (self.find(e), i))
+            .collect::<HashMap<S, usize>>();
+        let mut result : Vec<Vec<E>> = vec![];
+        result.resize_with(elems.len(), Default::default);
+
+        for elem in self.get_elements() {
+            let set = self.find(&elem);
+            let Some(idx) = elem_index.get(&set) else {
+                continue;
+            };
+            result[*idx].push(elem);
+        }
+
+        result
+    }
+}
+
+// Two union find implementations, the first with the underlying storage being
+// a Vec and the elements being IDs and the second with the underlying
+// storage being a HashSet (this is for if the sets we care about are sparse)
+define_id_type!(SetID);
+
+#[derive(Clone, Debug)]
+pub struct UnionFindVec<T> {
+    sets: Vec<UnionFindNode<SetID>>,
+    _phantom: PhantomData<T>,
+}
+
+impl<T> UnionFindVec<T> {
+    pub fn new() -> Self {
+        UnionFindVec {
+            sets: vec![],
+            _phantom: PhantomData,
+        }
+    }
+}
+
+impl<T: ID> UnionFind<T, SetID> for UnionFindVec<T> {
+    fn element_set(&mut self, e: &T) -> SetID {
+        SetID::new(e.idx())
+    }
+
+    fn get_set(&mut self, s: SetID) -> &mut UnionFindNode<SetID> {
+        let id = s.idx();
+        for i in self.sets.len()..=id {
+            self.sets.push(UnionFindNode::new(SetID::new(i)));
+        }
+        &mut self.sets[id]
+    }
+
+    fn get_elements(&self) -> Vec<T> {
+        (0..self.sets.len()).map(|i| T::new(i)).collect()
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct UnionFindSparse<T: Hash> {
+    elems: HashMap<T, SetID>,
+    sets:  Vec<UnionFindNode<SetID>>,
+}
+
+impl<T: Eq + Hash> UnionFindSparse<T> {
+    pub fn new() -> Self {
+        UnionFindSparse {
+            elems: HashMap::new(),
+            sets: vec![],
+        }
+    }
+}
+
+impl<T: Eq + Hash + Clone> UnionFind<T, SetID> for UnionFindSparse<T> {
+    fn element_set(&mut self, e: &T) -> SetID {
+        *self.elems.entry(e.clone())
+            .or_insert_with(|| {
+                let set = SetID::new(self.sets.len());
+                self.sets.push(UnionFindNode::new(set));
+                set
+            })
+    }
+
+    fn get_set(&mut self, s: SetID) -> &mut UnionFindNode<SetID> {
+        &mut self.sets[s.idx()]
+    }
+
+    fn get_elements(&self) -> Vec<T> {
+        self.elems.iter().map(|(t, _)| t.clone()).collect()
+    }
+}