From 7fc489b9b7998415c3b72a32a2b80516b8cd315c Mon Sep 17 00:00:00 2001
From: Ben Deane <ben@elbeno.com>
Date: Tue, 11 Feb 2025 10:23:54 -0700
Subject: [PATCH] :sparkles: Add `transform_reduce` on bitset

Problem:
- It is relatively common to want to do for_each on a bitset, but also capture
  whether or not the action succeeded at all, e.g. when handling a message with
  an indexed handler.

Solution:
- Add `transform_reduce` on bitset.

Note:
- Both `for_each` and `transform_reduce` are applicable variadically, however
  this is not yet implemented. But this is the reason that the bitset is the
  final argument.
- There are some differences (for the better) from the `std::transform_reduce`
  API.
  - The transform function comes before the reduction function.
    `std::transform_reduce` is confusing, because the transform happens first
    but is given second.
  - The type of the `init` value can be easily given in a template argument;
    this is helpful for avoiding warnings when passing `0` (an `int`) to
    accumulate potentially larger integral types.
---
 docs/bitset.adoc        |  9 +++++++++
 include/stdx/bitset.hpp | 31 +++++++++++++++++++++++++++++++
 test/bitset.cpp         | 14 ++++++++++++++
 3 files changed, 54 insertions(+)

diff --git a/docs/bitset.adoc b/docs/bitset.adoc
index d851574..0ed75e6 100644
--- a/docs/bitset.adoc
+++ b/docs/bitset.adoc
@@ -103,3 +103,12 @@ objects, `lowest_unset` is also provided:
 auto bs = stdx::bitset<8>{0b11'0111ul};
 auto i = bs.lowest_unset(); // i == 3
 ----
+
+`transform_reduce` is also provided for bitsets:
+[source,cpp]
+----
+auto bs = stdx::bitset<8>{0b1010'1010ul};
+auto result = transform_reduce([](auto i) { return i * 2 },
+                               std::plus{}, std::size_t{}, bs);
+// result is 1*2 + 3*2 + 5*2 + 7*2
+----
diff --git a/include/stdx/bitset.hpp b/include/stdx/bitset.hpp
index d81f8c8..56fd2b1 100644
--- a/include/stdx/bitset.hpp
+++ b/include/stdx/bitset.hpp
@@ -117,6 +117,24 @@ class bitset {
     template <typename F, auto M, typename... S>
     friend constexpr auto for_each(F &&f, bitset<M, S> const &...bs) -> F;
 
+    template <typename T, typename F, typename R>
+    constexpr auto transform_reduce(F &&f, R &&r, T init) const -> T {
+        std::size_t i = 0;
+        for (auto e : storage) {
+            while (e != 0) {
+                auto const offset = static_cast<std::size_t>(countr_zero(e));
+                e &= static_cast<elem_t>(~(bit << offset));
+                init = r(std::move(init), f(i + offset));
+            }
+            i += std::numeric_limits<elem_t>::digits;
+        }
+        return init;
+    }
+
+    template <typename T, typename F, typename R, auto M, typename... S>
+    friend constexpr auto transform_reduce(F &&f, R &&r, T init,
+                                           bitset<M, S> const &...bs) -> T;
+
   public:
     constexpr bitset() = default;
     constexpr explicit bitset(std::uint64_t value) {
@@ -421,6 +439,19 @@ constexpr auto for_each(F &&f, bitset<M, S> const &...bs) -> F {
     }
 }
 
+template <typename T, typename F, typename R, auto M, typename... S>
+constexpr auto transform_reduce(F &&f, R &&r, T init,
+                                bitset<M, S> const &...bs) -> T {
+    if constexpr (sizeof...(bs) == 1) {
+        return (bs.transform_reduce(std::forward<F>(f), std::forward<R>(r),
+                                    std::move(init)),
+                ...);
+    } else {
+        static_assert(stdx::always_false_v<F>, "unimplemented");
+        return init;
+    }
+}
+
 #if __cplusplus >= 202002L
 template <std::size_t N> bitset(ct_string<N>) -> bitset<N - 1>;
 #endif
diff --git a/test/bitset.cpp b/test/bitset.cpp
index d6e84f1..6d6fa07 100644
--- a/test/bitset.cpp
+++ b/test/bitset.cpp
@@ -331,6 +331,20 @@ TEMPLATE_TEST_CASE("for_each iterates in order lsb to msb", "[bitset]",
     CHECK(result == bs);
 }
 
+TEMPLATE_TEST_CASE("transform_reduce", "[bitset]", std::uint8_t, std::uint16_t,
+                   std::uint32_t, std::uint64_t) {
+    constexpr auto bs = stdx::bitset<8, TestType>{0b10101010ul};
+    int calls{};
+    auto const result = transform_reduce(
+        [&](auto i) {
+            ++calls;
+            return i == 3 and calls == 2;
+        },
+        std::logical_or{}, false, bs);
+    CHECK(result);
+    CHECK(calls == bs.count());
+}
+
 TEMPLATE_TEST_CASE("set range of bits (lsb, length)", "[bitset]", std::uint8_t,
                    std::uint16_t, std::uint32_t, std::uint64_t) {
     using namespace stdx::literals;