Expand Table, BitSet API's (#27465)

* Expand Table, BitSet API's

Add a population count to BitSet, use it to add a count() method to
Table to get the number of fields set.

Add a ForEach to Table so that it can be iterated

* Automated change: Fix sanity tests

* fix

* Automated change: Fix sanity tests

Co-authored-by: ctiller <ctiller@users.noreply.github.com>

BUILD

@@ -868,6 +868,7 @@ grpc_cc_library(
     public_hdrs = ["src/core/lib/gprpp/bitset.h"],
     deps = [
         "gpr_platform",
+        "useful",
     ],
 )
 

build_autogenerated.yaml

@@ -4669,6 +4669,7 @@ targets:
   build: test
   language: c++
   headers:
+  - src/core/lib/gpr/useful.h
   - src/core/lib/gprpp/bitset.h
   src:
   - test/core/gprpp/bitset_test.cc
@@ -6342,6 +6343,7 @@ targets:
   build: test
   language: c++
   headers:
+  - src/core/lib/gpr/useful.h
   - src/core/lib/gprpp/bitset.h
   - src/core/lib/gprpp/construct_destruct.h
   - src/core/lib/promise/detail/basic_join.h
@@ -7627,6 +7629,7 @@ targets:
   build: test
   language: c++
   headers:
+  - src/core/lib/gpr/useful.h
   - src/core/lib/gprpp/bitset.h
   - src/core/lib/gprpp/table.h
   src:
@@ -7788,6 +7791,7 @@ targets:
   build: test
   language: c++
   headers:
+  - src/core/lib/gpr/useful.h
   - src/core/lib/gprpp/bitset.h
   - src/core/lib/gprpp/construct_destruct.h
   - src/core/lib/promise/detail/basic_join.h

src/core/lib/gpr/useful.h

@@ -76,6 +76,17 @@ inline constexpr uint32_t BitCount(uint32_t i) {
           255);
 }
 
+inline constexpr uint32_t BitCount(uint64_t i) {
+  return BitCount(uint32_t(i)) + BitCount(uint32_t(i >> 32));
+}
+
+inline constexpr uint32_t BitCount(uint16_t i) { return BitCount(uint32_t(i)); }
+inline constexpr uint32_t BitCount(uint8_t i) { return BitCount(uint32_t(i)); }
+inline constexpr uint32_t BitCount(int64_t i) { return BitCount(uint64_t(i)); }
+inline constexpr uint32_t BitCount(int32_t i) { return BitCount(uint32_t(i)); }
+inline constexpr uint32_t BitCount(int16_t i) { return BitCount(uint16_t(i)); }
+inline constexpr uint32_t BitCount(int8_t i) { return BitCount(uint8_t(i)); }
+
 // This function uses operator< to implement a qsort-style comparison, whereby:
 // if a is smaller than b, a number smaller than 0 is returned.
 // if a is bigger than b, a number greater than 0 is returned.

src/core/lib/gprpp/bitset.h

@@ -19,6 +19,8 @@
 
 #include <utility>
 
+#include "src/core/lib/gpr/useful.h"
+
 #if __cplusplus > 201103l
 #define GRPC_BITSET_CONSTEXPR_MUTATOR constexpr
 #else
@@ -29,7 +31,7 @@ namespace grpc_core {
 
 // Given a bit count as an integer, vend as member type `Type` a type with
 // exactly that number of bits. Undefined if that bit count is not available.
-template <std::size_t kBits>
+template <size_t kBits>
 struct UintSelector;
 template <>
 struct UintSelector<8> {
@@ -49,7 +51,7 @@ struct UintSelector<64> {
 };
 
 // An unsigned integer of some number of bits.
-template <std::size_t kBits>
+template <size_t kBits>
 using Uint = typename UintSelector<kBits>::Type;
 
 // Given the total number of bits that need to be stored, choose the size of
@@ -58,7 +60,7 @@ using Uint = typename UintSelector<kBits>::Type;
 // size and performance
 // - the details will likely be tweaked into the future.
 // Once we get over 96 bits, we just use uint64_t for everything.
-constexpr std::size_t ChooseUnitBitsForBitSet(std::size_t total_bits) {
+constexpr size_t ChooseUnitBitsForBitSet(size_t total_bits) {
   return total_bits <= 8    ? 8
          : total_bits <= 16 ? 16
          : total_bits <= 24 ? 8
@@ -74,11 +76,10 @@ constexpr std::size_t ChooseUnitBitsForBitSet(std::size_t total_bits) {
 // kUnitBits. e.g. to store 72 bits in 8 bit chunks, we'd say BitSet<72, 8>.
 // Since most users shouldn't care about the size of unit used, we default
 // kUnitBits to whatever is selected by ChooseUnitBitsForBitSet
-template <std::size_t kTotalBits,
-          std::size_t kUnitBits = ChooseUnitBitsForBitSet(kTotalBits)>
+template <size_t kTotalBits,
+          size_t kUnitBits = ChooseUnitBitsForBitSet(kTotalBits)>
 class BitSet {
-  static constexpr std::size_t kUnits =
-      (kTotalBits + kUnitBits - 1) / kUnitBits;
+  static constexpr size_t kUnits = (kTotalBits + kUnitBits - 1) / kUnitBits;
 
  public:
   // Initialize to all bits false
@@ -113,7 +114,7 @@ class BitSet {
     if (kTotalBits % kUnitBits == 0) {
       // kTotalBits is a multiple of kUnitBits ==> we can just check for all
       // ones in each unit.
-      for (std::size_t i = 0; i < kUnits; i++) {
+      for (size_t i = 0; i < kUnits; i++) {
         if (units_[i] != all_ones()) return false;
       }
       return true;
@@ -121,7 +122,7 @@ class BitSet {
       // kTotalBits is not a multiple of kUnitBits ==> we need special handling
       // for checking partial filling of the last unit (since not all of its
       // bits are used!)
-      for (std::size_t i = 0; i < kUnits - 1; i++) {
+      for (size_t i = 0; i < kUnits - 1; i++) {
         if (units_[i] != all_ones()) return false;
       }
       return units_[kUnits - 1] == n_ones(kTotalBits % kUnitBits);
@@ -130,20 +131,27 @@ class BitSet {
 
   // Return true if *no* bits are set.
   bool none() const {
-    for (std::size_t i = 0; i < kUnits; i++) {
+    for (size_t i = 0; i < kUnits; i++) {
       if (units_[i] != 0) return false;
     }
     return true;
   }
 
+  // Return a count of how many bits are set.
+  uint32_t count() const {
+    uint32_t count = 0;
+    for (size_t i = 0; i < kUnits; i++) {
+      count += BitCount(units_[i]);
+    }
+    return count;
+  }
+
  private:
   // Given a bit index, return which unit it's stored in.
-  static constexpr std::size_t unit_for(std::size_t bit) {
-    return bit / kUnitBits;
-  }
+  static constexpr size_t unit_for(size_t bit) { return bit / kUnitBits; }
 
   // Given a bit index, return a mask to access that bit within it's unit.
-  static constexpr Uint<kUnitBits> mask_for(std::size_t bit) {
+  static constexpr Uint<kUnitBits> mask_for(size_t bit) {
     return Uint<kUnitBits>{1} << (bit % kUnitBits);
   }
 
@@ -153,7 +161,7 @@ class BitSet {
   }
 
   // Return a value with n bottom bits ones
-  static constexpr Uint<kUnitBits> n_ones(std::size_t n) {
+  static constexpr Uint<kUnitBits> n_ones(size_t n) {
     return n == kUnitBits ? all_ones() : (Uint<kUnitBits>(1) << n) - 1;
   }
 

src/core/lib/gprpp/table.h

@@ -290,6 +290,15 @@ class Table {
     }
   }
 
+  // Iterate through each set field in the table
+  template <typename F>
+  void ForEach(F f) const {
+    ForEachImpl(std::move(f), absl::make_index_sequence<sizeof...(Ts)>());
+  }
+
+  // Count the number of set fields in the table
+  size_t count() const { return present_bits_.count(); }
+
  private:
   // Bit field for which elements of the table are set (true) or un-set (false,
   // the default) -- one bit for each type in Ts.
@@ -354,6 +363,14 @@ class Table {
     }
   }
 
+  // Call (*f)(value) if that value is in the table.
+  template <size_t I, typename F>
+  void CallIf(F* f) const {
+    if (auto* p = get<I>()) {
+      (*f)(*p);
+    }
+  }
+
   // For each field (element I=0, 1, ...) if that field is present, call its
   // destructor.
   template <size_t... I>
@@ -377,6 +394,12 @@ class Table {
         {(MoveIf<or_clear, I>(std::forward<Table>(rhs)), 1)...});
   }
 
+  // For each field (element I=0, 1, ...) if that field is present, call f.
+  template <typename F, size_t... I>
+  void ForEachImpl(F f, absl::index_sequence<I...>) const {
+    table_detail::do_these_things({(CallIf<I>(&f), 1)...});
+  }
+
   // Bit field indicating which elements are set.
   GPR_NO_UNIQUE_ADDRESS PresentBits present_bits_;
   // The memory to store the elements themselves.

test/core/gpr/useful_test.cc

@@ -59,6 +59,7 @@ TEST(UsefulTest, BitOps) {
   EXPECT_EQ(grpc_core::ClearBit(&bitset, 3), 2);
   EXPECT_EQ(grpc_core::BitCount(bitset), 1);
   EXPECT_EQ(grpc_core::GetBit(bitset, 3), 0);
+  EXPECT_EQ(grpc_core::BitCount(std::numeric_limits<uint64_t>::max()), 64);
 }
 
 }  // namespace grpc_core

test/core/gprpp/bitset_test.cc

@@ -14,15 +14,17 @@
 
 #include "src/core/lib/gprpp/bitset.h"
 
+#include <random>
+
 #include <gtest/gtest.h>
 
 namespace grpc_core {
 namespace testing {
 
 // Stand in type to make the size to test a type
-template <std::size_t K>
+template <size_t K>
 struct Size {
-  static constexpr std::size_t kBits = K;
+  static constexpr size_t kBits = K;
 };
 
 using TestSizes = ::testing::Types<
@@ -53,27 +55,41 @@ TYPED_TEST(BitSetTest, NoneAtInit) {
 }
 
 TYPED_TEST(BitSetTest, OneBit) {
-  constexpr std::size_t kBits = TypeParam::kBits;
-  for (std::size_t i = 0; i < kBits; i++) {
+  constexpr size_t kBits = TypeParam::kBits;
+  for (size_t i = 0; i < kBits; i++) {
     BitSet<kBits> b;
     b.set(i);
     EXPECT_FALSE(b.none());
-    for (std::size_t j = 0; j < kBits; j++) {
+    for (size_t j = 0; j < kBits; j++) {
       EXPECT_EQ(b.is_set(j), i == j);
     }
   }
 }
 
 TYPED_TEST(BitSetTest, AllSet) {
-  constexpr std::size_t kBits = TypeParam::kBits;
+  constexpr size_t kBits = TypeParam::kBits;
   BitSet<kBits> b;
-  for (std::size_t i = 0; i < kBits; i++) {
+  for (size_t i = 0; i < kBits; i++) {
     EXPECT_FALSE(b.all());
     b.set(i);
   }
   EXPECT_TRUE(b.all());
 }
 
+TYPED_TEST(BitSetTest, Count) {
+  constexpr size_t kBits = TypeParam::kBits;
+  BitSet<kBits> b;
+  std::set<size_t> bits_set;
+  std::random_device rd;
+  std::uniform_int_distribution<size_t> dist(0, kBits - 1);
+  for (size_t i = 0; i < 4 * kBits; i++) {
+    size_t bit = dist(rd);
+    bits_set.insert(bit);
+    b.set(bit);
+    EXPECT_EQ(b.count(), bits_set.size());
+  }
+}
+
 }  // namespace testing
 }  // namespace grpc_core
 

test/core/gprpp/table_test.cc

@@ -112,6 +112,18 @@ TEST(Table, SameTypes) {
   EXPECT_EQ(t.get<2>(), nullptr);
 }
 
+TEST(Table, ForEach) {
+  Table<int, int, int> t;
+  t.set<0>(1);
+  t.set<1>(2);
+  t.set<2>(3);
+  int i = 1;
+  t.ForEach([&i](int x) {
+    EXPECT_EQ(x, i);
+    i++;
+  });
+}
+
 #if !defined(_MSC_VER)
 // Test suite proving this is memory efficient compared to
 // tuple<optional<Ts>...>