wavelet_tree.h

Go to the documentation of this file.

/*
 * nvbio
 * Copyright (c) 2011-2014, NVIDIA CORPORATION. All rights reserved.
 * 
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *    * Redistributions of source code must retain the above copyright
 *      notice, this list of conditions and the following disclaimer.
 *    * Redistributions in binary form must reproduce the above copyright
 *      notice, this list of conditions and the following disclaimer in the
 *      documentation and/or other materials provided with the distribution.
 *    * Neither the name of the NVIDIA CORPORATION nor the
 *      names of its contributors may be used to endorse or promote products
 *      derived from this software without specific prior written permission.
 * 
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
 * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
 * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#pragma once

#include <nvbio/strings/string.h>
#include <nvbio/basic/packed_vector.h>
#include <nvbio/basic/primitives.h>
#include <nvbio/basic/cuda/sort.h>
#include <thrust/sort.h>
#include <algorithm>
#include <stack>

namespace nvbio {




template <typename BitStreamIterator, typename IndexIterator, typename SymbolType = uint8>
struct WaveletTree
{
    // define the system tag
    typedef typename iterator_system<BitStreamIterator>::type       system_tag;
    typedef typename stream_traits<BitStreamIterator>::index_type   index_type;
    typedef SymbolType                                              symbol_type;
    typedef SymbolType                                              value_type;
    typedef BitStreamIterator                                       bit_iterator;
    typedef IndexIterator                                           index_iterator;
    typedef WaveletTree<BitStreamIterator,IndexIterator>            text_type;   // the text is the wavelet tree itself

    typedef typename vector_type<index_type,2>::type                range_type;
    typedef null_type                                               vector_type; // unsupported, would require knowing alphabet size

    NVBIO_HOST_DEVICE
    WaveletTree(
        const uint32            _symbol_size = 0,
        const index_type        _size        = 0,
        const BitStreamIterator _bits        = BitStreamIterator(),
        const IndexIterator     _nodes       = IndexIterator(),
        const IndexIterator     _occ         = IndexIterator()) :
        m_symbol_size( _symbol_size ),
        m_size  ( _size ),
        m_bits  ( _bits ),
        m_nodes ( _nodes ),
        m_occ   ( _occ ) {}

    NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
    void resize(const uint32 _size, const uint32 _symbol_size)
    {
        m_size        = _size;
        m_symbol_size = _symbol_size;
    }

    NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
    uint32 symbol_count() const { return 1u << m_symbol_size; }

    NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
    uint32 symbol_size() const { return m_symbol_size; }

    NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
    index_type size() const { return m_size; }

    NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
    BitStreamIterator bits() const { return m_bits; }

    NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
    IndexIterator splits() const { return m_nodes; }

    NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
    IndexIterator occ() const { return m_occ; }

    NVBIO_HOST_DEVICE
    index_type rank(const uint32 l, const uint32 node, const index_type node_begin, const index_type r, const uint8 b) const;

    NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
    SymbolType operator[] (const index_type i) const;

    NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
    SymbolType operator() (const index_type i) const { return this->operator[](i); }

    NVBIO_FORCEINLINE NVBIO_HOST_DEVICE       text_type& text()       { return *this; }
    NVBIO_FORCEINLINE NVBIO_HOST_DEVICE const text_type& text() const { return *this; }

    uint32              m_symbol_size;
    index_type          m_size;
    BitStreamIterator   m_bits;
    IndexIterator       m_nodes;
    IndexIterator       m_occ;
};

template <typename SystemTag, typename IndexType = uint32, typename SymbolType = uint8>
struct WaveletTreeStorage
{
    // define the system tag
    typedef SystemTag                                       system_tag;
    typedef IndexType                                       index_type;
    typedef SymbolType                                      symbol_type;

    typedef PackedVector<system_tag,1u,true,index_type>     bit_vector_type;
    typedef nvbio::vector<system_tag,index_type>            index_vector_type;
    typedef typename bit_vector_type::iterator                    bit_iterator;
    typedef typename bit_vector_type::const_iterator        const_bit_iterator;
    typedef typename index_vector_type::iterator                  index_iterator;
    typedef typename index_vector_type::const_iterator      const_index_iterator;

    typedef WaveletTree<      bit_iterator,      index_iterator>       plain_view_type;
    typedef WaveletTree<const_bit_iterator,const_index_iterator> const_plain_view_type;

    WaveletTreeStorage() :
        m_symbol_size( 0u ),
        m_size( 0u ) {}

    void resize(const uint32 _size, const uint32 _symbol_size)
    {
        m_size        = _size;
        m_symbol_size = _symbol_size;

        const uint32 n_symbols = 1u << _symbol_size;

        m_bits.resize( m_size * m_symbol_size );
        m_nodes.resize( n_symbols );
        m_occ.resize( n_symbols + util::divide_ri( m_size * m_symbol_size, 32u ) );
    }

    NVBIO_HOST_DEVICE
    uint32 symbol_size() const { return m_symbol_size; }

    NVBIO_HOST_DEVICE
    index_type size() const { return m_size; }

    bit_iterator bits() { return m_bits.begin(); }

    index_iterator splits() { return m_nodes.begin(); }

    index_iterator occ() { return m_occ.begin(); }

    const_bit_iterator bits() const { return m_bits.begin(); }

    const_index_iterator splits() const { return m_nodes.begin(); }

    const_index_iterator occ() const { return m_occ.begin(); }

    operator plain_view_type()
    {
        return plain_view_type(
            m_symbol_size,
            m_size,
            bits(),
            splits(),
            occ() );
    }
    operator const_plain_view_type() const
    {
        return const_plain_view_type(
            m_symbol_size,
            m_size,
            bits(),
            splits(),
            occ() );
    }

    uint32              m_symbol_size;
    index_type          m_size;
    bit_vector_type     m_bits;
    index_vector_type   m_nodes;
    index_vector_type   m_occ;
};


template <typename system_tag, typename string_iterator, typename index_type, typename symbol_type>
void setup(
    const index_type                                        string_len,
    const string_iterator&                                  string,
    WaveletTreeStorage<system_tag,index_type,symbol_type>&  out_tree);

template <typename BitStreamIterator, typename IndexIterator, typename IndexType, typename SymbolType>
NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
SymbolType text(const WaveletTree<BitStreamIterator,IndexIterator,SymbolType>& tree, const IndexType i);

template <typename BitStreamIterator, typename IndexIterator, typename SymbolType>
NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
typename WaveletTree<BitStreamIterator,IndexIterator,SymbolType>::index_type
rank(
    const          WaveletTree<BitStreamIterator,IndexIterator,SymbolType>&             tree,
    const typename WaveletTree<BitStreamIterator,IndexIterator,SymbolType>::index_type  i,
    const uint32                                                                        c);

template <typename BitStreamIterator, typename IndexIterator, typename SymbolType>
NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
typename WaveletTree<BitStreamIterator,IndexIterator,SymbolType>::range_type
rank(
    const          WaveletTree<BitStreamIterator,IndexIterator,SymbolType>&             tree,
    const typename WaveletTree<BitStreamIterator,IndexIterator,SymbolType>::range_type  range,
    const uint32                                                                        c);

template <typename SystemTag, typename IndexType, typename SymbolType>
typename WaveletTreeStorage<SystemTag,IndexType,SymbolType>::plain_view_type plain_view(WaveletTreeStorage<SystemTag,IndexType,SymbolType>& tree)
{
    return tree;
}
template <typename SystemTag, typename IndexType, typename SymbolType>
typename WaveletTreeStorage<SystemTag,IndexType,SymbolType>::const_plain_view_type plain_view(const WaveletTreeStorage<SystemTag,IndexType,SymbolType>& tree)
{
    return tree;
}


} // namespace nvbio

#include <nvbio/strings/wavelet_tree_inl.h>