Kyle Bayes
2b0629d75f
PiperOrigin-RevId: 688519096 Change-Id: I05ab1576703c2458968ef81915ac682a8beac391
75 lines
5.2 KiB
C
75 lines
5.2 KiB
C
// Copyright 2024 DeepMind Technologies Limited
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
#ifndef MUJOCO_SRC_ENGINE_ENGINE_SORT_H_
|
|
#define MUJOCO_SRC_ENGINE_ENGINE_SORT_H_
|
|
|
|
// threshold size of a run to do insertion sort on
|
|
#define _mjRUNSIZE 32
|
|
|
|
// insertion sort sub-macro that runs on a sub-array [start, ..., end)
|
|
#define _mjINSERTION_SORT(type, arr, start, end, cmp, context) \
|
|
{ \
|
|
for (int j = start + 1; j < end; j++) { \
|
|
type tmp = arr[j]; \
|
|
int k = j - 1; \
|
|
for (; k >= start && cmp(arr + k, &tmp, context) > 0; k--) { \
|
|
arr[k + 1] = arr[k]; \
|
|
} \
|
|
arr[k + 1] = tmp; \
|
|
} \
|
|
}
|
|
|
|
// sub-macro that merges two sub-sorted arrays [start, ..., mid), [mid, ..., end) together
|
|
#define _mjMERGE(type, arr, buf, start, mid, end, cmp, context) \
|
|
{ \
|
|
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_
|