diff --git a/upb/message/internal/extension.c b/upb/message/internal/extension.c index ca8a5865fd62..e4f866b67ee6 100644 --- 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); +} \ No newline at end of file diff --git a/upb/message/internal/extension.h b/upb/message/internal/extension.h index c7af38acc3b2..f6e346d19918 100644 --- 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 index 38e6e2a2351c..2ed22a23dba0 100644 --- 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; }