Merge pull request #3396 from teerthsharma:topo/linear-island-scratch

PiperOrigin-RevId: 951110709
Change-Id: I0c9c96365a5667172c1d676026ab797b7f8e8137

doc/changelog.rst

@@ -48,6 +48,8 @@ Engine
   as originally guarded by ``ngravcomp``. Since the engine uses these integers as flags (zero vs. non-zero), the new
   flags are honest boolean properties, writeable from the Python bindings at runtime. The field ``ngravcomp`` is
   deprecated and will be removed in a future release.
+- Replaced quadratic scratch in DFS flood-fill island discovery with a linear-memory Union-Find (disjoint set).
+  Contribution by :github:user:`teerthsharma`.
 
 .. admonition:: Breaking API changes
    :class: attention

src/engine/engine_island.c

@@ -16,7 +16,6 @@
 
 #include <stdio.h>
 #include <stddef.h>
-#include <string.h>
 
 #include <mujoco/mjdata.h>
 #include <mujoco/mjmodel.h>
@@ -82,7 +81,74 @@ static int arenaAllocIsland(const mjModel* m, mjData* d) {
 }
 
 
-//-------------------------- flood-fill and graph construction  ------------------------------------
+//-------------------------- union-find and flood-fill  --------------------------------------------
+
+// find the canonical root of an active tree and compress its path
+int mj_dsuRoot(int* parent, int tree) {
+  int root = tree;
+  while (parent[root] != root) {
+    root = parent[root];
+  }
+
+  while (parent[tree] != tree) {
+    int next = parent[tree];
+    parent[tree] = root;
+    tree = next;
+  }
+
+  return root;
+}
+
+
+// activate and union two incident trees; -1 denotes a static endpoint
+void mj_dsuMerge(int* parent, int tree1, int tree2) {
+  if (tree1 == -1 && tree2 == -1) {
+    mjERROR("self-incidence of the static tree");  // SHOULD NOT OCCUR
+    return;
+  }
+
+  if (tree1 == -1) tree1 = tree2;
+  if (tree2 == -1) tree2 = tree1;
+
+  if (parent[tree1] == -1) parent[tree1] = tree1;
+  if (parent[tree2] == -1) parent[tree2] = tree2;
+
+  if (parent[tree1] == parent[tree2]) return;
+
+  int root1 = mj_dsuRoot(parent, tree1);
+  int root2 = mj_dsuRoot(parent, tree2);
+  if (root1 < root2) {
+    parent[root2] = root1;
+  } else if (root2 < root1) {
+    parent[root1] = root2;
+  }
+}
+
+
+// assign deterministic island ids in ascending canonical-root order
+int mj_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++) {
+    if (parent[tree] == -1) {
+      island[tree] = -1;
+      continue;
+    }
+
+    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;
+}
+
 
 // find disjoint subgraphs ("islands") given sparse symmetric adjacency matrix
 //   arguments:
@@ -280,60 +346,28 @@ static void treeIterInit(const mjModel* m, const mjData* d, int i, mjTreeIter* i
 }
 
 
-// add 0, 1 or 2 edges to uncompressed CSR adjacency matrix
-//   increment rownnz using tree_tree to de-dupe; return number of edges added
-static int addEdge(int* rownnz, int* colind, mjtByte* tree_tree, int ntree, int tree1, int tree2) {
-  if (tree1 == -1 && tree2 == -1) {
-    mjERROR("self-edge of the static tree");  // SHOULD NOT OCCUR
-    return 0;
-  }
-
-  // handle static trees (treat as self-edge)
-  if (tree1 == -1) tree1 = tree2;
-  if (tree2 == -1) tree2 = tree1;
-
-  // skip if edge already present
-  if (tree_tree[tree1*ntree + tree2]) {
-    return 0;
-  }
-
-  // add edge
-  tree_tree[tree1*ntree + tree2] = 1;
-  colind[tree1*ntree + rownnz[tree1]++] = tree2;  // uncompressed format, rowadr is known
-
-  // add flipped edge (off-diagonal)
-  if (tree1 != tree2) {
-    tree_tree[tree2*ntree + tree1] = 1;
-    colind[tree2*ntree + rownnz[tree2]++] = tree1;  // uncompressed format, rowadr is known
-    return 2;
-  }
-
-  return 1;
+// return whether repeated scalar rows of this constraint require separate tree scans
+static int isFlexEquality(const mjModel* m, int efc_type, int efc_id) {
+  return efc_type == mjCNSTR_EQUALITY &&
+         (m->eq_type[efc_id] == mjEQ_FLEX ||
+          m->eq_type[efc_id] == mjEQ_FLEXVERT ||
+          m->eq_type[efc_id] == mjEQ_FLEXSTRAIN);
 }
 
 
-// find tree-tree edges (column indices), return total number of edges
-//   efc_tree: first nonegative tree index of each constraint
-static int findEdges(const mjModel* m, const mjData* d,
-                     int* rownnz, int* colind, mjtByte* tree_tree, int* efc_tree, int ntree) {
+// activate and union all trees with direct incidence in a constraint
+static const char* unionConstraintTrees(const mjModel* m, const mjData* d, int* parent,
+                                        int* efc_tree, int* err_i) {
   int nefc = d->nefc;
-  int nnz = 0;
   int efc_type = -1;
   int efc_id = -1;
 
-  // clear row nonzeros
-  mju_zeroInt(rownnz, ntree);
-
-  // iterate over constraints, compute tree-tree edges, assign efc_tree
+  // 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]) {
+    // row i is still in the same constraint: skip it
+    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 (!(efc_type == mjCNSTR_EQUALITY &&
-            (m->eq_type[efc_id] == mjEQ_FLEX ||
-             m->eq_type[efc_id] == mjEQ_FLEXVERT ||
-             m->eq_type[efc_id] == mjEQ_FLEXSTRAIN))) {
-        // copy tree assignment from previous constraint and continue
+      if (!isFlexEquality(m, efc_type, efc_id)) {
         efc_tree[i] = efc_tree[i-1];
         continue;
       }
@@ -349,25 +383,26 @@ static int findEdges(const mjModel* m, const mjData* d,
     int tree1 = treeNext(m, d, i, &iter);
     if (tree1 != -2) {
       int tree2 = treeNext(m, d, i, &iter);
-
       // assign tree to constraint, one of (tree1, tree2) must be non-negative
       efc_tree[i] = tree1 >= 0 ? tree1 : tree2;
       if (efc_tree[i] < 0) {
-        mjERROR("constraint %d is between two static bodies", i);  // SHOULD NOT OCCUR
+        *err_i = i;
+        return "constraint %d is between two static bodies";
       }
 
-      // add one edge or continue to search for more edges
+      // activate a singleton or union all trees in a multi-tree constraint
       if (tree2 == -2) {
-        nnz += addEdge(rownnz, colind, tree_tree, ntree, tree1, -1);
+        mj_dsuMerge(parent, tree1, -1);
       } else {
         while (tree2 != -2) {
-          nnz += addEdge(rownnz, colind, tree_tree, ntree, tree1, tree2);
+          mj_dsuMerge(parent, tree1, tree2);
           tree1 = tree2;
           tree2 = treeNext(m, d, i, &iter);
         }
       }
     } else {
-      mjERROR("no tree found for constraint %d", i);  // SHOULD NOT OCCUR
+      *err_i = i;
+      return "no tree found for constraint %d";
     }
   }
 
@@ -387,6 +422,7 @@ static int findEdges(const mjModel* m, const mjData* d,
     if (m->flex_bendingadr[f] < 0 && (sadr < 0 || m->flex_stiffness[sadr] == 0)) {
       continue;
     }
+
     int num, adr;
     const int* bodyid;
     if (m->flex_interp[f]) {
@@ -398,21 +434,22 @@ static int findEdges(const mjModel* m, const mjData* d,
       adr = m->flex_vertadr[f];
       bodyid = m->flex_vertbodyid;
     }
+
     int tree1 = -1;
     for (int j=0; j < num; j++) {
-      int treeid = m->body_treeid[bodyid[adr+j]];
-      if (treeid < 0 || treeid == tree1 || !d->tree_awake[treeid]) {
+      int tree2 = m->body_treeid[bodyid[adr+j]];
+      if (tree2 < 0 || tree2 == tree1 || !d->tree_awake[tree2]) {
         continue;
       }
       if (tree1 < 0) {
-        tree1 = treeid;
+        tree1 = tree2;
       } else {
-        nnz += addEdge(rownnz, colind, tree_tree, ntree, tree1, treeid);
+        mj_dsuMerge(parent, tree1, tree2);
       }
     }
   }
 
-  return nnz;
+  return NULL;
 }
 
 
@@ -431,29 +468,19 @@ void mj_island(const mjModel* m, mjData* d) {
 
   mj_markStack(d);
 
-  // dense tree-tree adjacency matrix
-  int ntree2 = ntree * ntree;
-  mjtByte* tree_tree = mjSTACKALLOC(d, ntree2, mjtByte);
-  memset(tree_tree, 0, ntree2);
-
-  // CSR representation of tree-tree adjacency matrix (uncompressed)
-  int* colind = mjSTACKALLOC(d, ntree2, int);
-  int* rownnz = mjSTACKALLOC(d, ntree, int);
-  int* rowadr = mjSTACKALLOC(d, ntree, int);
-  for (int r=0; r < ntree; r++) {
-    rowadr[r] = r * ntree;
-  }
-
-  // first non-negative tree index of each constraint, used later for computing efc_island
+  // union direct tree incidence and assign deterministic components
   int* efc_tree = mjSTACKALLOC(d, nefc, int);
-
-  // compute tree-tree adjacency matrix: fill rownnz and colind
-  int nnz = findEdges(m, d, rownnz, colind, tree_tree, efc_tree, ntree);
-
-  // discover islands
+  int* parent = mjSTACKALLOC(d, ntree, int);
+  mju_fillInt(parent, -1, ntree);
+  int err_i = -1;
+  const char* err_msg = unionConstraintTrees(m, d, parent, efc_tree, &err_i);
+  if (err_msg) {
+    mj_freeStack(d);
+    mjERROR(err_msg, err_i);
+  }
   int* tree_island = mjSTACKALLOC(d, ntree, int);
-  int* stack = mjSTACKALLOC(d, nnz, int);
-  d->nisland = mj_floodFill(tree_island, ntree, rownnz, rowadr, colind, stack);
+  int nidof;
+  d->nisland = mj_dsuAssign(tree_island, parent, m->tree_dofnum, ntree, &nidof);
 
   // no islands found: quick return
   if (!d->nisland) {
@@ -462,13 +489,6 @@ void mj_island(const mjModel* m, mjData* d) {
     return;
   }
 
-  // count nidof: total number of dofs in islands
-  int nidof = 0;
-  for (int i=0; i < ntree; i++) {
-    if (tree_island[i] >= 0) {
-      nidof += m->tree_dofnum[i];
-    }
-  }
   d->nidof = nidof;
 
   // allocate island arrays on arena
@@ -575,8 +595,8 @@ void mj_island(const mjModel* m, mjData* d) {
   mju_zeroInt(d->island_nf, nisland);
   mju_zeroInt(d->island_nefc, nisland);
   for (int i=0; i < nefc; i++) {
-    int island = tree_island[efc_tree[i]];
-    d->efc_island[i] = island;
+    d->efc_island[i] = tree_island[efc_tree[i]];
+    int island = d->efc_island[i];
     d->island_nefc[island]++;
     switch (d->efc_type[i]) {
       case mjCNSTR_EQUALITY:
@@ -605,17 +625,15 @@ void mj_island(const mjModel* m, mjData* d) {
     int ic = d->island_iefcadr[island] + island_nefc2[island]++;
     d->map_efc2iefc[c] = ic;
     d->map_iefc2efc[ic] = c;
+    d->iefc_type[ic] = d->efc_type[c];
+    d->iefc_id[ic] = d->efc_id[c];
+    d->iefc_frictionloss[ic] = d->efc_frictionloss[c];
+    d->iefc_D[ic] = d->efc_D[c];
+    d->iefc_R[ic] = d->efc_R[c];
   }
 
   // SHOULD NOT OCCUR
   if (!mju_compare(island_nefc2, d->island_nefc, nisland)) mjERROR("island_nefc miscount");
 
-  // copy position-dependent efc vectors required by solver
-  mju_gatherInt(d->iefc_type, d->efc_type, d->map_iefc2efc, nefc);
-  mju_gatherInt(d->iefc_id, d->efc_id, d->map_iefc2efc, nefc);
-  mju_gather(d->iefc_frictionloss, d->efc_frictionloss, d->map_iefc2efc, nefc);
-  mju_gather(d->iefc_D, d->efc_D, d->map_iefc2efc, nefc);
-  mju_gather(d->iefc_R, d->efc_R, d->map_iefc2efc, nefc);
-
   mj_freeStack(d);
 }

src/engine/engine_island.h

@@ -23,6 +23,11 @@
 extern "C" {
 #endif
 
+// disjoint-set roots are minimum tree indices; mj_dsuRoot requires parent[tree] >= 0
+MJAPI void mj_dsuMerge(int* parent, int tree1, int tree2);
+MJAPI int mj_dsuRoot(int* parent, int tree);
+MJAPI int mj_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/engine/engine_island_test.cc

@@ -16,6 +16,9 @@
 
 #include "src/engine/engine_island.h"
 
+#include <array>
+#include <cstddef>
+#include <cstdint>
 #include <string>
 #include <vector>
 
@@ -35,6 +38,301 @@ using ::testing::NotNull;
 using ::testing::Pointwise;
 using IslandTest = MujocoTest;
 
+
+
+TEST_F(IslandTest, DsuRootReturnsCanonicalRootAndCompressesPath) {
+  int parent[] = {0, 0, 1, 2, 3};
+
+  EXPECT_EQ(mj_dsuRoot(parent, 0), 0);
+  EXPECT_EQ(mj_dsuRoot(parent, 4), 0);
+  EXPECT_THAT(parent, ElementsAre(0, 0, 0, 0, 0));
+}
+
+TEST_F(IslandTest, DsuMergeActivatesEndpointsAndUsesMinimumRoot) {
+  int parent[] = {-1, -1, -1, -1, -1, -1};
+
+  mj_dsuMerge(parent, -1, 4);
+  mj_dsuMerge(parent, 3, -1);
+  mj_dsuMerge(parent, 5, 2);
+  mj_dsuMerge(parent, 4, 5);
+  mj_dsuMerge(parent, 3, 4);
+
+  EXPECT_THAT(parent, ElementsAre(-1, -1, 2, 2, 2, 2));
+  for (int tree = 2; tree < 6; ++tree) {
+    EXPECT_EQ(mj_dsuRoot(parent, tree), 2);
+  }
+  EXPECT_THAT(parent, ElementsAre(-1, -1, 2, 2, 2, 2));
+}
+
+TEST_F(IslandTest, DsuMergeRedundantAndReversedEdgesAreIdempotent) {
+  int parent[] = {-1, -1, -1, -1};
+  mj_dsuMerge(parent, 3, 1);
+  mj_dsuMerge(parent, 2, 1);
+  EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 1));
+
+  mj_dsuMerge(parent, 1, 3);
+  mj_dsuMerge(parent, 3, 1);
+  mj_dsuMerge(parent, 2, 2);
+  mj_dsuMerge(parent, -1, 2);
+  EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 1));
+}
+
+TEST_F(IslandTest, DsuMergeFastPathActivatesBeforeTestingParents) {
+  int self_parent[] = {-1, -1, -1};
+  mj_dsuMerge(self_parent, 1, 1);
+  EXPECT_THAT(self_parent, ElementsAre(-1, 1, -1));
+
+  int static_first[] = {-1, -1, -1};
+  mj_dsuMerge(static_first, -1, 2);
+  EXPECT_THAT(static_first, ElementsAre(-1, -1, 2));
+
+  int static_second[] = {-1, -1, -1};
+  mj_dsuMerge(static_second, 0, -1);
+  EXPECT_THAT(static_second, ElementsAre(0, -1, -1));
+}
+
+TEST_F(IslandTest, DsuMergeFastPathDistinguishesParentsFromRoots) {
+  int distinct_parent[] = {0, 0, 2, 2};
+  mj_dsuMerge(distinct_parent, 1, 3);
+  EXPECT_THAT(distinct_parent, ElementsAre(0, 0, 0, 2));
+
+  int shared_parent[] = {0, 0, 0, 3};
+  mj_dsuMerge(shared_parent, 1, 2);
+  EXPECT_THAT(shared_parent, ElementsAre(0, 0, 0, 3));
+
+  int long_paths[] = {0, 0, 1, 3, 3, 4};
+  mj_dsuMerge(long_paths, 2, 5);
+  EXPECT_THAT(long_paths, ElementsAre(0, 0, 0, 0, 3, 3));
+}
+
+TEST_F(IslandTest, DsuMergeFastPathPreservesCyclesDuplicatesAndForest) {
+  int parent[] = {-1, -1, -1, -1, -1, -1};
+  mj_dsuMerge(parent, 0, 1);
+  mj_dsuMerge(parent, 1, 2);
+  mj_dsuMerge(parent, 2, 0);
+  mj_dsuMerge(parent, 0, 2);
+  mj_dsuMerge(parent, 3, 4);
+  mj_dsuMerge(parent, 4, 5);
+  EXPECT_THAT(parent, ElementsAre(0, 0, 0, 3, 3, 3));
+
+  mj_dsuMerge(parent, 5, 0);
+  EXPECT_THAT(parent, ElementsAre(0, 0, 0, 0, 3, 3));
+}
+
+TEST_F(IslandTest, DsuMergeRejectsStaticSelfIncidence) {
+  int parent[] = {-1, 1, 1, 3};
+
+  EXPECT_EQ(MjuErrorMessageFrom(mj_dsuMerge)(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(mj_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(mj_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(mj_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(mj_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(mj_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, -1);
+  std::vector<int> island(kTreeCount, -2);
+  std::vector<int> tree_dofnum(kTreeCount);
+
+  int expected_nidof = 0;
+  for (int tree = kTreeCount - 1; tree > 0; --tree) {
+    mj_dsuMerge(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(mj_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, -1);
+    for (const auto& edge : edges) {
+      mj_dsuMerge(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 = mj_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
@@ -156,6 +454,230 @@ TEST_F(IslandTest, FloodFill3b) {
   EXPECT_THAT(island, ElementsAre(0, 1, 1, -1, 0, 0, 0));
 }
 
+TEST_F(IslandTest, ProductionStaticFirstAndRepeatedRows) {
+  static constexpr char xml[] = R"(
+<mujoco>
+  <option jacobian="sparse"><flag contact="disable" gravity="disable"/></option>
+  <worldbody>
+    <site name="world"/>
+    <body name="b0">
+      <inertial pos="0 0 0" mass="1" diaginertia="1 1 1"/>
+      <joint type="slide"/><site name="s0"/>
+    </body>
+    <body name="b1">
+      <inertial pos="0 0 0" mass="1" diaginertia="1 1 1"/>
+      <joint type="slide"/><site name="s1"/>
+    </body>
+    <body name="b2">
+      <inertial pos="0 0 0" mass="1" diaginertia="1 1 1"/>
+      <joint type="slide"/><site name="s2"/>
+    </body>
+  </worldbody>
+  <equality>
+    <connect site1="world" site2="s0"/>
+    <connect site1="s1" site2="s2"/>
+  </equality>
+</mujoco>
+)";
+  char error[1024] = {};
+  MjModelPtr model = LoadModelFromString(xml, error, sizeof(error));
+  ASSERT_THAT(model.get(), NotNull()) << error;
+  ASSERT_EQ(model->ntree, 3);
+  ASSERT_EQ(model->nv, 3);
+
+  // The first equality incidence is static first, then dynamic tree 0.
+  ASSERT_EQ(model->eq_objtype[0], mjOBJ_SITE);
+  int body1 = model->site_bodyid[model->eq_obj1id[0]];
+  int body2 = model->site_bodyid[model->eq_obj2id[0]];
+  EXPECT_EQ(model->body_treeid[body1], -1);
+  EXPECT_EQ(model->body_treeid[body2], 0);
+
+  MjDataPtr data = MakeData(model);
+  mj_fwdPosition(model.get(), data.get());
+
+  ASSERT_EQ(data->nefc, 6);
+  EXPECT_EQ(data->nisland, 2);
+  EXPECT_EQ(data->nidof, 3);
+  EXPECT_EQ(data->ne, 6);
+  EXPECT_EQ(data->nf, 0);
+  EXPECT_THAT(
+      AsVector(data->efc_type, data->nefc),
+      ElementsAre(mjCNSTR_EQUALITY, mjCNSTR_EQUALITY, mjCNSTR_EQUALITY,
+                  mjCNSTR_EQUALITY, mjCNSTR_EQUALITY, mjCNSTR_EQUALITY));
+  EXPECT_THAT(AsVector(data->efc_id, data->nefc),
+              ElementsAre(0, 0, 0, 1, 1, 1));
+  EXPECT_THAT(AsVector(data->tree_island, model->ntree), ElementsAre(0, 1, 1));
+  EXPECT_THAT(AsVector(data->island_ntree, data->nisland), ElementsAre(1, 2));
+  EXPECT_THAT(AsVector(data->island_itreeadr, data->nisland),
+              ElementsAre(0, 1));
+  EXPECT_THAT(AsVector(data->map_itree2tree, model->ntree),
+              ElementsAre(0, 1, 2));
+  EXPECT_THAT(AsVector(data->dof_island, model->nv), ElementsAre(0, 1, 1));
+  EXPECT_THAT(AsVector(data->island_nv, data->nisland), ElementsAre(1, 2));
+  EXPECT_THAT(AsVector(data->island_idofadr, data->nisland), ElementsAre(0, 1));
+  EXPECT_THAT(AsVector(data->island_dofadr, data->nisland), ElementsAre(0, 1));
+  EXPECT_THAT(AsVector(data->map_dof2idof, model->nv), ElementsAre(0, 1, 2));
+  EXPECT_THAT(AsVector(data->map_idof2dof, model->nv), ElementsAre(0, 1, 2));
+  EXPECT_THAT(AsVector(data->efc_island, data->nefc),
+              ElementsAre(0, 0, 0, 1, 1, 1));
+  EXPECT_THAT(AsVector(data->island_ne, data->nisland), ElementsAre(3, 3));
+  EXPECT_THAT(AsVector(data->island_nf, data->nisland), ElementsAre(0, 0));
+  EXPECT_THAT(AsVector(data->island_nefc, data->nisland), ElementsAre(3, 3));
+  EXPECT_THAT(AsVector(data->island_iefcadr, data->nisland), ElementsAre(0, 3));
+  EXPECT_THAT(AsVector(data->map_efc2iefc, data->nefc),
+              ElementsAre(0, 1, 2, 3, 4, 5));
+  EXPECT_THAT(AsVector(data->map_iefc2efc, data->nefc),
+              ElementsAre(0, 1, 2, 3, 4, 5));
+}
+
+TEST_F(IslandTest, ReportsConstraintBetweenTwoStaticBodies) {
+  static constexpr char xml[] = R"(
+<mujoco>
+  <option jacobian="sparse"><flag contact="disable" gravity="disable"/></option>
+  <worldbody>
+    <body>
+      <inertial pos="0 0 0" mass="1" diaginertia="1 1 1"/>
+      <joint type="slide" frictionloss="1"/>
+    </body>
+  </worldbody>
+</mujoco>
+)";
+  char error[1024] = {};
+  MjModelPtr model = LoadModelFromString(xml, error, sizeof(error));
+  ASSERT_THAT(model.get(), NotNull()) << error;
+  MjDataPtr data = MakeData(model);
+  mj_fwdPosition(model.get(), data.get());
+  ASSERT_GT(data->nefc, 0);
+  ASSERT_EQ(data->efc_type[0], mjCNSTR_FRICTION_DOF);
+
+  model->dof_treeid[data->efc_id[0]] = -1;
+
+  EXPECT_EQ(MjuErrorMessageFrom(mj_island)(model.get(), data.get()),
+            "constraint 0 is between two static bodies");
+}
+
+TEST_F(IslandTest, ProductionFlexEqualityRescansRows) {
+  static constexpr char xml[] = R"(
+<mujoco>
+  <option jacobian="sparse"><flag contact="disable" gravity="disable"/></option>
+  <worldbody>
+    <flexcomp name="f" type="grid" dim="1" count="3 1 1"
+              spacing=".05 .05 .05" radius=".01" mass="1">
+      <edge equality="true"/>
+      <contact internal="false" selfcollide="none"/>
+    </flexcomp>
+  </worldbody>
+</mujoco>
+)";
+  char error[1024] = {};
+  MjModelPtr model = LoadModelFromString(xml, error, sizeof(error));
+  ASSERT_THAT(model.get(), NotNull()) << error;
+  ASSERT_EQ(model->ntree, 3);
+  ASSERT_EQ(model->nv, 9);
+  ASSERT_EQ(model->neq, 1);
+  ASSERT_EQ(model->eq_type[0], mjEQ_FLEX);
+  ASSERT_TRUE(mj_isSparse(model.get()));
+
+  MjDataPtr data = MakeData(model);
+  mj_fwdPosition(model.get(), data.get());
+
+  ASSERT_EQ(data->nefc, 2);
+  auto row_trees = [&](int row) {
+    std::vector<int> trees;
+    for (int j=0; j < data->efc_J_rownnz[row]; j++) {
+      int dof = data->efc_J_colind[data->efc_J_rowadr[row] + j];
+      int tree = model->dof_treeid[dof];
+      if (trees.empty() || trees.back() != tree) {
+        trees.push_back(tree);
+      }
+    }
+    return trees;
+  };
+
+  // Rows share one flex equality id but have different tree incidence.
+  EXPECT_THAT(row_trees(0), ElementsAre(0, 1));
+  EXPECT_THAT(row_trees(1), ElementsAre(1, 2));
+  EXPECT_THAT(AsVector(data->efc_type, data->nefc),
+              ElementsAre(mjCNSTR_EQUALITY, mjCNSTR_EQUALITY));
+  EXPECT_THAT(AsVector(data->efc_id, data->nefc), ElementsAre(0, 0));
+  EXPECT_EQ(data->nisland, 1);
+  EXPECT_EQ(data->nidof, 9);
+  EXPECT_EQ(data->ne, 2);
+  EXPECT_EQ(data->nf, 0);
+  EXPECT_THAT(AsVector(data->tree_island, model->ntree), ElementsAre(0, 0, 0));
+  EXPECT_THAT(AsVector(data->island_ntree, data->nisland), ElementsAre(3));
+  EXPECT_THAT(AsVector(data->island_itreeadr, data->nisland), ElementsAre(0));
+  EXPECT_THAT(AsVector(data->map_itree2tree, model->ntree),
+              ElementsAre(0, 1, 2));
+  EXPECT_THAT(AsVector(data->dof_island, model->nv),
+              ElementsAre(0, 0, 0, 0, 0, 0, 0, 0, 0));
+  EXPECT_THAT(AsVector(data->island_nv, data->nisland), ElementsAre(9));
+  EXPECT_THAT(AsVector(data->island_idofadr, data->nisland), ElementsAre(0));
+  EXPECT_THAT(AsVector(data->island_dofadr, data->nisland), ElementsAre(0));
+  EXPECT_THAT(AsVector(data->map_dof2idof, model->nv),
+              ElementsAre(0, 1, 2, 3, 4, 5, 6, 7, 8));
+  EXPECT_THAT(AsVector(data->map_idof2dof, model->nv),
+              ElementsAre(0, 1, 2, 3, 4, 5, 6, 7, 8));
+  EXPECT_THAT(AsVector(data->efc_island, data->nefc), ElementsAre(0, 0));
+  EXPECT_THAT(AsVector(data->island_ne, data->nisland), ElementsAre(2));
+  EXPECT_THAT(AsVector(data->island_nf, data->nisland), ElementsAre(0));
+  EXPECT_THAT(AsVector(data->island_nefc, data->nisland), ElementsAre(2));
+  EXPECT_THAT(AsVector(data->island_iefcadr, data->nisland), ElementsAre(0));
+  EXPECT_THAT(AsVector(data->map_efc2iefc, data->nefc), ElementsAre(0, 1));
+  EXPECT_THAT(AsVector(data->map_iefc2efc, data->nefc), ElementsAre(0, 1));
+}
+
+TEST_F(IslandTest, BoundedArenaSupports1024Trees) {
+  constexpr int kTreeCount = 1024;
+  constexpr size_t kArenaBytes = 2 * 1024 * 1024;
+  std::string xml = R"(
+<mujoco>
+  <size memory="2M"/>
+  <option jacobian="sparse">
+    <flag contact="disable"/>
+  </option>
+  <worldbody>
+)";
+  xml.reserve(160 * kTreeCount);
+  for (int i=0; i < kTreeCount; i++) {
+    std::string name = std::to_string(i);
+    xml += "<body name=\"b" + name + "\">";
+    xml += "<inertial pos=\"0 0 0\" mass=\"1\" diaginertia=\"1 1 1\"/>";
+    xml += "<joint name=\"j" + name + "\" type=\"slide\"/>";
+    if (i < 2) {
+      xml += "<site name=\"s" + name + "\"/>";
+    }
+    xml += "</body>";
+  }
+  xml += R"(
+  </worldbody>
+  <equality>
+    <connect site1="s0" site2="s1"/>
+  </equality>
+</mujoco>
+)";
+
+  char error[1024] = {};
+  MjModelPtr model = LoadModelFromString(xml, error, sizeof(error));
+  ASSERT_THAT(model.get(), NotNull()) << error;
+  ASSERT_EQ(model->ntree, kTreeCount);
+  ASSERT_EQ(model->narena, kArenaBytes);
+  MjDataPtr data = MakeData(model);
+  ASSERT_THAT(data.get(), NotNull());
+
+  mj_fwdPosition(model.get(), data.get());
+
+  ASSERT_EQ(data->nefc, 3);
+  ASSERT_EQ(data->nisland, 1);
+  ASSERT_THAT(data->tree_island, NotNull());
+  EXPECT_EQ(data->tree_island[0], 0);
+  EXPECT_EQ(data->tree_island[1], 0);
+  for (int tree=2; tree < kTreeCount; tree++) {
+    EXPECT_EQ(data->tree_island[tree], -1);
+  }
+  EXPECT_LE(data->maxuse_arena, kArenaBytes);
+}
+
 static const char* const kAbacusPath = "engine/testdata/island/abacus.xml";
 
 TEST_F(IslandTest, Abacus) {
@@ -342,6 +864,16 @@ 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);