<commit_before>#ifndef _FM_INDEX_WAVELET_MATRIX_HPP
#define _FM_INDEX_WAVELET_MATRIX_HPP

#include "wavelet_tree.hpp"
#include <am/succinct/sdarray/SDArray.hpp>
#include <am/succinct/utils.hpp>
#include <am/interval_heap.hpp>

#include <deque>


NS_IZENELIB_AM_BEGIN

namespace succinct
{
namespace fm_index
{

template <class CharT>
class WaveletMatrix : public WaveletTree<CharT>
{
public:
    typedef CharT char_type;
    typedef WaveletMatrix<CharT> self_type;
    typedef std::vector<FilterList<self_type> *, boost::stl_allocator<FilterList<self_type> *> > aux_filter_list_type;

    WaveletMatrix(size_t alphabet_num, bool support_select, bool dense);
    ~WaveletMatrix();

    void build(const char_type *char_seq, size_t len);

    char_type access(size_t pos) const;
    char_type access(size_t pos, size_t &rank) const;

    size_t rank(char_type c, size_t pos) const;
    size_t select(char_type c, size_t rank) const;

    void intersect(
            const std::vector<std::pair<size_t, size_t> > &patterns,
            size_t thres,
            size_t max_count,
            std::vector<char_type> &results) const;

    void topKUnion(
            const range_list_type &patterns,
            size_t topK,
            std::vector<std::pair<double, char_type> > &results,
            boost::auto_alloc& alloc) const;

    void topKUnionWithFilters(
            const range_list_type &filters,
            const range_list_type &patterns,
            size_t topK,
            std::vector<std::pair<double, char_type> > &results,
            boost::auto_alloc& alloc) const;

    void topKUnionWithAuxFilters(
            const aux_filter_list_type &aux_filters,
            const range_list_type &patterns,
            size_t topK,
            std::vector<std::pair<double, char_type> > &results,
            boost::auto_alloc& alloc) const;

    size_t getOcc(char_type c) const;
    WaveletTreeNode *getRoot() const;

    size_t length() const;
    size_t allocSize() const;

    void save(std::ostream &ostr) const;
    void load(std::istream &istr);

private:
    void doIntersect_(
            const std::vector<std::pair<size_t, size_t> > &patterns,
            size_t thres,
            size_t max_count,
            size_t level,
            char_type symbol,
            std::vector<char_type> &results) const;

protected:
    sdarray::SDArray occ_;
    std::vector<size_t> zero_counts_;
    std::vector<WaveletTreeNode *> nodes_;
};

template <class CharT>
WaveletMatrix<CharT>::WaveletMatrix(uint64_t alphabet_num, bool support_select, bool dense)
    : WaveletTree<CharT>(alphabet_num, support_select, dense)
{
}

template <class CharT>
WaveletMatrix<CharT>::~WaveletMatrix()
{
    for (size_t i = 0; i < nodes_.size(); ++i)
    {
        if (nodes_[i]) delete nodes_[i];
    }
}

template <class CharT>
void WaveletMatrix<CharT>::build(const char_type *char_seq, size_t len)
{
    if (this->alphabet_num_ == 0) return;

    this->alphabet_bit_num_ = SuccinctUtils::log2(this->alphabet_num_ - 1);

    zero_counts_.resize(this->alphabet_bit_num_);

    nodes_.resize(this->alphabet_bit_num_);

    std::vector<size_t> prev_begin_pos(1), node_begin_pos(1);
    char_type bit_mask, subscript;

    for (size_t i = 0; i < this->alphabet_bit_num_; ++i)
    {
        node_begin_pos.clear();
        node_begin_pos.resize((1ULL << (i + 1)) + 1);

        nodes_[i] = new WaveletTreeNode(this->support_select_, this->dense_);
        nodes_[i]->resize(len);

        bit_mask = (char_type)1 << i;

        for (size_t j = 0; j < len; ++j)
        {
            subscript = char_seq[j] & (bit_mask - 1);

            if (char_seq[j] & bit_mask)
            {
                nodes_[i]->setBit(prev_begin_pos[subscript]++);
                ++node_begin_pos[subscript + bit_mask + 1];
            }
            else
            {
                nodes_[i]->unsetBit(prev_begin_pos[subscript]++);
                ++node_begin_pos[subscript + 1];
            }
        }

        nodes_[i]->build();

        prev_begin_pos = node_begin_pos;
        for (size_t j = 2; j < prev_begin_pos.size(); ++j)
        {
            prev_begin_pos[j] += prev_begin_pos[j - 1];
        }
        zero_counts_[i] = prev_begin_pos[bit_mask];
    }

    for (size_t i = 1; i <= this->alphabet_num_; ++i)
    {
        occ_.add(node_begin_pos[i]);
    }
    occ_.build();

    for (size_t i = 1; i < nodes_.size(); ++i)
    {
        nodes_[i]->parent_ = nodes_[i - 1];
        nodes_[i - 1]->left_ = nodes_[i];<commit_msg>reimplement constructing wavelet matrix as a radix-sort<commit_after>#ifndef _FM_INDEX_WAVELET_MATRIX_HPP
#define _FM_INDEX_WAVELET_MATRIX_HPP

#include "wavelet_tree.hpp"
#include <am/succinct/sdarray/SDArray.hpp>
#include <am/succinct/utils.hpp>
#include <am/interval_heap.hpp>

#include <deque>


NS_IZENELIB_AM_BEGIN

namespace succinct
{
namespace fm_index
{

template <class CharT>
class WaveletMatrix : public WaveletTree<CharT>
{
public:
    typedef CharT char_type;
    typedef WaveletMatrix<CharT> self_type;
    typedef std::vector<FilterList<self_type> *, boost::stl_allocator<FilterList<self_type> *> > aux_filter_list_type;

    WaveletMatrix(size_t alphabet_num, bool support_select, bool dense);
    ~WaveletMatrix();

    void build(const char_type *char_seq, size_t len);

    char_type access(size_t pos) const;
    char_type access(size_t pos, size_t &rank) const;

    size_t rank(char_type c, size_t pos) const;
    size_t select(char_type c, size_t rank) const;

    void intersect(
            const std::vector<std::pair<size_t, size_t> > &patterns,
            size_t thres,
            size_t max_count,
            std::vector<char_type> &results) const;

    void topKUnion(
            const range_list_type &patterns,
            size_t topK,
            std::vector<std::pair<double, char_type> > &results,
            boost::auto_alloc& alloc) const;

    void topKUnionWithFilters(
            const range_list_type &filters,
            const range_list_type &patterns,
            size_t topK,
            std::vector<std::pair<double, char_type> > &results,
            boost::auto_alloc& alloc) const;

    void topKUnionWithAuxFilters(
            const aux_filter_list_type &aux_filters,
            const range_list_type &patterns,
            size_t topK,
            std::vector<std::pair<double, char_type> > &results,
            boost::auto_alloc& alloc) const;

    size_t getOcc(char_type c) const;
    WaveletTreeNode *getRoot() const;

    size_t length() const;
    size_t allocSize() const;

    void save(std::ostream &ostr) const;
    void load(std::istream &istr);

private:
    void doIntersect_(
            const std::vector<std::pair<size_t, size_t> > &patterns,
            size_t thres,
            size_t max_count,
            size_t level,
            char_type symbol,
            std::vector<char_type> &results) const;

protected:
    sdarray::SDArray occ_;
    std::vector<size_t> zero_counts_;
    std::vector<WaveletTreeNode *> nodes_;
};

template <class CharT>
WaveletMatrix<CharT>::WaveletMatrix(uint64_t alphabet_num, bool support_select, bool dense)
    : WaveletTree<CharT>(alphabet_num, support_select, dense)
{
}

template <class CharT>
WaveletMatrix<CharT>::~WaveletMatrix()
{
    for (size_t i = 0; i < nodes_.size(); ++i)
    {
        if (nodes_[i]) delete nodes_[i];
    }
}

template <class CharT>
void WaveletMatrix<CharT>::build(const char_type *char_seq, size_t len)
{
    if (this->alphabet_num_ == 0) return;

    this->alphabet_bit_num_ = SuccinctUtils::log2(this->alphabet_num_ - 1);

    zero_counts_.resize(this->alphabet_bit_num_);

    nodes_.resize(this->alphabet_bit_num_);

    std::vector<size_t> prev_begin_pos(1), node_begin_pos(1);

    char_type *temp_seq = new char_type[len];
    char_type *curr = const_cast<char_type *>(char_seq);
    char_type *next = temp_seq;

    for (size_t i = 0; i < this->alphabet_bit_num_; ++i)
    {
        nodes_[i] = new WaveletTreeNode(this->support_select_, this->dense_);
        nodes_[i]->resize(len);

        size_t zero_count = 0;
        char_type bit_mask = (char_type)1 << i;

        for (size_t j = 0; j < len; ++j)
        {
            if (curr[j] & bit_mask)
            {
                nodes_[i]->setBit(j);
            }
            else
            {
                nodes_[i]->unsetBit(j);
                next[zero_count++] = curr[j];
            }
        }

        nodes_[i]->build();
        zero_counts_[i] = zero_count;

        for (size_t j = 0; j < len; ++j)
        {
            if (curr[j] & bit_mask)
            {
                next[zero_count++] = curr[j];
            }
        }
        std::swap(curr, next);
    }

    size_t curr_char = 0;
    size_t curr_count = 0;

    for (size_t i = 0; i < len; ++i)
    {
        if (curr[i] == curr_char)
        {
            ++curr_count;
        }
        else
        {
            occ_.add(curr_count);
            for (++curr_char; curr_char < curr[i]; ++curr_char)
            {
                occ_.add(0);
            }
            curr_count = 1;
        }
    }

    occ_.add(curr_count);
    for (++curr_char; curr_char <= 1LLU << this->alphabet_bit_num_; ++curr_char)
    {
        occ_.add(0);
    }

    occ_.build();

    delete temp_seq;

    for (size_t i = 1; i < nodes_.size(); ++i)
    {
        nodes_[i]->parent_ = nodes_[i - 1];
        nodes_[i - 1]->left_ = nodes_[i];
{'fraction_non_alphanumeric': 0.12016530677121967, 'fraction_numerical': 0.004874430433400445, 'mean_word_length': 2.5561416729464956, 'pattern_counts': {'":': 0, '<': 92, '<?xml version=': 0, '>': 100, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '363208', 'n_tokens_mistral': 3336, 'n_tokens_neox': 3206, 'n_words': 646}