Introduce non-canonical extensions in upb messages.

upb/message/copy.c

@@ -333,6 +333,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/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) {
+static upb_Extension* _upb_Message_GetOrCreateExtensionInternal(
+    struct upb_Message* msg, const upb_MiniTableExtension* e, upb_Arena* a,
+    bool non_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,18 @@ 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++] =
+      non_canonical ? upb_TaggedAuxPtr_MakeNonCanonicalExtension(ext)
+                    : upb_TaggedAuxPtr_MakeExtension(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, false);
+}
+
+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, true);
+}

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/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 not known to the message
+  // so it is treated like an unknown field, but it 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.

upb/mini_table/internal/extension.h

@@ -23,6 +23,8 @@ struct upb_MiniTableExtension {
   struct upb_MiniTableField UPB_PRIVATE(field);
 
   union upb_MiniTableSub UPB_PRIVATE(sub);  // NULL unless submsg or proto2 enum
+  // A known extendee schema for a canonical extension. Otherwise, the
+  // `extendee` info should be ignored/NULL if not a canonical extension.
   const struct upb_MiniTable* UPB_PRIVATE(extendee);
 };