Fasttable: apply UnknownField fast path to empty slots when permissible.

PiperOrigin-RevId: 908298171

Author: protobuf-github-bot

upb/wire/decode_fast/BUILD

@@ -9,7 +9,7 @@ load("@bazel_skylib//rules:common_settings.bzl", "bool_flag")
 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 @@ cc_library(
         "field_generic.c",
         "field_message.c",
         "field_string.c",
+        "field_unknown.c",
         "field_varint.c",
     ],
     hdrs = ["field_parsers.h"],
@@ -66,6 +67,7 @@ cc_library(
         "//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 @@ cc_test(
     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 @@ cc_binary(
         "@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",
+    ],
+)

upb/wire/decode_fast/combinations.h

@@ -216,6 +216,9 @@ UPB_INLINE upb_DecodeFast_Type upb_DecodeFast_GetType(uint32_t function_idx) {
   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_

upb/wire/decode_fast/field_parsers.h

@@ -43,6 +43,9 @@ const char* _upb_FastDecoder_FallbackToMiniTable(PARSE_PARAMS);
 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
 

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);
+}

upb/wire/decode_fast/function_array.c

@@ -27,9 +27,15 @@ static _upb_FieldParser* funcs[] = {UPB_DECODEFAST_FUNCTIONS(ADDR_OF_FUNC)};
 #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"

upb/wire/decode_fast/select.c

@@ -182,11 +182,14 @@ static bool upb_DecodeFast_GetFunctionData(const upb_MiniTable* m,
 
 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,19 +208,47 @@ int upb_DecodeFast_BuildTable(const upb_MiniTable* m,
     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;
       max = UPB_MAX(max, slot);
     }
   }
 
-  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 @@ const char* upb_DecodeFast_GetFunctionName(uint32_t function_idx) {
 #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];
 }

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 @@ TEST(SelectTest, FunctionIndicesMatch) {
                              << 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