BLI: use templates for disjoint set data structure

Differential Revision: https://developer.blender.org/D16472

3 files changed, 26 insertions, 24 deletions

diff --git a/source/blender/blenlib/BLI_disjoint_set.hh b/source/blender/blenlib/BLI_disjoint_set.hh
index d135aa90307..3adc609dad3 100644
--- a/source/blender/blenlib/BLI_disjoint_set.hh
+++ b/source/blender/blenlib/BLI_disjoint_set.hh
@@ -9,23 +9,24 @@
  */
 
 #include "BLI_array.hh"
+#include "BLI_index_range.hh"
 
 namespace blender {
 
-class DisjointSet {
+template<typename T = int64_t> class DisjointSet {
  private:
-  Array<int64_t> parents_;
-  Array<int64_t> ranks_;
+  Array<T> parents_;
+  Array<T> ranks_;
 
  public:
   /**
    * Create a new disjoint set with the given size. Initially, every element is in a separate set.
    */
-  DisjointSet(int64_t size) : parents_(size), ranks_(size, 0)
+  DisjointSet(const int64_t size) : parents_(size), ranks_(size, 0)
   {
     BLI_assert(size >= 0);
-    for (int64_t i = 0; i < size; i++) {
-      parents_[i] = i;
+    for (const int64_t i : IndexRange(size)) {
+      parents_[i] = T(i);
     }
   }
 
@@ -33,10 +34,10 @@ class DisjointSet {
    * Join the sets containing elements x and y. Nothing happens when they have been in the same set
    * before.
    */
-  void join(int64_t x, int64_t y)
+  void join(const T x, const T y)
   {
-    int64_t root1 = this->find_root(x);
-    int64_t root2 = this->find_root(y);
+    T root1 = this->find_root(x);
+    T root2 = this->find_root(y);
 
     /* x and y are in the same set already. */
     if (root1 == root2) {
@@ -57,29 +58,30 @@ class DisjointSet {
   /**
    * Return true when x and y are in the same set.
    */
-  bool in_same_set(int64_t x, int64_t y)
+  bool in_same_set(const T x, const T y)
   {
-    int64_t root1 = this->find_root(x);
-    int64_t root2 = this->find_root(y);
+    T root1 = this->find_root(x);
+    T root2 = this->find_root(y);
     return root1 == root2;
   }
 
   /**
    * Find the element that represents the set containing x currently.
    */
-  int64_t find_root(int64_t x)
+  T find_root(const T x)
   {
     /* Find root by following parents. */
-    int64_t root = x;
+    T root = x;
     while (parents_[root] != root) {
       root = parents_[root];
     }
 
     /* Compress path. */
-    while (parents_[x] != root) {
-      int64_t parent = parents_[x];
-      parents_[x] = root;
-      x = parent;
+    T to_root = x;
+    while (parents_[to_root] != root) {
+      const T parent = parents_[to_root];
+      parents_[to_root] = root;
+      to_root = parent;
     }
 
     return root;
diff --git a/source/blender/nodes/geometry/nodes/node_geo_input_mesh_island.cc b/source/blender/nodes/geometry/nodes/node_geo_input_mesh_island.cc
index 6b54828b042..d658b14092a 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_input_mesh_island.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_input_mesh_island.cc
@@ -35,7 +35,7 @@ class IslandFieldInput final : public bke::MeshFieldInput {
   {
     const Span<MEdge> edges = mesh.edges();
 
-    DisjointSet islands(mesh.totvert);
+    DisjointSet<int> islands(mesh.totvert);
     for (const int i : edges.index_range()) {
       islands.join(edges[i].v1, edges[i].v2);
     }
@@ -43,7 +43,7 @@ class IslandFieldInput final : public bke::MeshFieldInput {
     Array<int> output(mesh.totvert);
     VectorSet<int> ordered_roots;
     for (const int i : IndexRange(mesh.totvert)) {
-      const int64_t root = islands.find_root(i);
+      const int root = islands.find_root(i);
       output[i] = ordered_roots.index_of_or_add(root);
     }
 
@@ -81,14 +81,14 @@ class IslandCountFieldInput final : public bke::MeshFieldInput {
   {
     const Span<MEdge> edges = mesh.edges();
 
-    DisjointSet islands(mesh.totvert);
+    DisjointSet<int> islands(mesh.totvert);
     for (const int i : edges.index_range()) {
       islands.join(edges[i].v1, edges[i].v2);
     }
 
     Set<int> island_list;
     for (const int i_vert : IndexRange(mesh.totvert)) {
-      const int64_t root = islands.find_root(i_vert);
+      const int root = islands.find_root(i_vert);
       island_list.add(root);
     }
 
diff --git a/source/blender/nodes/geometry/nodes/node_geo_scale_elements.cc b/source/blender/nodes/geometry/nodes/node_geo_scale_elements.cc
index 5f1baa23511..3cfdafa76c0 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_scale_elements.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_scale_elements.cc
@@ -248,7 +248,7 @@ static Vector<ElementIsland> prepare_face_islands(const Mesh &mesh, const IndexM
   const Span<MLoop> loops = mesh.loops();
 
   /* Use the disjoint set data structure to determine which vertices have to be scaled together. */
-  DisjointSet disjoint_set(mesh.totvert);
+  DisjointSet<int> disjoint_set(mesh.totvert);
   for (const int poly_index : face_selection) {
     const MPoly &poly = polys[poly_index];
     const Span<MLoop> poly_loops = loops.slice(poly.loopstart, poly.totloop);
@@ -344,7 +344,7 @@ static Vector<ElementIsland> prepare_edge_islands(const Mesh &mesh, const IndexM
   const Span<MEdge> edges = mesh.edges();
 
   /* Use the disjoint set data structure to determine which vertices have to be scaled together. */
-  DisjointSet disjoint_set(mesh.totvert);
+  DisjointSet<int> disjoint_set(mesh.totvert);
   for (const int edge_index : edge_selection) {
     const MEdge &edge = edges[edge_index];
     disjoint_set.join(edge.v1, edge.v2);