upb: add upb_Message_ExtensionByIndex()

PiperOrigin-RevId: 592091162

upb/message/compat.c

@@ -17,10 +17,19 @@
 // Must be last.
 #include "upb/port/def.inc"
 
+const upb_Extension* upb_Message_ExtensionByIndex(const upb_Message* msg,
+                                                  size_t index) {
+  size_t count;
+  const upb_Extension* ext = UPB_PRIVATE(_upb_Message_Getexts)(msg, &count);
+
+  UPB_ASSERT(index < count);
+  return &ext[index];
+}
+
 const upb_Extension* upb_Message_FindExtensionByNumber(const upb_Message* msg,
                                                        uint32_t field_number) {
-  size_t count = 0;
-  const upb_Extension* ext = _upb_Message_Getexts(msg, &count);
+  size_t count;
+  const upb_Extension* ext = UPB_PRIVATE(_upb_Message_Getexts)(msg, &count);
 
   while (count--) {
     if (upb_MiniTableExtension_Number(ext->ext) == field_number) return ext;

upb/message/compat.h

@@ -24,6 +24,9 @@
 extern "C" {
 #endif
 
+const upb_Extension* upb_Message_ExtensionByIndex(const upb_Message* msg,
+                                                  size_t index);
+
 // Returns the extension with the given field number, or NULL on failure.
 const upb_Extension* upb_Message_FindExtensionByNumber(const upb_Message* msg,
                                                        uint32_t field_number);

upb/message/copy.c

@@ -19,14 +19,12 @@
 #include "upb/message/internal/array.h"
 #include "upb/message/internal/extension.h"
 #include "upb/message/internal/map.h"
-#include "upb/message/internal/message.h"
 #include "upb/message/map.h"
 #include "upb/message/message.h"
 #include "upb/message/tagged_ptr.h"
 #include "upb/mini_table/extension.h"
 #include "upb/mini_table/field.h"
 #include "upb/mini_table/internal/field.h"
-#include "upb/mini_table/internal/message.h"
 #include "upb/mini_table/internal/size_log2.h"
 #include "upb/mini_table/message.h"
 #include "upb/mini_table/sub.h"
@@ -257,7 +255,7 @@ upb_Message* _upb_Message_Copy(upb_Message* dst, const upb_Message* src,
   }
   // Clone extensions.
   size_t ext_count;
-  const upb_Extension* ext = _upb_Message_Getexts(src, &ext_count);
+  const upb_Extension* ext = UPB_PRIVATE(_upb_Message_Getexts)(src, &ext_count);
   for (size_t i = 0; i < ext_count; ++i) {
     const upb_Extension* msg_ext = &ext[i];
     const upb_MiniTableField* field = &msg_ext->ext->UPB_PRIVATE(field);

upb/message/internal/extension.c

@@ -18,10 +18,10 @@
 // Must be last.
 #include "upb/port/def.inc"
 
-const upb_Extension* _upb_Message_Getext(const upb_Message* msg,
-                                         const upb_MiniTableExtension* e) {
+const struct upb_Extension* _upb_Message_Getext(
+    const upb_Message* msg, const upb_MiniTableExtension* e) {
   size_t n;
-  const upb_Extension* ext = _upb_Message_Getexts(msg, &n);
+  const struct upb_Extension* ext = UPB_PRIVATE(_upb_Message_Getexts)(msg, &n);
 
   // For now we use linear search exclusively to find extensions.
   // If this becomes an issue due to messages with lots of extensions,
@@ -35,12 +35,12 @@ const upb_Extension* _upb_Message_Getext(const upb_Message* msg,
   return NULL;
 }
 
-const upb_Extension* _upb_Message_Getexts(const upb_Message* msg,
-                                          size_t* count) {
+const struct upb_Extension* UPB_PRIVATE(_upb_Message_Getexts)(
+    const upb_Message* msg, size_t* count) {
   const upb_Message_Internal* in = upb_Message_Getinternal(msg);
   if (in->internal) {
-    *count =
-        (in->internal->size - in->internal->ext_begin) / sizeof(upb_Extension);
+    *count = (in->internal->size - in->internal->ext_begin) /
+             sizeof(struct upb_Extension);
     return UPB_PTR_AT(in->internal, in->internal->ext_begin, void);
   } else {
     *count = 0;
@@ -48,16 +48,18 @@ const upb_Extension* _upb_Message_Getexts(const upb_Message* msg,
   }
 }
 
-upb_Extension* _upb_Message_GetOrCreateExtension(
+struct upb_Extension* _upb_Message_GetOrCreateExtension(
     upb_Message* msg, const upb_MiniTableExtension* e, upb_Arena* arena) {
-  upb_Extension* ext = (upb_Extension*)_upb_Message_Getext(msg, e);
+  struct upb_Extension* ext =
+      (struct upb_Extension*)_upb_Message_Getext(msg, e);
   if (ext) return ext;
-  if (!UPB_PRIVATE(_upb_Message_Realloc)(msg, sizeof(upb_Extension), arena))
+  if (!UPB_PRIVATE(_upb_Message_Realloc)(msg, sizeof(struct upb_Extension),
+                                         arena))
     return NULL;
   upb_Message_Internal* in = upb_Message_Getinternal(msg);
-  in->internal->ext_begin -= sizeof(upb_Extension);
+  in->internal->ext_begin -= sizeof(struct upb_Extension);
   ext = UPB_PTR_AT(in->internal, in->internal->ext_begin, void);
-  memset(ext, 0, sizeof(upb_Extension));
+  memset(ext, 0, sizeof(struct upb_Extension));
   ext->ext = e;
   return ext;
 }

upb/message/internal/extension.h

@@ -45,11 +45,11 @@ upb_Extension* _upb_Message_GetOrCreateExtension(
 
 // Returns an array of extensions for this message.
 // Note: the array is ordered in reverse relative to the order of creation.
-const upb_Extension* _upb_Message_Getexts(const upb_Message* msg,
-                                          size_t* count);
+const upb_Extension* UPB_PRIVATE(_upb_Message_Getexts)(const upb_Message* msg,
+                                                       size_t* count);
 
-// Returns an extension for the given field number, or NULL if no extension
-// exists for this field number.
+// 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_Message_Getext(const upb_Message* msg,
                                          const upb_MiniTableExtension* ext);
 

upb/message/message.c

@@ -72,6 +72,6 @@ void upb_Message_DeleteUnknown(upb_Message* msg, const char* data, size_t len) {
 
 size_t upb_Message_ExtensionCount(const upb_Message* msg) {
   size_t count;
-  _upb_Message_Getexts(msg, &count);
+  UPB_PRIVATE(_upb_Message_Getexts)(msg, &count);
   return count;
 }

upb/reflection/message.c

@@ -146,7 +146,7 @@ bool upb_Message_Next(const upb_Message* msg, const upb_MessageDef* m,
   if (ext_pool) {
     // Return any extensions that are set.
     size_t count;
-    const upb_Extension* ext = _upb_Message_Getexts(msg, &count);
+    const upb_Extension* ext = UPB_PRIVATE(_upb_Message_Getexts)(msg, &count);
     if (i - n < count) {
       ext += count - 1 - (i - n);
       memcpy(out_val, &ext->data, sizeof(*out_val));

upb/wire/encode.c

@@ -570,7 +570,8 @@ static void encode_message(upb_encstate* e, const upb_Message* msg,
      * these in field number order relative to normal fields or even to each
      * other. */
     size_t ext_count;
-    const upb_Extension* ext = _upb_Message_Getexts(msg, &ext_count);
+    const upb_Extension* ext =
+        UPB_PRIVATE(_upb_Message_Getexts)(msg, &ext_count);
     if (ext_count) {
       if (e->options & kUpb_EncodeOption_Deterministic) {
         _upb_sortedmap sorted;