Add Flex component.

PiperOrigin-RevId: 572830650
Change-Id: I6908228087b7b9683be3506c8d9cdc725ed5dcd5

src/engine/engine_util_sparse.c

@@ -368,6 +368,176 @@ void mju_combineSparseInc(mjtNum* dst, const mjtNum* src, int n, mjtNum a, mjtNu
 
 
 
+// dst += src, only at common non-zero indices
+void mju_addToSparseInc(mjtNum* dst, const mjtNum* src,
+                        int nnzdst, const int* inddst,
+                        int nnzsrc, const int* indsrc) {
+  if (!nnzdst || !nnzsrc) {
+    return;
+  }
+
+  int adrs = 0, adrd = 0, inds = indsrc[0], indd = inddst[0];
+  while (1) {
+    // common non-zero index
+    if (inds==indd) {
+      // add
+      dst[adrd] += src[adrs];
+
+      // advance src
+      if (++adrs<nnzsrc) {
+        inds = indsrc[adrs];
+      } else {
+        return;
+      }
+
+      // advance dst
+      if (++adrd<nnzdst) {
+        indd = inddst[adrd];
+      } else {
+        return;
+      }
+    }
+
+    // src non-zero index smaller: advance src
+    else if (inds<indd) {
+      if (++adrs<nnzsrc) {
+        inds = indsrc[adrs];
+      } else {
+        return;
+      }
+    }
+
+    // dst non-zero index smaller: advance dst
+    else {
+      if (++adrd<nnzdst) {
+        indd = inddst[adrd];
+      } else {
+        return;
+      }
+    }
+  }
+}
+
+
+
+// add to sparse matrix: dst = dst + scl*src, return nnz of result
+int mju_addToSparseMat(mjtNum* dst, const mjtNum* src, int n, int nrow, mjtNum scl,
+                       int dst_nnz, int src_nnz, int* dst_ind, const int* src_ind,
+                       mjtNum* buf, int* buf_ind) {
+  // check for identical pattern
+  if (dst_nnz==src_nnz) {
+    if (dst_nnz==0) {
+      return 0;
+    }
+    if (mju_compare(dst_ind, src_ind, dst_nnz)) {
+      // combine mjtNum data directly
+      mju_addToScl(dst, src, scl, nrow*dst_nnz);
+      return dst_nnz;
+    }
+  }
+
+  // prepare to merge scr and dst into buf^T
+  int si = 0, di = 0, nnz = 0;
+  int sadr = src_nnz ? src_ind[0] : n+1;
+  int dadr = dst_nnz ? dst_ind[0] : n+1;
+
+  // merge matrices
+  while (si<src_nnz || di<dst_nnz) {
+    // both
+    if (sadr==dadr) {
+      for (int k=0; k<nrow; k++) {
+        buf[nrow*nnz + k] = dst[di + k*dst_nnz] + scl*src[si + k*src_nnz];
+      }
+
+      buf_ind[nnz++] = sadr;
+      si++;
+      di++;
+      sadr = si<src_nnz ? src_ind[si] : n+1;
+      dadr = di<dst_nnz ? dst_ind[di] : n+1;
+    }
+
+    // dst only
+    else if (dadr<sadr) {
+      for (int k=0; k<nrow; k++) {
+        buf[nrow*nnz + k] = dst[di + k*dst_nnz];
+      }
+
+      buf_ind[nnz++] = dadr;
+      di++;
+      dadr = di<dst_nnz ? dst_ind[di] : n+1;
+    }
+
+    // src only
+    else {
+      for (int k=0; k<nrow; k++) {
+        buf[nrow*nnz + k] = scl*src[si + k*src_nnz];
+      }
+
+      buf_ind[nnz++] = sadr;
+      si++;
+      sadr = si<src_nnz ? src_ind[si] : n+1;
+    }
+  }
+
+  // copy transposed buf into dst
+  mju_transpose(dst, buf, nnz, nrow);
+  mju_copyInt(dst_ind, buf_ind, nnz);
+
+  return nnz;
+}
+
+
+
+// add(merge) two chains
+int mju_addChains(int* res, int n, int NV1, int NV2,
+                  const int* chain1, const int* chain2) {
+  // check for identical pattern
+  if (NV1==NV2) {
+    if (NV1==0) {
+      return 0;
+    }
+    if (mju_compare(chain1, chain2, NV1)) {
+      mju_copyInt(res, chain1, NV1);
+      return NV1;
+    }
+  }
+
+  // prepare to merge
+  int i1 = 0, i2 = 0, NV = 0;
+  int adr1 = NV1 ? chain1[0] : n+1;
+  int adr2 = NV2 ? chain2[0] : n+1;
+
+  // merge chains
+  while (i1<NV1 || i2<NV2) {
+    // both
+    if (adr1==adr2) {
+      res[NV++] = adr1;
+      i1++;
+      i2++;
+      adr1 = i1<NV1 ? chain1[i1] : n+1;
+      adr2 = i2<NV2 ? chain2[i2] : n+1;
+    }
+
+    // chain1 only
+    else if (adr1<adr2) {
+      res[NV++] = adr1;
+      i1++;
+      adr1 = i1<NV1 ? chain1[i1] : n+1;
+    }
+
+    // chain2 only
+    else {
+      res[NV++] = adr2;
+      i2++;
+      adr2 = i2<NV2 ? chain2[i2] : n+1;
+    }
+  }
+
+  return NV;
+}
+
+
+
 // compress layout of sparse matrix
 void mju_compressSparse(mjtNum* mat, int nr, int nc, int* rownnz, int* rowadr, int* colind) {
   rowadr[0] = 0;