From 4270f73aee8de940c0a449284d2c6568cf2c38ec Mon Sep 17 00:00:00 2001
From: Kevin King <kcking@google.com>
Date: Thu, 11 Jan 2024 09:12:31 -0800
Subject: [PATCH] Add `upb_Array_GetMutable` accessor

PiperOrigin-RevId: 597573362
---
 upb/message/accessors_test.cc | 23 +++++++++++++++++++++++
 upb/message/array.c           |  9 +++++++++
 upb/message/array.h           |  4 ++++
 3 files changed, 36 insertions(+)

diff --git a/upb/message/accessors_test.cc b/upb/message/accessors_test.cc
index 9db6eb5ddc..59a5e51ad7 100644
--- a/upb/message/accessors_test.cc
+++ b/upb/message/accessors_test.cc
@@ -49,6 +49,7 @@ const uint32_t kFieldOptionalBool = 13;
 const uint32_t kFieldOptionalString = 14;
 const uint32_t kFieldOptionalNestedMessage = 18;
 const uint32_t kFieldOptionalRepeatedInt32 = 31;
+const uint32_t kFieldOptionalRepeatedNestedMessage = 48;
 const uint32_t kFieldOptionalNestedMessageA = 1;
 const uint32_t kFieldOptionalOneOfUInt32 = 111;
 const uint32_t kFieldOptionalOneOfString = 113;
@@ -402,6 +403,28 @@ TEST(GeneratedCode, RepeatedScalar) {
   upb_Arena_Free(arena);
 }
 
+TEST(GeneratedCode, RepeatedMessage) {
+  upb_Arena* arena = upb_Arena_New();
+  protobuf_test_messages_proto2_TestAllTypesProto2* msg =
+      protobuf_test_messages_proto2_TestAllTypesProto2_new(arena);
+
+  const upb_MiniTableField* repeated_nested_message_field =
+      find_proto2_field(kFieldOptionalRepeatedNestedMessage);
+  upb_Message* nested_message = (upb_Message*)
+      protobuf_test_messages_proto2_TestAllTypesProto2_NestedMessage_new(arena);
+
+  upb_Array* array = upb_Message_GetOrCreateMutableArray(
+      UPB_UPCAST(msg), repeated_nested_message_field, arena);
+  upb_MessageValue new_value;
+  new_value.msg_val = nested_message;
+  EXPECT_TRUE(upb_Array_Append(array, new_value, arena));
+
+  EXPECT_EQ(nested_message, upb_Array_GetMutable(array, 0).msg);
+  EXPECT_EQ(nested_message, upb_Array_Get(array, 0).msg_val);
+
+  upb_Arena_Free(arena);
+}
+
 TEST(GeneratedCode, GetMutableMessage) {
   upb_Arena* arena = upb_Arena_New();
   protobuf_test_messages_proto2_TestAllTypesProto2* msg =
diff --git a/upb/message/array.c b/upb/message/array.c
index 45ce299dd0..79a1e58f84 100644
--- a/upb/message/array.c
+++ b/upb/message/array.c
@@ -40,6 +40,15 @@ upb_MessageValue upb_Array_Get(const upb_Array* arr, size_t i) {
   return ret;
 }
 
+upb_MutableMessageValue upb_Array_GetMutable(upb_Array* arr, size_t i) {
+  upb_MutableMessageValue ret;
+  char* data = _upb_array_ptr(arr);
+  const int lg2 = UPB_PRIVATE(_upb_Array_ElemSizeLg2)(arr);
+  UPB_ASSERT(i < arr->UPB_PRIVATE(size));
+  memcpy(&ret, data + (i << lg2), 1 << lg2);
+  return ret;
+}
+
 void upb_Array_Set(upb_Array* arr, size_t i, upb_MessageValue val) {
   char* data = _upb_array_ptr(arr);
   const int lg2 = UPB_PRIVATE(_upb_Array_ElemSizeLg2)(arr);
diff --git a/upb/message/array.h b/upb/message/array.h
index b217bfbcb8..684dfe0b30 100644
--- a/upb/message/array.h
+++ b/upb/message/array.h
@@ -32,6 +32,10 @@ UPB_API size_t upb_Array_Size(const upb_Array* arr);
 // Returns the given element, which must be within the array's current size.
 UPB_API upb_MessageValue upb_Array_Get(const upb_Array* arr, size_t i);
 
+// Returns a mutating pointer to the given element, which must be within the
+// array's current size.
+UPB_API upb_MutableMessageValue upb_Array_GetMutable(upb_Array* arr, size_t i);
+
 // Sets the given element, which must be within the array's current size.
 UPB_API void upb_Array_Set(upb_Array* arr, size_t i, upb_MessageValue val);