Optimize ReadPackedVarint

Count the number of varints before allocating in ReadPackedVarint.
This is pretty cheap, so we save more on avoiding realloctions.

BM_ParsePacked/1             524.2Mi ±   2%    532.2Mi ±   2%   +1.53% (p=0.041 n=6)
BM_ParsePacked/8             758.2Mi ±   2%   1163.7Mi ±   1%  +53.48% (p=0.002 n=6)
BM_ParsePacked/64            1.582Gi ±   2%    2.656Gi ±   2%  +67.94% (p=0.002 n=6)
BM_ParsePacked/512           2.014Gi ±   3%    3.377Gi ±   4%  +67.66% (p=0.002 n=6)
BM_ParsePacked/4k            2.009Gi ±   2%    3.358Gi ±   2%  +67.10% (p=0.002 n=6)
BM_ParsePacked/32k           1.882Gi ±   2%    3.151Gi ±   3%  +67.46% (p=0.002 n=6)
BM_ParsePacked/256k          1.338Gi ±  15%    2.408Gi ±   8%  +80.03% (p=0.002 n=6)

Memory usage is also 25% lower

BM_ParsePacked/1               192.0 ± 0%           192.0 ± 0%        ~ (p=1.000 n=6) ¹
BM_ParsePacked/8              2.048k ± 0%          1.568k ± 0%  -23.44% (p=0.002 n=6)
BM_ParsePacked/64             16.19k ± 0%          12.05k ± 0%  -25.59% (p=0.002 n=6)
BM_ParsePacked/512           130.69k ± 0%          98.69k ± 0%  -24.49% (p=0.002 n=6)
BM_ParsePacked/4k            1048.6k ± 0%          795.9k ± 0%  -24.10% (p=0.002 n=6)
BM_ParsePacked/32k            8.045M ± 0%          5.677M ± 0%  -29.43% (p=0.002 n=6)
BM_ParsePacked/256k           64.01M ± 0%          44.67M ± 0%  -30.21% (p=0.002 n=6)

PiperOrigin-RevId: 813430457

Author: TocarIP

src/google/protobuf/generated_message_tctable_lite.cc

@@ -1225,20 +1225,23 @@ PROTOBUF_ALWAYS_INLINE const char* TcParser::PackedVarint(
   // Since ctx->ReadPackedVarint does not use TailCall or Return, sync any
   // pending hasbits now:
   SyncHasbits(msg, hasbits, table);
-  auto* field = &RefAt<RepeatedField<FieldType>>(msg, data.offset());
-  return ctx->ReadPackedVarint(ptr, [field](uint64_t varint) {
-    FieldType val;
-    if (zigzag) {
-      if (sizeof(FieldType) == 8) {
-        val = WireFormatLite::ZigZagDecode64(varint);
-      } else {
-        val = WireFormatLite::ZigZagDecode32(varint);
-      }
-    } else {
-      val = varint;
-    }
-    field->Add(val);
-  });
+  auto& field = RefAt<RepeatedField<FieldType>>(msg, data.offset());
+  return ctx->ReadPackedVarintWithField(
+      ptr,
+      [field](uint64_t varint) {
+        FieldType val;
+        if (zigzag) {
+          if (sizeof(FieldType) == 8) {
+            val = WireFormatLite::ZigZagDecode64(varint);
+          } else {
+            val = WireFormatLite::ZigZagDecode32(varint);
+          }
+        } else {
+          val = varint;
+        }
+        return val;
+      },
+      field);
 }
 
 PROTOBUF_NOINLINE const char* TcParser::FastV8P1(PROTOBUF_TC_PARAM_DECL) {

src/google/protobuf/parse_context.h

@@ -260,6 +260,11 @@ class PROTOBUF_EXPORT EpsCopyInputStream {
   template <typename Add, typename SizeCb>
   [[nodiscard]] const char* ReadPackedVarint(const char* ptr, Add add,
                                              SizeCb size_callback);
+  // Same as above, but pass the field directly , so we can preallocate.
+  template <typename Convert, typename T>
+  [[nodiscard]] const char* ReadPackedVarintWithField(const char* ptr,
+                                                      Convert conv,
+                                                      RepeatedField<T>& out);
 
   uint32_t LastTag() const { return last_tag_minus_1_ + 1; }
   bool ConsumeEndGroup(uint32_t start_tag) {
@@ -1362,6 +1367,83 @@ const char* ReadPackedVarintArray(const char* ptr, const char* end, Add add) {
   return ptr;
 }
 
+template <typename Convert, typename T>
+const char* ReadPackedVarintArrayWithField(const char* ptr, const char* end,
+                                           Convert conv,
+                                           RepeatedField<T>& out) {
+  // If we have enough bytes, we will spend more cpu cycles growing repeated
+  // field, than parsing, so count the number of ints first and preallocate.
+  // Assume that varint are valid and just count the number of bytes with
+  // continuation bit not set. In a valid varint there is only 1 such byte.
+  if ((end - ptr) >= 16 && (out.Capacity() - out.size() < end - ptr)) {
+    int old_size = out.size();
+    int count = out.Capacity() - out.size();
+    // We are not guaranteed to have enough space for worst possible case,
+    // do an actual count and reserve.
+    if (count < end - ptr) {
+      count = std::count_if(ptr, end, [](char c) { return (c & 0x80) == 0; });
+      // We can overread, so if the last byte has a continuation bit set,
+      // we need to account for that.
+      if (end[-1] & 0x80) count++;
+      out.Reserve(old_size + count);
+    }
+    T* x = out.AddNAlreadyReserved(count);
+    ptr = ReadPackedVarintArray(ptr, end, [&](uint64_t varint) {
+      *x = conv(varint);
+      x++;
+    });
+    int new_size = x - out.data();
+    ABSL_DCHECK_LE(new_size, old_size + count);
+    // We may have overreserved if there was enough capacitiy.
+    // Or encountered malformed data, so set the actaul size to
+    // avoid exposing uninitialized memory.
+    out.Truncate(new_size);
+    return ptr;
+  } else {
+    return ReadPackedVarintArray(
+        ptr, end, [&](uint64_t varint) { out.Add(conv(varint)); });
+  }
+}
+
+template <typename Convert, typename T>
+const char* EpsCopyInputStream::ReadPackedVarintWithField(
+    const char* ptr, Convert conv, RepeatedField<T>& out) {
+  int size = ReadSize(&ptr);
+
+  GOOGLE_PROTOBUF_PARSER_ASSERT(ptr);
+  int chunk_size = static_cast<int>(buffer_end_ - ptr);
+  while (size > chunk_size) {
+    ptr = ReadPackedVarintArrayWithField(ptr, buffer_end_, conv, out);
+    if (ptr == nullptr) return nullptr;
+    int overrun = static_cast<int>(ptr - buffer_end_);
+    ABSL_DCHECK(overrun >= 0 && overrun <= kSlopBytes);
+    if (size - chunk_size <= kSlopBytes) {
+      // The current buffer contains all the information needed, we don't need
+      // to flip buffers. However we must parse from a buffer with enough space
+      // so we are not prone to a buffer overflow.
+      char buf[kSlopBytes + 10] = {};
+      std::memcpy(buf, buffer_end_, kSlopBytes);
+      ABSL_CHECK_LE(size - chunk_size, kSlopBytes);
+      auto end = buf + (size - chunk_size);
+      auto result = ReadPackedVarintArray(
+          buf + overrun, end, [&](uint64_t varint) { out.Add(conv(varint)); });
+      if (result == nullptr || result != end) return nullptr;
+      return buffer_end_ + (result - buf);
+    }
+    size -= overrun + chunk_size;
+    ABSL_DCHECK_GT(size, 0);
+    // We must flip buffers
+    if (limit_ <= kSlopBytes) return nullptr;
+    ptr = Next();
+    if (ptr == nullptr) return nullptr;
+    ptr += overrun;
+    chunk_size = static_cast<int>(buffer_end_ - ptr);
+  }
+  auto end = ptr + size;
+  ptr = ReadPackedVarintArrayWithField(ptr, end, conv, out);
+  return end == ptr ? ptr : nullptr;
+}
+
 template <typename Add, typename SizeCb>
 const char* EpsCopyInputStream::ReadPackedVarint(const char* ptr, Add add,
                                                  SizeCb size_callback) {