Fancify enum_set and introduce enum_iter_t

Allow iterating over the values of an enum class.

src/enum_set.h

@@ -1,26 +1,85 @@
 #pragma once
 
 #include <bitset>
+#include <cassert>
+#include <iterator>
+
+/// A type (to specialize) that provides a count for an enum.
+/// Example:
+///   template<> struct enum_info_t<MyEnum>
+///   { static constexpr auto count = MyEnum::COUNT; };
+template <typename T>
+struct enum_info_t {};
 
 template <typename T>
-class enum_set_t {
+class enum_set_t : private std::bitset<static_cast<size_t>(enum_info_t<T>::count)> {
    private:
-    using base_type_t = typename std::underlying_type<T>::type;
-    std::bitset<8 * sizeof(base_type_t)> bitmask{0};
-    static int index_of(T t) { return static_cast<base_type_t>(t); }
+    using super = std::bitset<static_cast<size_t>(enum_info_t<T>::count)>;
+    static size_t index_of(T t) { return static_cast<size_t>(t); }
+
+    explicit enum_set_t(unsigned long raw) : super(raw) {}
 
    public:
-    bool get(T t) const { return bitmask.test(index_of(t)); }
+    enum_set_t() = default;
 
-    void set(T t, bool v) {
-        if (v) {
-            bitmask.set(index_of(t));
-        } else {
-            bitmask.reset(index_of(t));
-        }
-    }
+    explicit enum_set_t(T v) { set(v); }
 
-    void set(T t) { bitmask.set(index_of(t)); }
+    static enum_set_t from_raw(unsigned long v) { return enum_set_t{v}; }
 
-    void clear(T t) { bitmask.reset(index_of(t)); }
+    unsigned long to_raw() const { return super::to_ulong(); }
+
+    bool get(T t) const { return super::test(index_of(t)); }
+
+    void set(T t, bool v = true) { super::set(index_of(t), v); }
+
+    void clear(T t) { super::reset(index_of(t)); }
+
+    bool none() const { return super::none(); }
+
+    bool any() const { return super::any(); }
+
+    bool operator==(const enum_set_t &rhs) const { return super::operator==(rhs); }
+
+    bool operator!=(const enum_set_t &rhs) const { return super::operator!=(rhs); }
+};
+
+/// A counting iterator for an enum class.
+/// This enumerates the values of an enum class from 0 up to (not including) count.
+/// Example:
+///   for (auto v : enum_iter_t<MyEnum>) {...}
+template <typename T>
+class enum_iter_t {
+    using base_type_t = typename std::underlying_type<T>::type;
+    struct iterator_t {
+        friend class enum_iter_t;
+        using iterator_category = std::forward_iterator_tag;
+        using value_type = T;
+        using difference_type = long;
+
+        explicit iterator_t(base_type_t v) : v_(v) {}
+
+        T operator*() const { return static_cast<T>(v_); }
+        const T *operator->() const { return static_cast<const T *>(v_); }
+
+        iterator_t &operator++() {
+            v_ += 1;
+            return *this;
+        }
+
+        iterator_t operator++(int) {
+            auto res = *this;
+            v_ += 1;
+            return res;
+        }
+
+        bool operator==(iterator_t rhs) const { return v_ == rhs.v_; }
+        bool operator!=(iterator_t rhs) const { return v_ != rhs.v_; }
+
+       private:
+        base_type_t v_{};
+    };
+
+   public:
+    iterator_t begin() const { return iterator_t{0}; }
+    iterator_t end() const { return iterator_t{static_cast<base_type_t>(enum_info_t<T>::count)}; }
 };

src/fish_tests.cpp

@@ -286,6 +286,38 @@ static void test_str_to_num() {
              L"converting invalid num to long did not fail");
 }
 
+enum class test_enum { alpha, beta, gamma, COUNT };
+
+template <>
+struct enum_info_t<test_enum> {
+    static constexpr auto count = test_enum::COUNT;
+};
+
+static void test_enum_set() {
+    say(L"Testing enum set");
+    enum_set_t<test_enum> es;
+    do_test(es.none());
+    do_test(!es.any());
+    do_test(es.to_raw() == 0);
+    do_test(es == enum_set_t<test_enum>::from_raw(0));
+    do_test(es != enum_set_t<test_enum>::from_raw(1));
+
+    es.set(test_enum::beta);
+    do_test(es.to_raw() == 2);
+    do_test(es == enum_set_t<test_enum>::from_raw(2));
+    do_test(es == enum_set_t<test_enum>{test_enum::beta});
+    do_test(es != enum_set_t<test_enum>::from_raw(3));
+    do_test(es.any());
+    do_test(!es.none());
+
+    unsigned idx = 0;
+    for (auto v : enum_iter_t<test_enum>{}) {
+        do_test(static_cast<unsigned>(v) == idx);
+        idx++;
+    }
+    do_test(static_cast<unsigned>(test_enum::COUNT) == idx);
+}
+
 /// Test sane escapes.
 static void test_unescape_sane() {
     const struct test_t {
@@ -5082,6 +5114,7 @@ int main(int argc, char **argv) {
     if (should_test_function("wcstring_tok")) test_wcstring_tok();
     if (should_test_function("env_vars")) test_env_vars();
     if (should_test_function("str_to_num")) test_str_to_num();
+    if (should_test_function("enum")) test_enum_set();
     if (should_test_function("highlighting")) test_highlighting();
     if (should_test_function("new_parser_ll2")) test_new_parser_ll2();
     if (should_test_function("new_parser_fuzzing"))

src/proc.h

@@ -153,6 +153,13 @@ enum class job_flag_t {
     JOB_CONTROL,
     /// Whether the job wants to own the terminal when in the foreground.
     TERMINAL,
+
+    JOB_FLAG_COUNT
+};
+
+template <>
+struct enum_info_t<job_flag_t> {
+    static constexpr auto count = job_flag_t::JOB_FLAG_COUNT;
 };
 
 typedef int job_id_t;