compact_ranked_bitvector.hpp

#pragma once

#include <stdexcept>
#include <succinct/bit_vector.hpp>
#include <succinct/broadword.hpp>

#include "global_parameters.hpp"
#include "util.hpp"

namespace quasi_succinct {

struct compact_ranked_bitvector {

    struct offsets {
        offsets(uint64_t base_offset,
            uint64_t universe,
            uint64_t n,
            global_parameters const& params)
            : universe(universe)
            , n(n)
            , log_rank1_sampling(params.rb_log_rank1_sampling)
            , log_sampling1(params.rb_log_sampling1)

            , rank1_sample_size(ceil_log2(n + 1))
            , pointer_size(ceil_log2(universe))
            , rank1_samples(universe >> params.rb_log_rank1_sampling)
            , pointers1(n >> params.rb_log_sampling1)

            , rank1_samples_offset(base_offset)
            , pointers1_offset(rank1_samples_offset + rank1_samples * rank1_sample_size)
            , bits_offset(pointers1_offset + pointers1 * pointer_size)
            , end(bits_offset + universe)
        {
        }

        uint64_t universe;
        uint64_t n;
        uint64_t log_rank1_sampling;
        uint64_t log_sampling1;

        uint64_t rank1_sample_size;
        uint64_t pointer_size;

        uint64_t rank1_samples;
        uint64_t pointers1;

        uint64_t rank1_samples_offset;
        uint64_t pointers1_offset;
        uint64_t bits_offset;
        uint64_t end;
    };

    static QS_FLATTEN_FUNC uint64_t
    bitsize(global_parameters const& params, uint64_t universe, uint64_t n)
    {
        return offsets(0, universe, n, params).end;
    }

    template <typename Iterator>
    static void write(succinct::bit_vector_builder& bvb,
        Iterator begin,
        uint64_t universe, uint64_t n,
        global_parameters const& params)
    {
        using succinct::util::ceil_div;

        uint64_t base_offset = bvb.size();
        offsets of(base_offset, universe, n, params);
        // initialize all the bits to 0
        bvb.zero_extend(of.end - base_offset);

        uint64_t offset;

        auto set_rank1_samples = [&](uint64_t begin, uint64_t end,
            uint64_t rank) {
            for (uint64_t sample = ceil_div(begin, uint64_t(1) << of.log_rank1_sampling);
                 (sample << of.log_rank1_sampling) < end;
                 ++sample) {
                if (!sample)
                    continue;
                offset = of.rank1_samples_offset + (sample - 1) * of.rank1_sample_size;
                assert(offset + of.rank1_sample_size <= of.pointers1_offset);
                bvb.set_bits(offset, rank, of.rank1_sample_size);
            }
        };

        uint64_t sample1_mask = (uint64_t(1) << of.log_sampling1) - 1;
        uint64_t last = 0;
        Iterator it = begin;
        for (size_t i = 0; i < n; ++i) {
            uint64_t v = *it++;
            if (i && v == last) {
                throw std::runtime_error("Duplicate element");
            }
            if (i && v < last) {
                throw std::runtime_error("Sequence is not sorted");
            }

            assert(!i || v > last);
            assert(v <= universe);

            bvb.set(of.bits_offset + v, 1);

            if (i && (i & sample1_mask) == 0) {
                uint64_t ptr1 = i >> of.log_sampling1;
                assert(ptr1 > 0);
                offset = of.pointers1_offset + (ptr1 - 1) * of.pointer_size;
                assert(offset + of.pointer_size <= of.bits_offset);
                bvb.set_bits(offset, v, of.pointer_size);
            }

            set_rank1_samples(last + 1, v + 1, i);
            last = v;
        }

        set_rank1_samples(last + 1, universe, n);
    }

    class enumerator {
    public:
        typedef std::pair<uint64_t, uint64_t> value_type; // (position, value)

        enumerator(succinct::bit_vector const& bv, uint64_t offset,
            uint64_t universe, uint64_t n,
            global_parameters const& params)
            : m_bv(&bv)
            , m_of(offset, universe, n, params)
            , m_position(size())
            , m_value(m_of.universe)
        {
        }

        value_type move(uint64_t position)
        {
            assert(position <= size());

            if (position == m_position) {
                return value();
            }

            // optimize small forward skips
            uint64_t skip = position - m_position;
            if (QS_LIKELY(position > m_position && skip <= linear_scan_threshold)) {
                m_position = position;
                if (QS_UNLIKELY(m_position == size())) {
                    m_value = m_of.universe;
                } else {
                    succinct::bit_vector::unary_enumerator he = m_enumerator;
                    for (size_t i = 0; i < skip; ++i) {
                        he.next();
                    }
                    m_value = he.position() - m_of.bits_offset;
                    m_enumerator = he;
                }

                return value();
            }

            return slow_move(position);
        }

        value_type next_geq(uint64_t lower_bound)
        {
            if (lower_bound == m_value) {
                return value();
            }

            uint64_t diff = lower_bound - m_value;
            if (QS_LIKELY(lower_bound > m_value
                    && diff <= linear_scan_threshold)) {
                // optimize small skips
                succinct::bit_vector::unary_enumerator he = m_enumerator;
                uint64_t val;
                do {
                    m_position += 1;
                    if (QS_LIKELY(m_position < size())) {
                        val = he.next() - m_of.bits_offset;
                    } else {
                        val = m_of.universe;
                        break;
                    }
                } while (val < lower_bound);

                m_value = val;
                m_enumerator = he;
                return value();
            } else {
                return slow_next_geq(lower_bound);
            }
        }

        value_type next()
        {
            m_position += 1;
            assert(m_position <= size());

            if (QS_LIKELY(m_position < size())) {
                m_value = read_next();
            } else {
                m_value = m_of.universe;
            }
            return value();
        }

        uint64_t size() const
        {
            return m_of.n;
        }

        uint64_t prev_value() const
        {
            if (m_position == 0) {
                return 0;
            }

            uint64_t pos = 0;
            if (QS_LIKELY(m_position < size())) {
                pos = m_bv->predecessor1(m_enumerator.position() - 1);
            } else {
                pos = m_bv->predecessor1(m_of.end - 1);
            }

            return pos - m_of.bits_offset;
        }

    private:
        value_type QS_NOINLINE slow_move(uint64_t position)
        {
            uint64_t skip = position - m_position;
            if (QS_UNLIKELY(position == size())) {
                m_position = position;
                m_value = m_of.universe;
                return value();
            }

            uint64_t to_skip;
            if (position > m_position
                && (skip >> m_of.log_sampling1) == 0) {
                to_skip = skip - 1;
            } else {
                uint64_t ptr = position >> m_of.log_sampling1;
                uint64_t ptr_pos = pointer1(ptr);

                m_enumerator = succinct::bit_vector::unary_enumerator(*m_bv, m_of.bits_offset + ptr_pos);
                to_skip = position - (ptr << m_of.log_sampling1);
            }

            m_enumerator.skip(to_skip);
            m_position = position;
            m_value = read_next();

            return value();
        }

        value_type QS_NOINLINE slow_next_geq(uint64_t lower_bound)
        {
            using succinct::broadword::popcount;

            if (QS_UNLIKELY(lower_bound >= m_of.universe)) {
                return move(size());
            }

            uint64_t skip = lower_bound - m_value;
            m_enumerator = succinct::bit_vector::unary_enumerator(*m_bv, m_of.bits_offset + lower_bound);

            uint64_t begin;
            if (lower_bound > m_value
                && (skip >> m_of.log_rank1_sampling) == 0) {
                begin = m_of.bits_offset + m_value;
            } else {
                uint64_t block = lower_bound >> m_of.log_rank1_sampling;
                m_position = rank1_sample(block);

                begin = m_of.bits_offset + (block << m_of.log_rank1_sampling);
            }

            uint64_t end = m_of.bits_offset + lower_bound;
            uint64_t begin_word = begin / 64;
            uint64_t begin_shift = begin % 64;
            uint64_t end_word = end / 64;
            uint64_t end_shift = end % 64;
            uint64_t word = (m_bv->data()[begin_word] >> begin_shift) << begin_shift;

            while (begin_word < end_word) {
                m_position += popcount(word);
                word = m_bv->data()[++begin_word];
            }
            if (end_shift) {
                m_position += popcount(word << (64 - end_shift));
            }

            if (m_position < size()) {
                m_value = read_next();
            } else {
                m_value = m_of.universe;
            }

            return value();
        }

        static const uint64_t linear_scan_threshold = 8;

        inline value_type value() const
        {
            return value_type(m_position, m_value);
        }

        inline uint64_t read_next()
        {
            return m_enumerator.next() - m_of.bits_offset;
        }

        inline uint64_t pointer(uint64_t offset, uint64_t i, uint64_t size) const
        {
            if (i == 0) {
                return 0;
            } else {
                return m_bv->get_word56(offset + (i - 1) * size)
                    & ((uint64_t(1) << size) - 1);
            }
        }

        inline uint64_t pointer1(uint64_t i) const
        {
            return pointer(m_of.pointers1_offset, i, m_of.pointer_size);
        }

        inline uint64_t rank1_sample(uint64_t i) const
        {
            return pointer(m_of.rank1_samples_offset, i,
                m_of.rank1_sample_size);
        }

        succinct::bit_vector const* m_bv;
        offsets m_of;

        uint64_t m_position;
        uint64_t m_value;
        succinct::bit_vector::unary_enumerator m_enumerator;
    };
};
}