Remove upb_ExtensionRegistry_AddAllLinkedExtensions API.

This is now an implementation detail of the public GeneratedRegistry.

PiperOrigin-RevId: 827277782

Author: mkruskal-google

hpb/extension.h

@@ -21,6 +21,7 @@
 #include "hpb/ptr.h"
 #include "upb/message/accessors.h"
 #include "upb/mini_table/extension_registry.h"
+#include "upb/mini_table/generated_registry.h"
 
 namespace hpb {
 // upb has a notion of an ExtensionRegistry. We expect most callers to use
@@ -66,17 +67,21 @@ class ExtensionRegistry {
 
  private:
 #if HPB_INTERNAL_BACKEND == HPB_INTERNAL_BACKEND_UPB
+  explicit ExtensionRegistry(upb_ExtensionRegistry* registry)
+      : registry_(registry) {}
+
   friend upb_ExtensionRegistry* ::hpb::internal::GetUpbExtensions(
       const ExtensionRegistry& extension_registry);
   upb_ExtensionRegistry* registry_;
 #endif
   // TODO: b/379100963 - Introduce ShutdownHpbLibrary
   static const ExtensionRegistry* NewGeneratedRegistry() {
 #if HPB_INTERNAL_BACKEND == HPB_INTERNAL_BACKEND_UPB
-    static hpb::Arena* global_arena = new hpb::Arena();
-    ExtensionRegistry* registry = new ExtensionRegistry(*global_arena);
-    upb_ExtensionRegistry_AddAllLinkedExtensions(registry->registry_);
-    return registry;
+    static const upb_GeneratedRegistryRef* registry_ref =
+        upb_GeneratedRegistry_Load();
+    // Const cast is safe because we're returning a const wrapper.
+    return new ExtensionRegistry(const_cast<upb_ExtensionRegistry*>(
+        upb_GeneratedRegistry_Get(registry_ref)));
 #elif HPB_INTERNAL_BACKEND == HPB_INTERNAL_BACKEND_CPP
     ExtensionRegistry* registry = new ExtensionRegistry();
     return registry;

upb/mini_table/extension_registry.c

@@ -15,8 +15,6 @@
 #include "upb/hash/str_table.h"
 #include "upb/mem/arena.h"
 #include "upb/mini_table/extension.h"
-#include "upb/mini_table/generated_registry.h"
-#include "upb/mini_table/internal/generated_registry.h"
 #include "upb/mini_table/message.h"
 
 // Must be last.
@@ -91,47 +89,6 @@ upb_ExtensionRegistryStatus upb_ExtensionRegistry_AddArray(
   return status;
 }
 
-const UPB_PRIVATE(upb_GeneratedExtensionListEntry) *
-    UPB_PRIVATE(upb_generated_extension_list) = NULL;
-
-bool upb_ExtensionRegistry_AddAllLinkedExtensions(upb_ExtensionRegistry* r) {
-  const UPB_PRIVATE(upb_GeneratedExtensionListEntry)* entry =
-      UPB_PRIVATE(upb_generated_extension_list);
-  while (entry != NULL) {
-    // Comparing pointers to different objects is undefined behavior, so we
-    // convert them to uintptr_t and compare their values.
-    uintptr_t begin = (uintptr_t)entry->start;
-    uintptr_t end = (uintptr_t)entry->stop;
-    uintptr_t current = begin;
-    while (current < end) {
-      const upb_MiniTableExtension* ext =
-          (const upb_MiniTableExtension*)current;
-      // Sentinels and padding introduced by the linker can result in zeroed
-      // entries, so simply skip them.
-      if (upb_MiniTableExtension_Number(ext) == 0) {
-        // MSVC introduces padding that might not be sized exactly the same as
-        // upb_MiniTableExtension, so we can't iterate by sizeof.  This is a
-        // valid thing for any linker to do, so it's safer to just always do it.
-        current += UPB_ALIGN_OF(upb_MiniTableExtension);
-        continue;
-      }
-
-      if (upb_ExtensionRegistry_Add(r, ext) !=
-          kUpb_ExtensionRegistryStatus_Ok) {
-        return false;
-      }
-      current += sizeof(upb_MiniTableExtension);
-    }
-    entry = entry->next;
-  }
-  return true;
-}
-
-const upb_ExtensionRegistry* upb_ExtensionRegistry_GetGenerated(
-    const upb_GeneratedRegistryRef* genreg) {
-  return genreg != NULL ? genreg->registry : NULL;
-}
-
 const upb_MiniTableExtension* upb_ExtensionRegistry_Lookup(
     const upb_ExtensionRegistry* r, const upb_MiniTable* t, uint32_t num) {
   char buf[EXTREG_KEY_SIZE];

upb/mini_table/extension_registry.h

@@ -13,7 +13,6 @@
 
 #include "upb/mem/arena.h"
 #include "upb/mini_table/extension.h"
-#include "upb/mini_table/generated_registry.h"
 #include "upb/mini_table/message.h"
 
 // Must be last.
@@ -80,28 +79,6 @@ UPB_API upb_ExtensionRegistryStatus upb_ExtensionRegistry_Add(
 upb_ExtensionRegistryStatus upb_ExtensionRegistry_AddArray(
     upb_ExtensionRegistry* r, const upb_MiniTableExtension** e, size_t count);
 
-// Adds all extensions linked into the binary into the registry.  The set of
-// linked extensions is assembled by the linker using linker arrays.  This
-// will likely not work properly if the extensions are split across multiple
-// shared libraries.
-//
-// Returns true if all extensions were added successfully, false on out of
-// memory or if any extensions were already present.
-//
-// This API is currently not available on MSVC (though it *is* available on
-// Windows using clang-cl).
-UPB_API bool upb_ExtensionRegistry_AddAllLinkedExtensions(
-    upb_ExtensionRegistry* r);
-
-// Returns the extension registry contained by a reference to the generated
-// registry.  The reference must be held for the lifetime of the registry.
-//
-// TODO This should actually be moved to generated_registry.h, but
-// can't because of the current location of
-// upb_ExtensionRegistry_AddAllLinkedExtensions.
-UPB_API const upb_ExtensionRegistry* upb_ExtensionRegistry_GetGenerated(
-    const upb_GeneratedRegistryRef* genreg);
-
 // Looks up the extension (if any) defined for message type |t| and field
 // number |num|. Returns the extension if found, otherwise NULL.
 UPB_API const upb_MiniTableExtension* upb_ExtensionRegistry_Lookup(

upb/mini_table/generated_registry.c

@@ -11,6 +11,7 @@
 
 #include "upb/mem/alloc.h"
 #include "upb/mem/arena.h"
+#include "upb/mini_table/extension.h"
 #include "upb/mini_table/extension_registry.h"
 #include "upb/mini_table/internal/generated_registry.h"  // IWYU pragma: keep
 #include "upb/port/atomic.h"
@@ -22,6 +23,9 @@
 #include <sched.h>
 #endif  // UPB_TSAN
 
+const UPB_PRIVATE(upb_GeneratedExtensionListEntry) *
+    UPB_PRIVATE(upb_generated_extension_list) = NULL;
+
 typedef struct upb_GeneratedRegistry {
   UPB_ATOMIC(upb_GeneratedRegistryRef*) ref;
   UPB_ATOMIC(int32_t) ref_count;
@@ -32,6 +36,40 @@ static upb_GeneratedRegistry* _upb_generated_registry(void) {
   return &r;
 }
 
+static bool _upb_GeneratedRegistry_AddAllLinkedExtensions(
+    upb_ExtensionRegistry* r) {
+  const UPB_PRIVATE(upb_GeneratedExtensionListEntry)* entry =
+      UPB_PRIVATE(upb_generated_extension_list);
+  while (entry != NULL) {
+    // Comparing pointers to different objects is undefined behavior, so we
+    // convert them to uintptr_t and compare their values.
+    uintptr_t begin = (uintptr_t)entry->start;
+    uintptr_t end = (uintptr_t)entry->stop;
+    uintptr_t current = begin;
+    while (current < end) {
+      const upb_MiniTableExtension* ext =
+          (const upb_MiniTableExtension*)current;
+      // Sentinels and padding introduced by the linker can result in zeroed
+      // entries, so simply skip them.
+      if (upb_MiniTableExtension_Number(ext) == 0) {
+        // MSVC introduces padding that might not be sized exactly the same as
+        // upb_MiniTableExtension, so we can't iterate by sizeof.  This is a
+        // valid thing for any linker to do, so it's safer to just always do it.
+        current += UPB_ALIGN_OF(upb_MiniTableExtension);
+        continue;
+      }
+
+      if (upb_ExtensionRegistry_Add(r, ext) !=
+          kUpb_ExtensionRegistryStatus_Ok) {
+        return false;
+      }
+      current += sizeof(upb_MiniTableExtension);
+    }
+    entry = entry->next;
+  }
+  return true;
+}
+
 // Constructs a new GeneratedRegistryRef, adding all linked extensions to the
 // registry or returning NULL on failure.
 static upb_GeneratedRegistryRef* _upb_GeneratedRegistry_New(void) {
@@ -47,7 +85,7 @@ static upb_GeneratedRegistryRef* _upb_GeneratedRegistry_New(void) {
   ref->arena = arena;
   ref->registry = extreg;
 
-  if (!upb_ExtensionRegistry_AddAllLinkedExtensions(extreg)) goto err;
+  if (!_upb_GeneratedRegistry_AddAllLinkedExtensions(extreg)) goto err;
 
   return ref;
 
@@ -138,3 +176,9 @@ void upb_GeneratedRegistry_Release(const upb_GeneratedRegistryRef* r) {
     }
   }
 }
+
+const upb_ExtensionRegistry* upb_GeneratedRegistry_Get(
+    const upb_GeneratedRegistryRef* r) {
+  if (r == NULL) return NULL;
+  return r->registry;
+}

upb/mini_table/generated_registry.h

@@ -8,6 +8,8 @@
 #ifndef UPB_MINI_TABLE_GENERATED_REGISTRY_H_
 #define UPB_MINI_TABLE_GENERATED_REGISTRY_H_
 
+#include "upb/mini_table/extension_registry.h"
+
 // Must be last.
 #include "upb/port/def.inc"
 
@@ -47,6 +49,13 @@ UPB_API const upb_GeneratedRegistryRef* upb_GeneratedRegistry_Load(void);
 // callers.
 UPB_API void upb_GeneratedRegistry_Release(const upb_GeneratedRegistryRef* r);
 
+// Returns the extension registry contained by a reference to the generated
+// registry.
+//
+// The reference must be held for the lifetime of the registry.
+UPB_API const upb_ExtensionRegistry* upb_GeneratedRegistry_Get(
+    const upb_GeneratedRegistryRef* r);
+
 #ifdef __cplusplus
 } /* extern "C" */
 #endif

upb/mini_table/generated_registry_test.cc

@@ -28,7 +28,7 @@ class GeneratedRegistryTest : public ::testing::Test {
   void SetUp() override {
     ref_ = upb_GeneratedRegistry_Load();
     ASSERT_NE(ref_, nullptr);
-    reg_ = upb_ExtensionRegistry_GetGenerated(ref_);
+    reg_ = upb_GeneratedRegistry_Get(ref_);
     ASSERT_NE(reg_, nullptr);
   }
 
@@ -92,6 +92,10 @@ TEST_F(GeneratedRegistryTest, ReleaseOnError) {
   upb_GeneratedRegistry_Release(nullptr);
 }
 
+TEST_F(GeneratedRegistryTest, GetOnError) {
+  EXPECT_EQ(upb_GeneratedRegistry_Get(nullptr), nullptr);
+}
+
 // On 32-bit systems, the stack size is smaller, so we can't run too many
 // concurrent threads without overflowing the stack.
 constexpr int kRaceIterations = sizeof(void*) < 8 ? 100 : 2000;
@@ -104,7 +108,7 @@ TEST(GeneratedRegistryRaceTest, Load) {
     threads.push_back(std::thread([&, i]() {
       barrier.Block();
       const upb_GeneratedRegistryRef* ref = upb_GeneratedRegistry_Load();
-      EXPECT_NE(upb_ExtensionRegistry_GetGenerated(ref), nullptr);
+      EXPECT_NE(upb_GeneratedRegistry_Get(ref), nullptr);
       refs[i] = ref;
     }));
   }
@@ -129,7 +133,7 @@ TEST(GeneratedRegistryRaceTest, Release) {
 
   for (int i = 0; i < kRaceIterations; ++i) {
     threads.push_back(std::thread([&, i]() {
-      EXPECT_NE(upb_ExtensionRegistry_GetGenerated(refs[i]), nullptr);
+      EXPECT_NE(upb_GeneratedRegistry_Get(refs[i]), nullptr);
       barrier.Block();
       upb_GeneratedRegistry_Release(refs[i]);
     }));
@@ -146,7 +150,7 @@ TEST(GeneratedRegistryRaceTest, LoadRelease) {
     threads.push_back(std::thread([&]() {
       barrier.Block();
       const upb_GeneratedRegistryRef* ref = upb_GeneratedRegistry_Load();
-      EXPECT_NE(upb_ExtensionRegistry_GetGenerated(ref), nullptr);
+      EXPECT_NE(upb_GeneratedRegistry_Get(ref), nullptr);
       upb_GeneratedRegistry_Release(ref);
     }));
   }
@@ -174,7 +178,7 @@ TEST(GeneratedRegistryRaceTest, ReleaseLastAndLoadMultiple) {
       threads.push_back(std::thread([&]() {
         barrier.Block();
         const upb_GeneratedRegistryRef* ref2 = upb_GeneratedRegistry_Load();
-        EXPECT_NE(upb_ExtensionRegistry_GetGenerated(ref2), nullptr);
+        EXPECT_NE(upb_GeneratedRegistry_Get(ref2), nullptr);
         upb_GeneratedRegistry_Release(ref2);
       }));
     }

upb/reflection/def_pool.c

@@ -545,5 +545,5 @@ bool _upb_DefPool_LoadDefInit(upb_DefPool* s, const _upb_DefPool_Init* init) {
 
 const upb_ExtensionRegistry* _upb_DefPool_GeneratedExtensionRegistry(
     const upb_DefPool* s) {
-  return upb_ExtensionRegistry_GetGenerated(s->generated_extreg);
+  return upb_GeneratedRegistry_Get(s->generated_extreg);
 }