diff --git a/upb/mini_table.c b/upb/mini_table.c index 5fa60ed6ec..4347b1515e 100644 --- a/upb/mini_table.c +++ b/upb/mini_table.c @@ -411,7 +411,7 @@ static void upb_MiniTable_SetField(upb_MtDecoder* d, uint8_t ch, } field->descriptortype = kUpb_EncodedToType[type]; if (upb_MiniTable_HasSub(type, msg_modifiers)) { - field->submsg_index = (*sub_count)++; + field->submsg_index = sub_count ? (*sub_count)++ : 0; } } @@ -528,15 +528,20 @@ static const char* upb_MtDecoder_DecodeOneofs(upb_MtDecoder* d, static const char* upb_MtDecoder_ParseModifier(upb_MtDecoder* d, const char* ptr, char first_ch, + uint16_t field_count, uint64_t* msg_modifiers) { uint32_t mod; ptr = upb_MiniTable_DecodeBase92Varint(d, ptr, first_ch, kUpb_EncodedValue_MinModifier, kUpb_EncodedValue_MaxModifier, &mod); - if (d->table->field_count == 0) { + if (field_count == 0) { + if (!d->table) { + upb_MtDecoder_ErrorFormat(d, "Extensions cannot have message modifiers"); + UPB_UNREACHABLE(); + } *msg_modifiers = mod; } else { - upb_MiniTable_Field* field = &d->fields[d->table->field_count - 1]; + upb_MiniTable_Field* field = &d->fields[field_count - 1]; upb_MtDecoder_ModifyField(d, mod, field); } @@ -549,38 +554,37 @@ static void upb_MtDecoder_AllocateSubs(upb_MtDecoder* d, uint32_t sub_count) { upb_MtDecoder_CheckOutOfMemory(d, d->table->subs); } -static void upb_MtDecoder_Parse(upb_MtDecoder* d, const char* ptr) { - // Buffer length is an upper bound on the number of fields. We will return - // what we don't use. - size_t first_size = d->end - ptr; - d->fields = upb_Arena_Malloc(d->arena, sizeof(*d->fields) * first_size); - upb_MtDecoder_CheckOutOfMemory(d, d->fields); - - d->table->field_count = 0; - d->table->fields = d->fields; - +static void upb_MtDecoder_Parse(upb_MtDecoder* d, const char* ptr, size_t len, + void* fields, size_t field_size, + uint16_t* field_count, uint32_t* sub_count) { uint64_t msg_modifiers = 0; uint32_t last_field_number = 0; - uint32_t sub_count = 0; - bool saw_skip = false; + bool need_dense_below = d->table != NULL; + + d->end = ptr + len; while (ptr < d->end) { char ch = *ptr++; if (ch <= kUpb_EncodedValue_MaxField) { - // Field type. - upb_MiniTable_Field* field = &d->fields[d->table->field_count++]; + upb_MiniTable_Field* field = fields; + *field_count += 1; + fields = (char*)fields + field_size; field->number = ++last_field_number; - upb_MiniTable_SetField(d, ch, field, msg_modifiers, &sub_count); + upb_MiniTable_SetField(d, ch, field, msg_modifiers, sub_count); } else if (kUpb_EncodedValue_MinModifier <= ch && ch <= kUpb_EncodedValue_MaxModifier) { - ptr = upb_MtDecoder_ParseModifier(d, ptr, ch, &msg_modifiers); + ptr = upb_MtDecoder_ParseModifier(d, ptr, ch, *field_count, &msg_modifiers); } else if (ch == kUpb_EncodedValue_End) { + if (!d->table) { + upb_MtDecoder_ErrorFormat(d, "Extensions cannot have oneofs."); + UPB_UNREACHABLE(); + } ptr = upb_MtDecoder_DecodeOneofs(d, ptr); } else if (kUpb_EncodedValue_MinSkip <= ch && ch <= kUpb_EncodedValue_MaxSkip) { - if (!saw_skip) { + if (need_dense_below) { d->table->dense_below = d->table->field_count; - saw_skip = true; + need_dense_below = false; } uint32_t skip; ptr = upb_MiniTable_DecodeBase92Varint(d, ptr, ch, @@ -591,14 +595,29 @@ static void upb_MtDecoder_Parse(upb_MtDecoder* d, const char* ptr) { } } - if (!saw_skip) { + if (need_dense_below) { d->table->dense_below = d->table->field_count; } +} + +static void upb_MtDecoder_ParseMessage(upb_MtDecoder* d, const char* data, + size_t len) { + // Buffer length is an upper bound on the number of fields. We will return + // what we don't use. + d->fields = upb_Arena_Malloc(d->arena, sizeof(*d->fields) * len); + upb_MtDecoder_CheckOutOfMemory(d, d->fields); + + uint32_t sub_count = 0; + d->table->field_count = 0; + d->table->fields = d->fields; + upb_MtDecoder_Parse(d, data, len, d->fields, sizeof(*d->fields), + &d->table->field_count, &sub_count); // Return unused memory from fields array. - assert(d->table->field_count <= first_size); - upb_Arena_Realloc(d->arena, d->fields, sizeof(*d->fields) * first_size, - sizeof(*d->fields) * d->table->field_count); + assert(d->table->field_count <= len); + d->fields = upb_Arena_Realloc(d->arena, d->fields, sizeof(*d->fields) * len, + sizeof(*d->fields) * d->table->field_count); + d->table->fields = d->fields; upb_MtDecoder_AllocateSubs(d, sub_count); } @@ -762,7 +781,6 @@ upb_MiniTable* upb_MiniTable_BuildWithBuf(const char* data, size_t len, upb_Arena* arena, void** buf, size_t* buf_size, upb_Status* status) { upb_MtDecoder decoder = { - .end = UPB_PTRADD(data, len), .platform = platform, .vec = { @@ -789,7 +807,7 @@ upb_MiniTable* upb_MiniTable_BuildWithBuf(const char* data, size_t len, decoder.table->table_mask = 0; decoder.table->required_count = 0; - upb_MtDecoder_Parse(&decoder, data); + upb_MtDecoder_ParseMessage(&decoder, data, len); upb_MtDecoder_AssignHasbits(decoder.table); upb_MtDecoder_SortLayoutItems(&decoder); upb_MtDecoder_AssignOffsets(&decoder); @@ -854,6 +872,37 @@ upb_MiniTable* upb_MiniTable_BuildMapEntry(upb_FieldType key_type, return ret; } +upb_MiniTable_Extension* upb_MiniTable_BuildExtensions(const char* data, + size_t len, + size_t* ext_count, + upb_Arena* arena, + upb_Status* status) { + upb_MtDecoder decoder = { + .arena = arena, + .status = status, + .table = NULL, + }; + + upb_MiniTable_Extension* exts; + + if (UPB_SETJMP(decoder.err)) { + exts = NULL; + *ext_count = 0; + goto done; + } + + uint16_t count = 0; + exts = upb_Arena_Malloc(arena, len); + upb_MtDecoder_CheckOutOfMemory(&decoder, exts); + upb_MtDecoder_Parse(&decoder, data, len, exts, sizeof(*exts), &count, NULL); + exts = upb_Arena_Realloc(arena, exts, sizeof(*exts) * len, + sizeof(*exts) * count); + +done: + *ext_count = count; + return exts; +} + upb_MiniTable* upb_MiniTable_Build(const char* data, size_t len, upb_MiniTablePlatform platform, upb_Arena* arena, upb_Status* status) { diff --git a/upb/mini_table.h b/upb/mini_table.h index 6923476d25..9cd65f7c99 100644 --- a/upb/mini_table.h +++ b/upb/mini_table.h @@ -125,6 +125,12 @@ void upb_MiniTable_SetSubEnum(upb_MiniTable* table, upb_MiniTable_Field* field, const upb_MiniTable_Enum* sub); +upb_MiniTable_Extension* upb_MiniTable_BuildExtensions(const char* data, + size_t len, + size_t* ext_count, + upb_Arena* arena, + upb_Status* status); + // Special-case functions for MessageSet layout and map entries. upb_MiniTable* upb_MiniTable_BuildMessageSet(upb_MiniTablePlatform platform, upb_Arena* arena); diff --git a/upb/mini_table_test.cc b/upb/mini_table_test.cc index 32adc2c50a..94d2814eda 100644 --- a/upb/mini_table_test.cc +++ b/upb/mini_table_test.cc @@ -147,7 +147,7 @@ TEST_P(MiniTableTest, AllScalarTypesOneof) { upb::Status status; upb_MiniTable* table = upb_MiniTable_Build( e.data().data(), e.data().size(), GetParam(), arena.ptr(), status.ptr()); - ASSERT_NE(nullptr, table); + ASSERT_NE(nullptr, table) << status.error_message(); EXPECT_EQ(count, table->field_count); absl::flat_hash_set offsets; for (int i = 0; i < 16; i++) { diff --git a/upb/upb.h b/upb/upb.h index 30dda6378f..969e757eee 100644 --- a/upb/upb.h +++ b/upb/upb.h @@ -225,6 +225,17 @@ UPB_INLINE void* upb_Arena_Malloc(upb_Arena* a, size_t size) { UPB_INLINE void* upb_Arena_Realloc(upb_Arena* a, void* ptr, size_t oldsize, size_t size) { + _upb_ArenaHead* h = (_upb_ArenaHead*)a; + char* ch_ptr = (char*)ptr; + oldsize = UPB_ALIGN_MALLOC(size); + size = UPB_ALIGN_MALLOC(size); + if (ch_ptr + oldsize == h->ptr) { + if (h->end - ch_ptr >= size) { + h->ptr = ch_ptr + size; + return ptr; + } + } + void* ret = upb_Arena_Malloc(a, size); if (ret && oldsize > 0) { diff --git a/upbc/protoc-gen-upb.cc b/upbc/protoc-gen-upb.cc index ffdb64202e..f282b9c508 100644 --- a/upbc/protoc-gen-upb.cc +++ b/upbc/protoc-gen-upb.cc @@ -370,14 +370,15 @@ upb_MiniTable_Field* FindField(upb_MiniTable* mt, uint32_t field_number) { return NULL; } -class FileLayout { +class FilePlatformLayout { public: - FileLayout(const protobuf::FileDescriptor* fd, upb_MiniTablePlatform platform) + FilePlatformLayout(const protobuf::FileDescriptor* fd, + upb_MiniTablePlatform platform) : platform_(platform) { ComputeLayout(fd); } - upb_MiniTable* GetTable(const protobuf::Descriptor* m) { + upb_MiniTable* GetTable(const protobuf::Descriptor* m) const { auto it = table_map_.find(m); assert(it != table_map_.end()); return it->second; @@ -478,6 +479,29 @@ class FileLayout { upb_MiniTablePlatform platform_; }; +class FileLayout { + public: + FileLayout(const protobuf::FileDescriptor* fd) + : descriptor_(fd), + layout32_(fd, kUpb_MiniTablePlatform_32Bit), + layout64_(fd, kUpb_MiniTablePlatform_64Bit) {} + + const protobuf::FileDescriptor* descriptor() const { return descriptor_; } + + const upb_MiniTable* GetMiniTable32(const protobuf::Descriptor* m) const { + return layout32_.GetTable(m); + } + + const upb_MiniTable* GetMiniTable64(const protobuf::Descriptor* m) const { + return layout64_.GetTable(m); + } + + private: + const protobuf::FileDescriptor* descriptor_; + FilePlatformLayout layout32_; + FilePlatformLayout layout64_; +}; + void DumpEnumValues(const protobuf::EnumDescriptor* desc, Output& output) { std::vector values; for (int i = 0; i < desc->value_count(); i++) { @@ -1312,15 +1336,19 @@ void WriteMessageField(const protobuf::FieldDescriptor* field, } // Writes a single message into a .upb.c source file. -void WriteMessage(const protobuf::Descriptor* message, Output& output, - bool fasttable_enabled) { +void WriteMessage(const protobuf::Descriptor* message, const FileLayout& layout, + Output& output, bool fasttable_enabled) { std::string msg_name = ToCIdent(message->full_name()); std::string fields_array_ref = "NULL"; std::string submsgs_array_ref = "NULL"; std::string subenums_array_ref = "NULL"; uint8_t dense_below = 0; const int dense_below_max = std::numeric_limits::max(); - MessageLayout layout(message); + const upb_MiniTable* mt_32 = layout.GetMiniTable32(message); + const upb_MiniTable* mt_64 = layout.GetMiniTable64(message); + (void)mt_32; + (void)mt_64; + MessageLayout msg_layout(message); SubLayoutArray sublayout_array(message); if (sublayout_array.total_count()) { @@ -1343,6 +1371,8 @@ void WriteMessage(const protobuf::Descriptor* message, Output& output, std::vector field_number_order = FieldNumberOrder(message); + assert(field_number_order.size() == mt_32->field_count); + assert(field_number_order.size() == mt_64->field_count); if (!field_number_order.empty()) { std::string fields_array_name = msg_name + "__fields"; fields_array_ref = "&" + fields_array_name + "[0]"; @@ -1365,7 +1395,7 @@ void WriteMessage(const protobuf::Descriptor* message, Output& output, sublayout_index = sublayout_array.GetIndex(field->enum_type()); } - WriteMessageField(field, layout, sublayout_index, output); + WriteMessageField(field, msg_layout, sublayout_index, output); } output("};\n\n"); } @@ -1374,7 +1404,7 @@ void WriteMessage(const protobuf::Descriptor* message, Output& output, uint8_t table_mask = -1; if (fasttable_enabled) { - table = FastDecodeTable(message, layout); + table = FastDecodeTable(message, msg_layout); } if (table.size() > 1) { @@ -1395,9 +1425,9 @@ void WriteMessage(const protobuf::Descriptor* message, Output& output, output("const upb_MiniTable $0 = {\n", MessageInit(message)); output(" $0,\n", submsgs_array_ref); output(" $0,\n", fields_array_ref); - output(" $0, $1, $2, $3, $4, $5,\n", GetSizeInit(layout.message_size()), + output(" $0, $1, $2, $3, $4, $5,\n", GetSizeInit(msg_layout.message_size()), field_number_order.size(), msgext, dense_below, table_mask, - layout.required_count()); + msg_layout.required_count()); if (!table.empty()) { output(" UPB_FASTTABLE_INIT({\n"); for (const auto& ent : table) { @@ -1464,14 +1494,15 @@ int WriteEnums(const protobuf::FileDescriptor* file, Output& output) { return this_file_enums.size(); } -int WriteMessages(const protobuf::FileDescriptor* file, Output& output, +int WriteMessages(const FileLayout& layout, Output& output, bool fasttable_enabled) { + const protobuf::FileDescriptor* file = layout.descriptor(); std::vector file_messages = SortedMessages(file); if (file_messages.empty()) return 0; for (auto message : file_messages) { - WriteMessage(message, output, fasttable_enabled); + WriteMessage(message, layout, output, fasttable_enabled); } output("static const upb_MiniTable *$0[$1] = {\n", kMessagesInit, @@ -1541,8 +1572,9 @@ int WriteExtensions(const protobuf::FileDescriptor* file, Output& output) { } // Writes a .upb.c source file. -void WriteSource(const protobuf::FileDescriptor* file, Output& output, +void WriteSource(const FileLayout& layout, Output& output, bool fasttable_enabled) { + const protobuf::FileDescriptor* file = layout.descriptor(); EmitFileWarning(file, output); output( @@ -1604,13 +1636,13 @@ bool Generator::Generate(const protobuf::FileDescriptor* file, } } - FileLayout layout(file, kUpb_MiniTablePlatform_64Bit); + FileLayout layout(file); Output h_output(context->Open(HeaderFilename(file))); WriteHeader(file, h_output); Output c_output(context->Open(SourceFilename(file))); - WriteSource(file, c_output, fasttable_enabled); + WriteSource(layout, c_output, fasttable_enabled); return true; }