Optimize counting varints when parsing packed repeated fields.

Count them in 64-bit blocks, using `& 0x8080808080808080` and
`absl::popcount()`.

PiperOrigin-RevId: 847794231

Author: QrczakMK

src/google/protobuf/parse_context.cc

@@ -19,6 +19,7 @@
 #include "absl/base/optimization.h"
 #include "absl/base/prefetch.h"
 #include "absl/log/absl_check.h"
+#include "absl/numeric/bits.h"
 #include "absl/strings/cord.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
@@ -767,6 +768,27 @@ template std::pair<const char*, bool> EpsCopyInputStream::DoneFallback<false>(
 template std::pair<const char*, bool> EpsCopyInputStream::DoneFallback<true>(
     int, int);
 
+int CountVarintsAssumingLargeArray(const char* ptr, const char* end) {
+  // The number of varints is the number of bytes with the highest bit clear.
+  // This is easier to compute as the total number of bytes, minus the number
+  // of bytes with the highest bit set.
+  int num_varints = end - ptr;
+  ABSL_DCHECK_GE(num_varints, int{sizeof(uint64_t)});
+
+  // Count in whole blocks, except for the last one.
+  const char* const limit = end - sizeof(uint64_t);
+  while (ptr < limit) {
+    num_varints -=
+        absl::popcount(EndianHelper<8>::Load(ptr) & 0x8080808080808080);
+    ptr += sizeof(uint64_t);
+  }
+
+  // Count in the last, possibly incomplete block.
+  return num_varints -
+         absl::popcount(EndianHelper<8>::Load(limit) &
+                        (0x8080808080808080 << ((ptr - limit) * 8)));
+}
+
 }  // namespace internal
 }  // namespace protobuf
 }  // namespace google

src/google/protobuf/parse_context.h

@@ -66,6 +66,11 @@ inline void WriteVarint(uint32_t num, uint64_t val, UnknownFieldSet* unknown);
 inline void WriteLengthDelimited(uint32_t num, absl::string_view val,
                                  UnknownFieldSet* unknown);
 
+// Counts the number of varints in the array, assuming that end - ptr >= 8.
+// If varints are valid only up to some point, then returns at least the number
+// of valid varints.
+int CountVarintsAssumingLargeArray(const char* ptr, const char* end);
+
 
 // The basic abstraction the parser is designed for is a slight modification
 // of the ZeroCopyInputStream (ZCIS) abstraction. A ZCIS presents a serialized
@@ -1389,13 +1394,13 @@ const char* EpsCopyInputStream::ReadPackedVarintArrayWithField(
   // 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)) {
+  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; });
+      count = CountVarintsAssumingLargeArray(ptr, end);
       // We can overread, so if the last byte has a continuation bit set,
       // we need to account for that.
       if (end[-1] & 0x80) count++;
@@ -1408,8 +1413,8 @@ const char* EpsCopyInputStream::ReadPackedVarintArrayWithField(
     });
     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
+    // We may have overreserved if there was enough capacity.
+    // Or encountered malformed data, so set the actual size to
     // avoid exposing uninitialized memory.
     out.Truncate(new_size);
     return ptr;