upb/wire/decode_fast/field_map.c

// Protocol Buffers - Google's data interchange format
// Copyright 2026 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 <stddef.h>
#include <stdint.h>
#include <string.h>

#include "upb/base/string_view.h"
#include "upb/hash/common.h"
#include "upb/message/internal/map.h"
#include "upb/message/internal/message.h"
#include "upb/message/map.h"
#include "upb/message/message.h"
#include "upb/mini_table/field.h"
#include "upb/mini_table/internal/message.h"
#include "upb/mini_table/internal/sub.h"
#include "upb/mini_table/message.h"
#include "upb/wire/decode.h"
#include "upb/wire/decode_fast/cardinality.h"
#include "upb/wire/decode_fast/combinations.h"
#include "upb/wire/decode_fast/data.h"
#include "upb/wire/decode_fast/dispatch.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/internal/eps_copy_input_stream.h"
#include "upb/wire/reader.h"
#include "upb/wire/types.h"
#include "utf8_range.h"

// Must be last.
#include "upb/port/def.inc"

typedef struct {
  upb_Map* map;
  const upb_MiniTable* sub_table;
} upb_DecodeFastMap;

// Resolves and initializes the map field on the target message.
// Traverses the MiniTable metadata only when the map value is a submessage,
// creates the upb_Map if not yet allocated, and synchronizes hasbits.
UPB_FORCEINLINE
bool upb_DecodeFast_GetMap(upb_Decoder* d, upb_Message* msg,
                           const upb_MiniTable* table, uint64_t data,
                           bool is_str_map, uint64_t* hasbits,
                           upb_DecodeFastMap* map_ctx,
                           upb_DecodeFastNext* next) {
  // Sync hasbits so we don't have to preserve them across the map entries.
  upb_DecodeFast_SetHasbits(msg, *hasbits);
  *hasbits = 0;

  const upb_MiniTable* sub_table = NULL;
  if (upb_DecodeFastData_ValIsMessage(data)) {
    uint32_t submsg_ofs = upb_DecodeFastData_GetSubofs(data) * 8;
    const upb_MiniTableSubInternal* sub = UPB_PTR_AT(
        table->UPB_ONLYBITS(fields), submsg_ofs, upb_MiniTableSubInternal);
    const upb_MiniTable* entry_table = sub->UPB_PRIVATE(submsg);
    // Map entries are synthetic messages co-generated with the parent message
    // in the same compilation unit, so they are never tree-shaken independently
    // of the parent message, and are strictly validated to have exactly 2
    // fields (field 1 = key, field 2 = value). Unlinked dynamic map tables are
    // filtered out prior to fast decode dispatch.
    UPB_ASSERT(entry_table);
    UPB_ASSERT(entry_table->UPB_PRIVATE(field_count) == 2);

    const upb_MiniTableField* val_field = &entry_table->UPB_PRIVATE(fields)[1];
    sub_table = upb_MiniTable_GetSubMessageTable(val_field);
    if (UPB_UNLIKELY(!sub_table)) {
      return UPB_DECODEFAST_EXIT(kUpb_DecodeFastNext_FallbackToMiniTable, next);
    }
  }

  uint16_t offset = upb_DecodeFastData_GetOffset(data);
  upb_Map** map_p = UPB_PTR_AT(msg, offset, upb_Map*);
  upb_Map* map = *map_p;
  if (UPB_UNLIKELY(!map)) {
    size_t key_size = is_str_map ? UPB_MAPTYPE_STRING
                                 : (upb_DecodeFastData_KeyIs32(data) ? 4 : 8);
    size_t val_size = upb_DecodeFastData_GetValSize(data);
    map = _upb_Map_New(&d->arena, key_size, val_size);
    if (UPB_UNLIKELY(!map)) {
      return UPB_DECODEFAST_ERROR(d, kUpb_DecodeStatus_OutOfMemory, next);
    }
    *map_p = map;
  }

  map_ctx->map = map;
  map_ctx->sub_table = sub_table;
  return true;
}

// Parses a single map entry (key, value) and inserts it directly into upb_Map.
UPB_FORCEINLINE
bool upb_DecodeFast_ParseMapEntry(upb_Decoder* d, const char** ptr,
                                  const upb_DecodeFastMap* map_ctx,
                                  uint64_t data, bool is_str_map,
                                  upb_DecodeFastNext* next) {
  UPB_ASSERT(map_ctx->map->UPB_PRIVATE(is_strtable) == is_str_map);

  const char* p = *ptr;
  int size;
  if (UPB_UNLIKELY(!upb_DecodeFast_DecodeSize(d, &p, &size, next))) {
    return false;
  }

  if (UPB_UNLIKELY(
          !upb_EpsCopyInputStream_SizeFitsWithoutSwap(EPS(d), p, size))) {
    return UPB_DECODEFAST_EXIT(kUpb_DecodeFastNext_FallbackToMiniTable, next);
  }

  const char* entry_end = p + size;
  if (UPB_UNLIKELY(p >= entry_end)) {
    return UPB_DECODEFAST_EXIT(kUpb_DecodeFastNext_FallbackToMiniTable, next);
  }

  // Parse Key (Field 1)
  UPB_PRIVATE(upb_EpsCopyInputStream_ConsumeBytes)(EPS(d), 1);
  uint8_t key_tag = *p++;
  uint8_t expected_key_tag = upb_DecodeFastData_GetKeyTag(data);
  if (is_str_map) {
    UPB_ASSUME(expected_key_tag == ((1 << 3) | kUpb_WireType_Delimited));
  }
  if (UPB_UNLIKELY(key_tag != expected_key_tag)) {
    return UPB_DECODEFAST_EXIT(kUpb_DecodeFastNext_FallbackToMiniTable, next);
  }

  UPB_PRIVATE(upb_EpsCopyInputStream_BoundsCheckedToEnd)(EPS(d), p, entry_end);

  upb_StringView strkey = {NULL, 0};
  uintptr_t intkey = 0;
  if (is_str_map) {
    int klen;
    if (UPB_UNLIKELY(!upb_DecodeFast_DecodeSize(d, &p, &klen, next))) {
      return false;
    }
    if (UPB_UNLIKELY(entry_end - p < klen || klen < 0)) {
      return UPB_DECODEFAST_EXIT(kUpb_DecodeFastNext_FallbackToMiniTable, next);
    }
    // Note: We do not need to copy string keys here even when aliasing is
    // disabled because upb_strtable_insert() always deep-copies keys into its
    // own arena-allocated upb_SizePrefixString.
    p = upb_EpsCopyInputStream_ReadStringAlwaysAlias(EPS(d), p, klen, &strkey);
    if (UPB_UNLIKELY(p == NULL)) {
      return UPB_DECODEFAST_EXIT(kUpb_DecodeFastNext_FallbackToMiniTable, next);
    }
    // Validate the key before looking at the value so that an entry with both
    // a bad key and a bad value reports the same error as the MiniTable
    // decoder, which validates the key as soon as it reads it.
    if (upb_DecodeFastData_KeyValidateUtf8(data) &&
        !utf8_range_IsValid(strkey.data, strkey.size)) {
      return UPB_DECODEFAST_ERROR(d, kUpb_DecodeStatus_BadUtf8, next);
    }
  } else {
    uint64_t k = 0;
    switch (expected_key_tag & 0x7) {
      case kUpb_WireType_Varint:
        p = upb_WireReader_ReadVarint(p, &k, EPS(d));
        if (UPB_UNLIKELY(upb_DecodeFastData_KeyIsZigZag(data))) {
          k = upb_DecodeFastData_KeyIs32(data) ? _upb_Decoder_ZigZagDecode32(k)
                                               : _upb_Decoder_ZigZagDecode64(k);
        }
        break;
      case kUpb_WireType_32Bit: {
        if (UPB_UNLIKELY(entry_end - p < 4)) {
          return UPB_DECODEFAST_EXIT(kUpb_DecodeFastNext_FallbackToMiniTable,
                                     next);
        }
        uint32_t k32;
        p = upb_WireReader_ReadFixed32(p, &k32, EPS(d));
        k = k32;
        break;
      }
      case kUpb_WireType_64Bit:
        if (UPB_UNLIKELY(entry_end - p < 8)) {
          return UPB_DECODEFAST_EXIT(kUpb_DecodeFastNext_FallbackToMiniTable,
                                     next);
        }
        p = upb_WireReader_ReadFixed64(p, &k, EPS(d));
        break;
      default:
        // upb_DecodeFast_TryFillMapEntry() only selects the int map parsers for
        // key types that use one of the wire types above.
        UPB_UNREACHABLE();
    }
    if (upb_DecodeFastData_KeyIs32(data)) {
      intkey = (uint32_t)k;
    } else {
      UPB_ASSERT(upb_DecodeFastData_GetKeySize(data) == 8);
      intkey = (uintptr_t)k;
    }
  }

  if (UPB_UNLIKELY(p >= entry_end)) {
    return UPB_DECODEFAST_EXIT(kUpb_DecodeFastNext_FallbackToMiniTable, next);
  }

  // Parse Value (Field 2)
  UPB_PRIVATE(upb_EpsCopyInputStream_BoundsCheckedToEnd)(EPS(d), p, entry_end);
  UPB_PRIVATE(upb_EpsCopyInputStream_ConsumeBytes)(EPS(d), 1);
  uint8_t val_tag = *p++;
  upb_value tabval = {0};

  uint8_t expected_val_tag = upb_DecodeFastData_GetValTag(data);
  uint8_t expected_val_wire_type = expected_val_tag & 0x7;
  // Fasttable won't contain entries for anything but these
  UPB_ASSERT(expected_val_wire_type == kUpb_WireType_Varint ||
             expected_val_wire_type == kUpb_WireType_32Bit ||
             expected_val_wire_type == kUpb_WireType_64Bit ||
             expected_val_wire_type == kUpb_WireType_Delimited);

  if (UPB_UNLIKELY(val_tag != expected_val_tag)) {
    return UPB_DECODEFAST_EXIT(kUpb_DecodeFastNext_FallbackToMiniTable, next);
  }

  // The value must have at least one byte; an empty value is malformed.
  if (UPB_UNLIKELY(p >= entry_end)) {
    return UPB_DECODEFAST_EXIT(kUpb_DecodeFastNext_FallbackToMiniTable, next);
  }

  UPB_PRIVATE(upb_EpsCopyInputStream_BoundsCheckedToEnd)(EPS(d), p, entry_end);
  switch (expected_val_wire_type) {
    case kUpb_WireType_Varint: {
      uint64_t v = 0;
      const char* p_next = upb_WireReader_ReadVarint(p, &v, EPS(d));
      if (UPB_UNLIKELY(p_next != entry_end)) {
        return UPB_DECODEFAST_EXIT(kUpb_DecodeFastNext_FallbackToMiniTable,
                                   next);
      }
      p = p_next;
      if (upb_DecodeFastData_ValIsBool(data)) {
        tabval.val = (v != 0);
      } else {
        if (UPB_UNLIKELY(upb_DecodeFastData_ValIsZigZag(data))) {
          v = upb_DecodeFastData_ValIs32(data) ? _upb_Decoder_ZigZagDecode32(v)
                                               : _upb_Decoder_ZigZagDecode64(v);
        }
        tabval.val = upb_DecodeFastData_ValIs32(data) ? (uint32_t)v : v;
      }
      break;
    }
    case kUpb_WireType_64Bit: {
      if (UPB_UNLIKELY(entry_end - p != 8)) {
        return UPB_DECODEFAST_EXIT(kUpb_DecodeFastNext_FallbackToMiniTable,
                                   next);
      }
      p = upb_WireReader_ReadFixed64(p, &tabval.val, EPS(d));
      break;
    }
    case kUpb_WireType_32Bit: {
      if (UPB_UNLIKELY(entry_end - p != 4)) {
        return UPB_DECODEFAST_EXIT(kUpb_DecodeFastNext_FallbackToMiniTable,
                                   next);
      }
      uint32_t v32;
      p = upb_WireReader_ReadFixed32(p, &v32, EPS(d));
      tabval.val = v32;
      break;
    }
    case kUpb_WireType_Delimited: {
      int vlen;
      if (UPB_UNLIKELY(!upb_DecodeFast_DecodeSize(d, &p, &vlen, next))) {
        return false;
      }
      if (UPB_UNLIKELY(entry_end - p != vlen)) {
        return UPB_DECODEFAST_EXIT(kUpb_DecodeFastNext_FallbackToMiniTable,
                                   next);
      }
      if (map_ctx->sub_table) {
        upb_Message* sub_msg = _upb_Message_New(map_ctx->sub_table, &d->arena);
        if (UPB_UNLIKELY(!sub_msg)) {
          return UPB_DECODEFAST_ERROR(d, kUpb_DecodeStatus_OutOfMemory, next);
        }
        upb_DecodeFast_MessageContext ctx = {map_ctx->sub_table, false,
                                             sub_msg};
        if (upb_EpsCopyInputStream_TryParseDelimitedFast(
                EPS(d), &p, vlen, &upb_DecodeFast_MessageData, &ctx)) {
          if (UPB_UNLIKELY(p == NULL)) {
            return UPB_DECODEFAST_ERROR(d, kUpb_DecodeStatus_Malformed, next);
          }
        } else {
          ptrdiff_t delta = upb_EpsCopyInputStream_PushLimit(EPS(d), p, vlen);
          p = upb_DecodeFast_MessageData(EPS(d), p, vlen, &ctx);
          if (UPB_UNLIKELY(p == NULL)) {
            return UPB_DECODEFAST_ERROR(d, kUpb_DecodeStatus_Malformed, next);
          }
          upb_EpsCopyInputStream_PopLimit(EPS(d), p, delta);
        }
        tabval = upb_value_ptr(sub_msg);
      } else {
        upb_StringView str_val;
        if (UPB_UNLIKELY(!_upb_Decoder_ReadString(
                d, &p, vlen, &str_val,
                upb_DecodeFastData_ValValidateUtf8(data)))) {
          return UPB_DECODEFAST_ERROR(d, kUpb_DecodeStatus_OutOfMemory, next);
        }
        if (UPB_UNLIKELY(!_upb_map_tovalue(&str_val, UPB_MAPTYPE_STRING,
                                           &tabval, &d->arena))) {
          return UPB_DECODEFAST_ERROR(d, kUpb_DecodeStatus_OutOfMemory, next);
        }
      }
      break;
    }
    default:
      UPB_UNREACHABLE();
  }

  bool replaced;
  bool ok =
      is_str_map
          ? upb_strtable_insert_or_replace(&map_ctx->map->t.strtable,
                                           strkey.data, strkey.size, tabval,
                                           &replaced, &d->arena)
          : upb_inttable_insert_or_replace(&map_ctx->map->t.inttable, intkey,
                                           tabval, &replaced, &d->arena);
  if (UPB_UNLIKELY(!ok)) {
    return UPB_DECODEFAST_ERROR(d, kUpb_DecodeStatus_OutOfMemory, next);
  }

  *ptr = p;
  return true;
}

// Estimates the number of contiguous map entries with single-byte lengths by
// skip-scanning ahead. If an entry has a length > 127 (multibyte varint),
// scanning stops early to avoid cache pollution on large entries.
UPB_FORCEINLINE
size_t upb_DecodeFast_EstimateMapSize(upb_Decoder* d, const char* ptr,
                                      uint16_t expected,
                                      upb_DecodeFast_TagSize tagsize) {
  size_t count = 0;
  const char* limit_ptr = d->input.limit_ptr;
  int saved_guaranteed =
      UPB_PRIVATE(upb_EpsCopyInputStream_GetBoundsCheckedBytes)(EPS(d));

  if (tagsize == kUpb_DecodeFast_Tag1Byte) {
    if (limit_ptr - ptr >= 2) {
      UPB_PRIVATE(upb_EpsCopyInputStream_BoundsCheckedToEnd)(EPS(d), ptr,
                                                             limit_ptr);
      while (limit_ptr - ptr >= 2) {
        UPB_PRIVATE(upb_EpsCopyInputStream_ConsumeBytes)(EPS(d), 2);
        uint16_t val;
        memcpy(&val, ptr, 2);
        if ((val & 0x80ff) != (uint8_t)expected) break;
        uint32_t len = val >> 8;
        uint32_t skip = 2 + len;
        if (limit_ptr - ptr < skip) break;
        UPB_PRIVATE(upb_EpsCopyInputStream_ConsumeBytes)(EPS(d), len);
        count++;
        ptr += skip;
      }
    }
  } else {
    UPB_ASSERT(tagsize == kUpb_DecodeFast_Tag2Byte);
    if (limit_ptr - ptr >= 3) {
      UPB_PRIVATE(upb_EpsCopyInputStream_BoundsCheckedToEnd)(EPS(d), ptr,
                                                             limit_ptr);
      while (limit_ptr - ptr >= 3) {
        UPB_PRIVATE(upb_EpsCopyInputStream_ConsumeBytes)(EPS(d), 4);
        uint32_t val = 0;
        memcpy(&val, ptr, 4);
        if ((val & 0x80ffff) != expected) break;
        uint32_t len = (val >> 16) & 0x7f;
        uint32_t skip = 3 + len;
        if (skip < 4 || limit_ptr - ptr < skip) break;
        UPB_PRIVATE(upb_EpsCopyInputStream_ConsumeBytes)(EPS(d), skip - 4);
        count++;
        ptr += skip;
      }
    }
  }

  UPB_PRIVATE(upb_EpsCopyInputStream_BoundsChecked)(EPS(d), saved_guaranteed);
  return count;
}

UPB_FORCEINLINE
void upb_DecodeFast_Map(upb_Decoder* d, const char** ptr, upb_Message* msg,
                        const upb_MiniTable* table, uint64_t* hasbits,
                        uint64_t data, uint64_t data2,
                        upb_DecodeFast_TagSize tagsize, bool is_str_map,
                        upb_DecodeFastNext* next) {
  UPB_ASSERT(is_str_map == upb_DecodeFastData_IsStrMap(data));
  UPB_ASSERT((tagsize == kUpb_DecodeFast_Tag2Byte) ==
             upb_DecodeFastData_IsTag2(data));
  uint16_t expected = upb_DecodeFastData_GetExpectedTag(data);
  uint16_t actual = upb_DecodeFastData2_GetOriginalTag(data2);
  if (UPB_UNLIKELY(!upb_DecodeFast_TagMatches(expected, actual, tagsize))) {
    UPB_DECODEFAST_EXIT(kUpb_DecodeFastNext_FallbackMismatchedSlot, next);
    return;
  }

  upb_DecodeFastMap map_ctx;
  if (UPB_UNLIKELY(!upb_DecodeFast_GetMap(d, msg, table, data, is_str_map,
                                          hasbits, &map_ctx, next))) {
    return;
  }

  // The sizes packed into `data` by upb_DecodeFast_TryFillMapEntry() describe
  // the same map that _upb_Decoder_CreateMap() sized from kSizeInMap[], so they
  // must agree. Use UPB_ASSERT rather than UPB_ASSUME here so the compiler does
  // not replace register bitfield extractions from `data` with memory loads
  // from `map_ctx.map`.
  UPB_ASSERT(map_ctx.map->UPB_PRIVATE(is_strtable) == is_str_map);
  UPB_ASSERT(map_ctx.map->key_size == upb_DecodeFastData_GetKeySize(data));
  UPB_ASSERT(map_ctx.map->val_size == upb_DecodeFastData_GetValSize(data));
  UPB_ASSERT(is_str_map ==
             (upb_DecodeFastData_GetKeySize(data) == UPB_MAPTYPE_STRING));

  int depth_cost = map_ctx.sub_table ? 2 : 1;
  if (d->depth < depth_cost) {
    // This isn't yet an error, because if the map's entries are malformed
    // (value field not actually present or wrong wire type) then it won't
    // actually recurse and thus won't exceed the max depth. But we don't want
    // to repeatedly check the depth limit on every entry, so let the fallback
    // decoder handle this.
    UPB_DECODEFAST_EXIT(kUpb_DecodeFastNext_FallbackToMiniTable, next);
    return;
  }
  d->depth -= depth_cost;

  size_t estimated_entries =
      upb_DecodeFast_EstimateMapSize(d, *ptr, expected, tagsize);
  if (estimated_entries > 0) {
    // The map's table discriminator is derived from the same entry MiniTable
    // that selected this parser, so it always agrees with our compile-time
    // is_str_map. Place the assume immediately before _upb_Map_Size() with no
    // intervening memory writes so EarlyCSE folds the strtable/inttable branch.
    UPB_ASSUME(map_ctx.map->UPB_PRIVATE(is_strtable) == is_str_map);
    if (UPB_UNLIKELY(!_upb_Map_Reserve(
            map_ctx.map, _upb_Map_Size(map_ctx.map) + estimated_entries,
            &d->arena))) {
      d->depth += depth_cost;
      UPB_DECODEFAST_ERROR(d, kUpb_DecodeStatus_OutOfMemory, next);
      return;
    }
  }

  const char* p = *ptr + upb_DecodeFast_TagSizeBytes(tagsize);
  while (1) {
    const char* entry_start = p - upb_DecodeFast_TagSizeBytes(tagsize);
    if (UPB_UNLIKELY(!upb_DecodeFast_ParseMapEntry(d, &p, &map_ctx, data,
                                                   is_str_map, next))) {
      *ptr = entry_start;
      break;
    }
    *ptr = p;
    _upb_Decoder_Trace(d, 'F');
    if (!upb_DecodeFast_TryMatchTag(d, p, expected, next, tagsize)) {
      break;
    }
    p += upb_DecodeFast_TagSizeBytes(tagsize);
  }

  d->depth += depth_cost;
}

#define F(key_type, is_str, tagsize)                               \
  UPB_NOINLINE UPB_PRESERVE_NONE upb_FastDecoder_Return            \
  upb_DecodeFast_##key_type##Map_##tagsize(UPB_PARSE_PARAMS) {     \
    upb_DecodeFastNext next = kUpb_DecodeFastNext_Dispatch;        \
    upb_DecodeFast_Map(d, &ptr, msg, table, &hasbits, data, data2, \
                       kUpb_DecodeFast_##tagsize, is_str, &next);  \
    UPB_DECODEFAST_NEXT(next);                                     \
  }

F(Int, false, Tag1Byte)
F(Int, false, Tag2Byte)
F(Str, true, Tag1Byte)
F(Str, true, Tag2Byte)

#undef F