Refactor constraint tree iteration and island assignment.

- Simplified constraint tree discovery with a unified iterator.
- Filling in a new `efc_tree` array to avoid calling tree discovery again in the constraint phase.
- Implemented dof skipping in the dense case for efficiency.

PiperOrigin-RevId: 875440387
Change-Id: Id44a5a6fe2a80edaaf30935506a023a307e21303

src/engine/engine_island.c

@@ -137,24 +137,42 @@ int mj_floodFill(int* island, int nr, const int* rownnz, const int* rowadr, cons
 }
 
 
-// return id of next tree in Jacobian row i that is different from tree, -1 if not found
+// state of iterator for finding trees involved in a constraint
-//   start search from *index
+typedef struct {
-//   write the index of the found tree to *index
+  int trees[2];   // pre-calculated trees (special-cased constraints); -2: empty/sentinel
-//   if J is (dense/sparse) *index is the (column/nonzero) index, respectively
+  int jac_idx;    // generic scan: current lookup index in Jacobian row; -1: scan disabled
-static int treeNext(const mjModel* m, const mjData* d, int tree, int i, int *index) {
+  int tree_prev;  // generic scan: previous tree in ongoing scan
-  int tree_next = -1;
+} mjTreeIter;
-  int j;  // local loop variable, saved to *index
+
+
+// return next tree of constraint i from iterator; -2: no more trees
+static int treeNext(const mjModel* m, const mjData* d, int i, mjTreeIter* iter) {
+  // handle special cases
+  if (iter->trees[0] != -2) {
+    // get first tree, queue up second tree, return first tree
+    int tree = iter->trees[0];
+    iter->trees[0] = iter->trees[1];
+    iter->trees[1] = -2;
+    return tree;
+  }
+
+  // special case mode complete
+  if (iter->jac_idx == -1) {
+    return -2;
+  }
+
+  // generic scan mode
+  int j;
+  int tree_next = -2;
 
   // sparse
   if (mj_isSparse(m)) {
     int rownnz = d->efc_J_rownnz[i];
-    int* colind = d->efc_J_colind + d->efc_J_rowadr[i];
+    const int* colind = d->efc_J_colind + d->efc_J_rowadr[i];
-
+    for (j = iter->jac_idx; j < rownnz; j++) {
-    // loop over remaining nonzeros, look for different tree
-    for (j=(*index); j < rownnz; j++) {
       int tree_j = m->dof_treeid[colind[j]];
-      if (tree_j != tree) {
+      if (tree_j != iter->tree_prev) {
-        // found different tree
+        // found new tree
         tree_next = tree_j;
         break;
       }
@@ -164,125 +182,107 @@ static int treeNext(const mjModel* m, const mjData* d, int tree, int i, int *ind
   // dense
   else {
     int nv = m->nv;
-
+    const mjtNum* J = d->efc_J + nv * i;
-    // scan row, look for different tree
+    for (j = iter->jac_idx; j < nv; j++) {
-    for (j=(*index); j < nv; j++) {
+      if (J[j]) {
-      if (d->efc_J[nv*i + j]) {
         int tree_j = m->dof_treeid[j];
-        if (tree_j != tree) {
+        if (tree_j != iter->tree_prev) {
-          // found different tree
+          // found new tree
           tree_next = tree_j;
           break;
         }
+
+        // skip to end of tree's dof block
+        j = m->tree_dofadr[tree_j] + m->tree_dofnum[tree_j] - 1;
       }
     }
   }
 
-  // save last index
+  // update iterator state
-  *index = j;
+  iter->jac_idx = j;
+  if (tree_next != -2) {
+    iter->tree_prev = tree_next;
+  }
 
   return tree_next;
 }
 
 
-// find first and possibly second nonegative tree ids in Jacobian row i
+// initialize tree iterator, handle special cases
-//   if row i is special-cased (no more trees), return -1
+static void treeIterInit(const mjModel* m, const mjData* d, int i, mjTreeIter* iter) {
-//   otherwise call treeNext, starting scan at index 0, return index
+  iter->trees[0] = -2;
-static int treeFirst(const mjModel* m, const mjData* d, int tree[2], int i) {
+  iter->trees[1] = -2;
+  iter->jac_idx = -1;
+  iter->tree_prev = -1;
+
   int efc_type = d->efc_type[i];
   int efc_id = d->efc_id[i];
 
-  // clear outputs
+  // ==== special cases: fill iter->trees where possible
-  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];
+    iter->trees[0] = m->dof_treeid[efc_id];
-    return -1;
   }
 
   // joint limit
-  if (efc_type == mjCNSTR_LIMIT_JOINT) {
+  else if (efc_type == mjCNSTR_LIMIT_JOINT) {
-    tree[0] = m->dof_treeid[m->jnt_dofadr[efc_id]];
+    iter->trees[0] = m->dof_treeid[m->jnt_dofadr[efc_id]];
-    return -1;
   }
 
   // contact
-  if (efc_type == mjCNSTR_CONTACT_FRICTIONLESS ||
+  else if (efc_type == mjCNSTR_CONTACT_FRICTIONLESS ||
-      efc_type == mjCNSTR_CONTACT_PYRAMIDAL ||
+           efc_type == mjCNSTR_CONTACT_PYRAMIDAL ||
-      efc_type == mjCNSTR_CONTACT_ELLIPTIC) {
+           efc_type == mjCNSTR_CONTACT_ELLIPTIC) {
     int g1 = d->contact[efc_id].geom[0];
     int g2 = d->contact[efc_id].geom[1];
 
-    // no shortcut for flex contacts (handled in the generic case)
+    // geom-geom contact
-    if (g1 >=0 && g2 >= 0) {
+    if (g1 >= 0 && g2 >= 0) {
-      tree[0] = m->body_treeid[m->geom_bodyid[g1]];
+      iter->trees[0] = m->body_treeid[m->geom_bodyid[g1]];
-      tree[1] = m->body_treeid[m->geom_bodyid[g2]];
+      iter->trees[1] = m->body_treeid[m->geom_bodyid[g2]];
-
+      if (iter->trees[0] < 0 && iter->trees[1] < 0) {
-      // handle static bodies
+        mjERROR("contact %d is between two static bodies", efc_id);  // SHOULD NOT OCCUR
-      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;
+    // no shortcut for flex contacts: enable generic scan
+    else {
+      iter->jac_idx = 0;
     }
   }
 
   // connect or weld constraints
-  if (efc_type == mjCNSTR_EQUALITY) {
+  else if (efc_type == mjCNSTR_EQUALITY &&
-    mjtEq eq_type = m->eq_type[efc_id];
+           (m->eq_type[efc_id] == mjEQ_CONNECT ||
-    if (eq_type == mjEQ_CONNECT || eq_type == mjEQ_WELD) {
+            m->eq_type[efc_id] == mjEQ_WELD)) {
-      int b1 = m->eq_obj1id[efc_id];
+    int b1 = m->eq_obj1id[efc_id];
-      int b2 = m->eq_obj2id[efc_id];
+    int b2 = m->eq_obj2id[efc_id];
 
-      // get body ids if using site semantics
+    // get body ids if using site semantics
-      if (m->eq_objtype[efc_id] == mjOBJ_SITE) {
+    if (m->eq_objtype[efc_id] == mjOBJ_SITE) {
-        b1 = m->site_bodyid[b1];
+      b1 = m->site_bodyid[b1];
-        b2 = m->site_bodyid[b2];
+      b2 = m->site_bodyid[b2];
-      }
+    }
 
-      tree[0] = m->body_treeid[b1];
+    // get trees
-      tree[1] = m->body_treeid[b2];
+    iter->trees[0] = m->body_treeid[b1];
-
+    iter->trees[1] = m->body_treeid[b2];
-      // handle static bodies
+    if (iter->trees[0] < 0 && iter->trees[1] < 0) {
-      if (tree[0] < 0) {
+      mjERROR("equality %d is between two static bodies", efc_id);  // SHOULD NOT OCCUR
-        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
+  // otherwise enable generic scan
-  int index = 0;
+  else {
-  tree[0] = treeNext(m, d, -1, i, &index);
+    iter->jac_idx = 0;
-
-  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
+// 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* edge, mjtByte* tree_tree, int ntree, int tree1, int tree2) {
+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;
@@ -299,12 +299,12 @@ static int addEdge(int* rownnz, int* edge, mjtByte* tree_tree, int ntree, int tr
 
   // add edge
   tree_tree[tree1*ntree + tree2] = 1;
-  edge[tree1*ntree + rownnz[tree1]++] = tree2;
+  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;
-    edge[tree2*ntree + rownnz[tree2]++] = tree1;
+    colind[tree2*ntree + rownnz[tree2]++] = tree1;  // uncompressed format, rowadr is known
     return 2;
   }
 
@@ -312,9 +312,10 @@ static int addEdge(int* rownnz, int* edge, mjtByte* tree_tree, int ntree, int tr
 }
 
 
-// find tree-tree edges (column indices), increment row nonzeros, return total number of edges
+// 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 ntree) {
+                     int* rownnz, int* colind, mjtByte* tree_tree, int* efc_tree, int ntree) {
   int nefc = d->nefc;
   int nnz = 0;
   int efc_type = -1;
@@ -323,6 +324,7 @@ static int findEdges(const mjModel* m, const mjData* d,
   // clear row nonzeros
   mju_zeroInt(rownnz, ntree);
 
+  // iterate over constraints, compute tree-tree edges, assign efc_tree
   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]) {
@@ -330,36 +332,41 @@ static int findEdges(const mjModel* m, const mjData* d,
       if (!(efc_type == mjCNSTR_EQUALITY &&
             (m->eq_type[efc_id] == mjEQ_FLEX ||
              m->eq_type[efc_id] == mjEQ_FLEXVERT))) {
+        // copy tree assignment from previous constraint and continue
+        efc_tree[i] = efc_tree[i-1];
         continue;
       }
     }
     efc_type = d->efc_type[i];
     efc_id = d->efc_id[i];
 
-    int tree[2];
+    // initialize tree iterator
-    int index = treeFirst(m, d, tree, i);
+    mjTreeIter iter;
-    int tree1 = tree[0];
+    treeIterInit(m, d, i, &iter);
-    int tree2 = tree[1];
 
-    // simple constraint (no more edges), add edge and continue
+    // iterate over trees involved in constraint i
-    if (index == -1) {
+    int tree1 = treeNext(m, d, i, &iter);
-      nnz += addEdge(rownnz, colind, tree_tree, ntree, tree1, tree2);
+    if (tree1 != -2) {
-      continue;
+      int tree2 = treeNext(m, d, i, &iter);
-    }
 
-    // add first edge
+      // assign tree to constraint, one of (tree1, tree2) must be non-negative
-    tree2 = treeNext(m, d, tree1, i, &index);
+      efc_tree[i] = tree1 >= 0 ? tree1 : tree2;
-    nnz += addEdge(rownnz, colind, tree_tree, ntree, tree1, tree2);
+      if (efc_tree[i] < 0) {
-
+        mjERROR("constraint %d is between two static bodies", i);  // SHOULD NOT OCCUR
-    // possibly more edges, scan Jacobian row
-    if (tree2 != -1) {
-      int tree3 = treeNext(m, d, tree2, i, &index);
-      while (tree3 > -1 && tree3 != tree2) {
-        tree1 = tree2;
-        tree2 = tree3;
-        nnz += addEdge(rownnz, colind, tree_tree, ntree, tree1, tree2);
-        tree3 = treeNext(m, d, tree2, i, &index);
       }
+
+      // add one edge or continue to search for more edges
+      if (tree2 == -2) {
+        nnz += addEdge(rownnz, colind, tree_tree, ntree, tree1, -1);
+      } else {
+        while (tree2 != -2) {
+          nnz += addEdge(rownnz, colind, tree_tree, ntree, tree1, tree2);
+          tree1 = tree2;
+          tree2 = treeNext(m, d, i, &iter);
+        }
+      }
+    } else {
+      mjERROR("no tree found for constraint %d", i);  // SHOULD NOT OCCUR
     }
   }
 
@@ -395,8 +402,11 @@ void mj_island(const mjModel* m, mjData* d) {
     rowadr[r] = r * ntree;
   }
 
+  // first non-negative tree index of each constraint, used later for computing efc_island
+  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, ntree);
+  int nnz = findEdges(m, d, rownnz, colind, tree_tree, efc_tree, ntree);
 
   // discover islands
   int* tree_island = mjSTACKALLOC(d, ntree, int);
@@ -532,9 +542,7 @@ 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 tree[2];
+    int island = tree_island[efc_tree[i]];
-    treeFirst(m, d, tree, i);
-    int island = tree_island[tree[0]];
     d->efc_island[i] = island;
     d->island_nefc[island]++;
     switch (d->efc_type[i]) {