[api] Count packed varints a machine word at a time (#20122)

esphome/components/api/proto.h

@@ -56,22 +56,25 @@ inline constexpr int64_t decode_zigzag64(uint64_t value) {
   return (value & 1) ? static_cast<int64_t>(~(value >> 1)) : static_cast<int64_t>(value >> 1);
 }
 
-/// Count number of varints in a packed buffer
-inline uint16_t count_packed_varints(const uint8_t *data, size_t len) {
-  uint16_t count = 0;
-  while (len > 0) {
-    // Skip varint bytes until we find one without continuation bit
-    while (len > 0 && (*data & 0x80)) {
-      data++;
-      len--;
-    }
-    if (len > 0) {
-      data++;
-      len--;
-      count++;
+/// Count varints in a packed buffer: len minus bytes with the continuation bit, summed a word at a time.
+/// Word is a template parameter so tests can cover the 32-bit path on a 64-bit host.
+template<typename Word = size_t> inline uint16_t count_packed_varints(const uint8_t *data, size_t len) {
+  constexpr size_t word_size = sizeof(Word);
+  constexpr Word lane_ones = ~Word{0} / 0xFF;  // 0x01..01
+  const uint8_t *end = data + len;
+  size_t continuations = 0;
+  while (data != end) {
+    // Unaligned word loads fault on Xtensa
+    if ((reinterpret_cast<uintptr_t>(data) & (word_size - 1)) == 0 && static_cast<size_t>(end - data) >= word_size) {
+      Word word;
+      memcpy(&word, __builtin_assume_aligned(data, word_size), word_size);
+      continuations += (((word >> 7) & lane_ones) * lane_ones) >> (word_size * 8 - 8);
+      data += word_size;
+    } else {
+      continuations += *data++ >> 7;
     }
   }
-  return count;
+  return static_cast<uint16_t>(len - continuations);
 }
 
 /// Encode a varint directly into a pre-allocated buffer.

tests/components/api/test_count_packed_varints.cpp

@@ -0,0 +1,84 @@
+#include <gtest/gtest.h>
+
+#include <cstdint>
+#include <random>
+#include <vector>
+
+#include "esphome/components/api/proto.h"
+
+namespace esphome::api::testing {
+
+// The original byte at a time implementation.
+static uint16_t reference_count(const uint8_t *data, size_t len) {
+  uint16_t count = 0;
+  while (len > 0) {
+    while (len > 0 && (*data & 0x80)) {
+      data++;
+      len--;
+    }
+    if (len > 0) {
+      data++;
+      len--;
+      count++;
+    }
+  }
+  return count;
+}
+
+TEST(CountPackedVarints, EmptyBuffer) {
+  const uint8_t data[1] = {0x00};
+  EXPECT_EQ(count_packed_varints(data, 0), 0);
+}
+
+TEST(CountPackedVarints, SingleByteVarints) {
+  const uint8_t data[] = {0x00, 0x01, 0x7F};
+  EXPECT_EQ(count_packed_varints(data, sizeof(data)), 3);
+}
+
+TEST(CountPackedVarints, MultiByteVarints) {
+  // 3 varints: 2 bytes, 3 bytes, 1 byte
+  const uint8_t data[] = {0x80, 0x01, 0x80, 0x80, 0x01, 0x05};
+  EXPECT_EQ(count_packed_varints(data, sizeof(data)), 3);
+}
+
+TEST(CountPackedVarints, TruncatedTrailingVarintIsNotCounted) {
+  const uint8_t data[] = {0x05, 0x80, 0x80, 0x80};
+  EXPECT_EQ(count_packed_varints(data, sizeof(data)), 1);
+}
+
+TEST(CountPackedVarints, AllContinuationBytes) {
+  std::vector<uint8_t> data(5000, 0x80);
+  EXPECT_EQ(count_packed_varints(data.data(), data.size()), 0);
+}
+
+TEST(CountPackedVarints, EveryStartOffsetAndLength) {
+  // Cover every alignment and length around word boundaries.
+  std::mt19937 rng(42);  // NOLINT(cert-msc32-c,cert-msc51-cpp,bugprone-random-generator-seed) reproducible
+  std::vector<uint8_t> buf(300);
+  for (auto &byte : buf)
+    byte = static_cast<uint8_t>(rng() & 0xFF);
+  for (size_t offset = 0; offset < 16; offset++) {
+    for (size_t len = 0; len + offset <= buf.size(); len++) {
+      const uint16_t expected = reference_count(buf.data() + offset, len);
+      EXPECT_EQ(count_packed_varints<uint32_t>(buf.data() + offset, len), expected)
+          << "offset=" << offset << " len=" << len;
+      EXPECT_EQ(count_packed_varints<uint64_t>(buf.data() + offset, len), expected)
+          << "offset=" << offset << " len=" << len;
+    }
+  }
+}
+
+TEST(CountPackedVarints, LongBuffer) {
+  std::mt19937 rng(7);  // NOLINT(cert-msc32-c,cert-msc51-cpp,bugprone-random-generator-seed) reproducible
+  std::vector<uint8_t> buf(5000);
+  for (auto &byte : buf)
+    byte = static_cast<uint8_t>((rng() % 4 == 0) ? (0x80 | (rng() & 0x7F)) : (rng() & 0x7F));
+  for (size_t offset = 0; offset < 8; offset++) {
+    const size_t len = buf.size() - offset;
+    const uint16_t expected = reference_count(buf.data() + offset, len);
+    EXPECT_EQ(count_packed_varints<uint32_t>(buf.data() + offset, len), expected) << "offset=" << offset;
+    EXPECT_EQ(count_packed_varints<uint64_t>(buf.data() + offset, len), expected) << "offset=" << offset;
+  }
+}
+
+}  // namespace esphome::api::testing