Add an ArrayList container type.

PiperOrigin-RevId: 558887266
Change-Id: I201978bfa5e6cfa2f4e495feaecf61374baf7b7b

src/engine/engine_util_container.c

@@ -0,0 +1,93 @@
+// Copyright 2023 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.
+
+#include "engine/engine_util_container.h"
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+#include <mujoco/mujoco.h>
+#include "engine/engine_io.h"
+
+// stack allocate and initialize new mjArrayList
+mjArrayList* mju_arrayListCreate(mjData* d, size_t element_size, size_t initial_capacity) {
+  mjArrayList* array_list = (mjArrayList*) mj_stackAllocBytes(d, sizeof(mjArrayList));
+  array_list->d = d;
+  array_list->element_size = element_size;
+  array_list->capacity = initial_capacity;
+  array_list->size = 0;
+  array_list->next_segment = NULL;
+
+  // allocate array list buffer
+  array_list->buffer = (void*) mj_stackAllocBytes(d, element_size * initial_capacity);
+  return array_list;
+}
+
+
+
+// returns total number of elements in mjArrayList
+size_t mju_arrayListSize(mjArrayList* array_list) {
+  mjArrayList* cursor = array_list;
+  size_t array_list_size = 0;
+  while (cursor) {
+    array_list_size += cursor->size;
+    cursor = cursor->next_segment;
+  }
+  return array_list_size;
+}
+
+
+
+// copies one element into an mjArrayList
+void mju_arrayListAdd(mjArrayList* array_list, void* element) {
+  mjArrayList* cursor = array_list;
+
+  // find a non full segment or add a new segment
+  while (cursor->size == cursor->capacity) {
+    if (cursor->next_segment == NULL) {
+      // add a new segment with twice the capacity of the last segment
+      cursor->next_segment = mju_arrayListCreate(
+          cursor->d, cursor->element_size, 2 * cursor->capacity);
+    }
+    cursor = cursor->next_segment;
+  }
+  // copy element into segment
+  memcpy(cursor->buffer + cursor->element_size * cursor->size, element, cursor->element_size);
+  ++cursor->size;
+}
+
+
+
+// returns pointer to element at index, NULL if out of bounds
+void* mju_arrayListAt(mjArrayList* array_list, size_t index) {
+  // if the index is larger than the current capacity, then it is in a later segment
+  mjArrayList* cursor = array_list;
+  size_t total_capacity = 0;
+  while (cursor != NULL && index >= total_capacity + cursor->capacity) {
+    total_capacity += cursor->capacity;
+    cursor = cursor->next_segment;
+  }
+
+  if (!cursor) {
+    return NULL;
+  }
+
+  if (index - total_capacity >= cursor->size) {
+    return NULL;
+  }
+
+  return cursor->buffer +
+      (cursor->element_size * (index - total_capacity));
+}
+

src/engine/engine_util_container.h

@@ -0,0 +1,64 @@
+// Copyright 2023 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_UTIL_CONTAINER_H_
+#define MUJOCO_SRC_ENGINE_ENGINE_UTIL_CONTAINER_H_
+
+#include <mujoco/mjdata.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+//---------------------------------- mjArrayList ---------------------------------------------------
+struct mjArrayList_ {
+  // pointer to mjData to allow for growth of the list
+  mjData* d;
+
+  // size of element
+  size_t element_size;
+
+  // maximum number of elements
+  size_t capacity;
+
+  // number of elements in this list
+  size_t size;
+
+  // pointer to the next segment of the array list, NULL if last
+  struct mjArrayList_* next_segment;
+
+  // buffer for data in this segment stored in d->arena
+  void* buffer;
+};
+typedef struct mjArrayList_ mjArrayList;
+
+// stack allocate and initialize new mjArrayList
+mjArrayList* mju_arrayListCreate(mjData* d, size_t element_size, size_t initial_capacity);
+
+// returns total number of elements in mjArrayList
+size_t mju_arrayListSize(mjArrayList* array_list);
+
+// copies an element into an mjArrayList
+void mju_arrayListAdd(mjArrayList* array_list, void* element);
+
+// returns a pointer to the element at the specified location of the arraylist
+// NULL returned if index is not in the mjArrayList
+void* mju_arrayListAt(mjArrayList* array_list, size_t index);
+
+#ifdef __cplusplus
+}
+#endif
+
+
+#endif  // MUJOCO_SRC_ENGINE_ENGINE_UTIL_CONTAINER_H_

test/engine/engine_util_container_test.cc

@@ -0,0 +1,60 @@
+// Copyright 2023 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.
+#include "src/engine/engine_util_container.h"
+
+#include <array>
+
+#include <mujoco/mjdata.h>
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include <mujoco/mujoco.h>
+#include "test/fixture.h"
+
+namespace mujoco {
+namespace {
+
+TEST(TestMjArrayList, TestMjArrayListSingleThreaded) {
+  constexpr char xml[] = "<mujoco/>";
+
+  std::array<char, 1024> error;
+  mjModel* m = LoadModelFromString(xml, error.data(), error.size());
+  ASSERT_THAT(m, testing::NotNull()) << "Failed to load model: "
+      << error.data();
+  mjData* d = mj_makeData(m);
+  mjMARKSTACK;
+  mjArrayList* array_list = mju_arrayListCreate(d, sizeof(int), 10);
+
+  for (int i = 0; i < 35; ++i) {
+    mju_arrayListAdd(array_list, &i);
+  }
+  EXPECT_EQ(mju_arrayListSize(array_list), 35);
+
+  // Approximately (3 * sizeof(int) + 3 * sizeof(mjArrayList)) / sizeof(mjtNum)
+  // However there is padding for alignment/etc.
+  EXPECT_EQ(d->maxuse_stack, 54);
+
+  for (int i = 0; i < 35; ++i) {
+    EXPECT_EQ(*(int*)mju_arrayListAt(array_list, i), i);
+  }
+
+  EXPECT_TRUE(mju_arrayListAt(array_list, 100) == NULL);
+
+
+  mjFREESTACK;
+  mj_deleteData(d);
+  mj_deleteModel(m);
+}
+
+}  // namespace
+}  // namespace mujoco