Support non-canonical extensions in message copy/freeze/encode operations.

upb/message/BUILD

@@ -186,13 +186,15 @@ cc_library(
     features = UPB_DEFAULT_FEATURES,
     visibility = ["//visibility:public"],
     deps = [
+        ":copy",
         ":internal",
         ":message",
         "//upb/base",
         "//upb/mem",
         "//upb/mini_table",
         "//upb/port",
         "//upb/wire",
+        "//upb/wire:encoder",
         "//upb/wire:eps_copy_input_stream",
         "//upb/wire:reader",
     ],
@@ -381,6 +383,7 @@ cc_test(
         "//upb/mini_descriptor",
         "//upb/mini_descriptor:internal",
         "//upb/mini_table",
+        "//upb/port",
         "//upb/test:test_messages_proto2_upb_proto",
         "//upb/test:test_messages_proto3_upb_proto",
         "//upb/test:test_proto_upb_minitable",

upb/message/copy.c

@@ -245,12 +245,19 @@ upb_Message* _upb_Message_Copy(upb_Message* dst, const upb_Message* src,
 
   for (size_t i = 0; i < in->size; i++) {
     upb_TaggedAuxPtr tagged_ptr = in->aux_data[i];
-    if (upb_TaggedAuxPtr_IsExtension(tagged_ptr)) {
-      // Clone extension
-      const upb_Extension* msg_ext = upb_TaggedAuxPtr_Extension(tagged_ptr);
+    if (upb_TaggedAuxPtr_IsExtension(tagged_ptr) ||
+        upb_TaggedAuxPtr_IsNonCanonicalExtension(tagged_ptr)) {
+      // Clone a canonical or non-canonical extension
+      bool canonical = upb_TaggedAuxPtr_IsExtension(tagged_ptr);
+      const upb_Extension* msg_ext =
+          canonical ? upb_TaggedAuxPtr_Extension(tagged_ptr)
+                    : upb_TaggedAuxPtr_NonCanonicalExtension(tagged_ptr);
       const upb_MiniTableField* field = &msg_ext->ext->UPB_PRIVATE(field);
-      upb_Extension* dst_ext = UPB_PRIVATE(_upb_Message_GetOrCreateExtension)(
-          dst, msg_ext->ext, arena);
+      upb_Extension* dst_ext =
+          canonical ? UPB_PRIVATE(_upb_Message_GetOrCreateExtension)(
+                          dst, msg_ext->ext, arena)
+                    : UPB_PRIVATE(_upb_Message_CreateNonCanonicalExtension)(
+                          dst, msg_ext->ext, arena);
       if (!dst_ext) return NULL;
       if (upb_MiniTableField_IsScalar(field)) {
         if (!upb_Clone_ExtensionValue(msg_ext->ext, msg_ext, dst_ext, arena)) {
@@ -268,7 +275,7 @@ upb_Message* _upb_Message_Copy(upb_Message* dst, const upb_Message* src,
         dst_ext->data.array_val = cloned_array;
       }
     } else if (upb_TaggedAuxPtr_IsUnknown(tagged_ptr)) {
-      // Clone unknown
+      // Clone a raw unknown field
       upb_StringView* unknown = upb_TaggedAuxPtr_UnknownData(tagged_ptr);
       // Make a copy into destination arena.
       if (!UPB_PRIVATE(_upb_Message_AddUnknown)(
@@ -333,6 +340,15 @@ bool upb_Message_ShallowCopy(upb_Message* dst, const upb_Message* src,
         dst_in->aux_data[i] = upb_TaggedAuxPtr_MakeUnknownDataAliased(dst_sv);
         break;
       }
+      case kUpb_TaggedAuxType_NonCanonicalExtension: {
+        const upb_Extension* msg_ext = aux.extension;
+        upb_Extension* dst_ext = upb_Arena_Malloc(arena, sizeof(upb_Extension));
+        if (!dst_ext) return false;
+        *dst_ext = *msg_ext;
+        dst_in->aux_data[i] =
+            upb_TaggedAuxPtr_MakeNonCanonicalExtension(dst_ext);
+        break;
+      }
     }
   }
 

upb/message/copy_test.cc

@@ -27,6 +27,8 @@
 #include "upb/base/upcast.h"
 #include "upb/mem/arena.h"
 #include "upb/message/accessors.h"
+#include "upb/message/internal/accessors.h"
+#include "upb/message/internal/extension.h"
 #include "upb/message/internal/message.h"
 #include "upb/message/map.h"
 #include "upb/message/message.h"
@@ -533,4 +535,53 @@ TEST(GeneratedCode, ShallowCloneMessage) {
   upb_Arena_Free(arena);
 }
 
+TEST(GeneratedCode, DeepCloneMessageNonCanonicalExtensions) {
+  upb_Arena* source_arena = upb_Arena_New();
+  upb_test_ModelWithExtensions* msg =
+      upb_test_ModelWithExtensions_new(source_arena);
+  upb_test_ModelExtension1* ext1 = upb_test_ModelExtension1_new(source_arena);
+  upb_test_ModelExtension1_set_str(ext1,
+                                   upb_StringView_FromString("LifecycleValue"));
+
+  // Attach as non-canonical extension
+  UPB_PRIVATE(_upb_Message_SetNonCanonicalExtension)(
+      UPB_UPCAST(msg), upb_test_ModelExtension1_model_ext_ext, &ext1,
+      source_arena);
+
+  // Deep clone msg to clone
+  upb_Arena* arena = upb_Arena_New();
+  upb_test_ModelWithExtensions* clone =
+      (upb_test_ModelWithExtensions*)upb_Message_DeepClone(
+          UPB_UPCAST(msg), &upb_0test__ModelWithExtensions_msg_init, arena);
+  ASSERT_NE(clone, nullptr);
+
+  // Mutate original
+  upb_test_ModelExtension1_set_str(ext1, upb_StringView_FromString("Mutated"));
+  upb_Arena_Free(source_arena);
+
+  // Check if clone has the non-canonical extension and it's unmodified
+  upb_MessageUnknown data;
+  uintptr_t iter = kUpb_Message_UnknownBegin;
+  bool has_non_canonical = false;
+  const upb_Extension* ext_found = nullptr;
+  while (upb_Message_NextUnknown2(UPB_UPCAST(clone), &data, &iter)) {
+    if (data.type == kUpb_MessageUnknownType_NonCanonicalExtension) {
+      has_non_canonical = true;
+      ext_found = (const upb_Extension*)data.value.extension;
+    }
+  }
+  EXPECT_TRUE(has_non_canonical);
+  ASSERT_NE(ext_found, nullptr);
+
+  const upb_test_ModelExtension1* cloned_ext =
+      (const upb_test_ModelExtension1*)ext_found->data.msg_val;
+  EXPECT_TRUE(
+      upb_StringView_IsEqual(upb_test_ModelExtension1_str(cloned_ext),
+                             upb_StringView_FromString("LifecycleValue")));
+
+  upb_Arena_Free(arena);
+}
+
 }  // namespace
+
+#include "upb/port/undef.inc"

upb/message/internal/accessors.h

@@ -331,6 +331,19 @@ UPB_API_INLINE bool upb_Message_SetExtension(struct upb_Message* msg,
   return true;
 }
 
+UPB_API_INLINE bool UPB_PRIVATE(_upb_Message_SetNonCanonicalExtension)(
+    struct upb_Message* msg, const upb_MiniTableExtension* e, const void* val,
+    upb_Arena* a) {
+  UPB_ASSERT(!upb_Message_IsFrozen(msg));
+  UPB_ASSERT(a);
+  upb_Extension* ext =
+      UPB_PRIVATE(_upb_Message_CreateNonCanonicalExtension)(msg, e, a);
+  if (!ext) return false;
+  UPB_PRIVATE(_upb_MiniTableField_DataCopy)
+  (&e->UPB_PRIVATE(field), &ext->data, val);
+  return true;
+}
+
 // Sets the value of the given field in the given msg. The return value is true
 // if the operation completed successfully, or false if memory allocation
 // failed.

upb/message/internal/extension.c

@@ -37,8 +37,9 @@ const upb_Extension* UPB_PRIVATE(_upb_Message_Getext)(
   return NULL;
 }
 
-upb_Extension* UPB_PRIVATE(_upb_Message_GetOrCreateExtension)(
-    struct upb_Message* msg, const upb_MiniTableExtension* e, upb_Arena* a) {
+UPB_INLINE upb_Extension* _upb_Message_GetOrCreateExtensionInternal(
+    struct upb_Message* msg, const upb_MiniTableExtension* e, upb_Arena* a,
+    bool canonical) {
   UPB_ASSERT(!upb_Message_IsFrozen(msg));
   upb_Extension* ext = (upb_Extension*)UPB_PRIVATE(_upb_Message_Getext)(msg, e);
   if (ext) return ext;
@@ -49,6 +50,20 @@ upb_Extension* UPB_PRIVATE(_upb_Message_GetOrCreateExtension)(
   if (!ext) return NULL;
   memset(ext, 0, sizeof(upb_Extension));
   ext->ext = e;
-  in->aux_data[in->size++] = upb_TaggedAuxPtr_MakeExtension(ext);
+  in->aux_data[in->size++] =
+      canonical ? upb_TaggedAuxPtr_MakeExtension(ext)
+                : upb_TaggedAuxPtr_MakeNonCanonicalExtension(ext);
   return ext;
 }
+
+upb_Extension* UPB_PRIVATE(_upb_Message_GetOrCreateExtension)(
+    struct upb_Message* msg, const upb_MiniTableExtension* e, upb_Arena* a) {
+  return _upb_Message_GetOrCreateExtensionInternal(msg, e, a,
+                                                   /*canonical=*/true);
+}
+
+upb_Extension* UPB_PRIVATE(_upb_Message_CreateNonCanonicalExtension)(
+    struct upb_Message* msg, const upb_MiniTableExtension* e, upb_Arena* a) {
+  return _upb_Message_GetOrCreateExtensionInternal(msg, e, a,
+                                                   /*canonical=*/false);
+}

upb/message/internal/extension.h

@@ -45,6 +45,14 @@ UPB_NODISCARD upb_Extension* UPB_PRIVATE(_upb_Message_GetOrCreateExtension)(
     struct upb_Message* msg, const upb_MiniTableExtension* ext,
     upb_Arena* arena);
 
+// Adds the given non-canonical extension data to the given message.
+// |ext| is copied into the message instance.
+// This logically replaces any previously-added extension with this number.
+UPB_NODISCARD upb_Extension* UPB_PRIVATE(
+    _upb_Message_CreateNonCanonicalExtension)(struct upb_Message* msg,
+                                              const upb_MiniTableExtension* ext,
+                                              upb_Arena* arena);
+
 // Returns an extension for a message with a given mini table,
 // or NULL if no extension exists with this mini table.
 const upb_Extension* UPB_PRIVATE(_upb_Message_Getext)(

upb/message/internal/map_sorter.h

@@ -90,7 +90,7 @@ bool _upb_mapsorter_pushmap(_upb_mapsorter* s, upb_FieldType key_type,
                             const struct upb_Map* map, _upb_sortedmap* sorted);
 
 bool _upb_mapsorter_pushexts(_upb_mapsorter* s, const upb_Message_Internal* in,
-                             _upb_sortedmap* sorted);
+                             _upb_sortedmap* sorted, bool include_noncanonical);
 
 #ifdef __cplusplus
 } /* extern "C" */

upb/message/internal/message.h

@@ -46,7 +46,8 @@ typedef struct upb_TaggedAuxPtr {
   // Two lowest bits form a tag:
   // 00 - non-aliased unknown data
   // 10 - aliased unknown data
-  // 01 - extension
+  // 01 - canonical extension
+  // 11 - unknown as a non-canonical extension
   //
   // The main semantic difference between aliased and non-aliased unknown data
   // is that non-aliased unknown data can be assumed to have the following
@@ -64,11 +65,19 @@ typedef struct upb_TaggedAuxPtr {
   // For aliased unknown data, this layout is _not_ guaranteed, since the
   // pointer to the StringView can be anywhere in the allocation, and the
   // StringView may point to non-data memory.
+  //
+  // For a non-canonical extension, its schema is known but not
+  // the one expected by the message, so it should be treated like an unknown
+  // field, but is stored as an extension for efficiency.
   uintptr_t ptr;
 } upb_TaggedAuxPtr;
 
 UPB_INLINE bool upb_TaggedAuxPtr_IsExtension(upb_TaggedAuxPtr ptr) {
-  return ptr.ptr & 1;
+  return (ptr.ptr & 3) == 1;
+}
+
+UPB_INLINE bool upb_TaggedAuxPtr_IsNonCanonicalExtension(upb_TaggedAuxPtr ptr) {
+  return (ptr.ptr & 3) == 3;
 }
 
 UPB_INLINE bool upb_TaggedAuxPtr_IsUnknown(upb_TaggedAuxPtr ptr) {
@@ -84,13 +93,20 @@ UPB_INLINE upb_Extension* upb_TaggedAuxPtr_Extension(upb_TaggedAuxPtr ptr) {
   return (upb_Extension*)(ptr.ptr & ~3ULL);
 }
 
+UPB_INLINE upb_Extension* upb_TaggedAuxPtr_NonCanonicalExtension(
+    upb_TaggedAuxPtr ptr) {
+  UPB_ASSERT(upb_TaggedAuxPtr_IsNonCanonicalExtension(ptr));
+  return (upb_Extension*)(ptr.ptr & ~3ULL);
+}
+
 UPB_INLINE upb_StringView* upb_TaggedAuxPtr_UnknownData(upb_TaggedAuxPtr ptr) {
   UPB_ASSERT(!upb_TaggedAuxPtr_IsExtension(ptr));
   return (upb_StringView*)(ptr.ptr & ~3ULL);
 }
 
 typedef enum {
   kUpb_TaggedAuxType_Extension,
+  kUpb_TaggedAuxType_NonCanonicalExtension,
   kUpb_TaggedAuxType_Unknown,
   kUpb_TaggedAuxType_AliasedUnknown
 } upb_TaggedAuxType;
@@ -108,10 +124,13 @@ UPB_INLINE upb_TaggedAuxType upb_TaggedAux_Get(upb_TaggedAuxPtr ptr,
   } else if (upb_TaggedAuxPtr_IsUnknownAliased(ptr)) {
     data->unknown_data = *upb_TaggedAuxPtr_UnknownData(ptr);
     return kUpb_TaggedAuxType_AliasedUnknown;
-  } else {
-    UPB_ASSERT(upb_TaggedAuxPtr_IsUnknown(ptr));
+  } else if (upb_TaggedAuxPtr_IsUnknown(ptr)) {
     data->unknown_data = *upb_TaggedAuxPtr_UnknownData(ptr);
     return kUpb_TaggedAuxType_Unknown;
+  } else {
+    UPB_ASSERT(upb_TaggedAuxPtr_IsNonCanonicalExtension(ptr));
+    data->extension = upb_TaggedAuxPtr_NonCanonicalExtension(ptr);
+    return kUpb_TaggedAuxType_NonCanonicalExtension;
   }
 }
 
@@ -128,6 +147,13 @@ upb_TaggedAuxPtr_MakeExtension(const upb_Extension* e) {
   return ptr;
 }
 
+UPB_INLINE upb_TaggedAuxPtr
+upb_TaggedAuxPtr_MakeNonCanonicalExtension(const upb_Extension* e) {
+  upb_TaggedAuxPtr ptr;
+  ptr.ptr = (uintptr_t)e | 3;
+  return ptr;
+}
+
 // This tag means that the original allocation for this field starts with the
 // string view and ends with the end of the content referenced by the string
 // view.
@@ -234,7 +260,8 @@ UPB_NODISCARD UPB_INLINE struct upb_Message* _upb_Message_New(
   return msg;
 }
 
-// Discards the unknown fields for this message only.
+// Discards the unknown (including non-canonical extensions) for this message
+// only.
 void _upb_Message_DiscardUnknown_shallow(struct upb_Message* msg);
 
 UPB_NODISCARD UPB_NOINLINE bool UPB_PRIVATE(_upb_Message_AddUnknownSlowPath)(
@@ -342,12 +369,6 @@ UPB_INLINE bool upb_Message_NextUnknown(const struct upb_Message* msg,
   return false;
 }
 
-UPB_INLINE bool upb_Message_HasUnknown(const struct upb_Message* msg) {
-  upb_StringView data;
-  uintptr_t iter = kUpb_Message_UnknownBegin;
-  return upb_Message_NextUnknown(msg, &data, &iter);
-}
-
 UPB_INLINE bool upb_Message_NextExtension(const struct upb_Message* msg,
                                           const upb_MiniTableExtension** out_e,
                                           upb_MessageValue* out_v,

upb/message/map_sorter.c

@@ -164,10 +164,15 @@ static int _upb_mapsorter_cmpext(const void* _a, const void* _b) {
 }
 
 bool _upb_mapsorter_pushexts(_upb_mapsorter* s, const upb_Message_Internal* in,
-                             _upb_sortedmap* sorted) {
+                             _upb_sortedmap* sorted,
+                             bool include_noncanonical) {
   size_t count = 0;
   for (size_t i = 0; i < in->size; i++) {
-    count += upb_TaggedAuxPtr_IsExtension(in->aux_data[i]);
+    bool is_any_extension =
+        upb_TaggedAuxPtr_IsExtension(in->aux_data[i]) ||
+        (include_noncanonical &&
+         upb_TaggedAuxPtr_IsNonCanonicalExtension(in->aux_data[i]));
+    count += is_any_extension;
   }
   if (!_upb_mapsorter_resize(s, sorted, count)) return false;
   if (count == 0) return true;
@@ -177,6 +182,9 @@ bool _upb_mapsorter_pushexts(_upb_mapsorter* s, const upb_Message_Internal* in,
     upb_TaggedAuxPtr tagged_ptr = in->aux_data[i];
     if (upb_TaggedAuxPtr_IsExtension(tagged_ptr)) {
       *entry++ = upb_TaggedAuxPtr_Extension(tagged_ptr);
+    } else if (include_noncanonical &&
+               upb_TaggedAuxPtr_IsNonCanonicalExtension(tagged_ptr)) {
+      *entry++ = upb_TaggedAuxPtr_NonCanonicalExtension(tagged_ptr);
     }
   }
   qsort(&s->entries[sorted->start], count, sizeof(*s->entries),