Introduced a generated extension registry using linker arrays.

PiperOrigin-RevId: 640196225

upb/mini_table/extension_registry.c

@@ -70,6 +70,24 @@ failure:
   return false;
 }
 
+#ifdef UPB_LINKARR_DECLARE
+
+UPB_LINKARR_DECLARE(upb_AllExts, const upb_MiniTableExtension*);
+
+bool upb_ExtensionRegistry_AddAllLinkedExtensions(upb_ExtensionRegistry* r) {
+  const upb_MiniTableExtension* const* start = UPB_LINKARR_START(upb_AllExts);
+  const upb_MiniTableExtension* const* stop = UPB_LINKARR_STOP(upb_AllExts);
+  for (const upb_MiniTableExtension* const* p = start; p < stop; p++) {
+    // Windows can introduce zero padding, so we have to skip zeroes.
+    if (*p != 0) {
+      if (!upb_ExtensionRegistry_Add(r, *p)) return false;
+    }
+  }
+  return true;
+}
+
+#endif  // UPB_LINKARR_DECLARE
+
 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

@@ -71,6 +71,23 @@ bool upb_ExtensionRegistry_AddArray(upb_ExtensionRegistry* r,
                                     const upb_MiniTableExtension** e,
                                     size_t count);
 
+#ifdef UPB_LINKARR_DECLARE
+
+// 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);
+
+#endif  // UPB_LINKARR_DECLARE
+
 // 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/port/def.inc

@@ -350,3 +350,79 @@ void __asan_unpoison_memory_region(void const volatile *addr, size_t size);
 error UPB_TRACING_ENABLED Tracing should be disabled in production builds
 #endif
 #endif
+
+// Linker arrays combine elements from multiple translation units into a single
+// array that can be iterated over at runtime.
+//
+// It is an alternative to pre-main "registration" functions.
+//
+// Usage:
+//
+//   // In N translation units.
+//   UPB_LINKARR_APPEND(foo_array) static int elems[3] = {1, 2, 3};
+//
+//   // At runtime:
+//   UPB_LINKARR_DECLARE(foo_array, int);
+//
+//   void f() {
+//     const int* start = UPB_LINKARR_START(foo_array);
+//     const int* stop = UPB_LINKARR_STOP(foo_array);
+//     for (const int* p = start; p < stop; p++) {
+//       // Windows can introduce zero padding, so we have to skip zeroes.
+//       if (*p != 0) {
+//         vec.push_back(*p);
+//       }
+//     }
+//   }
+
+#if defined(__ELF__) || defined(__wasm__)
+
+#define UPB_LINKARR_APPEND(name) \
+  __attribute__((retain, used, section("linkarr_" #name)))
+#define UPB_LINKARR_DECLARE(name, type)     \
+  extern type const __start_linkarr_##name; \
+  extern type const __stop_linkarr_##name;  \
+  UPB_LINKARR_APPEND(name) type UPB_linkarr_internal_empty_##name[1] = {0}
+#define UPB_LINKARR_START(name) (&__start_linkarr_##name)
+#define UPB_LINKARR_STOP(name) (&__stop_linkarr_##name)
+
+#elif defined(__MACH__)
+
+/* As described in: https://stackoverflow.com/a/22366882 */
+#define UPB_LINKARR_APPEND(name) \
+  __attribute__((retain, used, section("__DATA,la_" #name)))
+#define UPB_LINKARR_DECLARE(name, type)           \
+  extern type const __start_linkarr_##name __asm( \
+      "section$start$__DATA$la_" #name);          \
+  extern type const __stop_linkarr_##name __asm(  \
+      "section$end$__DATA$"                       \
+      "la_" #name);                               \
+  UPB_LINKARR_APPEND(name) type UPB_linkarr_internal_empty_##name[1] = {0}
+#define UPB_LINKARR_START(name) (&__start_linkarr_##name)
+#define UPB_LINKARR_STOP(name) (&__stop_linkarr_##name)
+
+#elif defined(_MSC_VER) && defined(__clang__)
+
+/* See:
+ *   https://devblogs.microsoft.com/oldnewthing/20181107-00/?p=100155
+ *   https://devblogs.microsoft.com/oldnewthing/20181108-00/?p=100165
+ *   https://devblogs.microsoft.com/oldnewthing/20181109-00/?p=100175 */
+
+// Usage of __attribute__ here probably means this is Clang-specific, and would
+// not work on MSVC.
+#define UPB_LINKARR_APPEND(name) \
+  __declspec(allocate("la_" #name "$j")) __attribute__((retain, used))
+#define UPB_LINKARR_DECLARE(name, type)                               \
+  __declspec(allocate("la_" #name "$a")) type __start_linkarr_##name; \
+  __declspec(allocate("la_" #name "$z")) type __stop_linkarr_##name;  \
+  UPB_LINKARR_APPEND(name) type UPB_linkarr_internal_empty_##name[1] = {0}
+#define UPB_LINKARR_START(name) (&__start_linkarr_##name)
+#define UPB_LINKARR_STOP(name) (&__stop_linkarr_##name)
+
+#else
+
+// Linker arrays are not supported on this platform.  Make appends a no-op but
+// don't define the other macros.
+#define UPB_LINKARR_APPEND(name)
+
+#endif

upb/port/undef.inc

@@ -56,3 +56,7 @@
 #undef UPB_USE_C11_ATOMICS
 #undef UPB_PRIVATE
 #undef UPB_ONLYBITS
+#undef UPB_LINKARR_DECLARE
+#undef UPB_LINKARR_APPEND
+#undef UPB_LINKARR_START
+#undef UPB_LINKARR_STOP

upb_generator/protoc-gen-upb_minitable.cc

@@ -519,6 +519,7 @@ int WriteExtensions(const DefPoolPair& pools, upb::FileDefPtr file,
 
   output(
       "\n"
+      "UPB_LINKARR_APPEND(upb_AllExts)\n"
       "static const upb_MiniTableExtension *$0[$1] = {\n",
       kExtensionsInit, exts.size());