From 5f3e22a45f53c269bed4eb1871c8aa1d50e64c0a Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Tue, 29 Aug 2023 05:15:58 -0700
Subject: [PATCH] Refactor a loop in engine_solver.c

PiperOrigin-RevId: 560997065
Change-Id: I26540e4f750eb18d9497bf35396adbb15cc3a122
---
 src/engine/engine_solver.c | 60 ++++++++++++++++++++------------------
 1 file changed, 32 insertions(+), 28 deletions(-)

diff --git a/src/engine/engine_solver.c b/src/engine/engine_solver.c
index c4dcf5b9..a3ee385f 100644
--- a/src/engine/engine_solver.c
+++ b/src/engine/engine_solver.c
@@ -929,39 +929,43 @@ static void CGeval(const mjModel* m, mjData* d, mjCGContext* ctx, mjCGPnt* p) {
   mjtNum quadTotal[3];
   mju_copy3(quadTotal, ctx->quadGauss);
 
-  // equality
-  for (int i=0; i < ne; i++) {
-    mju_addTo3(quadTotal, ctx->quad+3*i);
-  }
-
-  // friction
-  for (int i=ne; i < ne+nf; i++) {
-    // search point, friction loss, bound (Rf)
-    mjtNum start = ctx->Jaref[i], dir = ctx->Jv[i];
-    mjtNum x = start + alpha*dir;
-    mjtNum f = d->efc_frictionloss[i];
-    mjtNum Rf = d->efc_R[i]*f;
-
-    // -bound < x < bound : quadratic
-    if (-Rf < x && x < Rf) {
+  // process constraints
+  for (int i=0; i < nefc; i++) {
+    // equality
+    if (i < ne) {
       mju_addTo3(quadTotal, ctx->quad+3*i);
+      continue;
     }
 
-    // x < -bound : linear negative
-    else if (x <= -Rf) {
-      mjtNum qf[3] = {f*(-0.5*Rf-start), -f*dir, 0};
-      mju_addTo3(quadTotal, qf);
+
+    // friction
+    if (i < ne + nf) {
+      // search point, friction loss, bound (Rf)
+      mjtNum start = ctx->Jaref[i], dir = ctx->Jv[i];
+      mjtNum x = start + alpha*dir;
+      mjtNum f = d->efc_frictionloss[i];
+      mjtNum Rf = d->efc_R[i]*f;
+
+      // -bound < x < bound : quadratic
+      if (-Rf < x && x < Rf) {
+        mju_addTo3(quadTotal, ctx->quad+3*i);
+      }
+
+      // x < -bound : linear negative
+      else if (x <= -Rf) {
+        mjtNum qf[3] = {f*(-0.5*Rf-start), -f*dir, 0};
+        mju_addTo3(quadTotal, qf);
+      }
+
+      // bound < x : linear positive
+      else {
+        mjtNum qf[3] = {f*(-0.5*Rf+start), f*dir, 0};
+        mju_addTo3(quadTotal, qf);
+      }
+      continue;
     }
 
-    // bound < x : linear positive
-    else {
-      mjtNum qf[3] = {f*(-0.5*Rf+start), f*dir, 0};
-      mju_addTo3(quadTotal, qf);
-    }
-  }
-
-  // limit and contact
-  for (int i=ne+nf; i < nefc; i++) {
+    // limit and contact
     if (d->efc_type[i] == mjCNSTR_CONTACT_ELLIPTIC) {         // elliptic cone
       // extract contact info
       mjContact* con = d->contact + d->efc_id[i];