Benchmark and expose disjoint-set islands

Signed-off-by: teerthsharma <teerths57@gmail.com>

src/engine/engine_island.c

@@ -84,7 +84,7 @@ static int arenaAllocIsland(const mjModel* m, mjData* d) {
 //-------------------------- flood-fill and graph construction  ------------------------------------
 
 // find the canonical root of an active tree and compress its path
-static int dsuFind(int* parent, int tree) {
+static inline int dsuFind(int* parent, int tree) {
   int root = tree;
   while (parent[root] != root) {
     root = parent[root];
@@ -101,13 +101,13 @@ static int dsuFind(int* parent, int tree) {
 
 
 // initialize all trees as inactive
-static void dsuInit(int* parent, int ntree) {
+static inline void dsuInit(int* parent, int ntree) {
   mju_fillInt(parent, -1, ntree);
 }
 
 
 // activate and union two incident trees; -1 denotes a static endpoint
-static void dsuUnion(int* parent, int tree1, int tree2) {
+static inline void dsuUnion(int* parent, int tree1, int tree2) {
   if (tree1 == -1 && tree2 == -1) {
     mjERROR("self-incidence of the static tree");  // SHOULD NOT OCCUR
     return;
@@ -119,6 +119,8 @@ static void dsuUnion(int* parent, int tree1, int tree2) {
   if (parent[tree1] == -1) parent[tree1] = tree1;
   if (parent[tree2] == -1) parent[tree2] = tree2;
 
+  if (parent[tree1] == parent[tree2]) return;
+
   int root1 = dsuFind(parent, tree1);
   int root2 = dsuFind(parent, tree2);
   if (root1 < root2) {
@@ -130,7 +132,8 @@ static void dsuUnion(int* parent, int tree1, int tree2) {
 
 
 // assign deterministic island ids in ascending canonical-root order
-static int dsuAssign(int* island, int* parent, const int* tree_dofnum, int ntree, int* nidof) {
+static inline int dsuAssign(int* island, int* parent, const int* tree_dofnum, int ntree,
+                            int* nidof) {
   int nisland = 0;
   *nidof = 0;
   for (int tree=0; tree < ntree; tree++) {
@@ -139,14 +142,41 @@ static int dsuAssign(int* island, int* parent, const int* tree_dofnum, int ntree
       continue;
     }
 
-    int root = dsuFind(parent, tree);
-    island[tree] = root == tree ? nisland++ : island[root];
+    if (parent[tree] == tree) {
+      island[tree] = nisland++;
+    } else {
+      // Union always links the larger root to the smaller root. Since trees are visited in
+      // ascending order, this predecessor has already been compressed and assigned an island.
+      parent[tree] = parent[parent[tree]];
+      island[tree] = island[parent[tree]];
+    }
     *nidof += tree_dofnum[tree];
   }
 
   return nisland;
 }
 
+
+// exported private wrappers for direct unit tests and benchmarks
+int _mjPRIVATE_dsuFind(int* parent, int tree) {
+  return dsuFind(parent, tree);
+}
+
+
+void _mjPRIVATE_dsuInit(int* parent, int ntree) {
+  dsuInit(parent, ntree);
+}
+
+
+void _mjPRIVATE_dsuUnion(int* parent, int tree1, int tree2) {
+  dsuUnion(parent, tree1, tree2);
+}
+
+
+int _mjPRIVATE_dsuAssign(int* island, int* parent, const int* tree_dofnum, int ntree, int* nidof) {
+  return dsuAssign(island, parent, tree_dofnum, ntree, nidof);
+}
+
 // find disjoint subgraphs ("islands") given sparse symmetric adjacency matrix
 //   arguments:
 //     island  (nr)   - island index assigned to vertex, -1 if vertex has no edges
@@ -353,7 +383,7 @@ static int isFlexEquality(const mjModel* m, int efc_type, int efc_id) {
 
 
 // activate and union all trees with direct incidence in a constraint
-static void unionConstraintTrees(const mjModel* m, const mjData* d, int* parent) {
+static void unionConstraintTrees(const mjModel* m, const mjData* d, int* parent, int* efc_tree) {
   int nefc = d->nefc;
   int efc_type = -1;
   int efc_id = -1;
@@ -361,9 +391,10 @@ static void unionConstraintTrees(const mjModel* m, const mjData* d, int* parent)
   // iterate over constraints and union incident trees
   for (int i=0; i < nefc; i++) {
     // row i is still in the same constraint: skip it
-    if (efc_type == d->efc_type[i] && efc_id == d->efc_id[i]) {
+    if (i > 0 && efc_type == d->efc_type[i] && efc_id == d->efc_id[i]) {
       // unless it is a flex equality, where the tree pattern changes per dof
       if (!isFlexEquality(m, efc_type, efc_id)) {
+        efc_tree[i] = efc_tree[i-1];
         continue;
       }
     }
@@ -378,6 +409,7 @@ static void unionConstraintTrees(const mjModel* m, const mjData* d, int* parent)
     int tree1 = treeNext(m, d, i, &iter);
     if (tree1 != -2) {
       int tree2 = treeNext(m, d, i, &iter);
+      efc_tree[i] = tree1 == -1 ? tree2 : tree1;
 
       // activate a singleton or union all trees in a multi-tree constraint
       if (tree2 == -2) {
@@ -393,35 +425,41 @@ static void unionConstraintTrees(const mjModel* m, const mjData* d, int* parent)
       mjERROR("no tree found for constraint %d", i);  // SHOULD NOT OCCUR
     }
   }
-}
 
-// assign each constraint from its first non-negative incident tree
-static void assignConstraintIslands(const mjModel* m, mjData* d, const int* tree_island) {
-  int efc_type = -1;
-  int efc_id = -1;
-
-  for (int i=0; i < d->nefc; i++) {
-    // reuse assignment for repeated scalar rows, except flex equality rows
-    if (efc_type == d->efc_type[i] && efc_id == d->efc_id[i] &&
-        !isFlexEquality(m, efc_type, efc_id)) {
-      d->efc_island[i] = d->efc_island[i-1];
+  // Flex stiffness couples all vertices (nodes for interpolated flexes) without a constraint
+  // row representing the coupling. Union the awake dynamic trees of each stiffness-active flex.
+  for (int f=0; f < m->nflex; f++) {
+    if (m->flex_rigid[f] || m->flex_dim[f] < 2) {
+      continue;
+    }
+    int sadr = m->flex_stiffnessadr[f];
+    if (m->flex_bendingadr[f] < 0 && (sadr < 0 || m->flex_stiffness[sadr] == 0)) {
       continue;
     }
-    efc_type = d->efc_type[i];
-    efc_id = d->efc_id[i];
 
-    mjTreeIter iter;
-    treeIterInit(m, d, i, &iter);
-
-    int tree;
-    do {
-      tree = treeNext(m, d, i, &iter);
-    } while (tree == -1);
-
-    if (tree == -2) {
-      mjERROR("no dynamic tree found for constraint %d", i);  // SHOULD NOT OCCUR
+    int num, adr;
+    const int* bodyid;
+    if (m->flex_interp[f]) {
+      num = m->flex_nodenum[f];
+      adr = m->flex_nodeadr[f];
+      bodyid = m->flex_nodebodyid;
     } else {
-      d->efc_island[i] = tree_island[tree];
+      num = m->flex_vertnum[f];
+      adr = m->flex_vertadr[f];
+      bodyid = m->flex_vertbodyid;
+    }
+
+    int tree1 = -1;
+    for (int j=0; j < num; j++) {
+      int tree2 = m->body_treeid[bodyid[adr+j]];
+      if (tree2 < 0 || tree2 == tree1 || !d->tree_awake[tree2]) {
+        continue;
+      }
+      if (tree1 < 0) {
+        tree1 = tree2;
+      } else {
+        dsuUnion(parent, tree1, tree2);
+      }
     }
   }
 }
@@ -443,9 +481,10 @@ void mj_island(const mjModel* m, mjData* d) {
   mj_markStack(d);
 
   // union direct tree incidence and assign deterministic components
+  int* efc_tree = mjSTACKALLOC(d, nefc, int);
   int* parent = mjSTACKALLOC(d, ntree, int);
   dsuInit(parent, ntree);
-  unionConstraintTrees(m, d, parent);
+  unionConstraintTrees(m, d, parent, efc_tree);
   int* tree_island = mjSTACKALLOC(d, ntree, int);
   int nidof;
   d->nisland = dsuAssign(tree_island, parent, m->tree_dofnum, ntree, &nidof);
@@ -558,14 +597,12 @@ void mj_island(const mjModel* m, mjData* d) {
 
   // ------------------------------------- constraints ---------------------------------------------
 
-  // compute efc_island from first non-negative tree of each constraint
-  assignConstraintIslands(m, d, tree_island);
-
   // compute efc_island, island_{ne,nf,nefc}
   mju_zeroInt(d->island_ne, nisland);
   mju_zeroInt(d->island_nf, nisland);
   mju_zeroInt(d->island_nefc, nisland);
   for (int i=0; i < nefc; i++) {
+    d->efc_island[i] = tree_island[efc_tree[i]];
     int island = d->efc_island[i];
     d->island_nefc[island]++;
     switch (d->efc_type[i]) {

src/engine/engine_island.h

@@ -23,6 +23,12 @@
 extern "C" {
 #endif
 
+MJAPI int _mjPRIVATE_dsuFind(int* parent, int tree);
+MJAPI void _mjPRIVATE_dsuInit(int* parent, int ntree);
+MJAPI void _mjPRIVATE_dsuUnion(int* parent, int tree1, int tree2);
+MJAPI int _mjPRIVATE_dsuAssign(int* island, int* parent,
+                               const int* tree_dofnum, int ntree, int* nidof);
+
 
 // find disjoint subgraphs ("islands") given sparse symmetric adjacency matrix
 MJAPI int mj_floodFill(int* island, int nr, const int* rownnz, const int* rowadr, const int* colind,

test/benchmark/CMakeLists.txt

@@ -79,6 +79,12 @@ mujoco_test(
   ADDITIONAL_LINK_LIBRARIES benchmark::benchmark absl::core_headers
 )
 
+mujoco_test(
+  island_benchmark_test
+  MAIN_TARGET benchmark::benchmark_main
+  ADDITIONAL_LINK_LIBRARIES benchmark::benchmark absl::core_headers
+)
+
 mujoco_test(
   engine_util_sparse_benchmark_test
   MAIN_TARGET benchmark::benchmark_main

test/benchmark/island_benchmark_test.cc

@@ -0,0 +1,374 @@
+// Copyright 2026 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Benchmarks island discovery on deterministic geodesic Rips graphs.  The
+// corpus spans the connectivity transition of points sampled on S^2 and adds
+// MuJoCo-relevant static and repeated incidences.  Corpus construction and
+// validation are deliberately outside the timed region.
+
+#include <benchmark/benchmark.h>
+
+#include <algorithm>
+#include <array>
+#include <cmath>
+#include <cstdint>
+#include <limits>
+#include <queue>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "src/engine/engine_island.h"
+
+namespace mujoco {
+namespace {
+
+struct Edge {
+  int first;
+  int second;
+};
+
+struct GraphCase {
+  std::string name;
+  int node_count;
+  std::vector<Edge> incidences;
+  std::vector<int> expected_partition;
+  int active_nodes;
+  int expected_components;
+  int pre_bridge_components;
+  bool bridge_added;
+  std::uint64_t expected_checksum;
+};
+
+struct Point {
+  double x;
+  double y;
+  double z;
+};
+
+std::uint64_t SplitMix64(std::uint64_t& state) {
+  state += 0x9e3779b97f4a7c15ULL;
+  std::uint64_t value = state;
+  value = (value ^ (value >> 30)) * 0xbf58476d1ce4e5b9ULL;
+  value = (value ^ (value >> 27)) * 0x94d049bb133111ebULL;
+  return value ^ (value >> 31);
+}
+
+double Uniform01(std::uint64_t& state) {
+  return static_cast<double>(SplitMix64(state) >> 11) * 0x1.0p-53;
+}
+
+std::vector<Point> SampleSphere(int count, std::uint64_t seed) {
+  constexpr double kTwoPi = 6.283185307179586476925286766559;
+  std::vector<Point> points;
+  points.reserve(count);
+  for (int i = 0; i < count; ++i) {
+    const double z = 2.0 * Uniform01(seed) - 1.0;
+    const double angle = kTwoPi * Uniform01(seed);
+    const double radial = std::sqrt(std::max(0.0, 1.0 - z*z));
+    points.push_back({radial * std::cos(angle), radial * std::sin(angle), z});
+  }
+  return points;
+}
+
+double Dot(const Point& a, const Point& b) {
+  return a.x*b.x + a.y*b.y + a.z*b.z;
+}
+
+std::vector<Edge> RipsEdges(const std::vector<Point>& points, double target_degree) {
+  const double probability = target_degree / (points.size() - 1);
+  const double radius = 2.0 * std::asin(std::sqrt(probability));
+  const double minimum_dot = std::cos(radius);
+  std::vector<Edge> edges;
+  for (int i = 0; i < static_cast<int>(points.size()); ++i) {
+    for (int j = i + 1; j < static_cast<int>(points.size()); ++j) {
+      if (Dot(points[i], points[j]) >= minimum_dot) {
+        edges.push_back({i, j});
+      }
+    }
+  }
+  return edges;
+}
+
+std::vector<int> CanonicalPartition(int node_count, const std::vector<Edge>& edges) {
+  std::vector<std::vector<int>> adjacency(node_count);
+  std::vector<bool> active(node_count, false);
+  for (const Edge& edge : edges) {
+    active[edge.first] = true;
+    active[edge.second] = true;
+    if (edge.first != edge.second) {
+      adjacency[edge.first].push_back(edge.second);
+      adjacency[edge.second].push_back(edge.first);
+    }
+  }
+
+  std::vector<int> partition(node_count, -1);
+  std::queue<int> pending;
+  for (int start = 0; start < node_count; ++start) {
+    if (!active[start] || partition[start] != -1) {
+      continue;
+    }
+    partition[start] = start;
+    pending.push(start);
+    while (!pending.empty()) {
+      const int node = pending.front();
+      pending.pop();
+      for (int neighbor : adjacency[node]) {
+        if (partition[neighbor] == -1) {
+          partition[neighbor] = start;
+          pending.push(neighbor);
+        }
+      }
+    }
+  }
+  return partition;
+}
+
+int CountComponents(const std::vector<int>& partition) {
+  int count = 0;
+  for (int node = 0; node < static_cast<int>(partition.size()); ++node) {
+    count += partition[node] == node;
+  }
+  return count;
+}
+
+std::uint64_t PartitionChecksum(const std::vector<int>& partition) {
+  std::uint64_t hash = 1469598103934665603ULL;
+  for (int value : partition) {
+    hash ^= static_cast<std::uint32_t>(value);
+    hash *= 1099511628211ULL;
+  }
+  return hash;
+}
+
+std::vector<int> CanonicalizeLabels(const std::vector<int>& labels) {
+  std::vector<int> minimum(labels.size(), std::numeric_limits<int>::max());
+  for (int node = 0; node < static_cast<int>(labels.size()); ++node) {
+    if (labels[node] >= 0) {
+      minimum[labels[node]] = std::min(minimum[labels[node]], node);
+    }
+  }
+  std::vector<int> canonical(labels.size(), -1);
+  for (int node = 0; node < static_cast<int>(labels.size()); ++node) {
+    if (labels[node] >= 0) {
+      canonical[node] = minimum[labels[node]];
+    }
+  }
+  return canonical;
+}
+
+void DeterministicShuffle(std::vector<Edge>& edges, std::uint64_t seed) {
+  for (std::size_t i = edges.size(); i > 1; --i) {
+    const std::size_t j = SplitMix64(seed) % i;
+    std::swap(edges[i - 1], edges[j]);
+  }
+}
+
+bool AddCriticalBridge(const std::vector<Point>& points, std::vector<Edge>& edges) {
+  const std::vector<int> partition = CanonicalPartition(points.size(), edges);
+  if (CountComponents(partition) < 2) {
+    return false;
+  }
+
+  double best_dot = -2.0;
+  Edge bridge{-1, -1};
+  for (int i = 0; i < static_cast<int>(points.size()); ++i) {
+    for (int j = i + 1; j < static_cast<int>(points.size()); ++j) {
+      if (partition[i] >= 0 && partition[j] >= 0 && partition[i] != partition[j] &&
+          Dot(points[i], points[j]) > best_dot) {
+        best_dot = Dot(points[i], points[j]);
+        bridge = {i, j};
+      }
+    }
+  }
+  if (bridge.first >= 0) {
+    edges.push_back(bridge);
+    return true;
+  }
+  return false;
+}
+
+GraphCase MakeCase(std::string name, int node_count, double target_degree,
+                   std::uint64_t seed, bool critical_bridge, bool static_rows,
+                   bool repeated_rows) {
+  const std::vector<Point> points = SampleSphere(node_count, seed);
+  std::vector<Edge> edges = RipsEdges(points, target_degree);
+  const int pre_bridge_components =
+      critical_bridge ? CountComponents(CanonicalPartition(node_count, edges)) : -1;
+  bool bridge_added = false;
+  if (critical_bridge) {
+    bridge_added = AddCriticalBridge(points, edges);
+  }
+
+  const std::vector<Edge> unique_edges = edges;
+  if (static_rows) {
+    for (int node = 0; node < node_count; node += 17) {
+      edges.push_back({node, node});
+    }
+  }
+  if (repeated_rows) {
+    for (std::size_t i = 0; i < unique_edges.size(); i += 11) {
+      edges.push_back(unique_edges[i]);
+      edges.push_back({unique_edges[i].second, unique_edges[i].first});
+    }
+  }
+  DeterministicShuffle(edges, seed ^ 0xd1b54a32d192ed03ULL);
+
+  std::vector<int> expected = CanonicalPartition(node_count, edges);
+  const int active_nodes = std::count_if(expected.begin(), expected.end(),
+                                         [](int component) { return component >= 0; });
+  const int components = CountComponents(expected);
+  const std::uint64_t checksum = PartitionChecksum(expected);
+  return {std::move(name), node_count, std::move(edges), std::move(expected), active_nodes,
+          components, pre_bridge_components, bridge_added, checksum};
+}
+
+const std::vector<GraphCase>& Corpus() {
+  static const std::vector<GraphCase> corpus = {
+      MakeCase("StableSparse_S2Rips_64", 64, 2.0, 0x33960001ULL, false, false, false),
+      MakeCase("CriticalBridge_S2Rips_256", 256, 0.75 * std::log(256.0),
+               0x33960002ULL, true, false, false),
+      MakeCase("SupercriticalDense_S2Rips_256", 256, 2.0 * std::ceil(std::log(256.0)),
+               0x33960003ULL, false, false, false),
+      MakeCase("GroundedStaticRepeated_S2Rips_256", 256,
+               2.0 * std::ceil(std::log(256.0)), 0x33960004ULL, false, true, true),
+      MakeCase("StableRepeated_S2Rips_1024", 1024, 2.0, 0x33960005ULL,
+               false, false, true),
+      MakeCase("CriticalLarge_S2Rips_1024", 1024, std::ceil(std::log(1024.0)),
+               0x33960006ULL, true, false, false),
+  };
+  return corpus;
+}
+
+struct FloodFillWorkspace {
+  explicit FloodFillWorkspace(int node_count)
+      : adjacency(node_count * node_count), rownnz(node_count), rowadr(node_count),
+        colind(node_count * node_count), stack(node_count * node_count + node_count),
+        island(node_count) {}
+
+  std::vector<unsigned char> adjacency;
+  std::vector<int> rownnz;
+  std::vector<int> rowadr;
+  std::vector<int> colind;
+  std::vector<int> stack;
+  std::vector<int> island;
+};
+
+int RunFloodFill(const GraphCase& graph, FloodFillWorkspace& work) {
+  const int n = graph.node_count;
+  std::fill(work.adjacency.begin(), work.adjacency.end(), 0);
+  std::fill(work.rownnz.begin(), work.rownnz.end(), 0);
+  for (const Edge& edge : graph.incidences) {
+    work.adjacency[edge.first*n + edge.second] = 1;
+    work.adjacency[edge.second*n + edge.first] = 1;
+  }
+
+  int address = 0;
+  for (int row = 0; row < n; ++row) {
+    work.rowadr[row] = address;
+    for (int column = 0; column < n; ++column) {
+      if (work.adjacency[row*n + column]) {
+        work.colind[address++] = column;
+        ++work.rownnz[row];
+      }
+    }
+  }
+  return mj_floodFill(work.island.data(), n, work.rownnz.data(), work.rowadr.data(),
+                      work.colind.data(), work.stack.data());
+}
+
+struct DsuWorkspace {
+  explicit DsuWorkspace(int node_count)
+      : parent(node_count), island(node_count), dof_count(node_count, 1) {}
+
+  std::vector<int> parent;
+  std::vector<int> island;
+  std::vector<int> dof_count;
+};
+
+int RunDsu(const GraphCase& graph, DsuWorkspace& work) {
+  _mjPRIVATE_dsuInit(work.parent.data(), graph.node_count);
+  for (const Edge& edge : graph.incidences) {
+    _mjPRIVATE_dsuUnion(work.parent.data(), edge.first, edge.second);
+  }
+  int dof_count = 0;
+  return _mjPRIVATE_dsuAssign(work.island.data(), work.parent.data(), work.dof_count.data(),
+                              graph.node_count, &dof_count);
+}
+
+bool Validate(const GraphCase& graph) {
+  FloodFillWorkspace flood(graph.node_count);
+  DsuWorkspace dsu(graph.node_count);
+  const int flood_components = RunFloodFill(graph, flood);
+  const int dsu_components = RunDsu(graph, dsu);
+  const bool bridge_valid = graph.pre_bridge_components < 0 ||
+                            (graph.bridge_added &&
+                             graph.pre_bridge_components == graph.expected_components + 1);
+  return bridge_valid && flood_components == graph.expected_components &&
+         dsu_components == graph.expected_components &&
+         CanonicalizeLabels(flood.island) == graph.expected_partition &&
+         CanonicalizeLabels(dsu.island) == graph.expected_partition &&
+         PartitionChecksum(graph.expected_partition) == graph.expected_checksum;
+}
+
+void BM_FloodFill(benchmark::State& state, const GraphCase* graph) {
+  if (!Validate(*graph)) {
+    state.SkipWithError("invalid S2-Rips graph fixture");
+    return;
+  }
+  FloodFillWorkspace work(graph->node_count);
+  state.SetLabel("edges=" + std::to_string(graph->incidences.size()) +
+                 " active=" + std::to_string(graph->active_nodes) +
+                 " components=" + std::to_string(graph->expected_components) +
+                 " pre_bridge=" + std::to_string(graph->pre_bridge_components) +
+                 " bridge_added=" + std::to_string(graph->bridge_added) +
+                 " checksum=" + std::to_string(graph->expected_checksum));
+  for (auto _ : state) {
+    int components = RunFloodFill(*graph, work);
+    benchmark::DoNotOptimize(components);
+    benchmark::ClobberMemory();
+  }
+  state.SetItemsProcessed(state.iterations() * graph->incidences.size());
+}
+
+void BM_Dsu(benchmark::State& state, const GraphCase* graph) {
+  if (!Validate(*graph)) {
+    state.SkipWithError("invalid S2-Rips graph fixture");
+    return;
+  }
+  DsuWorkspace work(graph->node_count);
+  state.SetLabel("edges=" + std::to_string(graph->incidences.size()) +
+                 " active=" + std::to_string(graph->active_nodes) +
+                 " components=" + std::to_string(graph->expected_components) +
+                 " pre_bridge=" + std::to_string(graph->pre_bridge_components) +
+                 " bridge_added=" + std::to_string(graph->bridge_added) +
+                 " checksum=" + std::to_string(graph->expected_checksum));
+  for (auto _ : state) {
+    int components = RunDsu(*graph, work);
+    benchmark::DoNotOptimize(components);
+    benchmark::ClobberMemory();
+  }
+  state.SetItemsProcessed(state.iterations() * graph->incidences.size());
+}
+
+const bool kRegistered = [] {
+  for (const GraphCase& graph : Corpus()) {
+    benchmark::RegisterBenchmark(("Island/FloodFill/" + graph.name).c_str(), BM_FloodFill, &graph);
+    benchmark::RegisterBenchmark(("Island/DSU/" + graph.name).c_str(), BM_Dsu, &graph);
+  }
+  return true;
+}();
+
+}  // namespace
+}  // namespace mujoco

test/engine/engine_island_test.cc

@@ -16,6 +16,8 @@
 
 #include "src/engine/engine_island.h"
 
+#include <array>
+#include <cstdint>
 #include <string>
 #include <vector>
 
@@ -35,6 +37,310 @@ using ::testing::NotNull;
 using ::testing::Pointwise;
 using IslandTest = MujocoTest;
 
+TEST_F(IslandTest, DsuInitHandlesEmptyAndNonemptyRanges) {
+  int parent[] = {8, 6, 7, 5};
+
+  _mjPRIVATE_dsuInit(parent, 0);
+  EXPECT_THAT(parent, ElementsAre(8, 6, 7, 5));
+
+  _mjPRIVATE_dsuInit(parent, 4);
+  EXPECT_THAT(parent, ElementsAre(-1, -1, -1, -1));
+}
+
+TEST_F(IslandTest, DsuFindReturnsCanonicalRootAndCompressesPath) {
+  int parent[] = {0, 0, 1, 2, 3};
+
+  EXPECT_EQ(_mjPRIVATE_dsuFind(parent, 0), 0);
+  EXPECT_EQ(_mjPRIVATE_dsuFind(parent, 4), 0);
+  EXPECT_THAT(parent, ElementsAre(0, 0, 0, 0, 0));
+}
+
+TEST_F(IslandTest, DsuUnionActivatesEndpointsAndUsesMinimumRoot) {
+  int parent[] = {-1, -1, -1, -1, -1, -1};
+
+  _mjPRIVATE_dsuUnion(parent, -1, 4);
+  _mjPRIVATE_dsuUnion(parent, 3, -1);
+  _mjPRIVATE_dsuUnion(parent, 5, 2);
+  _mjPRIVATE_dsuUnion(parent, 4, 5);
+  _mjPRIVATE_dsuUnion(parent, 3, 4);
+
+  EXPECT_THAT(parent, ElementsAre(-1, -1, 2, 2, 2, 2));
+  for (int tree = 2; tree < 6; ++tree) {
+    EXPECT_EQ(_mjPRIVATE_dsuFind(parent, tree), 2);
+  }
+  EXPECT_THAT(parent, ElementsAre(-1, -1, 2, 2, 2, 2));
+}
+
+TEST_F(IslandTest, DsuUnionRedundantAndReversedEdgesAreIdempotent) {
+  int parent[] = {-1, -1, -1, -1};
+  _mjPRIVATE_dsuUnion(parent, 3, 1);
+  _mjPRIVATE_dsuUnion(parent, 2, 1);
+  EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 1));
+
+  _mjPRIVATE_dsuUnion(parent, 1, 3);
+  _mjPRIVATE_dsuUnion(parent, 3, 1);
+  _mjPRIVATE_dsuUnion(parent, 2, 2);
+  _mjPRIVATE_dsuUnion(parent, -1, 2);
+  EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 1));
+}
+
+TEST_F(IslandTest, DsuUnionFastPathActivatesBeforeTestingParents) {
+  int self_parent[] = {-1, -1, -1};
+  _mjPRIVATE_dsuUnion(self_parent, 1, 1);
+  EXPECT_THAT(self_parent, ElementsAre(-1, 1, -1));
+
+  int static_first[] = {-1, -1, -1};
+  _mjPRIVATE_dsuUnion(static_first, -1, 2);
+  EXPECT_THAT(static_first, ElementsAre(-1, -1, 2));
+
+  int static_second[] = {-1, -1, -1};
+  _mjPRIVATE_dsuUnion(static_second, 0, -1);
+  EXPECT_THAT(static_second, ElementsAre(0, -1, -1));
+}
+
+TEST_F(IslandTest, DsuUnionFastPathDistinguishesParentsFromRoots) {
+  int distinct_parent[] = {0, 0, 2, 2};
+  _mjPRIVATE_dsuUnion(distinct_parent, 1, 3);
+  EXPECT_THAT(distinct_parent, ElementsAre(0, 0, 0, 2));
+
+  int shared_parent[] = {0, 0, 0, 3};
+  _mjPRIVATE_dsuUnion(shared_parent, 1, 2);
+  EXPECT_THAT(shared_parent, ElementsAre(0, 0, 0, 3));
+
+  int long_paths[] = {0, 0, 1, 3, 3, 4};
+  _mjPRIVATE_dsuUnion(long_paths, 2, 5);
+  EXPECT_THAT(long_paths, ElementsAre(0, 0, 0, 0, 3, 3));
+}
+
+TEST_F(IslandTest, DsuUnionFastPathPreservesCyclesDuplicatesAndForest) {
+  int parent[] = {-1, -1, -1, -1, -1, -1};
+  _mjPRIVATE_dsuUnion(parent, 0, 1);
+  _mjPRIVATE_dsuUnion(parent, 1, 2);
+  _mjPRIVATE_dsuUnion(parent, 2, 0);
+  _mjPRIVATE_dsuUnion(parent, 0, 2);
+  _mjPRIVATE_dsuUnion(parent, 3, 4);
+  _mjPRIVATE_dsuUnion(parent, 4, 5);
+  EXPECT_THAT(parent, ElementsAre(0, 0, 0, 3, 3, 3));
+
+  _mjPRIVATE_dsuUnion(parent, 5, 0);
+  EXPECT_THAT(parent, ElementsAre(0, 0, 0, 0, 3, 3));
+}
+
+TEST_F(IslandTest, DsuUnionRejectsStaticSelfIncidence) {
+  int parent[] = {-1, 1, 1, 3};
+
+  EXPECT_EQ(MjuErrorMessageFrom(_mjPRIVATE_dsuUnion)(parent, -1, -1),
+            "self-incidence of the static tree");
+  EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 3));
+}
+
+TEST_F(IslandTest, DsuAssignHandlesEmptyAndInactiveInputs) {
+  int island[] = {71};
+  int parent[] = {72};
+  const int tree_dofnum[] = {73};
+  int nidof = -1;
+
+  EXPECT_EQ(_mjPRIVATE_dsuAssign(island, parent, tree_dofnum, 0, &nidof), 0);
+  EXPECT_EQ(nidof, 0);
+  EXPECT_THAT(island, ElementsAre(71));
+  EXPECT_THAT(parent, ElementsAre(72));
+
+  parent[0] = -1;
+  EXPECT_EQ(_mjPRIVATE_dsuAssign(island, parent, tree_dofnum, 1, &nidof), 0);
+  EXPECT_EQ(nidof, 0);
+  EXPECT_THAT(island, ElementsAre(-1));
+  EXPECT_THAT(parent, ElementsAre(-1));
+}
+
+TEST_F(IslandTest, DsuAssignLabelsComponentsAndCountsOnlyActiveDofs) {
+  int parent[] = {-1, 1, 1, 2, 4, 4, 6};
+  const int tree_dofnum[] = {1000, 0, 3, 5, 7, 11, 13};
+  int island[] = {9, 9, 9, 9, 9, 9, 9};
+  int nidof = -1;
+
+  EXPECT_EQ(_mjPRIVATE_dsuAssign(island, parent, tree_dofnum, 7, &nidof), 3);
+  EXPECT_EQ(nidof, 39);
+  EXPECT_THAT(island, ElementsAre(-1, 0, 0, 0, 1, 1, 2));
+  EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 1, 4, 4, 6));
+}
+
+TEST_F(IslandTest, DsuAssignCompressesAscendingMultiHopForest) {
+  int parent[] = {-1, 1, 1, 2, 4, 4, 5, 7, 7, 8};
+  const int tree_dofnum[] = {99, 0, 2, 3, 0, 5, 7, 11, 0, 13};
+  int island[] = {9, 9, 9, 9, 9, 9, 9, 9, 9, 9};
+  int nidof = -1;
+
+  EXPECT_EQ(_mjPRIVATE_dsuAssign(island, parent, tree_dofnum, 10, &nidof), 3);
+  EXPECT_EQ(nidof, 41);
+  EXPECT_THAT(island, ElementsAre(-1, 0, 0, 0, 1, 1, 1, 2, 2, 2));
+  EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 1, 4, 4, 4, 7, 7, 7));
+}
+
+TEST_F(IslandTest, DsuAssignCompresses4096NodeAdversarialChain) {
+  constexpr int kTreeCount = 4096;
+  std::vector<int> parent(kTreeCount);
+  std::vector<int> island(kTreeCount, -2);
+  std::vector<int> tree_dofnum(kTreeCount);
+  parent[0] = 0;
+  int expected_nidof = 0;
+  for (int tree = 1; tree < kTreeCount; ++tree) {
+    parent[tree] = tree - 1;
+    tree_dofnum[tree] = tree % 5;
+    expected_nidof += tree_dofnum[tree];
+  }
+
+  int nidof = -1;
+  EXPECT_EQ(_mjPRIVATE_dsuAssign(island.data(), parent.data(), tree_dofnum.data(),
+                                kTreeCount, &nidof),
+            1);
+  EXPECT_EQ(nidof, expected_nidof);
+  for (int tree = 0; tree < kTreeCount; ++tree) {
+    EXPECT_EQ(island[tree], 0);
+    EXPECT_EQ(parent[tree], 0);
+  }
+}
+
+TEST_F(IslandTest, DsuHandlesLongConnectedBoundaryCase) {
+  constexpr int kTreeCount = 4096;
+  std::vector<int> parent(kTreeCount);
+  std::vector<int> island(kTreeCount, -2);
+  std::vector<int> tree_dofnum(kTreeCount);
+  _mjPRIVATE_dsuInit(parent.data(), kTreeCount);
+
+  int expected_nidof = 0;
+  for (int tree = kTreeCount - 1; tree > 0; --tree) {
+    _mjPRIVATE_dsuUnion(parent.data(), tree, tree - 1);
+  }
+  for (int tree = 0; tree < kTreeCount; ++tree) {
+    tree_dofnum[tree] = tree % 7;
+    expected_nidof += tree_dofnum[tree];
+  }
+
+  int nidof = -1;
+  EXPECT_EQ(_mjPRIVATE_dsuAssign(island.data(), parent.data(), tree_dofnum.data(),
+                                kTreeCount, &nidof),
+            1);
+  EXPECT_EQ(nidof, expected_nidof);
+  for (int tree = 0; tree < kTreeCount; ++tree) {
+    EXPECT_EQ(island[tree], 0);
+    EXPECT_EQ(parent[tree], 0);
+  }
+}
+
+TEST_F(IslandTest, DsuRandomizedDifferentialAgainstGraphTraversal) {
+  constexpr uint32_t kSeed = 0x5eed3396u;
+  constexpr int kTrials = 2000;
+  uint32_t state = kSeed;
+  auto next = [&state]() {
+    state = state * 1664525u + 1013904223u;
+    return state;
+  };
+
+  for (int trial = 0; trial < kTrials; ++trial) {
+    const int ntree = 1 + next() % 64;
+    const int nedge = next() % 192;
+    std::vector<std::array<int, 2>> edges;
+    edges.reserve(nedge);
+    for (int edge = 0; edge < nedge; ++edge) {
+      int tree1;
+      int tree2;
+      switch (next() % 8) {
+        case 0:
+          tree1 = -1;
+          tree2 = next() % ntree;
+          break;
+        case 1:
+          tree1 = next() % ntree;
+          tree2 = -1;
+          break;
+        case 2:
+          tree1 = next() % ntree;
+          tree2 = tree1;
+          break;
+        case 3:
+          if (!edges.empty()) {
+            const auto& previous = edges[next() % edges.size()];
+            tree1 = previous[0];
+            tree2 = previous[1];
+            break;
+          }
+          [[fallthrough]];
+        case 4:
+          if (!edges.empty()) {
+            const auto& previous = edges[next() % edges.size()];
+            tree1 = previous[1];
+            tree2 = previous[0];
+            break;
+          }
+          [[fallthrough]];
+        default:
+          tree1 = next() % ntree;
+          tree2 = next() % ntree;
+          break;
+      }
+      edges.push_back({tree1, tree2});
+    }
+
+    std::vector<int> parent(ntree);
+    _mjPRIVATE_dsuInit(parent.data(), ntree);
+    for (const auto& edge : edges) {
+      _mjPRIVATE_dsuUnion(parent.data(), edge[0], edge[1]);
+    }
+
+    std::vector<int> active(ntree);
+    std::vector<std::vector<int>> adjacent(ntree);
+    for (const auto& edge : edges) {
+      if (edge[0] >= 0) active[edge[0]] = 1;
+      if (edge[1] >= 0) active[edge[1]] = 1;
+      if (edge[0] >= 0 && edge[1] >= 0) {
+        adjacent[edge[0]].push_back(edge[1]);
+        adjacent[edge[1]].push_back(edge[0]);
+      }
+    }
+
+    std::vector<int> expected_island(ntree, -1);
+    std::vector<int> expected_parent(ntree, -1);
+    int expected_nisland = 0;
+    for (int start = 0; start < ntree; ++start) {
+      if (!active[start] || expected_island[start] != -1) continue;
+      std::vector<int> pending = {start};
+      std::vector<int> component;
+      expected_island[start] = expected_nisland;
+      while (!pending.empty()) {
+        const int tree = pending.back();
+        pending.pop_back();
+        component.push_back(tree);
+        for (int neighbor : adjacent[tree]) {
+          if (expected_island[neighbor] == -1) {
+            expected_island[neighbor] = expected_nisland;
+            pending.push_back(neighbor);
+          }
+        }
+      }
+      for (int tree : component) expected_parent[tree] = start;
+      ++expected_nisland;
+    }
+
+    std::vector<int> tree_dofnum(ntree);
+    int expected_nidof = 0;
+    for (int tree = 0; tree < ntree; ++tree) {
+      tree_dofnum[tree] = next() % 8;
+      if (active[tree]) expected_nidof += tree_dofnum[tree];
+    }
+    std::vector<int> island(ntree, -2);
+    int nidof = -1;
+    const int nisland = _mjPRIVATE_dsuAssign(
+        island.data(), parent.data(), tree_dofnum.data(), ntree, &nidof);
+
+    SCOPED_TRACE(::testing::Message() << "seed=" << kSeed << " trial=" << trial
+                                      << " ntree=" << ntree << " nedge=" << nedge);
+    EXPECT_EQ(nisland, expected_nisland);
+    EXPECT_EQ(nidof, expected_nidof);
+    EXPECT_EQ(island, expected_island);
+    EXPECT_EQ(parent, expected_parent);
+  }
+}
+
 TEST_F(IslandTest, FloodFillSingleton) {
   // adjacency matrix for the graph  0   1   2
   //                                 U       U
@@ -532,6 +838,13 @@ TEST_F(IslandTest, IslandEfc) {
   EXPECT_EQ(data->nf, 2);
   EXPECT_EQ(data->nl, 1);
   EXPECT_EQ(data->nefc, 30);
+  EXPECT_THAT(AsVector(data->island_ne, data->nisland), ElementsAre(1, 0, 0, 6));
+  EXPECT_THAT(AsVector(data->island_nf, data->nisland), ElementsAre(0, 1, 1, 0));
+  EXPECT_THAT(AsVector(data->island_nefc, data->nisland), ElementsAre(6, 17, 1, 6));
+  EXPECT_THAT(AsVector(data->efc_island, data->nefc),
+              ElementsAre(0, 3, 3, 3, 3, 3, 3, 1, 2, 0,
+                          0, 0, 0, 0, 1, 1, 1, 1, 1, 1,
+                          1, 1, 1, 1, 1, 1, 1, 1, 1, 1));
 
   mj_deleteData(data);
   mj_deleteModel(model);