diff --git a/upb/mini_table/decode.c b/upb/mini_table/decode.c index baaf1fcccb..02d70132cc 100644 --- a/upb/mini_table/decode.c +++ b/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; +} diff --git a/upb/mini_table/decode.h b/upb/mini_table/decode.h index 385bda4661..31e1192394 100644 --- a/upb/mini_table/decode.h +++ b/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 diff --git a/upbc/BUILD b/upbc/BUILD index d95d4175d7..ae561a5261 100644 --- a/upbc/BUILD +++ b/upbc/BUILD @@ -225,6 +225,7 @@ cc_library( "//:base", "//:json", "//:mem", + "//:mini_table", "//:port", "//:reflection", ], diff --git a/upbc/code_generator_request.c b/upbc/code_generator_request.c index 8aad9ac957..34f3ca1318 100644 --- a/upbc/code_generator_request.c +++ b/upbc/code_generator_request.c @@ -30,6 +30,7 @@ #include #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); diff --git a/upbc/code_generator_request.proto b/upbc/code_generator_request.proto index 07a8a02f00..0b55e69cfb 100644 --- a/upbc/code_generator_request.proto +++ b/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 mini_descriptors = 2; + // upb-specific info for the messages/enums/extensions in the request, keyed + // by the fully qualified names. + map upb_info = 2; }