fm_index.hpp

Go to the documentation of this file.

// SPDX-FileCopyrightText: 2006-2024 Knut Reinert & Freie Universität Berlin
// SPDX-FileCopyrightText: 2016-2024 Knut Reinert & MPI für molekulare Genetik
// SPDX-License-Identifier: BSD-3-Clause
 
#pragma once
 
#include <algorithm>
#include <filesystem>
#include <ranges>
 
#include <sdsl/suffix_trees.hpp>
 
#include <seqan3/alphabet/views/to_rank.hpp>
#include <seqan3/core/range/type_traits.hpp>
#include <seqan3/search/fm_index/concept.hpp>
#include <seqan3/search/fm_index/detail/fm_index_cursor.hpp>
#include <seqan3/search/fm_index/fm_index_cursor.hpp>
 
namespace seqan3::detail
{
struct fm_index_validator
{
    template <semialphabet alphabet_t, text_layout text_layout_mode_, std::ranges::range text_t>
    static void validate(text_t && text)
    {
        if constexpr (text_layout_mode_ == text_layout::single)
        {
            static_assert(std::ranges::bidirectional_range<text_t>, "The text must model bidirectional_range.");
            static_assert(std::convertible_to<range_innermost_value_t<text_t>, alphabet_t>,
                          "The alphabet of the text collection must be convertible to the alphabet of the index.");
            static_assert(range_dimension_v<text_t> == 1, "The input cannot be a text collection.");
 
            if (std::ranges::empty(text))
                throw std::invalid_argument("The text to index cannot be empty.");
        }
        else
        {
            static_assert(std::ranges::bidirectional_range<text_t>,
                          "The text collection must model bidirectional_range.");
            static_assert(std::ranges::bidirectional_range<std::ranges::range_reference_t<text_t>>,
                          "The elements of the text collection must model bidirectional_range.");
            static_assert(std::convertible_to<range_innermost_value_t<text_t>, alphabet_t>,
                          "The alphabet of the text collection must be convertible to the alphabet of the index.");
            static_assert(range_dimension_v<text_t> == 2, "The input must be a text collection.");
 
            if (std::ranges::empty(text))
                throw std::invalid_argument("The text collection to index cannot be empty.");
        }
        static_assert(alphabet_size<range_innermost_value_t<text_t>> <= 256, "The alphabet is too big.");
    }
};
} // namespace seqan3::detail
 
namespace seqan3
{
 
// forward declarations
template <typename index_t>
class fm_index_cursor;
 
template <typename index_t>
class bi_fm_index_cursor;
 
namespace detail
{
template <semialphabet alphabet_t, text_layout text_layout_mode_, detail::sdsl_index sdsl_index_type_>
class reverse_fm_index;
}
 
using sdsl_wt_index_type =
    sdsl::csa_wt<sdsl::wt_blcd<sdsl::bit_vector, // Wavelet tree type
                               sdsl::rank_support_v<>,
                               sdsl::select_support_scan<>,
                               sdsl::select_support_scan<0>>,
                 16,                           // Sampling rate of the suffix array
                 10'000'000,                   // Sampling rate of the inverse suffix array
                 sdsl::sa_order_sa_sampling<>, // How to sample positions in the suffix array (text VS SA sampling)
                 sdsl::isa_sampling<>,         // How to sample positons in the inverse suffix array
                 sdsl::plain_byte_alphabet>;   // How to represent the alphabet
 
using default_sdsl_index_type = sdsl_wt_index_type;
 
template <semialphabet alphabet_t,
          text_layout text_layout_mode_,
          detail::sdsl_index sdsl_index_type_ = default_sdsl_index_type>
class fm_index
{
private:
    using sdsl_index_type = sdsl_index_type_;
    using sdsl_char_type = typename sdsl_index_type::alphabet_type::char_type;
    using sdsl_sigma_type = typename sdsl_index_type::alphabet_type::sigma_type;
 
    friend class detail::reverse_fm_index<alphabet_t, text_layout_mode_, sdsl_index_type_>;
 
    sdsl_index_type index;
 
    sdsl::sd_vector<> text_begin;
    sdsl::select_support_sd<1> text_begin_ss;
    sdsl::rank_support_sd<1> text_begin_rs;
 
    template <typename output_it_t, typename sequence_t>
    static output_it_t copy_sequence_ranks_shifted_by_one(output_it_t output_it, sequence_t && sequence)
    {
        constexpr size_t sigma = alphabet_size<alphabet_t>;
        constexpr size_t max_sigma = text_layout_mode_ == text_layout::single ? 256u : 255u;
 
        constexpr auto warn_if_rank_out_of_range = [](uint8_t const rank)
        {
            if (rank >= max_sigma - 1) // same as rank + 1 >= max_sigma but without overflow
                throw std::out_of_range("The input text cannot be indexed, because for full"
                                        "character alphabets the last one/two values are reserved"
                                        "(single sequence/collection).");
        };
 
        return std::ranges::transform(sequence,
                                      output_it,
                                      [&warn_if_rank_out_of_range](auto const & chr)
                                      {
                                          uint8_t const rank = seqan3::to_rank(chr);
                                          if constexpr (sigma >= max_sigma)
                                              warn_if_rank_out_of_range(rank);
                                          else
                                              (void)warn_if_rank_out_of_range;
                                          return rank + 1;
                                      })
            .out;
    }
 
    template <std::ranges::range text_t>
        requires (text_layout_mode_ == text_layout::single)
    void construct(text_t && text)
    {
        detail::fm_index_validator::validate<alphabet_t, text_layout_mode_>(text);
 
        // TODO:
        // * check what happens in sdsl when constructed twice!
        // * choose between in-memory/external and construction algorithms
        // * sdsl construction currently only works for int_vector, std::string and char *, not ranges in general
        // uint8_t largest_char = 0;
        sdsl::int_vector<8> tmp_text(std::ranges::distance(text));
 
        // copy ranks into tmp_text
        copy_sequence_ranks_shifted_by_one(std::ranges::begin(tmp_text), text | std::views::reverse);
 
        sdsl::construct_im(index, tmp_text, 0);
 
        // TODO: would be nice but doesn't work since it's private and the public member references are const
        // index.m_C.resize(largest_char);
        // index.m_C.shrink_to_fit();
        // index.m_sigma = largest_char;
    }
 
    template <std::ranges::range text_t>
        requires (text_layout_mode_ == text_layout::collection)
    void construct(text_t && text, bool reverse = false)
    {
        detail::fm_index_validator::validate<alphabet_t, text_layout_mode_>(text);
 
        std::vector<size_t> text_sizes;
 
        for (auto && t : text)
            text_sizes.push_back(std::ranges::distance(t));
 
        size_t const number_of_texts{text_sizes.size()};
 
        // text size including delimiters
        size_t const text_size = std::accumulate(text_sizes.begin(), text_sizes.end(), number_of_texts);
 
        if (number_of_texts == text_size)
            throw std::invalid_argument("A text collection that only contains empty texts cannot be indexed.");
 
        constexpr auto sigma = alphabet_size<alphabet_t>;
 
        // Instead of creating a bitvector of size `text_size`, setting the bits to 1 and then compressing it, we can
        // use the `sd_vector_builder(text_size, number_of_ones)` because we know the parameters and the 1s we want to
        // set are in a strictly increasing order. This inplace construction of the compressed vector saves memory.
        sdsl::sd_vector_builder builder(text_size, number_of_texts);
        size_t prefix_sum{0};
 
        for (auto && size : text_sizes)
        {
            builder.set(prefix_sum);
            prefix_sum += size + 1;
        }
 
        text_begin = sdsl::sd_vector<>(builder);
        text_begin_ss = sdsl::select_support_sd<1>(&text_begin);
        text_begin_rs = sdsl::rank_support_sd<1>(&text_begin);
 
        // last text in collection needs no delimiter if we have more than one text in the collection
        sdsl::int_vector<8> tmp_text(text_size - (number_of_texts > 1));
 
        constexpr uint8_t delimiter = sigma >= 255 ? 255 : sigma + 1;
 
        auto copy_join_with = [](auto output_it, auto && collection)
        {
            // this is basically std::views::join() with a delimiter
            auto collection_it = std::ranges::begin(collection);
            auto const collection_sentinel = std::ranges::end(collection);
            if (collection_it == collection_sentinel)
                return;
 
            output_it = copy_sequence_ranks_shifted_by_one(output_it, *collection_it);
            ++collection_it;
 
            for (; collection_it != collection_sentinel; ++collection_it)
            {
                *output_it = delimiter;
                ++output_it;
                output_it = copy_sequence_ranks_shifted_by_one(output_it, *collection_it);
            }
        };
 
        // copy ranks into tmp_text
        copy_join_with(std::ranges::begin(tmp_text), text);
 
        if (!reverse)
        {
            // we need at least one delimiter
            if (number_of_texts == 1)
                tmp_text.back() = delimiter;
 
            std::ranges::reverse(tmp_text);
        }
        else
        {
            // If only one text is in the text collection, we still need one delimiter at the end to be able to
            // conduct rank and select queries when locating hits in the index.
            // Also, tmp_text looks like [text|0], but after reversing we need [txet|0] to be able to add the delimiter.
            if (number_of_texts == 1)
            {
                std::ranges::reverse(tmp_text.begin(), tmp_text.end() - 1);
                tmp_text.back() = delimiter;
            }
            else
            {
                std::ranges::reverse(tmp_text);
            }
        }
 
        sdsl::construct_im(index, tmp_text, 0);
    }
 
public:
    static constexpr text_layout text_layout_mode = text_layout_mode_;
 
    using alphabet_type = alphabet_t;
    using size_type = typename sdsl_index_type::size_type;
    using cursor_type = fm_index_cursor<fm_index>;
 
    template <typename bi_fm_index_t>
    friend class bi_fm_index_cursor;
 
    template <typename fm_index_t>
    friend class fm_index_cursor;
 
    template <typename fm_index_t>
    friend struct detail::fm_index_cursor_node;
 
    fm_index() = default; 
 
    fm_index(fm_index const & rhs) :
        index{rhs.index},
        text_begin{rhs.text_begin},
        text_begin_ss{rhs.text_begin_ss},
        text_begin_rs{rhs.text_begin_rs}
    {
        text_begin_ss.set_vector(&text_begin);
        text_begin_rs.set_vector(&text_begin);
    }
 
    fm_index(fm_index && rhs) :
        index{std::move(rhs.index)},
        text_begin{std::move(rhs.text_begin)},
        text_begin_ss{std::move(rhs.text_begin_ss)},
        text_begin_rs{std::move(rhs.text_begin_rs)}
    {
        text_begin_ss.set_vector(&text_begin);
        text_begin_rs.set_vector(&text_begin);
    }
 
    fm_index & operator=(fm_index rhs)
    {
        index = std::move(rhs.index);
        text_begin = std::move(rhs.text_begin);
        text_begin_ss = std::move(rhs.text_begin_ss);
        text_begin_rs = std::move(rhs.text_begin_rs);
 
        text_begin_ss.set_vector(&text_begin);
        text_begin_rs.set_vector(&text_begin);
 
        return *this;
    }
 
    ~fm_index() = default; 
 
    template <std::ranges::bidirectional_range text_t>
    explicit fm_index(text_t && text)
    {
        construct(std::forward<text_t>(text));
    }
 
    size_type size() const noexcept
    {
        return index.size();
    }
 
    bool empty() const noexcept
    {
        return size() == 0;
    }
 
    bool operator==(fm_index const & rhs) const noexcept
    {
        // (void) rhs;
        return (index == rhs.index) && (text_begin == rhs.text_begin);
    }
 
    bool operator!=(fm_index const & rhs) const noexcept
    {
        return !(*this == rhs);
    }
 
    cursor_type cursor() const noexcept
    {
        return {*this};
    }
 
    template <cereal_archive archive_t>
    void CEREAL_SERIALIZE_FUNCTION_NAME(archive_t & archive)
    {
        archive(index);
        archive(text_begin);
        archive(text_begin_ss);
        text_begin_ss.set_vector(&text_begin);
        archive(text_begin_rs);
        text_begin_rs.set_vector(&text_begin);
 
        auto sigma = alphabet_size<alphabet_t>;
        archive(sigma);
        if (sigma != alphabet_size<alphabet_t>)
        {
            throw std::logic_error{"The fm_index was built over an alphabet of size " + std::to_string(sigma)
                                   + " but it is being read into an fm_index with an alphabet of size "
                                   + std::to_string(alphabet_size<alphabet_t>) + "."};
        }
 
        bool tmp = text_layout_mode;
        archive(tmp);
        if (tmp != text_layout_mode)
        {
            throw std::logic_error{std::string{"The fm_index was built over a "}
                                   + (tmp ? "text collection" : "single text")
                                   + " but it is being read into an fm_index expecting a "
                                   + (text_layout_mode ? "text collection." : "single text.")};
        }
    }
};
 
template <std::ranges::range text_t>
fm_index(text_t &&) -> fm_index<range_innermost_value_t<text_t>, text_layout{range_dimension_v<text_t> != 1}>;
} // namespace seqan3
 
namespace seqan3::detail
{
 
template <semialphabet alphabet_t,
          text_layout text_layout_mode,
          detail::sdsl_index sdsl_index_type = default_sdsl_index_type>
class reverse_fm_index : public fm_index<alphabet_t, text_layout_mode, sdsl_index_type>
{
private:
    template <std::ranges::range text_t>
    void construct_(text_t && text)
    {
        if constexpr (text_layout_mode == text_layout::single)
        {
            auto reverse_text = text | std::views::reverse;
            this->construct(reverse_text);
        }
        else
        {
            auto reverse_text = text | views::deep{std::views::reverse} | std::views::reverse;
            this->construct(reverse_text, true);
        }
    }
 
public:
    using fm_index<alphabet_t, text_layout_mode, sdsl_index_type>::fm_index;
 
    template <std::ranges::bidirectional_range text_t>
    explicit reverse_fm_index(text_t && text)
    {
        construct_(std::forward<text_t>(text));
    }
};
 
} // namespace seqan3::detail