Replace mjQUICKSORT with faster, native sorting function. fixes #1638

PiperOrigin-RevId: 688519096
Change-Id: I05ab1576703c2458968ef81915ac682a8beac391

src/engine/engine_collision_driver.c

@@ -224,10 +224,8 @@ int mj_isElemActive(const mjModel* m, int f, int e) {
 //----------------------------- collision detection entry point ------------------------------------
 
 // compare contact pairs by their geom/elem/vert IDs
-quicksortfunc(contactcompare, context, el1, el2) {
+static inline int contactcompare(const mjContact* c1, const mjContact* c2, void* context) {
   const mjModel* m = (const mjModel*) context;
-  mjContact* c1 = (mjContact*)el1;
-  mjContact* c2 = (mjContact*)el2;
 
   // get colliding object ids
   int con1_obj1 = c1->geom[0] >= 0 ? c1->geom[0] : (c1->elem[0] >= 0 ? c1->elem[0] : c1->vert[0]);
@@ -258,6 +256,9 @@ quicksortfunc(contactcompare, context, el1, el2) {
   return 0;
 }
 
+// define contactSort function for sorting contacts
+mjSORT(contactSort, mjContact, contactcompare)
+
 
 
 // main collision function
@@ -364,8 +365,14 @@ void mj_collision(const mjModel* m, mjData* d) {
       int ncon_after = d->ncon;
 
       // sort contacts
-      mjQUICKSORT(d->contact + ncon_before, ncon_after - ncon_before,
+      int n = ncon_after - ncon_before;
-                  sizeof(mjContact), contactcompare, (void*) m);
+      if (n > 1) {
+        mj_markStack(d);
+        mjContact* buf = (mjContact*)mj_stackAllocByte(d, n * sizeof(mjContact),
+                                                       _Alignof(mjContact));
+        contactSort(d->contact + ncon_before, buf, n, (void*)m);
+        mj_freeStack(d);
+      }
     }
 
     // process bodyflex pair: all-to-all
@@ -1006,10 +1013,7 @@ typedef struct _mjtSAP mjtSAP;
 
 
 // comparison function for SAP
-quicksortfunc(SAPcompare, context, el1, el2) {
+static inline int SAPcmp(mjtSAP* obj1, mjtSAP* obj2, void* context) {
-  mjtSAP* obj1 = (mjtSAP*)el1;
-  mjtSAP* obj2 = (mjtSAP*)el2;
-
   if (obj1->value < obj2->value) {
     return -1;
   } else if (obj1->value == obj2->value) {
@@ -1019,6 +1023,8 @@ quicksortfunc(SAPcompare, context, el1, el2) {
   }
 }
 
+// define SAPsort function for sorting SAP sorting
+mjSORT(SAPsort, mjtSAP, SAPcmp)
 
 
 // given list of axis-aligned bounding boxes in AAMM (xmin[3], xmax[3]) format,
@@ -1043,7 +1049,8 @@ static int mj_SAP(mjData* d, const mjtNum* aamm, int n, int axis, int* pair, int
   }
 
   // sort along specified axis
-  mjQUICKSORT(sortbuf, 2*n, sizeof(mjtSAP), SAPcompare, 0);
+  mjtSAP* buf = (mjtSAP*) mj_stackAllocByte(d, 2*n*sizeof(mjtSAP), _Alignof(mjtSAP));
+  SAPsort(sortbuf, buf, 2*n, NULL);
 
   // define the other two axes
   int axisA, axisB;
@@ -1133,19 +1140,18 @@ static void updateCov(mjtNum cov[9], const mjtNum vec[3], const mjtNum cen[3]) {
 
 
 // comparison function for unsigned ints
-quicksortfunc(uintcompare, context, el1, el2) {
+static inline int uintcmp(int* i, int* j, void* context) {
-  unsigned int n1 = *(unsigned int*)el1;
+  if ((unsigned) *i < (unsigned) *j) {
-  unsigned int n2 = *(unsigned int*)el2;
-
-  if (n1 < n2) {
     return -1;
-  } else if (n1 == n2) {
+  } else if (*i == *j) {
     return 0;
   } else {
     return 1;
   }
 }
 
+// define bfsort function for sorting bodyflex pairs
+mjSORT(bfsort, int, uintcmp)
 
 
 // broadphase collision detector
@@ -1281,8 +1287,9 @@ int mj_broadphase(const mjModel* m, mjData* d, int* bfpair, int maxpair) {
 endbroad:
 
   // sort bodyflex pairs by signature
-  if (npair) {
+  if (npair > 1) {
-    mjQUICKSORT(bfpair, npair, sizeof(int), uintcompare, 0);
+    int* buf = mj_stackAllocInt(d, npair);
+    bfsort(bfpair, buf, npair, NULL);
   }
 
   mj_freeStack(d);

src/engine/engine_collision_sdf.c

@@ -306,11 +306,10 @@ static void undoTransformation(const mjModel* m, const mjData* d, int g,
 //---------------------------- narrow phase -----------------------------------------------
 
 // comparison function for contact sorting
-quicksortfunc(distcompare, dist, i1, i2) {
+static inline int distcmp(int* i, int* j, void* context) {
-  mjtNum d1 = ((mjtNum*)dist)[*(int*)i1];
+  mjtNum d1 = ((mjtNum*)context)[*i];
-  mjtNum d2 = ((mjtNum*)dist)[*(int*)i2];
+  mjtNum d2 = ((mjtNum*)context)[*j];
-
+    if (d1 < d2) {
-  if (d1 < d2) {
     return -1;
   } else if (d1 == d2) {
     return 0;
@@ -319,6 +318,9 @@ quicksortfunc(distcompare, dist, i1, i2) {
   }
 }
 
+// define distSort function for contact sorting
+mjSORT(distSort, int, distcmp)
+
 // check if the collision point already exists
 static int isknown(const mjtNum* points, const mjtNum x[3], int cnt) {
   for (int i = 0; i < cnt; i++) {
@@ -642,7 +644,10 @@ int mjc_MeshSDF(const mjModel* m, const mjData* d, mjContact* con, int g1, int g
   }
 
   // sort contacts using depth
-  mjQUICKSORT(index, ncandidate, sizeof(int), distcompare, dist);
+  if (ncandidate > 1) {
+    int buf[MAXMESHPNT];
+    distSort(index, buf, ncandidate, dist);
+  }
 
   // add only the first mjMAXCONPAIR pairs
   for (int i=0; i < mju_min(ncandidate, mjMAXCONPAIR); i++) {

src/engine/engine_sort.h

@@ -15,46 +15,60 @@
 #ifndef MUJOCO_SRC_ENGINE_ENGINE_SORT_H_
 #define MUJOCO_SRC_ENGINE_ENGINE_SORT_H_
 
-#if !defined(__cplusplus)
+// threshold size of a run to do insertion sort on
-#include <stddef.h>
+#define _mjRUNSIZE 32
-#include <stdlib.h>
 
-// sorting functions using q_sort_s/r
+// insertion sort sub-macro that runs on a sub-array [start, ..., end)
-#ifdef _WIN32
+#define _mjINSERTION_SORT(type, arr, start, end, cmp, context)                                     \
-#define mjQUICKSORT(buf, elnum, elsz, func, context) \
+{                                                                                                  \
-  qsort_s(buf, elnum, elsz, func, context)
+  for (int j = start + 1; j < end; j++) {                                                          \
-#define quicksortfunc(name, context, el1, el2) \
+    type tmp = arr[j];                                                                             \
-  static int name(void* context, const void* el1, const void* el2)
+    int k = j - 1;                                                                                 \
-#else  // assumes POSIX
+    for (; k >= start && cmp(arr + k, &tmp, context) > 0; k--) {                                   \
-#ifdef __APPLE__
+      arr[k + 1] = arr[k];                                                                         \
-#define mjQUICKSORT(buf, elnum, elsz, func, context) \
+    }                                                                                              \
-  qsort_r(buf, elnum, elsz, context, func)
+    arr[k + 1] = tmp;                                                                              \
-#define quicksortfunc(name, context, el1, el2) \
+  }                                                                                                \
-  static int name(void* context, const void* el1, const void* el2)
-#else  // non-Apple
-#define mjQUICKSORT(buf, elnum, elsz, func, context) \
-  qsort_r(buf, elnum, elsz, func, context)
-#define quicksortfunc(name, context, el1, el2) \
-  static int name(const void* el1, const void* el2, void* context)
-#endif
-#endif
-#else
-#include <algorithm>
-#include <cstddef>
-#include <cstdlib>
-
-// sorting function using std::sort
-template <typename T>
-void mjQUICKSORT(T* buf, size_t elnum, size_t elsz,
-                 int (*compare)(const void* a, const void* b, void* c),
-                 void* context) {
-  std::sort(buf, buf + elnum, [compare, context](const T& a, const T& b) {
-    return compare(&a, &b, context) < 0;
-  });
 }
 
-#define quicksortfunc(name, context, el1, el2) \
+// sub-macro that merges two sub-sorted arrays [start, ..., mid), [mid, ..., end) together
-  static int name(const void* el1, const void* el2, void* context)
+#define _mjMERGE(type, arr, buf, start, mid, end, cmp, context)                                    \
-#endif
+{                                                                                                  \
+  int len1 = mid - start, len2 = end - mid;                                                        \
+  type* left = buf, *right = buf + len1;                                                           \
+  for (int i = 0; i < len1; i++) left[i] = arr[start + i];                                         \
+  for (int i = 0; i < len2; i++) right[i] = arr[mid + i];                                          \
+  int i = 0, j = 0, k = start;                                                                     \
+  while (i < len1 && j < len2) {                                                                   \
+    if (cmp(left + i, right + j, context) <= 0) {                                                  \
+       arr[k++] = left[i++];                                                                       \
+    } else {                                                                                       \
+      arr[k++] = right[j++];                                                                       \
+    }                                                                                              \
+  }                                                                                                \
+  while (i < len1) arr[k++] = left[i++];                                                           \
+  while (j < len2) arr[k++] = right[j++];                                                          \
+}
+
+// defines an inline stable sorting function via tiled merge sorting (timsort)
+// function is of form:
+// void name(type* arr, type* buf, int n, void* context)
+// where arr is the array of size n to be sorted inplace and buf is a buffer of size n.
+#define mjSORT(name, type, cmp)                                                                    \
+  static inline void name(type* arr, type* buf, int n, void* context) {                            \
+    for (int start = 0; start < n; start += _mjRUNSIZE) {                                          \
+      int end = (start + _mjRUNSIZE < n) ? start + _mjRUNSIZE : n;                                 \
+      _mjINSERTION_SORT(type, arr, start, end, cmp, context);                                      \
+    }                                                                                              \
+    for (int len = _mjRUNSIZE; len < n; len *= 2) {                                                \
+      for (int start = 0; start < n; start += 2*len) {                                             \
+        int mid = start + len;                                                                     \
+        int end = (start + 2*len < n) ? start + 2*len : n;                                         \
+        if (mid < end) {                                                                           \
+          _mjMERGE(type, arr, buf, start, mid, end, cmp, context);                                 \
+        }                                                                                          \
+      }                                                                                            \
+    }                                                                                              \
+  }
 
 #endif  // MUJOCO_SRC_ENGINE_ENGINE_SORT_H_

src/render/render_gl3.c

@@ -766,10 +766,10 @@ static void setView(int view, mjrRect viewport, const mjvScene* scn, const mjrCo
 
 
 // comparison function for geom sorting
-quicksortfunc(geomcompare, context, el1, el2) {
+static inline int geomcmp(int* i, int* j, void* context) {
   mjvGeom* geom = (mjvGeom*) context;
-  float d1 = geom[*(int*)el1].camdist;
+  float d1 = geom[*i].camdist;
-  float d2 = geom[*(int*)el2].camdist;
+  float d2 = geom[*j].camdist;
 
   if (d1 < d2) {
     return -1;
@@ -780,6 +780,9 @@ quicksortfunc(geomcompare, context, el1, el2) {
   }
 }
 
+// define geomSort function for sorting geoms
+mjSORT(geomSort, int, geomcmp)
+
 
 
 // adjust light n position and direction
@@ -907,7 +910,11 @@ void mjr_render(mjrRect viewport, mjvScene* scn, const mjrContext* con) {
   }
 
   // sort transparent geoms according to distance to camera
-  mjQUICKSORT(scn->geomorder, nt, sizeof(int), geomcompare, scn->geoms);
+  if (nt > 1) {
+    int *buf = (int*) mju_malloc(nt * sizeof(int));
+    geomSort(scn->geomorder, buf, nt, scn->geoms);
+    mju_free(buf);
+  }
 
   // allow only one reflective geom
   int j = 0;

test/engine/engine_sort_test.cc

@@ -14,11 +14,12 @@
 
 #include "src/engine/engine_sort.h"
 
-#include <cstddef>
+#include <algorithm>
+#include <array>
+#include <utility>
 #include <vector>
 
 #include <gtest/gtest.h>
-#include "src/engine/engine_sort.h"
 #include "test/fixture.h"
 
 namespace mujoco {
@@ -26,23 +27,57 @@ namespace {
 
 using EngineSortTest = MujocoTest;
 
+constexpr int kPermutationSize = 5;  // 5! = 120 checks
+
+constexpr int factorial() {
+  int n = 1;
+  for (int i = 1; i <= kPermutationSize; ++i) {
+    n *= i;
+  }
+  return n;
+}
+
 struct IntStruct {
   int value;
 };
 
-quicksortfunc(int_compare, context, x, y) {
+// computes next permutation in lexicographical order in place
-  int a = *(int*)x;
+// returns false if there is no next permutation
-  int b = *(int*)y;
+bool NextPermutation(std::array<int, kPermutationSize>& arr) {
-  if (a < b) {
+  int j;
+  for (j = arr.size() - 2; j >= 0; --j) {
+    if (arr[j] < arr[j + 1]) {
+      break;
+    }
+  }
+
+  // last permutation i.e. [n, n - 1, ..., 3, 2, 1]
+  if (j < 0) {
+    return false;
+  }
+
+  for (int l = arr.size() - 1; l > -1; --l) {
+    if (arr[j] < arr[l]) {
+      std::swap(arr[j], arr[l]);
+      std::reverse(arr.begin() + j + 1, arr.end());
+      return true;
+    }
+  }
+  return true;
+}
+
+int IntCompare(int* i, int* j, void* context) {
+  if (*i < *j) {
     return -1;
-  } else if (a == b) {
+  } else if (*i == *j) {
     return 0;
   } else {
     return 1;
   }
 }
+mjSORT(IntSort, int, IntCompare)
 
-quicksortfunc(intstruct_compare, context, x, y) {
+int IntStructCompare(const IntStruct* x, const IntStruct* y, void* context) {
   IntStruct* a = (IntStruct*)x;
   IntStruct* b = (IntStruct*)y;
   if (a->value < b->value) {
@@ -53,18 +88,40 @@ quicksortfunc(intstruct_compare, context, x, y) {
     return 1;
   }
 }
+mjSORT(IntStructSort, IntStruct, IntStructCompare)
 
 TEST_F(EngineSortTest, Sort) {
-  // test int
+  int total = 0;
-  std::vector<int> x = {1, 3, 2};
+  std::array<int, kPermutationSize> initial, buf;
-  mjQUICKSORT(x.data(), x.size(), sizeof(int), int_compare, x.data());
+  for (int i = 0; i < kPermutationSize; ++i) {
-  EXPECT_EQ(x[0], 1);
+    initial[i] = i + 1;
-  EXPECT_EQ(x[1], 2);
+  }
-  EXPECT_EQ(x[2], 3);
+  std::array<int, kPermutationSize> arr = initial;
+  do {
+    ++total;
+    std::array<int, kPermutationSize> sorted_arr = arr;
+    IntSort(sorted_arr.data(), buf.data(), sorted_arr.size(), nullptr);
+    EXPECT_EQ(sorted_arr, initial);
+  } while (NextPermutation(arr));
+  ASSERT_EQ(total, factorial());
+}
 
-  // test custom struct with mjQUICKSORT
+TEST_F(EngineSortTest, LargeSort) {
+  std::array<int, 2500> arr, buf;
+  int n = 1;
+  for (int i = arr.size() - 1; i >= 0; --i) {
+    arr[i] = n++;
+  }
+  IntSort(arr.data(), buf.data(), arr.size(), nullptr);
+  for (int i = 0; i < arr.size(); ++i) {
+    EXPECT_EQ(arr[i], i + 1);
+  }
+}
+
+TEST_F(EngineSortTest, SortStruct) {
   std::vector<IntStruct> y = {{1}, {3}, {2}};
-  mjQUICKSORT(y.data(), y.size(), sizeof(IntStruct), intstruct_compare, NULL);
+  std::vector<IntStruct> buf(3);
+  IntStructSort(y.data(), buf.data(), y.size(), nullptr);
   EXPECT_EQ(y[0].value, 1);
   EXPECT_EQ(y[1].value, 2);
   EXPECT_EQ(y[2].value, 3);