Fasttable: apply UnknownField fast path to empty slots when permissible. PiperOrigin-RevId: 908298171
diff --git a/upb/wire/decode_fast/BUILD b/upb/wire/decode_fast/BUILD index 5ad2647..9ebe6ed 100644 --- a/upb/wire/decode_fast/BUILD +++ b/upb/wire/decode_fast/BUILD
@@ -9,7 +9,7 @@ load("@rules_cc//cc:cc_binary.bzl", "cc_binary") load("@rules_cc//cc:cc_test.bzl", "cc_test") load("@rules_cc//cc:defs.bzl", "cc_library") -load("//upb/bazel:copts.bzl", "UPB_DEFAULT_COPTS", "UPB_DEFAULT_FEATURES") +load("//upb/bazel:copts.bzl", "UPB_DEFAULT_COPTS", "UPB_DEFAULT_CPPOPTS", "UPB_DEFAULT_FEATURES") package(default_applicable_licenses = ["//:license"]) @@ -39,6 +39,7 @@ "field_generic.c", "field_message.c", "field_string.c", + "field_unknown.c", "field_varint.c", ], hdrs = ["field_parsers.h"], @@ -66,6 +67,7 @@ "//upb/message:types", "//upb/mini_table", "//upb/port", + "//upb/wire", "//upb/wire:decoder", "//upb/wire:eps_copy_input_stream", "//upb/wire:reader", @@ -119,6 +121,9 @@ deps = [ ":combinations", ":select", + "//upb/mem", + "//upb/wire/test_util:field_types", + "//upb/wire/test_util:make_mini_table", "@abseil-cpp//absl/base:core_headers", "@googletest//:gtest", "@googletest//:gtest_main", @@ -182,3 +187,23 @@ "@abseil-cpp//absl/flags:parse", ], ) + +cc_test( + name = "test_unknown", + srcs = ["test_unknown.cc"], + copts = UPB_DEFAULT_CPPOPTS, + features = UPB_DEFAULT_FEATURES, + deps = [ + ":combinations", + "//upb/base", + "//upb/mem", + "//upb/message", + "//upb/port", + "//upb/wire", + "//upb/wire/test_util:field_types", + "//upb/wire/test_util:make_mini_table", + "//upb/wire/test_util:wire_message", + "@googletest//:gtest", + "@googletest//:gtest_main", + ], +)
diff --git a/upb/wire/decode_fast/combinations.h b/upb/wire/decode_fast/combinations.h index 58a1e04..cdb57cd 100644 --- a/upb/wire/decode_fast/combinations.h +++ b/upb/wire/decode_fast/combinations.h
@@ -216,6 +216,9 @@ UPB_DECODEFAST_COMBINATION_IS_ENABLED(type, card, size) #endif +// A special value for function_idx to indicate that the field is unknown. +#define kUpb_DecodeFast_Unknown (UINT32_MAX - 1) + #include "upb/port/undef.inc" #endif // UPB_WIRE_INTERNAL_DECODE_FAST_COMBINATIONS_H_
diff --git a/upb/wire/decode_fast/field_parsers.h b/upb/wire/decode_fast/field_parsers.h index a166f87..57f5d81 100644 --- a/upb/wire/decode_fast/field_parsers.h +++ b/upb/wire/decode_fast/field_parsers.h
@@ -43,6 +43,9 @@ UPB_PRESERVE_NONE const char* _upb_FastDecoder_DecodeGeneric(PARSE_PARAMS); +UPB_PRESERVE_NONE +const char* _upb_FastDecoder_DecodeUnknown(PARSE_PARAMS); + #undef F #undef PARSE_PARAMS
diff --git a/upb/wire/decode_fast/field_unknown.c b/upb/wire/decode_fast/field_unknown.c new file mode 100644 index 0000000..e4c5416 --- /dev/null +++ b/upb/wire/decode_fast/field_unknown.c
@@ -0,0 +1,123 @@ +// Protocol Buffers - Google's data interchange format +// Copyright 2025 Google LLC. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file or at +// https://developers.google.com/open-source/licenses/bsd + +#include <stdint.h> + +#include "upb/base/string_view.h" +#include "upb/message/internal/message.h" +#include "upb/message/message.h" +#include "upb/wire/decode.h" +#include "upb/wire/decode_fast/cardinality.h" +#include "upb/wire/decode_fast/field_parsers.h" +#include "upb/wire/eps_copy_input_stream.h" +#include "upb/wire/internal/decoder.h" +#include "upb/wire/reader.h" + +// Must be last. +#include "upb/port/def.inc" + +UPB_FORCEINLINE bool _upb_FastDecoder_DoDecodeUnknown( + struct upb_Decoder* d, const char** ptr, upb_Message* msg, intptr_t table, + uint64_t hasbits, uint64_t data, upb_DecodeFastNext* ret) { + const char* start = *ptr; + + // Important: if the branch is correctly predicted, the ptr incremen is + // treated as constant and subsequent loads will not have a data dependency on + // the branch. + if (UPB_LIKELY((data & 0x80) == 0)) { + *ptr += 1; + // Ensure the field number is not 0. + // Use bitwise op to only examine first byte minus additional tag data. + if (UPB_UNLIKELY((data & 0xF8) == 0)) { + return UPB_DECODEFAST_ERROR(d, kUpb_DecodeStatus_Malformed, ret); + } + } else if ((data & 0x8000) == 0) { + *ptr += 2; + // Ensure the field number is not 0. + // Use bitwise op to limit to first two bytes, and ignore continuation bit & + // additional tag data. + if (UPB_UNLIKELY((data & 0x7f78) == 0)) { + return UPB_DECODEFAST_ERROR(d, kUpb_DecodeStatus_Malformed, ret); + } + } else { + // Tag >=2048 + return UPB_DECODEFAST_EXIT(kUpb_DecodeFastNext_FallbackToMiniTable, ret); + } + + uint32_t wire_type = data & 0x07; + + if (UPB_UNLIKELY(wire_type == kUpb_WireType_EndGroup || + wire_type == kUpb_WireType_StartGroup)) { + // FastDecoder doesn't handle group fields, but it can be used to decode a + // message that is itself a group. When decoding a group, the end of the + // message is marked by an EndGroup tag. Since EndGroup tags are not in + // the MiniTable, they are routed to the unknown field handler. We must + // intercept them here to properly terminate the message. + return UPB_DECODEFAST_EXIT(kUpb_DecodeFastNext_FallbackToMiniTable, ret); + } + + upb_EpsCopyInputStream_StartCapture(&d->input, start); + + switch (wire_type) { + case kUpb_WireType_Varint: + *ptr = upb_WireReader_SkipVarint(*ptr, &d->input); + if (UPB_UNLIKELY(!ptr)) { + return UPB_DECODEFAST_ERROR(d, kUpb_DecodeStatus_Malformed, ret); + } + break; + case kUpb_WireType_32Bit: + UPB_PRIVATE(upb_EpsCopyInputStream_ConsumeBytes)(&d->input, 4); + *ptr += 4; + break; + case kUpb_WireType_64Bit: + UPB_PRIVATE(upb_EpsCopyInputStream_ConsumeBytes)(&d->input, 8); + *ptr += 8; + break; + case kUpb_WireType_Delimited: { + int size; + *ptr = upb_WireReader_ReadSize(*ptr, &size, &d->input); + if (UPB_UNLIKELY(!ptr || !upb_EpsCopyInputStream_CheckSize(&d->input, + *ptr, size))) { + return UPB_DECODEFAST_ERROR(d, kUpb_DecodeStatus_Malformed, ret); + } + *ptr += size; + break; + } + default: + return UPB_DECODEFAST_ERROR(d, kUpb_DecodeStatus_Malformed, ret); + } + + upb_StringView sv; + if (UPB_UNLIKELY(!upb_EpsCopyInputStream_EndCapture(&d->input, *ptr, &sv))) { + return UPB_DECODEFAST_ERROR(d, kUpb_DecodeStatus_Malformed, ret); + } + + upb_AddUnknownMode mode = kUpb_AddUnknown_Copy; + if (d->options & kUpb_DecodeOption_AliasString) { + if (sv.data != d->input.buffer_start) { + mode = kUpb_AddUnknown_AliasAllowMerge; + } else { + mode = kUpb_AddUnknown_Alias; + } + } + + if (!UPB_PRIVATE(_upb_Message_AddUnknown)(msg, sv.data, sv.size, &d->arena, + mode)) { + return UPB_DECODEFAST_ERROR(d, kUpb_DecodeStatus_OutOfMemory, ret); + } + + data = 0; + return true; +} + +UPB_PRESERVE_NONE const char* _upb_FastDecoder_DecodeUnknown( + struct upb_Decoder* d, const char* ptr, upb_Message* msg, intptr_t table, + uint64_t hasbits, uint64_t data) { + upb_DecodeFastNext ret = kUpb_DecodeFastNext_Dispatch; + _upb_FastDecoder_DoDecodeUnknown(d, &ptr, msg, table, hasbits, data, &ret); + UPB_DECODEFAST_NEXT(ret); +}
diff --git a/upb/wire/decode_fast/function_array.c b/upb/wire/decode_fast/function_array.c index 4f6f7d2..b30fe88 100644 --- a/upb/wire/decode_fast/function_array.c +++ b/upb/wire/decode_fast/function_array.c
@@ -27,9 +27,15 @@ #undef ADDR_OF_FUNC _upb_FieldParser* upb_DecodeFast_GetFunctionPointer(uint32_t function_idx) { - if (function_idx == UINT32_MAX) return &_upb_FastDecoder_DecodeGeneric; - UPB_ASSERT(function_idx < UPB_ARRAY_SIZE(funcs)); - return funcs[function_idx]; + switch (function_idx) { + case UINT32_MAX: + return &_upb_FastDecoder_DecodeGeneric; + case kUpb_DecodeFast_Unknown: + return &_upb_FastDecoder_DecodeUnknown; + default: + UPB_ASSERT(function_idx < UPB_ARRAY_SIZE(funcs)); + return funcs[function_idx]; + } } #include "upb/port/undef.inc"
diff --git a/upb/wire/decode_fast/select.c b/upb/wire/decode_fast/select.c index 883349c..ceb168d 100644 --- a/upb/wire/decode_fast/select.c +++ b/upb/wire/decode_fast/select.c
@@ -182,11 +182,14 @@ static bool upb_DecodeFast_TryFillEntry(const upb_MiniTable* m, const upb_MiniTableField* field, + bool* out_supported_tag_size, upb_DecodeFast_TableEntry* entry) { UPB_ASSERT(!upb_MiniTableField_IsExtension(field)); uint16_t tag; upb_DecodeFast_TagSize tag_size; - return upb_DecodeFast_GetEncodedTag(field, &tag, &tag_size) && + *out_supported_tag_size = + upb_DecodeFast_GetEncodedTag(field, &tag, &tag_size); + return *out_supported_tag_size && upb_DecodeFast_GetFunctionIndex(m, field, tag_size, &entry->function_idx) && UPB_DECODEFAST_ISENABLED( @@ -205,11 +208,22 @@ table[i].function_data = 0; } + // Fasttable only handles fields with tag size of 1 or 2 bytes. If all known + // fields with such tag sizes are covered, and the message is non-extensible, + // we can short circuit misses to unknown field handling + bool all_supported_tag_size_fields_map_to_assigned_slots = true; int max = 0; for (size_t i = 0, n = upb_MiniTable_FieldCount(m); i < n; i++) { const upb_MiniTableField* field = upb_MiniTable_GetFieldByIndex(m, i); + bool supported_tag_size; upb_DecodeFast_TableEntry entry; - if (!upb_DecodeFast_TryFillEntry(m, field, &entry)) continue; + if (!upb_DecodeFast_TryFillEntry(m, field, &supported_tag_size, &entry)) { + if (supported_tag_size) { + // Check if this tag collides + all_supported_tag_size_fields_map_to_assigned_slots = false; + } + continue; + } int slot = upb_DecodeFastData_GetTableSlot(entry.function_data); if (table[slot].function_idx == UINT32_MAX) { table[slot] = entry; @@ -217,7 +231,24 @@ } } - return max == 0 ? 0 : upb_RoundUpToPowerOfTwo(max + 1); + int table_size = max == 0 ? 0 : upb_RoundUpToPowerOfTwo(max + 1); + + // If the message is not extendable, we can swap the generic handler for a + // fast unknown field handler in remaining open slots. + // The fast unknown handler only covers 1/2 byte tags and falls back for >2 + // bytes; thus, we do not need to check for total exhaustiveness in field + // coverage, only for 1/2 byte tags. + if (all_supported_tag_size_fields_map_to_assigned_slots && + m->UPB_PRIVATE(ext) == kUpb_ExtMode_NonExtendable) { + for (int i = 0; i < table_size; i++) { + if (table[i].function_idx == UINT32_MAX) { + table[i].function_idx = kUpb_DecodeFast_Unknown; + table[i].function_data = 0; + } + } + } + + return table_size; } uint8_t upb_DecodeFast_GetTableMask(int table_size) { @@ -237,6 +268,9 @@ #undef FUNCSTR if (function_idx == UINT32_MAX) return "_upb_FastDecoder_DecodeGeneric"; + if (function_idx == kUpb_DecodeFast_Unknown) { + return "_upb_FastDecoder_DecodeUnknown"; + } UPB_ASSERT(function_idx < UPB_ARRAY_SIZE(names)); return names[function_idx]; }
diff --git a/upb/wire/decode_fast/select_test.cc b/upb/wire/decode_fast/select_test.cc index 2a167c0..cb4f8b1 100644 --- a/upb/wire/decode_fast/select_test.cc +++ b/upb/wire/decode_fast/select_test.cc
@@ -1,13 +1,26 @@ +// Protocol Buffers - Google's data interchange format +// Copyright 2025 Google LLC. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file or at +// https://developers.google.com/open-source/licenses/bsd + #include "upb/wire/decode_fast/select.h" #include <cstddef> #include <cstdint> #include <cstdio> +#include <cstring> #include <gtest/gtest.h> #include "absl/base/macros.h" +#include "upb/mem/arena.hpp" #include "upb/wire/decode_fast/combinations.h" +#include "upb/wire/test_util/field_types.h" +#include "upb/wire/test_util/make_mini_table.h" +namespace upb { +namespace test { namespace { TEST(SelectTest, FunctionIndicesMatch) { @@ -22,6 +35,25 @@ << upb_DecodeFast_GetFunctionName(numbers[i]); } #undef IDX -} // namespace +} + +TEST(SelectTest, UnknownSlotsAssignedForNonExtendable) { + upb::Arena mt_arena; + auto [mt, field] = MiniTable::MakeSingleFieldTable<field_types::Int32>( + 1, kUpb_DecodeFast_Scalar, mt_arena.ptr()); + + upb_DecodeFast_TableEntry table[32]; + int size = upb_DecodeFast_BuildTable(mt, table); + + // Field 1 maps to slot 1: (1 << 3) | 0 (varint wire type) = 8. + // Bits 3-7 are 00001, so slot 1. + // Table size will be 2 (power of 2 >= max+1). + EXPECT_EQ(size, 2); + EXPECT_NE(table[1].function_idx, UINT32_MAX); + EXPECT_NE(table[1].function_idx, kUpb_DecodeFast_Unknown); + EXPECT_EQ(table[0].function_idx, kUpb_DecodeFast_Unknown); +} } // namespace +} // namespace test +} // namespace upb \ No newline at end of file
diff --git a/upb/wire/decode_fast/test_unknown.cc b/upb/wire/decode_fast/test_unknown.cc new file mode 100644 index 0000000..81a9816 --- /dev/null +++ b/upb/wire/decode_fast/test_unknown.cc
@@ -0,0 +1,101 @@ +// Protocol Buffers - Google's data interchange format +// Copyright 2025 Google LLC. All rights reserved. +// +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file or at +// https://developers.google.com/open-source/licenses/bsd + +// This test requires fasttable support to work properly. +// Run with `--//third_party/upb:fasttable_enabled=true`. + +#include <cstdint> +#include <ostream> +#include <string> + +#include <gtest/gtest.h> +#include "upb/base/string_view.h" +#include "upb/mem/arena.hpp" +#include "upb/message/message.h" +#include "upb/wire/decode.h" +#include "upb/wire/decode_fast/combinations.h" +#include "upb/wire/test_util/field_types.h" +#include "upb/wire/test_util/make_mini_table.h" +#include "upb/wire/test_util/wire_message.h" + +// Must be last. +#include "upb/port/def.inc" + +namespace upb { +namespace test { +namespace { + +struct UnknownFieldTestCase { + uint32_t field_number; + wire_types::WireValue value; + std::string name; + + template <typename Sink> + friend void AbslStringify(Sink& sink, const UnknownFieldTestCase& tc) { + sink.Append("{field_number: "); + sink.Append(std::to_string(tc.field_number)); + sink.Append(", name: \""); + sink.Append(tc.name); + sink.Append("\"}"); + } +}; + +class UnknownFieldTest : public testing::TestWithParam<UnknownFieldTestCase> {}; + +TEST_P(UnknownFieldTest, UnknownFieldFastPath) { + const auto& test_case = GetParam(); + uint32_t unknown_field_num = test_case.field_number; + char trace_buf[64]; + upb::Arena msg_arena; + upb::Arena mt_arena; + + // Create a MiniTable with field 1. Any other field number will be unknown. + auto [mt, field] = MiniTable::MakeSingleFieldTable<field_types::Int32>( + 1, kUpb_DecodeFast_Scalar, mt_arena.ptr()); + upb_Message* msg = upb_Message_New(mt, msg_arena.ptr()); + + std::string payload = ToBinaryPayload( + wire_types::WireMessage{{unknown_field_num, test_case.value}}); + + upb_DecodeStatus result = + upb_DecodeWithTrace(payload.data(), payload.size(), msg, mt, nullptr, 0, + msg_arena.ptr(), trace_buf, sizeof(trace_buf)); + ASSERT_EQ(result, kUpb_DecodeStatus_Ok) << upb_DecodeStatus_String(result); + + // Verify it was saved as an unknown field. + ASSERT_TRUE(upb_Message_HasUnknown(msg)); + + // Verify the contents of the unknown field. + uintptr_t iter = kUpb_Message_UnknownBegin; + upb_StringView unknown_data; + std::string captured_unknown; + while (upb_Message_NextUnknown(msg, &unknown_data, &iter)) { + captured_unknown.append(unknown_data.data, unknown_data.size); + } + EXPECT_EQ(captured_unknown, payload); + +#if !defined(NDEBUG) && defined(UPB_ENABLE_FASTTABLE) + // Because it was parsed on the fast path, we should not see fallback ('<') + // or mini table ('M') in the trace output. + EXPECT_FALSE(absl::StrContains(trace_buf, "<")); + EXPECT_FALSE(absl::StrContains(trace_buf, "M")); +#endif +} + +INSTANTIATE_TEST_SUITE_P( + AllWireTypes, UnknownFieldTest, + testing::Values( + UnknownFieldTestCase{2, wire_types::Varint{123}, "Varint"}, + UnknownFieldTestCase{2, wire_types::Delimited{"Hello World"}, + "Delimited"}, + UnknownFieldTestCase{100, wire_types::Fixed64{0x1234567890ABCDEF}, + "Fixed64"}, + UnknownFieldTestCase{16, wire_types::Fixed32{0x12345678}, "Fixed32"})); + +} // namespace +} // namespace test +} // namespace upb