Added new APIs for linking a MiniTable all at one time

The new API upb_MiniTable_Link() links all sub-messages and sub-enums at a single time, by accepting an array of sub-tables and sub-enums.  The order of these sub-tables can be queried using a separate function `upb_MiniTable_GetSubList()`, and this information is added to `CodeGeneratorRequest` as part of the upb-specific info.

PiperOrigin-RevId: 513970874

upb/mini_table/decode.c

@@ -1062,3 +1062,64 @@ bool upb_MiniTable_SetSubEnum(upb_MiniTable* table, upb_MiniTableField* field,
   table_sub->subenum = sub;
   return true;
 }
+
+uint32_t upb_MiniTable_GetSubList(const upb_MiniTable* mt,
+                                  const upb_MiniTableField** subs) {
+  uint32_t msg_count = 0;
+  uint32_t enum_count = 0;
+
+  for (int i = 0; i < mt->field_count; i++) {
+    const upb_MiniTableField* f = &mt->fields[i];
+    if (upb_MiniTableField_CType(f) == kUpb_CType_Message) {
+      *subs = f;
+      ++subs;
+      msg_count++;
+    }
+  }
+
+  for (int i = 0; i < mt->field_count; i++) {
+    const upb_MiniTableField* f = &mt->fields[i];
+    if (upb_MiniTableField_CType(f) == kUpb_CType_Enum) {
+      *subs = f;
+      ++subs;
+      enum_count++;
+    }
+  }
+
+  return (msg_count << 16) | enum_count;
+}
+
+// The list of sub_tables and sub_enums must exactly match the number and order
+// of sub-message fields and sub-enum fields given by upb_MiniTable_GetSubList()
+// above.
+bool upb_MiniTable_Link(upb_MiniTable* mt, const upb_MiniTable** sub_tables,
+                        size_t sub_table_count,
+                        const upb_MiniTableEnum** sub_enums,
+                        size_t sub_enum_count) {
+  uint32_t msg_count = 0;
+  uint32_t enum_count = 0;
+
+  for (int i = 0; i < mt->field_count; i++) {
+    upb_MiniTableField* f = (upb_MiniTableField*)&mt->fields[i];
+    if (upb_MiniTableField_CType(f) == kUpb_CType_Message) {
+      const upb_MiniTable* sub = sub_tables[msg_count++];
+      if (msg_count > sub_table_count) return false;
+      if (sub != NULL) {
+        if (!upb_MiniTable_SetSubMessage(mt, f, sub)) return false;
+      }
+    }
+  }
+
+  for (int i = 0; i < mt->field_count; i++) {
+    upb_MiniTableField* f = (upb_MiniTableField*)&mt->fields[i];
+    if (upb_MiniTableField_CType(f) == kUpb_CType_Enum) {
+      const upb_MiniTableEnum* sub = sub_enums[enum_count++];
+      if (enum_count > sub_table_count) return false;
+      if (sub != NULL) {
+        if (!upb_MiniTable_SetSubEnum(mt, f, sub)) return false;
+      }
+    }
+  }
+
+  return true;
+}

upb/mini_table/decode.h

@@ -140,6 +140,33 @@ upb_MiniTable* upb_MiniTable_BuildWithBuf(const char* data, size_t len,
                                           upb_Arena* arena, void** buf,
                                           size_t* buf_size, upb_Status* status);
 
+// Returns a list of fields that require linking at runtime, to connect the
+// MiniTable to its sub-messages and sub-enums.  The list of fields will be
+// written to the `subs` array, which must have been allocated by the caller
+// and must be large enough to hold a list of all fields in the message.
+//
+// The order of the fields returned by this function is significant: it matches
+// the order expected by upb_MiniTable_Link() below.
+//
+// The return value packs the sub-message count and sub-enum count into a single
+// integer like so:
+//  return (msg_count << 16) | enum_count;
+UPB_API uint32_t upb_MiniTable_GetSubList(const upb_MiniTable* mt,
+                                          const upb_MiniTableField** subs);
+
+// Links a message to its sub-messages and sub-enums.  The caller must pass
+// arrays of sub-tables and sub-enums, in the same length and order as is
+// returned by upb_MiniTable_GetSubList() above.  However, individual elements
+// of the sub_tables may be NULL if those sub-messages were tree shaken.
+//
+// Returns false if either array is too short, or if any of the tables fails
+// to link.
+UPB_API bool upb_MiniTable_Link(upb_MiniTable* mt,
+                                const upb_MiniTable** sub_tables,
+                                size_t sub_table_count,
+                                const upb_MiniTableEnum** sub_enums,
+                                size_t sub_enum_count);
+
 #ifdef __cplusplus
 } /* extern "C" */
 #endif

upbc/BUILD

@@ -225,6 +225,7 @@ cc_library(
         "//:base",
         "//:json",
         "//:mem",
+        "//:mini_table",
         "//:port",
         "//:reflection",
     ],

upbc/code_generator_request.c

@@ -30,6 +30,7 @@
 #include <inttypes.h>
 
 #include "google/protobuf/compiler/plugin.upb.h"
+#include "upb/mini_table/decode.h"
 #include "upb/reflection/def.h"
 
 // Must be last.
@@ -68,11 +69,22 @@ static void upbc_State_Init(upbc_State* s) {
   if (!s->out) upbc_Error(s, __func__, "could not allocate request");
 }
 
-static void upbc_State_Emit(upbc_State* s, const char* name,
-                            upb_StringView encoding) {
+static upb_StringView upbc_State_StrDup(upbc_State* s, const char* str) {
+  upb_StringView from = upb_StringView_FromString(str);
+  char* to = upb_Arena_Malloc(s->arena, from.size);
+  if (!to) upbc_Error(s, __func__, "Out of memory");
+  memcpy(to, from.data, from.size);
+  return upb_StringView_FromDataAndSize(to, from.size);
+}
+
+static void upbc_State_AddMiniDescriptor(upbc_State* s, const char* name,
+                                         upb_StringView encoding) {
   const upb_StringView key = upb_StringView_FromString(name);
-  bool ok = upbc_CodeGeneratorRequest_mini_descriptors_set(s->out, key,
-                                                           encoding, s->arena);
+  upbc_CodeGeneratorRequest_UpbInfo* info =
+      upbc_CodeGeneratorRequest_UpbInfo_new(s->arena);
+  if (!info) upbc_Error(s, __func__, "Out of memory");
+  upbc_CodeGeneratorRequest_UpbInfo_set_mini_descriptor(info, encoding);
+  bool ok = upbc_CodeGeneratorRequest_upb_info_set(s->out, key, info, s->arena);
   if (!ok) upbc_Error(s, __func__, "could not set mini descriptor in map");
 }
 
@@ -86,7 +98,7 @@ static void upbc_Scrape_Enum(upbc_State* s, const upb_EnumDef* e) {
   bool ok = upb_EnumDef_MiniDescriptorEncode(e, s->arena, &desc);
   if (!ok) upbc_Error(s, __func__, "could not encode enum");
 
-  upbc_State_Emit(s, upb_EnumDef_FullName(e), desc);
+  upbc_State_AddMiniDescriptor(s, upb_EnumDef_FullName(e), desc);
 }
 
 static void upbc_Scrape_Extension(upbc_State* s, const upb_FieldDef* f) {
@@ -94,7 +106,7 @@ static void upbc_Scrape_Extension(upbc_State* s, const upb_FieldDef* f) {
   bool ok = upb_FieldDef_MiniDescriptorEncode(f, s->arena, &desc);
   if (!ok) upbc_Error(s, __func__, "could not encode extension");
 
-  upbc_State_Emit(s, upb_FieldDef_FullName(f), desc);
+  upbc_State_AddMiniDescriptor(s, upb_FieldDef_FullName(f), desc);
 }
 
 static void upbc_Scrape_FileEnums(upbc_State* s, const upb_FileDef* f) {
@@ -168,12 +180,54 @@ static void upbc_Scrape_NestedMessages(upbc_State* s, const upb_MessageDef* m) {
   }
 }
 
+static void upbc_Scrape_MessageSubs(upbc_State* s,
+                                    upbc_CodeGeneratorRequest_UpbInfo* info,
+                                    const upb_MessageDef* m) {
+  const upb_MiniTableField** fields =
+      malloc(upb_MessageDef_FieldCount(m) * sizeof(*fields));
+  const upb_MiniTable* mt = upb_MessageDef_MiniTable(m);
+  uint32_t counts = upb_MiniTable_GetSubList(mt, fields);
+  uint32_t msg_count = counts >> 16;
+  uint32_t enum_count = counts & 0xffff;
+
+  for (uint32_t i = 0; i < msg_count; i++) {
+    const upb_FieldDef* f =
+        upb_MessageDef_FindFieldByNumber(m, fields[i]->number);
+    if (!f) upbc_Error(s, __func__, "Missing f");
+    const upb_MessageDef* sub = upb_FieldDef_MessageSubDef(f);
+    if (!sub) upbc_Error(s, __func__, "Missing sub");
+    upb_StringView name = upbc_State_StrDup(s, upb_MessageDef_FullName(sub));
+    upbc_CodeGeneratorRequest_UpbInfo_add_sub_message(info, name, s->arena);
+  }
+
+  for (uint32_t i = 0; i < enum_count; i++) {
+    const upb_FieldDef* f =
+        upb_MessageDef_FindFieldByNumber(m, fields[msg_count + i]->number);
+    if (!f) upbc_Error(s, __func__, "Missing f (2)");
+    const upb_EnumDef* sub = upb_FieldDef_EnumSubDef(f);
+    if (!sub) upbc_Error(s, __func__, "Missing sub (2)");
+    upb_StringView name = upbc_State_StrDup(s, upb_EnumDef_FullName(sub));
+    upbc_CodeGeneratorRequest_UpbInfo_add_sub_enum(info, name, s->arena);
+  }
+
+  free(fields);
+}
+
 static void upbc_Scrape_Message(upbc_State* s, const upb_MessageDef* m) {
   upb_StringView desc;
   bool ok = upb_MessageDef_MiniDescriptorEncode(m, s->arena, &desc);
   if (!ok) upbc_Error(s, __func__, "could not encode message");
 
-  upbc_State_Emit(s, upb_MessageDef_FullName(m), desc);
+  upbc_CodeGeneratorRequest_UpbInfo* info =
+      upbc_CodeGeneratorRequest_UpbInfo_new(s->arena);
+  if (!info) upbc_Error(s, __func__, "Out of memory");
+  upbc_CodeGeneratorRequest_UpbInfo_set_mini_descriptor(info, desc);
+
+  upbc_Scrape_MessageSubs(s, info, m);
+
+  const upb_StringView key = upbc_State_StrDup(s, upb_MessageDef_FullName(m));
+  ok = upbc_CodeGeneratorRequest_upb_info_set(s->out, key, info, s->arena);
+  if (!ok) upbc_Error(s, __func__, "could not set mini descriptor in map");
 
   upbc_Scrape_NestedEnums(s, m);
   upbc_Scrape_NestedExtensions(s, m);

upbc/code_generator_request.proto

@@ -5,9 +5,22 @@ package upbc;
 import "src/google/protobuf/compiler/plugin.proto";
 
 message CodeGeneratorRequest {
+  message UpbInfo {
+    optional string mini_descriptor = 1;
+
+    // An ordered list of fully qualified sub-message names whose upb_MiniTable
+    // should be passed to upb_MiniTable_Link().
+    repeated string sub_message = 3;
+
+    // An ordered list of fully qualified sub-enum names whose upb_MiniTableEnum
+    // should be passed to upb_MiniTable_Link().
+    repeated string sub_enum = 4;
+  }
+
   // The pb sent by protoc to its plugins.
   optional google.protobuf.compiler.CodeGeneratorRequest request = 1;
 
-  // Mini descriptors for the above pb, keyed by the fully qualified names.
-  map<string, string> mini_descriptors = 2;
+  // upb-specific info for the messages/enums/extensions in the request, keyed
+  // by the fully qualified names.
+  map<string, UpbInfo> upb_info = 2;
 }