From 1aa375ef9afadb6096e88777f6d81e53aaaec6cb Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Mon, 4 Sep 2023 13:55:56 -0700
Subject: [PATCH] Allow RHS vector in `mj_mulM_island` to use uncompressed
 memory.

PiperOrigin-RevId: 562605681
Change-Id: If34bceed9eade59957ac9fd9b6a41fb736b58024
---
 src/engine/engine_support.c        | 18 ++++++++++++++----
 src/engine/engine_support.h        |  2 +-
 test/engine/engine_support_test.cc | 17 ++++++++++++++++-
 3 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c
index c58e5e06..14cd04d5 100644
--- a/src/engine/engine_support.c
+++ b/src/engine/engine_support.c
@@ -885,7 +885,8 @@ void mj_mulM(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec)
 
 
 // multiply vector by inertia matrix for one dof island
-void mj_mulM_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec, int island) {
+void mj_mulM_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec,
+                    int island, int flg_vecunc) {
   // if no island, call regular function
   if (island < 0) {
     mj_mulM(m, d, res, vec);
@@ -913,7 +914,11 @@ void mj_mulM_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum
     int adr = Madr[i];
 
     // diagonal
-    res[k] = M[adr]*vec[k];
+    if (flg_vecunc) {
+      res[k] = M[adr]*vec[i];
+    } else {
+      res[k] = M[adr]*vec[k];
+    }
 
     // simple dof: continue
     if (simplenum[i]) {
@@ -925,8 +930,13 @@ void mj_mulM_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum
     while (j >= 0) {
       adr++;
       int l = islandind[j];
-      res[k] += M[adr]*vec[l];
-      res[l] += M[adr]*vec[k];
+      if (flg_vecunc) {
+        res[k] += M[adr]*vec[j];
+        res[l] += M[adr]*vec[i];
+      } else {
+        res[k] += M[adr]*vec[l];
+        res[l] += M[adr]*vec[k];
+      }
 
       // advance to parent
       j = parentid[j];
diff --git a/src/engine/engine_support.h b/src/engine/engine_support.h
index 3b10fbe7..acacbeca 100644
--- a/src/engine/engine_support.h
+++ b/src/engine/engine_support.h
@@ -110,7 +110,7 @@ MJAPI void mj_mulM(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum*
 
 // multiply vector by inertia matrix for one dof island
 MJAPI void mj_mulM_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec,
-                          int island);
+                          int island, int flg_vecunc);
 
 // multiply vector by (inertia matrix)^(1/2)
 MJAPI void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec);
diff --git a/test/engine/engine_support_test.cc b/test/engine/engine_support_test.cc
index 1d7b8511..10c94734 100644
--- a/test/engine/engine_support_test.cc
+++ b/test/engine/engine_support_test.cc
@@ -493,8 +493,23 @@ TEST_F(SupportTest, MulMIsland) {
       vec_i[j] = vec[dofind[j]];
     }
 
+    // === compressed: use vec_i
+
     // multiply by Jacobian, for this island
-    mj_mulM_island(model, data, Mvec_i, vec_i, i);
+    int flg_vecunc = 0;
+    mj_mulM_island(model, data, Mvec_i, vec_i, i, flg_vecunc);
+
+    // expect corresponding values to match
+    for (int j=0; j < dofnum; j++) {
+      EXPECT_THAT(Mvec_i[j], DoubleNear(Mvec[dofind[j]], 1e-12));
+    }
+
+    // === uncompressed: use vec
+    mju_zero(Mvec_i, dofnum);  // clear output
+
+    // multiply by Jacobian, for this island
+    flg_vecunc = 1;
+    mj_mulM_island(model, data, Mvec_i, vec, i, flg_vecunc);
 
     // expect corresponding values to match
     for (int j=0; j < dofnum; j++) {