Revert "Move union find to utils"

This reverts commit 3d8e6e8a.

hercules_opt/src/array_to_prod.rs

@@ -6,6 +6,7 @@ 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
@@ -199,7 +200,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 = UnionFindVec::new();
+    let mut analysis = UnionFind::new();
     for node_idx in 0..num_nodes {
         let node_id = NodeID::new(node_idx);
         if editor.get_type(types[node_idx]).is_array() {
@@ -214,7 +215,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);
                     }
                 }
                 _ => {}
@@ -265,3 +266,87 @@ 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;
 
@@ -569,161 +567,3 @@ pub fn is_one(editor: &FunctionEditor, id: NodeID) -> bool {
             .unwrap_or(false)
         || nodes[id.idx()].is_undef()
 }
-
-/*
- * 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()
-    }
-}