bi_fm_index_cursor.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 <array>
#include <ranges>
 
#include <sdsl/suffix_trees.hpp>
 
#include <seqan3/alphabet/adaptation/char.hpp>
#include <seqan3/alphabet/adaptation/uint.hpp>
#include <seqan3/alphabet/concept.hpp>
#include <seqan3/core/range/type_traits.hpp>
#include <seqan3/search/fm_index/fm_index.hpp>
#include <seqan3/search/fm_index/fm_index_cursor.hpp>
#include <seqan3/utility/views/slice.hpp>
 
namespace seqan3
{
 
template <typename index_t>
class bi_fm_index_cursor
{
public:
    using index_type = index_t;
 
    using size_type = typename index_type::size_type;
 
    using fwd_cursor = fm_index_cursor<fm_index<typename index_type::alphabet_type,
                                                index_type::text_layout_mode,
                                                typename index_type::sdsl_index_type>>;
 
private:
    using sdsl_char_type = typename index_type::sdsl_char_type;
    using sdsl_index_type = typename index_t::sdsl_index_type;
    using sdsl_sigma_type = typename index_type::sdsl_sigma_type;
    using index_alphabet_type = typename index_t::alphabet_type;
    using locate_result_value_type = std::pair<size_type, size_type>;
    using locate_result_type = std::vector<locate_result_value_type>;
 
    index_type const * index{nullptr};
 
    size_type fwd_lb{};
    size_type fwd_rb{};
    size_type rev_lb{};
    size_type rev_rb{};
    //\}
 
    sdsl_sigma_type sigma{};
 
    // parent_* and _last_char only have to be stored for the (unidirectional) cursor that has been used last for
    // extend_right() or cycle_back() resp. extend_left() or cycle_front(), (i.e. either fwd or rev). Thus there is no
    // need to store it twice. Once the cursor is switched, the information becomes invalid anyway.
 
    size_type parent_lb{};
    size_type parent_rb{};
    sdsl_char_type _last_char{};
    //\}
 
    size_type depth{}; // equal for both cursors. only stored once
 
    // supports assertions to check whether cycle_back() resp. cycle_front() is called on the same direction as the last
    // extend_right([...]) resp. extend_left([...])
#ifndef NDEBUG
    //        cycle_back() and cycle_front()
    bool fwd_cursor_last_used = false;
#endif
 
    size_type offset() const noexcept
    {
        assert(index->size() > query_length());
        return index->size() - query_length() - 1; // since the string is reversed during construction
    }
 
    template <typename csa_t>
        requires (std::same_as<csa_t, typename index_type::sdsl_index_type>
                  || std::same_as<csa_t, typename index_type::rev_sdsl_index_type>) bool
    bidirectional_search(csa_t const & csa,
                         sdsl_char_type const c,
                         size_type & l_fwd,
                         size_type & r_fwd,
                         size_type & l_bwd,
                         size_type & r_bwd) const noexcept
    {
        assert((l_fwd <= r_fwd) && (r_fwd < csa.size()));
        assert(r_fwd + 1 >= l_fwd);
        assert(r_bwd + 1 - l_bwd == r_fwd + 1 - l_fwd);
 
        size_type _l_fwd, _r_fwd, _l_bwd, _r_bwd;
 
        size_type cc = c;
        if constexpr (!std::same_as<index_alphabet_type, sdsl::plain_byte_alphabet>)
        {
            cc = csa.char2comp[c];
            if (cc == 0 && c > 0) // [[unlikely]]
                return false;
        }
 
        size_type const c_begin = csa.C[cc];
        if (r_fwd + 1 - l_fwd == csa.size()) // [[unlikely]]
        {
            _l_fwd = c_begin;
            _l_bwd = c_begin;
            _r_fwd = csa.C[cc + 1] - 1;
            _r_bwd = _r_fwd;
            // if we use not the plain_byte_alphabet, we could return always return true here
        }
        else
        {
            auto const r_s_b = csa.wavelet_tree.lex_count(l_fwd, r_fwd + 1, c);
            size_type const rank_l = std::get<0>(r_s_b);
            size_type const s = std::get<1>(r_s_b), b = std::get<2>(r_s_b);
            size_type const rank_r = r_fwd - l_fwd - s - b + rank_l;
            _l_fwd = c_begin + rank_l;
            _r_fwd = c_begin + rank_r;
            _l_bwd = l_bwd + s;
            _r_bwd = r_bwd - b;
        }
 
        if (_r_fwd >= _l_fwd)
        {
            l_fwd = _l_fwd;
            r_fwd = _r_fwd;
            l_bwd = _l_bwd;
            r_bwd = _r_bwd;
            assert(r_fwd + 1 >= l_fwd);
            assert(r_bwd + 1 - l_bwd == r_fwd + 1 - l_fwd);
            return true;
        }
        return false;
    }
 
    template <typename csa_t>
        requires (std::same_as<csa_t, typename index_type::sdsl_index_type>
                  || std::same_as<csa_t, typename index_type::rev_sdsl_index_type>) bool
    bidirectional_search_cycle(csa_t const & csa,
                               sdsl_char_type const c,
                               size_type const l_parent,
                               size_type const r_parent,
                               size_type & l_fwd,
                               size_type & r_fwd,
                               size_type & l_bwd,
                               size_type & r_bwd) const noexcept
    {
        assert((l_parent <= r_parent) && (r_parent < csa.size()));
 
        size_type c_begin;
        if constexpr (std::same_as<index_alphabet_type, sdsl::plain_byte_alphabet>)
            c_begin = csa.C[c]; // TODO: check whether this can be removed
        else
            c_begin = csa.C[csa.char2comp[c]];
 
        auto const r_s_b = csa.wavelet_tree.lex_count(l_parent, r_parent + 1, c);
        size_type const s = std::get<1>(r_s_b), b = std::get<2>(r_s_b), rank_l = std::get<0>(r_s_b),
                        rank_r = r_parent - l_parent - s - b + rank_l;
 
        size_type const _l_fwd = c_begin + rank_l;
        size_type const _r_fwd = c_begin + rank_r;
        size_type const _l_bwd = r_bwd + 1;
        size_type const _r_bwd = r_bwd + 1 + rank_r - rank_l;
 
        if (_r_fwd >= _l_fwd)
        {
            l_fwd = _l_fwd;
            r_fwd = _r_fwd;
            l_bwd = _l_bwd;
            r_bwd = _r_bwd;
            assert(r_fwd + 1 >= l_fwd);
            assert(r_bwd + 1 - l_bwd == r_fwd + 1 - l_fwd);
            return true;
        }
        return false;
    }
 
public:
    //        Default construction is necessary to make this class semi-regular and e.g., to allow construction of
    //        std::array of cursors.
    bi_fm_index_cursor() noexcept = default;                                       
    bi_fm_index_cursor(bi_fm_index_cursor const &) noexcept = default;             
    bi_fm_index_cursor & operator=(bi_fm_index_cursor const &) noexcept = default; 
    bi_fm_index_cursor(bi_fm_index_cursor &&) noexcept = default;                  
    bi_fm_index_cursor & operator=(bi_fm_index_cursor &&) noexcept = default;      
    ~bi_fm_index_cursor() = default;                                               
 
    bi_fm_index_cursor(index_t const & _index) noexcept :
        index(&_index),
        fwd_lb(0),
        fwd_rb(_index.size() - 1),
        rev_lb(0),
        rev_rb(_index.size() - 1),
        sigma(_index.fwd_fm.index.sigma - index_t::text_layout_mode),
        depth(0)
    {}
    //\}
 
    bool operator==(bi_fm_index_cursor const & rhs) const noexcept
    {
        assert(index != nullptr);
        // equal SA interval implies equal parent node information (or both are root nodes)
        assert(!(fwd_lb == rhs.fwd_lb && fwd_rb == rhs.fwd_rb && depth == rhs.depth) || (depth == 0)
               || (parent_lb == rhs.parent_lb && parent_rb == rhs.parent_rb && _last_char == rhs._last_char));
 
        return std::tie(fwd_lb, fwd_rb, depth) == std::tie(rhs.fwd_lb, rhs.fwd_rb, rhs.depth);
    }
 
    bool operator!=(bi_fm_index_cursor const & rhs) const noexcept
    {
        assert(index != nullptr);
 
        return !(*this == rhs);
    }
 
    bool extend_right() noexcept
    {
#ifndef NDEBUG
        fwd_cursor_last_used = true;
#endif
 
        assert(index != nullptr);
 
        size_type new_parent_lb = fwd_lb, new_parent_rb = fwd_rb;
 
        sdsl_char_type c = 1; // NOTE: start with 0 or 1 depending on implicit_sentintel
        while (c < sigma
               && !bidirectional_search(index->fwd_fm.index,
                                        index->fwd_fm.index.comp2char[c],
                                        fwd_lb,
                                        fwd_rb,
                                        rev_lb,
                                        rev_rb))
        {
            ++c;
        }
 
        if (c != sigma)
        {
            parent_lb = new_parent_lb;
            parent_rb = new_parent_rb;
 
            _last_char = c;
            ++depth;
 
            return true;
        }
        return false;
    }
 
    bool extend_left() noexcept
    {
#ifndef NDEBUG
        fwd_cursor_last_used = false;
#endif
 
        assert(index != nullptr);
 
        size_type new_parent_lb = rev_lb, new_parent_rb = rev_rb;
 
        sdsl_char_type c = 1; // NOTE: start with 0 or 1 depending on implicit_sentintel
        while (c < sigma
               && !bidirectional_search(index->rev_fm.index,
                                        index->rev_fm.index.comp2char[c],
                                        rev_lb,
                                        rev_rb,
                                        fwd_lb,
                                        fwd_rb))
        {
            ++c;
        }
 
        if (c != sigma)
        {
            parent_lb = new_parent_lb;
            parent_rb = new_parent_rb;
 
            _last_char = c;
            ++depth;
 
            return true;
        }
        return false;
    }
 
    template <typename char_t>
        requires std::convertible_to<char_t, index_alphabet_type> bool
    extend_right(char_t const c) noexcept
    {
#ifndef NDEBUG
        fwd_cursor_last_used = true;
#endif
 
        assert(index != nullptr);
        // The rank cannot exceed 255 for single text and 254 for text collections as they are reserved as sentinels
        // for the indexed text.
        assert(seqan3::to_rank(static_cast<index_alphabet_type>(c))
               < ((index_type::text_layout_mode == text_layout::single) ? 255 : 254));
 
        size_type new_parent_lb = fwd_lb, new_parent_rb = fwd_rb;
 
        auto c_char = seqan3::to_rank(static_cast<index_alphabet_type>(c)) + 1;
        if (bidirectional_search(index->fwd_fm.index, c_char, fwd_lb, fwd_rb, rev_lb, rev_rb))
        {
            parent_lb = new_parent_lb;
            parent_rb = new_parent_rb;
 
            _last_char = c_char;
            ++depth;
 
            return true;
        }
        return false;
    }
 
    template <typename char_type>
        requires seqan3::detail::is_char_adaptation_v<char_type> bool
    extend_right(char_type const * cstring) noexcept
    {
        return extend_right(std::basic_string_view<char_type>{cstring});
    }
 
    template <typename char_t>
        requires std::convertible_to<char_t, index_alphabet_type> bool
    extend_left(char_t const c) noexcept
    {
#ifndef NDEBUG
        fwd_cursor_last_used = false;
#endif
 
        assert(index != nullptr);
        // The rank cannot exceed 255 for single text and 254 for text collections as they are reserved as sentinels
        // for the indexed text.
        assert(seqan3::to_rank(static_cast<index_alphabet_type>(c))
               < ((index_type::text_layout_mode == text_layout::single) ? 255 : 254));
 
        size_type new_parent_lb = rev_lb, new_parent_rb = rev_rb;
 
        auto c_char = seqan3::to_rank(static_cast<index_alphabet_type>(c)) + 1;
        if (bidirectional_search(index->rev_fm.index, c_char, rev_lb, rev_rb, fwd_lb, fwd_rb))
        {
            parent_lb = new_parent_lb;
            parent_rb = new_parent_rb;
 
            _last_char = c_char;
            ++depth;
 
            return true;
        }
        return false;
    }
 
    template <typename char_type>
        requires seqan3::detail::is_char_adaptation_v<char_type> bool
    extend_left(char_type const * cstring) noexcept
    {
        return extend_left(std::basic_string_view<char_type>{cstring});
    }
 
    template <std::ranges::range seq_t>
    bool extend_right(seq_t && seq) noexcept
    {
        static_assert(std::ranges::forward_range<seq_t>, "The query must model forward_range.");
        static_assert(std::convertible_to<range_innermost_value_t<seq_t>, index_alphabet_type>,
                      "The alphabet of the sequence must be convertible to the alphabet of the index.");
 
        assert(index != nullptr);
 
        auto first = std::ranges::begin(seq);
        auto last = std::ranges::end(seq);
 
#ifndef NDEBUG
        fwd_cursor_last_used = (first != last); // only if seq was not empty
#endif
 
        size_type _fwd_lb = fwd_lb, _fwd_rb = fwd_rb, _rev_lb = rev_lb, _rev_rb = rev_rb;
        size_type new_parent_lb = parent_lb, new_parent_rb = parent_rb;
        sdsl_char_type c = _last_char;
        size_t len{0};
 
        for (auto it = first; it != last; ++len, ++it)
        {
            // The rank cannot exceed 255 for single text and 254 for text collections as they are reserved as sentinels
            // for the indexed text.
            assert(seqan3::to_rank(static_cast<index_alphabet_type>(*it))
                   < ((index_type::text_layout_mode == text_layout::single) ? 255 : 254));
 
            c = seqan3::to_rank(static_cast<index_alphabet_type>(*it)) + 1;
 
            new_parent_lb = _fwd_lb;
            new_parent_rb = _fwd_rb;
            if (!bidirectional_search(index->fwd_fm.index, c, _fwd_lb, _fwd_rb, _rev_lb, _rev_rb))
                return false;
        }
 
        fwd_lb = _fwd_lb;
        fwd_rb = _fwd_rb;
        rev_lb = _rev_lb;
        rev_rb = _rev_rb;
 
        parent_lb = new_parent_lb;
        parent_rb = new_parent_rb;
 
        _last_char = c;
        depth += len;
 
        return true;
    }
 
    template <std::ranges::range seq_t>
    bool extend_left(seq_t && seq) noexcept
    {
        static_assert(std::ranges::bidirectional_range<seq_t>, "The query must model bidirectional_range.");
        static_assert(std::convertible_to<range_innermost_value_t<seq_t>, index_alphabet_type>,
                      "The alphabet of the sequence must be convertible to the alphabet of the index.");
        assert(index != nullptr);
 
        auto rev_seq = std::views::reverse(seq);
        auto first = std::ranges::begin(rev_seq);
        auto last = std::ranges::end(rev_seq);
 
#ifndef NDEBUG
        if (first != last) // only if seq was not empty
            fwd_cursor_last_used = false;
#endif
 
        size_type _fwd_lb = fwd_lb, _fwd_rb = fwd_rb, _rev_lb = rev_lb, _rev_rb = rev_rb;
        size_type new_parent_lb = parent_lb, new_parent_rb = parent_rb;
        sdsl_char_type c = _last_char;
        size_t len{0};
 
        for (auto it = first; it != last; ++len, ++it)
        {
            // The rank cannot exceed 255 for single text and 254 for text collections as they are reserved as sentinels
            // for the indexed text.
            assert(seqan3::to_rank(static_cast<index_alphabet_type>(*it))
                   < ((index_type::text_layout_mode == text_layout::single) ? 255 : 254));
 
            c = seqan3::to_rank(static_cast<index_alphabet_type>(*it)) + 1;
 
            new_parent_lb = _rev_lb;
            new_parent_rb = _rev_rb;
            if (!bidirectional_search(index->rev_fm.index, c, _rev_lb, _rev_rb, _fwd_lb, _fwd_rb))
                return false;
        }
 
        fwd_lb = _fwd_lb;
        fwd_rb = _fwd_rb;
        rev_lb = _rev_lb;
        rev_rb = _rev_rb;
 
        parent_lb = new_parent_lb;
        parent_rb = new_parent_rb;
        _last_char = c;
        depth += len;
 
        return true;
    }
 
    bool cycle_back() noexcept
    {
#ifndef NDEBUG
        // cycle_back() can only be used if the last extension was to the right.
        assert(fwd_cursor_last_used);
#endif
 
        assert(index != nullptr && query_length() > 0);
 
        sdsl_char_type c = _last_char + 1;
 
        while (c < sigma
               && !bidirectional_search_cycle(index->fwd_fm.index,
                                              index->fwd_fm.index.comp2char[c],
                                              parent_lb,
                                              parent_rb,
                                              fwd_lb,
                                              fwd_rb,
                                              rev_lb,
                                              rev_rb))
        {
            ++c;
        }
 
        if (c != sigma)
        {
            _last_char = c;
 
            return true;
        }
        return false;
    }
 
    bool cycle_front() noexcept
    {
#ifndef NDEBUG
        // cycle_front() can only be used if the last extension was to the left.
        assert(!fwd_cursor_last_used);
#endif
 
        assert(index != nullptr && query_length() > 0);
 
        sdsl_char_type c = _last_char + 1;
        while (c < sigma
               && !bidirectional_search_cycle(index->rev_fm.index,
                                              index->rev_fm.index.comp2char[c],
                                              parent_lb,
                                              parent_rb,
                                              rev_lb,
                                              rev_rb,
                                              fwd_lb,
                                              fwd_rb))
        {
            ++c;
        }
 
        if (c != sigma)
        {
            _last_char = c;
 
            return true;
        }
        return false;
    }
 
    size_type last_rank() noexcept
    {
        assert(index != nullptr && query_length() > 0);
 
        return index->fwd_fm.index.comp2char[_last_char] - 1; // text is not allowed to contain ranks of 0
    }
 
    size_type query_length() const noexcept
    {
        assert(index != nullptr);
        // depth == 0 -> root node
        assert(depth != 0 || (fwd_lb == rev_lb && fwd_rb == rev_rb && fwd_lb == 0 && fwd_rb == index->size() - 1));
 
        return depth;
    }
 
    fwd_cursor to_fwd_cursor() const noexcept
    {
        assert(index != nullptr);
 
        fwd_cursor cur{index->fwd_fm};
        cur.parent_lb = parent_lb;
        cur.parent_rb = parent_rb;
        cur.node = {fwd_lb, fwd_rb, depth, _last_char};
 
#ifndef NDEBUG
        if (!fwd_cursor_last_used)
        {
            // invalidate parent interval
            cur.parent_lb = 1;
            cur.parent_rb = 0;
        }
#endif
 
        return cur;
    }
 
    template <std::ranges::range text_t>
    auto path_label(text_t && text) const noexcept
        requires (index_t::text_layout_mode == text_layout::single)
    {
        static_assert(std::ranges::input_range<text_t>, "The text must model input_range.");
        static_assert(range_dimension_v<text_t> == 1, "The input cannot be a text collection.");
        static_assert(std::same_as<range_innermost_value_t<text_t>, index_alphabet_type>,
                      "The alphabet types of the given text and index differ.");
        assert(index != nullptr);
 
        size_type const query_begin = offset() - index->fwd_fm.index[fwd_lb];
        return text | views::slice(query_begin, query_begin + query_length());
    }
 
    template <std::ranges::range text_t>
    auto path_label(text_t && text) const noexcept
        requires (index_t::text_layout_mode == text_layout::collection)
    {
        static_assert(std::ranges::input_range<text_t>, "The text collection must model input_range.");
        static_assert(range_dimension_v<text_t> == 2, "The input must be a text collection.");
        static_assert(std::same_as<range_innermost_value_t<text_t>, index_alphabet_type>,
                      "The alphabet types of the given text and index differ.");
        assert(index != nullptr);
 
        // Position of query in concatenated text.
        size_type const location = offset() - index->fwd_fm.index[fwd_lb];
 
        // The rank represents the number of start positions of the individual sequences/texts in the collection
        // before position `location + 1` and thereby also the number of delimiters.
        size_type const rank = index->fwd_fm.text_begin_rs.rank(location + 1);
        assert(rank > 0);
        size_type const text_id = rank - 1;
 
        // The start location of the `text_id`-th text in the sequence (position of the `rank`-th 1 in the bitvector).
        size_type const start_location = index->fwd_fm.text_begin_ss.select(rank);
        // Substract lengths of previous sequences.
        size_type const query_begin = location - start_location;
 
        // Take subtext, slice query out of it
        return text[text_id] | views::slice(query_begin, query_begin + query_length());
    }
 
    size_type count() const noexcept
    {
        assert(index != nullptr && (1 + fwd_rb - fwd_lb == 1 + rev_rb - rev_lb));
 
        return 1 + fwd_rb - fwd_lb;
    }
 
    locate_result_type locate() const
        requires (index_t::text_layout_mode == text_layout::single)
    {
        assert(index != nullptr);
 
        locate_result_type occ{};
        occ.reserve(count());
        for (size_type i = 0; i < count(); ++i)
        {
            occ.emplace_back(0, offset() - index->fwd_fm.index[fwd_lb + i]);
        }
        return occ;
    }
 
    std::vector<std::pair<size_type, size_type>> locate() const
        requires (index_t::text_layout_mode == text_layout::collection)
    {
        assert(index != nullptr);
 
        std::vector<std::pair<size_type, size_type>> occ;
        occ.reserve(count());
        for (size_type i = 0; i < count(); ++i)
        {
            size_type loc = offset() - index->fwd_fm.index[fwd_lb + i];
            size_type sequence_rank = index->fwd_fm.text_begin_rs.rank(loc + 1);
            size_type sequence_position = loc - index->fwd_fm.text_begin_ss.select(sequence_rank);
            occ.emplace_back(sequence_rank - 1, sequence_position);
        }
        return occ;
    }
 
    auto lazy_locate() const
        requires (index_t::text_layout_mode == text_layout::single)
    {
        assert(index != nullptr);
 
        return std::views::iota(fwd_lb, fwd_lb + count())
             | std::views::transform(
                   [*this, _offset = offset()](auto sa_pos)
                   {
                       return locate_result_value_type{0u, _offset - index->fwd_fm.index[sa_pos]};
                   });
    }
 
    auto lazy_locate() const
        requires (index_t::text_layout_mode == text_layout::collection)
    {
        assert(index != nullptr);
 
        return std::views::iota(fwd_lb, fwd_lb + count())
             | std::views::transform(
                   [*this, _offset = offset()](auto sa_pos)
                   {
                       auto loc = _offset - index->fwd_fm.index[sa_pos];
                       size_type sequence_rank = index->fwd_fm.text_begin_rs.rank(loc + 1);
                       size_type sequence_position = loc - index->fwd_fm.text_begin_ss.select(sequence_rank);
                       return locate_result_value_type{sequence_rank - 1, sequence_position};
                   });
    }
 
    template <cereal_archive archive_t>
    void CEREAL_SERIALIZE_FUNCTION_NAME(archive_t & archive)
    {
        archive(fwd_lb);
        archive(fwd_rb);
        archive(rev_lb);
        archive(rev_rb);
        archive(sigma);
        archive(parent_lb);
        archive(parent_rb);
        archive(_last_char);
        archive(depth);
    }
};
 
} // namespace seqan3