Promote extensions in place of the unknown field they were promoted from

PiperOrigin-RevId: 708100356
protocolbuffers · Dec 20, 2024 · a673dcc · a673dcc
1 parent e890d15
commit a673dcc
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 6 deletions.
diff --git a/upb/message/internal/extension.c b/upb/message/internal/extension.c
@@ -52,3 +52,15 @@ upb_Extension* UPB_PRIVATE(_upb_Message_GetOrCreateExtension)(
   in->aux_data[in->size++] = upb_TaggedAuxPtr_MakeExtension(ext);
   return ext;
 }
+
+void upb_Message_ReplaceUnknownWithExtension(struct upb_Message* msg,
+                                             uintptr_t iter,
+                                             const upb_Extension* ext) {
+  UPB_ASSERT(iter != 0);
+  upb_Message_Internal* in = UPB_PRIVATE(_upb_Message_GetInternal)(msg);
+  UPB_ASSERT(in);
+  size_t index = iter - 1;
+  upb_TaggedAuxPtr tagged_ptr = in->aux_data[index];
+  UPB_ASSERT(upb_TaggedAuxPtr_IsUnknown(tagged_ptr));
+  in->aux_data[index] = upb_TaggedAuxPtr_MakeExtension(ext);
+}
diff --git a/upb/message/internal/extension.h b/upb/message/internal/extension.h
@@ -63,6 +63,10 @@ UPB_INLINE bool UPB_PRIVATE(_upb_Extension_IsEmpty)(const upb_Extension* ext) {
   UPB_UNREACHABLE();
 }
 
+// Replaces the unknown field at iter with the provided extension.
+void upb_Message_ReplaceUnknownWithExtension(struct upb_Message* msg,
+                                             uintptr_t iter,
+                                             const upb_Extension* ext);
 #ifdef __cplusplus
 } /* extern "C" */
 #endif

diff --git a/upb/message/promote.c b/upb/message/promote.c
@@ -101,17 +101,14 @@ upb_GetExtension_Status upb_Message_GetOrPromoteExtension(
   }
   upb_Message* extension_msg = parse_result.message;
   // Add to extensions.
-  upb_Extension* ext =
-      UPB_PRIVATE(_upb_Message_GetOrCreateExtension)(msg, ext_table, arena);
+  upb_Extension* ext = upb_Arena_Malloc(arena, sizeof(upb_Extension));
   if (!ext) {
     return kUpb_GetExtension_OutOfMemory;
   }
+  ext->ext = ext_table;
   ext->data.msg_val = extension_msg;
+  upb_Message_ReplaceUnknownWithExtension(msg, result.iter, ext);
   value->msg_val = extension_msg;
-  // Our storage may have been reallocated by adding the extension, but our
-  // iterators are stable.
-  upb_StringView data = upb_StringView_FromDataAndSize(result.ptr, result.len);
-  upb_Message_DeleteUnknown(msg, &data, &result.iter);
   return kUpb_GetExtension_Ok;
 }