Fix bugs in engine_island.c

PiperOrigin-RevId: 562519702
Change-Id: I3075e539ccc1cef92485152721a9344f180aa4b1

src/engine/engine_island.c

@@ -160,13 +160,12 @@ static int countMaxEdge(const mjModel* m, const mjData* d) {
 
 
 // return id of next tree in Jacobian row i that is different from tree, -1 if not found
-//   write the index of the found tree to *index if given
+//   start search from *index
-//   start search from *index if given, otherwise 0
+//   write the index of the found tree to *index
-//   if J is (dense/sparse) *index is the (column/nonzro) index, respectively
+//   if J is (dense/sparse) *index is the (column/nonzero) index, respectively
 static int treeNext(const mjModel* m, const mjData* d, int tree, int i, int *index) {
   int tree_next = -1;
-  int j0 = index ? *index : 0;  // start searching at *index if given, otherwise 0
+  int j;  // local loop variable, saved to *index
-  int j;  // loop variable, saved to *index
 
   // sparse
   if (mj_isSparse(m)) {
@@ -174,7 +173,7 @@ static int treeNext(const mjModel* m, const mjData* d, int tree, int i, int *ind
     int* colind = d->efc_J_colind + d->efc_J_rowadr[i];
 
     // loop over remaining nonzeros, look for different tree
-    for (j=j0; j < rownnz; j++) {
+    for (j=(*index); j < rownnz; j++) {
       int tree_j = m->dof_treeid[colind[j]];
       if (tree_j != tree) {
         // found different tree
@@ -189,7 +188,7 @@ static int treeNext(const mjModel* m, const mjData* d, int tree, int i, int *ind
     int nv = m->nv;
 
     // scan row, look for different tree
-    for (j=j0; j < nv; j++) {
+    for (j=(*index); j < nv; j++) {
       if (d->efc_J[nv*i + j]) {
         int tree_j = m->dof_treeid[j];
         if (tree_j != tree) {
@@ -202,13 +201,94 @@ static int treeNext(const mjModel* m, const mjData* d, int tree, int i, int *ind
   }
 
   // save last index
-  if (index) *index = j;
+  *index = j;
 
   return tree_next;
 }
 
 
 
+// find first and possibly second nonegative tree ids in Jacobian row i
+//   if row i is special-cased (no more trees), return -1
+//   otherwise call treeNext, starting scan at index 0, return index
+static int treeFirst(const mjModel* m, const mjData* d, int tree[2], int i) {
+  int efc_type = d->efc_type[i];
+  int efc_id = d->efc_id[i];
+
+  // clear outputs
+  tree[0] = -1;
+  tree[1] = -1;
+
+  // ==== fast handling of special cases
+
+  // joint friction
+  if (efc_type == mjCNSTR_FRICTION_DOF) {
+    tree[0] = m->dof_treeid[efc_id];
+    return -1;
+  }
+
+  // joint limit
+  if (efc_type == mjCNSTR_LIMIT_JOINT) {
+    tree[0] = m->dof_treeid[m->jnt_dofadr[efc_id]];
+    return -1;
+  }
+
+  // contact
+  if (efc_type == mjCNSTR_CONTACT_FRICTIONLESS ||
+      efc_type == mjCNSTR_CONTACT_PYRAMIDAL ||
+      efc_type == mjCNSTR_CONTACT_ELLIPTIC) {
+    tree[0] = m->body_treeid[m->geom_bodyid[d->contact[efc_id].geom1]];
+    tree[1] = m->body_treeid[m->geom_bodyid[d->contact[efc_id].geom2]];
+
+    // handle static bodies
+    if (tree[0] < 0) {
+      if (tree[1] < 0) {
+        mjERROR("contact %d is between two static bodies", efc_id);  // SHOULD NOT OCCUR
+      } else {
+        int tmp = tree[0];
+        tree[0] = tree[1];
+        tree[1] = tmp;
+      }
+    }
+
+    return -1;
+  }
+
+  // connect or weld constraints
+  if (efc_type == mjCNSTR_EQUALITY) {
+    mjtEq eq_type = m->eq_type[efc_id];
+    if (eq_type == mjEQ_CONNECT || eq_type == mjEQ_WELD) {
+      tree[0] = m->body_treeid[m->eq_obj1id[efc_id]];
+      tree[1] = m->body_treeid[m->eq_obj2id[efc_id]];
+
+      // handle static bodies
+      if (tree[0] < 0) {
+        if (tree[1] < 0) {
+          mjERROR("equality %d is between two static bodies", efc_id);  // SHOULD NOT OCCUR
+        } else {
+          int tmp = tree[0];
+          tree[0] = tree[1];
+          tree[1] = tmp;
+        }
+      }
+
+      return -1;
+    }
+  }
+
+  // ==== generic case: scan Jacobian
+  int index = 0;
+  tree[0] = treeNext(m, d, -1, i, &index);
+
+  if (tree[0] < 0) {
+    mjERROR("no tree found for constraint %d", i);  // SHOULD NOT OCCUR
+  }
+
+  return index;
+}
+
+
+
 // add 0 edges, 1 self-edge or 2 flipped edges to array, increment treenedge
 //   return current number of edges
 static int addEdge(int* treenedge, int* edge, int nedge, int tree1, int tree2, int nedge_max) {
@@ -273,7 +353,6 @@ static int findEdges(const mjModel* m, const mjData* d, int* treenedge, int* edg
   int nefc = d->nefc;
   int efc_type = -1;
   int efc_id = -1;
-  int tree1, tree2;
 
   // clear treenedge
   memset(treenedge, 0, m->ntree*sizeof(int));
@@ -287,60 +366,34 @@ static int findEdges(const mjModel* m, const mjData* d, int* treenedge, int* edg
     efc_type = d->efc_type[i];
     efc_id = d->efc_id[i];
 
-    // ==== fast handling of special cases
+    int tree[2];
+    int index = treeFirst(m, d, tree, i);
+    int tree1 = tree[0];
+    int tree2 = tree[1];
 
-    // joint friction
+    // no more edges to find, add and continue
-    if (efc_type == mjCNSTR_FRICTION_DOF) {
+    if (index == -1) {
-      tree1 = m->dof_treeid[efc_id];
+      nedge = addEdge(treenedge, edge, nedge, tree1, tree2 == -1 ? tree1 : tree2, nedge_max);
-      nedge = addEdge(treenedge, edge, nedge, tree1, tree1, nedge_max);
       continue;
     }
 
-    // joint limit
+    // possibly more edges, scan Jacobian row
-    if (efc_type == mjCNSTR_LIMIT_JOINT) {
+    else {
-      tree1 = m->dof_treeid[m->jnt_dofadr[efc_id]];
+      tree2 = treeNext(m, d, tree1, i, &index);
-      nedge = addEdge(treenedge, edge, nedge, tree1, tree1, nedge_max);
-      continue;
-    }
 
-    // contact
+      if (tree2 == -1) {
-    if (efc_type == mjCNSTR_CONTACT_FRICTIONLESS ||
+        // 1 tree found: add self-edge
-        efc_type == mjCNSTR_CONTACT_PYRAMIDAL ||
+        nedge = addEdge(treenedge, edge, nedge, tree1, tree1, nedge_max);
-        efc_type == mjCNSTR_CONTACT_ELLIPTIC) {
+      } else {
-      tree1 = m->body_treeid[m->geom_bodyid[d->contact[efc_id].geom1]];
+        // 2 trees found: add edge, keep scanning and adding until no more trees
-      tree2 = m->body_treeid[m->geom_bodyid[d->contact[efc_id].geom2]];
-      nedge = addEdge(treenedge, edge, nedge, tree1, tree2, nedge_max);
-      continue;
-    }
-
-    // connect or weld constraints
-    if (efc_type == mjCNSTR_EQUALITY) {
-      mjtEq eq_type = m->eq_type[efc_id];
-      if (eq_type == mjEQ_CONNECT || eq_type == mjEQ_WELD) {
-        tree1 = m->body_treeid[m->eq_obj1id[efc_id]];
-        tree2 = m->body_treeid[m->eq_obj2id[efc_id]];
         nedge = addEdge(treenedge, edge, nedge, tree1, tree2, nedge_max);
-        continue;
+        int tree3 = treeNext(m, d, tree2, i, &index);
-      }
+        while (tree3 > -1 && tree3 != tree2) {
-    }
+          tree1 = tree2;
-
+          tree2 = tree3;
-    // ==== generic case: scan Jacobian
+          nedge = addEdge(treenedge, edge, nedge, tree1, tree2, nedge_max);
-    int index = 0;
+          tree3 = treeNext(m, d, tree2, i, &index);
-    tree1 = treeNext(m, d, -1, i, &index);
+        }
-    tree2 = treeNext(m, d, tree1, i, &index);
-
-    if (tree2 == -1) {
-      // 1 tree found: add self-edge
-      nedge = addEdge(treenedge, edge, nedge, tree1, tree1, nedge_max);
-    } else {
-      // 2 trees found: add edge, keep scanning and adding until no more trees
-      nedge = addEdge(treenedge, edge, nedge, tree1, tree2, nedge_max);
-      int tree3 = treeNext(m, d, tree2, i, &index);
-      while (tree3 > -1 && tree3 != tree2) {
-        tree1 = tree2;
-        tree2 = tree3;
-        nedge = addEdge(treenedge, edge, nedge, tree1, tree2, nedge_max);
-        tree3 = treeNext(m, d, tree2, i, &index);
       }
     }
   }
@@ -452,7 +505,9 @@ void mj_island(const mjModel* m, mjData* d) {
   // compute efc_island, island_efcnum
   memset(d->island_efcnum, 0, nisland*sizeof(int));
   for (int i=0; i < nefc; i++) {
-    int island = tree_island[treeNext(m, d, -1, i, NULL)];
+    int tree[2];
+    treeFirst(m, d, tree, i);
+    int island = tree_island[tree[0]];
     d->efc_island[i] = island;
     d->island_efcnum[island]++;
   }

test/engine/engine_island_test.cc

@@ -350,33 +350,33 @@ TEST_F(IslandTest, IslandEfc) {
     mj_step(model, data);
   }
 
-  // sizes
-  int nv      = model->nv;
-  int nefc    = data->nefc;
-  int nisland = data->nisland;
-
   // expect island structure to correspond to comment at top of xml
-  EXPECT_EQ(nisland, 3);
+  EXPECT_EQ(data->nisland, 4);
-  EXPECT_EQ(data->ne, 1);
+  EXPECT_EQ(data->ne, 4);
   EXPECT_EQ(data->nf, 2);
   EXPECT_EQ(data->nl, 1);
-  EXPECT_EQ(nefc, 24);
+  EXPECT_EQ(data->nefc, 27);
-  EXPECT_THAT(AsVector(data->dof_island, nv),
+
-              ElementsAre(0, -1, 0, 0, 0, 1, 2, 0, 1, 1, 1, 1, 1, 1));
+  mj_deleteData(data);
-  EXPECT_THAT(AsVector(data->island_dofnum, nisland), ElementsAre(5, 7, 1));
+  mj_deleteModel(model);
-  EXPECT_THAT(AsVector(data->island_dofadr, nisland), ElementsAre(0, 5, 12));
+}
-  EXPECT_THAT(AsVector(data->island_dofind, nv),
+
-              ElementsAre(0, 2, 3, 4, 7, 5, 8, 9, 10, 11, 12, 13, 6, -1));
+TEST_F(IslandTest, IslandEfcElliptic) {
-  EXPECT_THAT(AsVector(data->dof_islandind, nv),
+  const std::string xml_path = GetTestDataFilePath(kIlslandEfcPath);
-              ElementsAre(0, -1, 1, 2, 3, 0, 0, 4, 1, 2, 3, 4, 5, 6));
+  mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0);
-  EXPECT_THAT(AsVector(data->efc_island, nefc),
+  mjData* data = mj_makeData(model);
-              ElementsAre(0, 1, 2, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
+
-                          1, 1, 1, 1, 1, 1));
+  model->opt.cone = mjCONE_ELLIPTIC;
-  EXPECT_THAT(AsVector(data->island_efcnum, nisland), ElementsAre(6, 17, 1));
+  while (data->time < 0.2) {
-  EXPECT_THAT(AsVector(data->island_efcadr, nisland), ElementsAre(0, 6, 23));
+    mj_step(model, data);
-  EXPECT_THAT(AsVector(data->island_efcind, nefc),
+  }
-              ElementsAre(0, 3, 4, 5, 6, 7, 1, 8, 9, 10, 11, 12, 13, 14, 15, 16,
+  mj_forward(model, data);
-                          17, 18, 19, 20, 21, 22, 23, 2));
+
+  EXPECT_EQ(data->nisland, 4);
+  EXPECT_EQ(data->ne, 4);
+  EXPECT_EQ(data->nf, 2);
+  EXPECT_EQ(data->nl, 1);
+  EXPECT_EQ(data->nefc, 22);
 
   mj_deleteData(data);
   mj_deleteModel(model);

test/engine/testdata/island/island_efc.xml

@@ -1,10 +1,10 @@
 <!--
   test model for constraint islanding:
-- 1 equality
+- 4 equalities (1:joint, 3:connect)
 - 1 limit
 - 2 friction constraints
 - 1 unconstrained dof
-- 3 islands, 2 with non-contiguous dofs
+- 4 islands, 2 with non-contiguous dofs
 -->
 
 <mujoco>
@@ -64,9 +64,15 @@
       <freejoint/>
       <geom type="box" size=".03 .03 .03" pos="0.01 0.01 0.01"/>
     </body>
+
+    <body pos="-.5 0 .3" name="connect">
+      <freejoint/>
+      <geom type="box" size=".05 .05 .05"/>
+    </body>
   </worldbody>
 
   <equality>
     <joint joint1="hinge1" joint2="hinge2"/>
+    <connect body1="connect" body2="world" anchor="-.05 -.05 .05"/>
   </equality>
 </mujoco>