From dae74030fee08430580eed7393b56d44d7639041 Mon Sep 17 00:00:00 2001
From: Nimrod Gileadi <nimrod@google.com>
Date: Tue, 22 Aug 2023 04:03:09 -0700
Subject: [PATCH] Handle initial_capacity = 0 in mju_arrayListCreate.

Add const to ArrayList function arguments where appropriate.

PiperOrigin-RevId: 559064386
Change-Id: I8763974a60a05f34dadd56e953cefa40fe78df42
---
 src/engine/engine_util_container.c        |  9 +++---
 src/engine/engine_util_container.h        |  6 ++--
 test/engine/engine_util_container_test.cc | 36 +++++++++++++++++++----
 3 files changed, 39 insertions(+), 12 deletions(-)

diff --git a/src/engine/engine_util_container.c b/src/engine/engine_util_container.c
index eacd443a..69778969 100644
--- a/src/engine/engine_util_container.c
+++ b/src/engine/engine_util_container.c
@@ -23,6 +23,7 @@
 // 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));
+  initial_capacity = mjMAX(1, initial_capacity);
   array_list->d = d;
   array_list->element_size = element_size;
   array_list->capacity = initial_capacity;
@@ -37,8 +38,8 @@ mjArrayList* mju_arrayListCreate(mjData* d, size_t element_size, size_t initial_
 
 
 // returns total number of elements in mjArrayList
-size_t mju_arrayListSize(mjArrayList* array_list) {
-  mjArrayList* cursor = array_list;
+size_t mju_arrayListSize(const mjArrayList* array_list) {
+  const mjArrayList* cursor = array_list;
   size_t array_list_size = 0;
   while (cursor) {
     array_list_size += cursor->size;
@@ -70,9 +71,9 @@ void mju_arrayListAdd(mjArrayList* array_list, void* element) {
 
 
 // returns pointer to element at index, NULL if out of bounds
-void* mju_arrayListAt(mjArrayList* array_list, size_t index) {
+void* mju_arrayListAt(const 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;
+  const mjArrayList* cursor = array_list;
   size_t total_capacity = 0;
   while (cursor != NULL && index >= total_capacity + cursor->capacity) {
     total_capacity += cursor->capacity;
diff --git a/src/engine/engine_util_container.h b/src/engine/engine_util_container.h
index 30a506fa..67dcafd5 100644
--- a/src/engine/engine_util_container.h
+++ b/src/engine/engine_util_container.h
@@ -15,6 +15,8 @@
 #ifndef MUJOCO_SRC_ENGINE_ENGINE_UTIL_CONTAINER_H_
 #define MUJOCO_SRC_ENGINE_ENGINE_UTIL_CONTAINER_H_
 
+#include <stddef.h>
+
 #include <mujoco/mjdata.h>
 
 #ifdef __cplusplus
@@ -47,14 +49,14 @@ typedef struct mjArrayList_ 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);
+size_t mju_arrayListSize(const 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);
+void* mju_arrayListAt(const mjArrayList* array_list, size_t index);
 
 #ifdef __cplusplus
 }
diff --git a/test/engine/engine_util_container_test.cc b/test/engine/engine_util_container_test.cc
index 4c20eb8e..d813ed25 100644
--- a/test/engine/engine_util_container_test.cc
+++ b/test/engine/engine_util_container_test.cc
@@ -24,13 +24,12 @@
 namespace mujoco {
 namespace {
 
-TEST(TestMjArrayList, TestMjArrayListSingleThreaded) {
-  constexpr char xml[] = "<mujoco/>";
+using testing::NotNull;
 
+TEST(TestMjArrayList, TestMjArrayListSingleThreaded) {
   std::array<char, 1024> error;
-  mjModel* m = LoadModelFromString(xml, error.data(), error.size());
-  ASSERT_THAT(m, testing::NotNull()) << "Failed to load model: "
-      << error.data();
+  mjModel* m = LoadModelFromString("<mujoco/>", error.data(), error.size());
+  ASSERT_THAT(m, NotNull()) << "Failed to load model: " << error.data();
   mjData* d = mj_makeData(m);
   mjMARKSTACK;
   mjArrayList* array_list = mju_arrayListCreate(d, sizeof(int), 10);
@@ -48,8 +47,33 @@ TEST(TestMjArrayList, TestMjArrayListSingleThreaded) {
     EXPECT_EQ(*(int*)mju_arrayListAt(array_list, i), i);
   }
 
-  EXPECT_TRUE(mju_arrayListAt(array_list, 100) == NULL);
+  EXPECT_EQ(mju_arrayListAt(array_list, 35), nullptr);
+  EXPECT_EQ(mju_arrayListAt(array_list, 100), nullptr);
 
+  mjFREESTACK;
+  mj_deleteData(d);
+  mj_deleteModel(m);
+}
+
+TEST(TestMjArrayList, ZeroInitialCapacity) {
+  mjModel* m = LoadModelFromString("<mujoco/>", nullptr, 0);
+  ASSERT_THAT(m, NotNull()) << "Failed to load model";
+  mjData* d = mj_makeData(m);
+  mjMARKSTACK;
+  mjArrayList* array_list =
+      mju_arrayListCreate(d, sizeof(double), /*initial_capacity=*/0);
+  EXPECT_EQ(mju_arrayListSize(array_list), 0);
+
+  for (int i = 0; i < 35; ++i) {
+    double value = i;
+    mju_arrayListAdd(array_list, &value);
+  }
+  EXPECT_EQ(mju_arrayListSize(array_list), 35);
+
+  for (int i = 0; i < 35; ++i) {
+    EXPECT_EQ(*(double*)mju_arrayListAt(array_list, i), i);
+  }
+  EXPECT_EQ(mju_arrayListAt(array_list, 35), nullptr);
 
   mjFREESTACK;
   mj_deleteData(d);