src/google/protobuf/generated_enum_util_test.cc

// Protocol Buffers - Google's data interchange format
// Copyright 2023 Google Inc.  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 "google/protobuf/generated_enum_util.h"

#include <algorithm>
#include <array>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <ostream>
#include <string>
#include <variant>
#include <vector>

#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/container/btree_set.h"
#include "absl/flags/marshalling.h"
#include "absl/strings/str_format.h"
#include "absl/types/span.h"
#include "google/protobuf/descriptor.h"
#include "google/protobuf/unittest.pb.h"
#include "google/protobuf/unittest_lite.pb.h"


// Must be included last.
#include "google/protobuf/port_def.inc"

using testing::_;
using testing::ElementsAre;
using testing::Gt;
using testing::HasSubstr;
using testing::IsEmpty;
using testing::SizeIs;
using testing::VariantWith;

namespace google {
namespace protobuf {
namespace internal {
namespace {

TEST(GenerateEnumDataTest, DebugChecks) {
#if GTEST_HAS_DEATH_TEST
  // Not unique
  EXPECT_DEBUG_DEATH((void)GenerateEnumData({1, 1}), "sorted_and_unique");
  // Not sorted
  EXPECT_DEBUG_DEATH((void)GenerateEnumData({2, 1}), "sorted_and_unique");
#endif
}

uint32_t Make32(uint16_t a, uint16_t b) { return a | (b << 16); }
std::array<uint16_t, 2> Unmake32(uint32_t v) {
  return {static_cast<uint16_t>(v), static_cast<uint16_t>(v >> 16)};
}

struct Header {
  int16_t sequence_start;
  uint16_t sequence_length;
  uint16_t bitmap_length;
  uint16_t ordered_length;

  std::string ToString() const {
    return absl::StrFormat("(%d,%d,%d,%d)", sequence_start, sequence_length,
                           bitmap_length, ordered_length);
  }
  friend std::ostream& operator<<(std::ostream& os, Header header) {
    return os << header.ToString();
  }
};

MATCHER_P4(HeaderHas, sequence_start, sequence_length, bitmap_length,
           ordered_length, "") {
  return testing::ExplainMatchResult(sequence_start, arg.sequence_start,
                                     result_listener) &&
         testing::ExplainMatchResult(sequence_length, arg.sequence_length,
                                     result_listener) &&
         testing::ExplainMatchResult(bitmap_length, arg.bitmap_length,
                                     result_listener) &&
         testing::ExplainMatchResult(ordered_length, arg.ordered_length,
                                     result_listener);
}

Header ExtractHeader(absl::Span<const uint32_t> data) {
  return {
      static_cast<int16_t>(Unmake32(data[0])[0]),
      Unmake32(data[0])[1],
      Unmake32(data[1])[0],
      Unmake32(data[1])[1],
  };
}

TEST(GenerateEnumDataTest, BitmapSpaceOptimizationWorks) {
  std::vector<int32_t> values = {0};

  auto encoded = GenerateEnumData(values);
  EXPECT_THAT(ExtractHeader(encoded), HeaderHas(0, 1, 0, 0));
  EXPECT_THAT(encoded, SizeIs(2));

  // Adding one large value puts it on the fallback.
  values.push_back(100);
  encoded = GenerateEnumData(values);
  EXPECT_THAT(ExtractHeader(encoded), HeaderHas(0, 1, 0, 1));
  EXPECT_THAT(encoded, SizeIs(3));

  // Adding a second one still prefers the fallback.
  values.push_back(101);
  encoded = GenerateEnumData(values);
  EXPECT_THAT(ExtractHeader(encoded), HeaderHas(0, 1, 0, 2));
  EXPECT_THAT(encoded, SizeIs(4));

  // Adding two more now makes bitmap more efficient, so they are collapsed
  // to it. Because we can fit the bitmap in 128 bits, which is the same as the
  // ints.
  values.push_back(102);
  values.push_back(103);
  encoded = GenerateEnumData(values);
  EXPECT_THAT(ExtractHeader(encoded), HeaderHas(0, 1, 128, 0));
  EXPECT_THAT(encoded, SizeIs(6));

  // Add one value that falls into the existing bitmap, nothing changes.
  values.push_back(104);
  encoded = GenerateEnumData(values);
  EXPECT_THAT(ExtractHeader(encoded), HeaderHas(0, 1, 128, 0));
  EXPECT_THAT(encoded, SizeIs(6));

  // Add one value that is in the next 32 bits. It should grow the bitmap.
  values.push_back(130);
  encoded = GenerateEnumData(values);
  EXPECT_THAT(ExtractHeader(encoded), HeaderHas(0, 1, 160, 0));
  EXPECT_THAT(encoded, SizeIs(7));

  // Add one value far away, it should go into fallback.
  values.push_back(200);
  encoded = GenerateEnumData(values);
  EXPECT_THAT(ExtractHeader(encoded), HeaderHas(0, 1, 160, 1));
  EXPECT_THAT(encoded, SizeIs(8));

  // Another in the next 32-bit block will still make them fallback.
  values.push_back(230);
  encoded = GenerateEnumData(values);
  EXPECT_THAT(ExtractHeader(encoded), HeaderHas(0, 1, 160, 2));
  EXPECT_THAT(encoded, SizeIs(9));

  // One more in that same block should collapse them to bitmap.
  values.push_back(231);
  encoded = GenerateEnumData(values);
  EXPECT_THAT(ExtractHeader(encoded), HeaderHas(0, 1, 256, 0));
  EXPECT_THAT(encoded, SizeIs(10));
}

template <typename T>
std::variant<T, std::string> TryParseFlag(absl::string_view str) {
  T value;
  std::string error;
  if (!absl::ParseFlag(str, &value, &error)) {
    return error;
  }
  return value;
}

TEST(ProtoEnumTest, HasAbseilFlagSupport) {
  using T = proto2_unittest::TestAllTypes;
  using E = typename T::NestedEnum;

  // Exact match.
  EXPECT_THAT(TryParseFlag<E>("FOO"), VariantWith<E>(T::FOO));
  // Lowercase match.
  EXPECT_THAT(TryParseFlag<E>("baz"), VariantWith<E>(T::BAZ));
  // Numeric value
  EXPECT_THAT(TryParseFlag<E>("2"), VariantWith<E>(T::BAR));

  // Error
  EXPECT_THAT(TryParseFlag<E>("xxx"),
              VariantWith<std::string>(
                  HasSubstr("Invalid value 'xxx' for enum "
                            "'proto2_unittest.TestAllTypes.NestedEnum'. "
                            "Supported values are: FOO, BAR, BAZ, NEG.")));

  EXPECT_EQ("FOO", absl::UnparseFlag(T::FOO));
  EXPECT_EQ("100", absl::UnparseFlag(static_cast<E>(100)));
}

TEST(ProtoEnumTest, HasAbseilFlagSupportWithLiteEnums) {
  using T = proto2_unittest::TestAllTypesLite;
  using E = typename T::NestedEnum;

  // Exact match.
  EXPECT_THAT(TryParseFlag<E>("FOO"), VariantWith<E>(T::FOO));
  // Lowercase match.
  EXPECT_THAT(TryParseFlag<E>("baz"), VariantWith<E>(T::BAZ));
  // Numeric value
  EXPECT_THAT(TryParseFlag<E>("2"), VariantWith<E>(T::BAR));

  // We don't check errors because we don't generate them for LITE.
  EXPECT_THAT(TryParseFlag<E>("xxx"), VariantWith<std::string>(_));

  EXPECT_EQ("FOO", absl::UnparseFlag(T::FOO));
  EXPECT_EQ("100", absl::UnparseFlag(static_cast<E>(100)));
}

TEST(ProtoEnumTest, HasAbseilFlagSupportForVectorEnum) {
  using T = proto2_unittest::TestAllTypes;
  using E = typename T::NestedEnum;
  using V = std::vector<E>;

  // Empty
  EXPECT_THAT(TryParseFlag<V>(""), VariantWith<V>(V{}));
  EXPECT_EQ("", absl::UnparseFlag(V{}));
  // Single element
  EXPECT_THAT(TryParseFlag<V>("FOO"), VariantWith<V>(V{T::FOO}));
  EXPECT_EQ("FOO", absl::UnparseFlag(V{T::FOO}));
  // Many elements
  EXPECT_THAT(TryParseFlag<V>("BAZ,2,foo"),
              VariantWith<V>(V{T::BAZ, T::BAR, T::FOO}));
  EXPECT_EQ("BAZ,BAR,FOO", absl::UnparseFlag(V{T::BAZ, T::BAR, T::FOO}));

  // Error
  EXPECT_THAT(TryParseFlag<V>("FOO,xxx"),
              VariantWith<std::string>(
                  HasSubstr("Invalid value 'xxx' for enum "
                            "'proto2_unittest.TestAllTypes.NestedEnum'. "
                            "Supported values are: FOO, BAR, BAZ, NEG.")));
}

TEST(ProtoEnumTest, HasAbseilFlagSupportForVectorEnumWithLiteEnums) {
  using T = proto2_unittest::TestAllTypesLite;
  using E = typename T::NestedEnum;
  using V = std::vector<E>;

  // Empty
  EXPECT_THAT(TryParseFlag<V>(""), VariantWith<V>(V{}));
  EXPECT_EQ("", absl::UnparseFlag(V{}));
  // Single element
  EXPECT_THAT(TryParseFlag<V>("FOO"), VariantWith<V>(V{T::FOO}));
  EXPECT_EQ("FOO", absl::UnparseFlag(V{T::FOO}));
  // Many elements
  EXPECT_THAT(TryParseFlag<V>("BAZ,2,foo"),
              VariantWith<V>(V{T::BAZ, T::BAR, T::FOO}));
  EXPECT_EQ("BAZ,BAR,FOO", absl::UnparseFlag(V{T::BAZ, T::BAR, T::FOO}));

  // Error
  EXPECT_THAT(TryParseFlag<V>("FOO,xxx"), VariantWith<std::string>(_));
}

void GatherValidValues(absl::Span<const uint32_t> data, int32_t min,
                       int32_t max, absl::btree_set<int32_t>& out) {
  if (min >= max) return;
  for (int32_t i = min;; ++i) {
    if (ValidateEnum(i, data.begin())) out.insert(i);
    // We check the top limit before ++i to avoid overflows
    if (i == max) break;
  }
}

std::vector<int32_t> GetValidValues(absl::Span<const uint32_t> data,
                                    int32_t min, int32_t max) {
  // Btree to keep them sorted. Makes testing easier.
  absl::btree_set<int32_t> s;
  GatherValidValues(data, min, max, s);
  return std::vector<int32_t>(s.begin(), s.end());
}

TEST(ValidateEnumTest, SequentialRangeTest) {
  EXPECT_THAT(GetValidValues({0, 0}, -100, 100), ElementsAre());
  EXPECT_THAT(GetValidValues(
                  {// sequence start=3, length=3
                   Make32(5, 3),
                   // no bitmap, no fallback
                   Make32(0, 0)},
                  -100, 100),
              ElementsAre(5, 6, 7));
  EXPECT_THAT(GetValidValues(
                  {// sequence start=-2, length=10
                   Make32(-2, 10),
                   // no bitmap, no fallback
                   Make32(0, 0)},
                  -100, 100),
              ElementsAre(-2, -1, 0, 1, 2, 3, 4, 5, 6, 7));
}

TEST(ValidateEnumTest, BitmapRangeTest) {
  EXPECT_THAT(GetValidValues(
                  {// no sequence
                   Make32(0, 0),
                   // bitmap of 32 bits, no fallback
                   Make32(32, 0),
                   // bitmap
                   0b10011010101},
                  -100, 100),
              ElementsAre(0, 2, 4, 6, 7, 10));
  EXPECT_THAT(GetValidValues(
                  {// no sequence
                   Make32(0, 0),
                   // bitmap of 64 bits, no fallback
                   Make32(64, 0),
                   // bitmap
                   (1 << 4) | (1 << 21), 1 << 6},
                  -100, 100),
              ElementsAre(4, 21, 32 + 6));
}

TEST(ValidateEnumTest, GenerateEnumDataSequential) {
  EXPECT_THAT(GenerateEnumData({0, 1, 2, 3}), ElementsAre(
                                                  // sequence start=0, length=4
                                                  Make32(0, 4),
                                                  // no bitmap, no fallback.
                                                  Make32(0, 0)));
  EXPECT_THAT(GenerateEnumData({-2, -1, 0, 1, 2, 3}),
              ElementsAre(
                  // sequence start=-2, length=6
                  Make32(-2, 6),
                  // no bitmap, no fallback.
                  Make32(0, 0)));
}

void TestRoundTrip(absl::Span<const int32_t> values, int line) {
  auto encoded = GenerateEnumData(values);

  absl::btree_set<int32_t> s;

  // We test that all elements in `values` exist in the encoded data, and also
  // test a range of other values to verify that they do not exist in the
  // encoded data.

  // We keep track of the max seen to avoid testing the same values many times.
  int64_t max_seen = std::numeric_limits<int64_t>::min();
  const auto gather_valid_values_around = [&](int32_t v) {
    int32_t min = std::max({
        static_cast<int64_t>(v) - 100,
        static_cast<int64_t>(std::numeric_limits<int32_t>::min()),
        max_seen,
    });
    int32_t max =
        std::min(static_cast<int64_t>(v) + 100,
                 static_cast<int64_t>(std::numeric_limits<int32_t>::max()));
    max_seen = std::max(max_seen, int64_t{max});
    GatherValidValues(encoded, min, max, s);
  };

  // We look at a few values around the expected ones.
  // We could in theory test the whole int32_t domain, but that takes too long
  // to run.
  for (int32_t v : values) {
    gather_valid_values_around(v);
  }
  // Also gather some around 0, just to add more coverage, specially when
  // `values` is empty.
  gather_valid_values_around(0);

  // Skip the checks below if we are correct because they are expensive.
  if (std::equal(s.begin(), s.end(), values.begin(), values.end())) return;

  std::vector<int32_t> false_negatives;
  for (int32_t v : values) {
    if (!ValidateEnum(v, encoded.data())) false_negatives.push_back(v);
    s.erase(v);
  }
  const auto& false_positives = s;
  const auto print_data = [&] {
    auto header = ExtractHeader(encoded);
    return absl::StrFormat("line=%d header=%s", line, header.ToString());
  };
  EXPECT_THAT(false_negatives, IsEmpty())
      << "Missing values from the input. " << print_data()
      << "\nEncoded: " << testing::PrintToString(encoded);
  EXPECT_THAT(false_positives, IsEmpty())
      << "Found values not in input. " << print_data()
      << "\nEncoded: " << testing::PrintToString(encoded);
}

TEST(ValidateEnumTest, GenerateEnumDataSequentialWithOverflow) {
  std::vector<int32_t> values;
  for (int32_t i = -33000; i < 33000; ++i) {
    values.push_back(i);
  }
  const auto data = GenerateEnumData(values);
  EXPECT_THAT(ExtractHeader(data),
              HeaderHas(
                  // The sequence starts at the minimum possible value,
                  std::numeric_limits<int16_t>::min(),
                  // and it is as long as possible.
                  0xFFFF,
                  // we have some values in the bitmap
                  Gt(0),
                  // we have some in the fallback
                  Gt(0)));

  TestRoundTrip(values, __LINE__);
}

TEST(ValidateEnumTest, GenerateEnumDataBitmap) {
  EXPECT_THAT(GenerateEnumData({0, 1, 2, 4, 8, 16, 32}),
              ElementsAre(Make32(0, 3), Make32(32, 0),
                          0b100000000000000010000000100010));
  TestRoundTrip({}, __LINE__);
  TestRoundTrip({0, 1, 2, 4, 8, 16}, __LINE__);
  TestRoundTrip({0, 1, 2, 4, 8, 16, 32, 64, 128, 256}, __LINE__);
  TestRoundTrip({10000, 10001, 10002, 10004, 10006, 10008, 10010}, __LINE__);
  TestRoundTrip({std::numeric_limits<int32_t>::min(), -123123, -123, 213,
                 213213, std::numeric_limits<int32_t>::max()},
                __LINE__);
}

TEST(ValidateEnumTest, GenerateEnumDataBitmapWithOverflow) {
  std::vector<int32_t> values;
  // We step by 10 to guarantee each new value is more cost effective to add to
  // the bitmap, which would cause an overflow of the 16-bit bitmap size if we
  // didn't prevent it in the generator.
  for (int32_t i = -33000; i < 33000; i += 10) {
    values.push_back(i);
  }
  const auto data = GenerateEnumData(values);

  EXPECT_THAT(ExtractHeader(data),
              HeaderHas(_, _,
                        // we reached the maximum size for the bitmap.
                        0x10000 - 32,
                        // we have some in the fallback
                        Gt(0)));

  TestRoundTrip(values, __LINE__);
}

TEST(ValidateEnumTest, GenerateEnumDataWithOverflowOnBoth) {
  std::vector<int32_t> values;
  for (int32_t i = -33000; i < 100000; ++i) {
    values.push_back(i);
  }
  const auto data = GenerateEnumData(values);

  EXPECT_THAT(ExtractHeader(data),
              HeaderHas(
                  // The sequence starts at the minimum possible value,
                  std::numeric_limits<int16_t>::min(),
                  // and it is as long as possible.
                  0xFFFF,
                  // we reached the maximum size for the bitmap.
                  0x10000 - 32,
                  // we have some in the fallback
                  Gt(0)));

  TestRoundTrip(values, __LINE__);
}



}  // namespace
}  // namespace internal
}  // namespace protobuf
}  // namespace google

#include "google/protobuf/port_undef.inc"