From 23ddf02bffe073655458f92f712c196a4dbcc29d Mon Sep 17 00:00:00 2001
From: Ben Deane <ben@elbeno.com>
Date: Tue, 11 Feb 2025 21:12:11 -0700
Subject: [PATCH] :art: Preserve enum values in bitset `for_each`

When using a `bitset` with an `enum`, it makes sense that the function passed to
`for_each` should take an `enum` value.
---
 include/stdx/bitset.hpp | 12 ++++++++----
 test/bitset.cpp         | 11 +++++++++++
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/include/stdx/bitset.hpp b/include/stdx/bitset.hpp
index 56fd2b1..50ab7b2 100644
--- a/include/stdx/bitset.hpp
+++ b/include/stdx/bitset.hpp
@@ -101,13 +101,16 @@ class bitset {
         return lhs;
     }
 
+    using iter_arg_t = conditional_t<std::is_enum_v<decltype(Size)>,
+                                     decltype(Size), std::size_t>;
+
     template <typename F> constexpr auto for_each(F &&f) const -> F {
         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));
-                f(i + offset);
+                f(static_cast<iter_arg_t>(i + offset));
             }
             i += std::numeric_limits<elem_t>::digits;
         }
@@ -124,7 +127,8 @@ class bitset {
             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));
+                init =
+                    r(std::move(init), f(static_cast<iter_arg_t>(i + offset)));
             }
             i += std::numeric_limits<elem_t>::digits;
         }
@@ -440,8 +444,8 @@ 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 {
+[[nodiscard]] 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)),
diff --git a/test/bitset.cpp b/test/bitset.cpp
index 6d6fa07..e383ee0 100644
--- a/test/bitset.cpp
+++ b/test/bitset.cpp
@@ -474,6 +474,17 @@ TEST_CASE("use bitset with enum struct (place_bits construct)", "[bitset]") {
     static_assert(bs.to_natural() == 1);
 }
 
+TEST_CASE("use bitset with enum struct (for_each)", "[bitset]") {
+    constexpr auto bs = stdx::bitset<Bits::MAX>{stdx::all_bits};
+    for_each([](Bits) {}, bs);
+}
+
+TEST_CASE("use bitset with enum struct (transform_reduce)", "[bitset]") {
+    constexpr auto bs = stdx::bitset<Bits::MAX>{stdx::all_bits};
+    CHECK(transform_reduce([](Bits) { return true; }, std::logical_or{}, false,
+                           bs));
+}
+
 #if __cplusplus >= 202002L
 TEST_CASE("construct with a ct_string", "[bitset]") {
     using namespace stdx::literals;