Fix GetMutableMessage not updating presence and storing new instance.

Add array resize API to support GetMutableRepeatedField api with size argument.

PiperOrigin-RevId: 452075516

upb/mini_table.c

@@ -316,6 +316,9 @@ char* upb_MtDataEncoder_PutEnumValue(upb_MtDataEncoder* e, char* ptr,
   uint32_t delta = val - in->state.enum_state.last_written_value;
   if (delta >= 5 && in->state.enum_state.present_values_mask) {
     ptr = upb_MtDataEncoder_FlushDenseEnumMask(e, ptr);
+    if (!ptr) {
+      return NULL;
+    }
     delta -= 5;
   }
 

upb/mini_table_accessors.c

@@ -58,6 +58,32 @@ size_t upb_MiniTable_Field_GetSize(const upb_MiniTable_Field* f) {
   return upb_IsRepeatedOrMap(f) ? sizeof(void*) : sizes[f->descriptortype];
 }
 
+// Maps descriptor type to elem_size_lg2.
+int upb_MiniTable_Field_CTypeLg2Size(const upb_MiniTable_Field* f) {
+  static const uint8_t sizes[] = {
+      -1,             /* invalid descriptor type */
+      3,              /* DOUBLE */
+      2,              /* FLOAT */
+      3,              /* INT64 */
+      3,              /* UINT64 */
+      2,              /* INT32 */
+      3,              /* FIXED64 */
+      2,              /* FIXED32 */
+      0,              /* BOOL */
+      UPB_SIZE(3, 4), /* STRING */
+      UPB_SIZE(2, 3), /* GROUP */
+      UPB_SIZE(2, 3), /* MESSAGE */
+      UPB_SIZE(3, 4), /* BYTES */
+      2,              /* UINT32 */
+      2,              /* ENUM */
+      2,              /* SFIXED32 */
+      3,              /* SFIXED64 */
+      2,              /* SINT32 */
+      3,              /* SINT64 */
+  };
+  return sizes[f->descriptortype];
+}
+
 bool upb_MiniTable_HasField(const upb_Message* msg,
                             const upb_MiniTable_Field* field) {
   if (_upb_MiniTable_Field_InOneOf(field)) {
@@ -84,6 +110,13 @@ void upb_MiniTable_ClearField(upb_Message* msg,
   memset(mem, 0, upb_MiniTable_Field_GetSize(field));
 }
 
+void* upb_MiniTable_ResizeArray(upb_Message* msg,
+                                const upb_MiniTable_Field* field, size_t len,
+                                upb_Arena* arena) {
+  return _upb_Array_Resize_accessor2(
+      msg, field->offset, len, upb_MiniTable_Field_CTypeLg2Size(field), arena);
+}
+
 typedef struct {
   const char* ptr;
   uint64_t val;

upb/mini_table_accessors.h

@@ -194,7 +194,8 @@ UPB_INLINE const upb_Message* upb_MiniTable_GetMessage(
 UPB_INLINE void upb_MiniTable_SetMessage(upb_Message* msg,
                                          const upb_MiniTable_Field* field,
                                          upb_Message* sub_message) {
-  UPB_ASSERT(field->descriptortype == kUpb_FieldType_Message);
+  UPB_ASSERT(field->descriptortype == kUpb_FieldType_Message ||
+             field->descriptortype == kUpb_FieldType_Group);
   _upb_MiniTable_SetPresence(msg, field);
   *UPB_PTR_AT(msg, field->offset, const upb_Message*) = sub_message;
 }
@@ -202,11 +203,14 @@ UPB_INLINE void upb_MiniTable_SetMessage(upb_Message* msg,
 UPB_INLINE upb_Message* upb_MiniTable_GetMutableMessage(
     upb_Message* msg, const upb_MiniTable* mini_table,
     const upb_MiniTable_Field* field, upb_Arena* arena) {
-  UPB_ASSERT(field->descriptortype == kUpb_FieldType_Message);
+  UPB_ASSERT(field->descriptortype == kUpb_FieldType_Message ||
+             field->descriptortype == kUpb_FieldType_Group);
   upb_Message* sub_message = *UPB_PTR_AT(msg, field->offset, upb_Message*);
   if (!sub_message) {
     sub_message =
         _upb_Message_New(mini_table->subs[field->submsg_index].submsg, arena);
+    *UPB_PTR_AT(msg, field->offset, upb_Message*) = sub_message;
+    _upb_MiniTable_SetPresence(msg, field);
   }
   return sub_message;
 }
@@ -221,6 +225,9 @@ UPB_INLINE upb_Array* upb_MiniTable_GetMutableArray(
   return (upb_Array*)*UPB_PTR_AT(msg, field->offset, upb_Array*);
 }
 
+void* upb_MiniTable_ResizeArray(upb_Message* msg,
+                                const upb_MiniTable_Field* field, size_t len,
+                                upb_Arena* arena);
 typedef enum {
   kUpb_GetExtension_Ok,
   kUpb_GetExtension_NotPresent,

upb/mini_table_accessors_test.cc

@@ -357,6 +357,25 @@ TEST(GeneratedCode, RepeatedScalar) {
   upb_Arena_Free(arena);
 }
 
+TEST(GeneratedCode, GetMutableMessage) {
+  upb_Arena* arena = upb_Arena_New();
+  protobuf_test_messages_proto2_TestAllTypesProto2* msg =
+      protobuf_test_messages_proto2_TestAllTypesProto2_new(arena);
+  // Message.
+  const upb_MiniTable_Field* optional_message_field =
+      find_proto2_field(kFieldOptionalNestedMessage);
+  upb_Message* msg1 = upb_MiniTable_GetMutableMessage(
+      msg, &protobuf_test_messages_proto2_TestAllTypesProto2_msginit,
+      optional_message_field, arena);
+  upb_Message* msg2 = upb_MiniTable_GetMutableMessage(
+      msg, &protobuf_test_messages_proto2_TestAllTypesProto2_msginit,
+      optional_message_field, arena);
+  // Verify that newly constructed sub message is stored in msg.
+  EXPECT_EQ(msg1, msg2);
+
+  upb_Arena_Free(arena);
+}
+
 TEST(GeneratedCode, Extensions) {
   upb_Arena* arena = upb_Arena_New();
   upb_test_ModelWithExtensions* msg = upb_test_ModelWithExtensions_new(arena);