IterableBitSet.hh

Go to the documentation of this file.

#ifndef ITERABLEBITSET_HH
#define ITERABLEBITSET_HH
 
#include "ranges.hh"
#include "xrange.hh"
#include <array>
#include <bit>
#include <cassert>
#include <concepts>
#include <cstddef>
#include <cstdint>
#include <type_traits>
 
template<size_t N>
class IterableBitSet
{
        using WordType = std::conditional_t<(N > 32), uint64_t,
                         std::conditional_t<(N > 16), uint32_t,
                         std::conditional_t<(N >  8), uint16_t,
                                                      uint8_t>>>;
        static constexpr size_t BITS_PER_WORD = 8 * sizeof(WordType);
        static constexpr size_t NUM_WORDS = (N + BITS_PER_WORD - 1) / BITS_PER_WORD;
 
public:
        [[nodiscard]] bool empty() const
        {
                return ranges::all_of(words, [](auto w) { return w == 0; });
        }
 
        void set(size_t pos)
        {
                assert(pos < N);
                if constexpr (NUM_WORDS > 1) {
                        auto w = pos / BITS_PER_WORD;
                        auto m = WordType(1) << (pos % BITS_PER_WORD);
                        words[w] |= m;
                } else {
                        words[0] |= WordType(1) << pos;
                }
        }
 
        void setPosN(size_t pos, size_t n)
        {
                setRange(pos, pos + n);
        }
 
        void setRange(size_t begin, size_t end)
        {
                assert(begin <= end);
                assert(end <= N);
 
                if (begin == end) return;
 
                if constexpr (NUM_WORDS > 1) {
                        auto firstWord = begin / BITS_PER_WORD;
                        auto lastWord = (end - 1) / BITS_PER_WORD;
                        if (firstWord != lastWord) {
                                auto firstPos = begin % BITS_PER_WORD;
                                auto firstMask = WordType(-1) << firstPos;
                                words[firstWord++] |= firstMask;
 
                                while (firstWord != lastWord) {
                                        words[firstWord++] = WordType(-1);
                                }
 
                                auto lastPos = end % BITS_PER_WORD;
                                auto lastMask = WordType(-1) >> (lastPos ? (BITS_PER_WORD - lastPos) : 0);
                                words[firstWord] |= lastMask;
                        } else {
                                auto firstPos = begin % BITS_PER_WORD;
                                auto mask = WordType(-1) << firstPos;
                                auto lastPos = end % BITS_PER_WORD;
                                if (lastPos) {
                                        auto lastMask = WordType(-1) >> (BITS_PER_WORD - lastPos);
                                        mask &= lastMask;
                                }
                                words[firstWord] |= mask;
                        }
                } else {
                        assert(end);
                        auto mask1 = WordType(-1) << begin;
                        auto mask2 = WordType(-1) >> (BITS_PER_WORD - end);
                        words[0] |= mask1 & mask2;
                }
        }
 
        void foreachSetBit(std::invocable<size_t> auto op) const
        {
                for (auto i : xrange(NUM_WORDS)) {
                        auto w = words[i];
                        while (w) {
                                op(i * BITS_PER_WORD + std::countr_zero(w));
                                w &= w - 1; // clear least significant set bit
                        }
                }
        }
 
private:
        std::array<WordType, NUM_WORDS> words = {};
};
 
#endif