indexlr.hpp

#ifndef BTLLIB_INDEXLR_HPP
#define BTLLIB_INDEXLR_HPP
 
#include "btllib/bloom_filter.hpp"
#include "btllib/nthash.hpp"
#include "btllib/order_queue.hpp"
#include "btllib/seq_reader.hpp"
#include "btllib/status.hpp"
#include "btllib/util.hpp"
 
#include <algorithm>
#include <atomic>
#include <cstdlib>
#include <cstring>
#include <functional>
#include <iostream>
#include <limits>
#include <memory>
#include <string>
#include <thread>
#include <vector>
 
namespace btllib {
 
class Indexlr
{
 
public:
  /* Has to be a struct and not an enum because:
   * 1) Non-class enums are not name qualified and can collide
   * 2) class enums can't be implicitly converted into integers
   */
  struct Flag
  {
    static const unsigned NO_ID = 1;
    static const unsigned BX = 2;
    static const unsigned SEQ = 4;
    static const unsigned FILTER_IN = 8;
    static const unsigned FILTER_OUT = 16;
    static const unsigned SHORT_MODE = 32;
    static const unsigned LONG_MODE = 64;
    static const unsigned QUAL = 128;
  };
 
  bool output_id() const { return bool(~flags & Flag::NO_ID); }
  bool output_bx() const { return bool(flags & Flag::BX); }
  bool output_seq() const { return bool(flags & Flag::SEQ); }
  bool output_qual() const { return bool(flags & Flag::QUAL); }
  bool filter_in() const { return bool(flags & Flag::FILTER_IN); }
  bool filter_out() const { return bool(flags & Flag::FILTER_OUT); }
  bool short_mode() const { return bool(flags & Flag::SHORT_MODE); }
  bool long_mode() const { return bool(flags & Flag::LONG_MODE); }
 
  struct Minimizer
  {
    Minimizer() = default;
 
    Minimizer(uint64_t min_hash,
              uint64_t out_hash,
              size_t pos,
              bool forward,
              std::string seq)
      : min_hash(min_hash)
      , out_hash(out_hash)
      , pos(pos)
      , forward(forward)
      , seq(std::move(seq))
    {
    }
 
    Minimizer(uint64_t min_hash,
              uint64_t out_hash,
              size_t pos,
              bool forward,
              std::string seq,
              std::string qual)
      : min_hash(min_hash)
      , out_hash(out_hash)
      , pos(pos)
      , forward(forward)
      , seq(std::move(seq))
      , qual(std::move(qual))
    {
    }
    Minimizer(uint64_t min_hash,
              uint64_t out_hash,
              size_t pos,
              bool forward,
              std::string seq,
              std::vector<uint64_t> all_hashes)
      : min_hash(min_hash)
      , out_hash(out_hash)
      , pos(pos)
      , forward(forward)
      , seq(std::move(seq))
      , all_hashes(std::move(all_hashes))
    {
    }
 
    Minimizer(uint64_t min_hash,
              uint64_t out_hash,
              size_t pos,
              bool forward,
              std::string seq,
              std::string qual,
              std::vector<uint64_t> all_hashes)
      : min_hash(min_hash)
      , out_hash(out_hash)
      , pos(pos)
      , forward(forward)
      , seq(std::move(seq))
      , qual(std::move(qual))
      , all_hashes(std::move(all_hashes))
    {
    }
 
    uint64_t min_hash = 0, out_hash = 0;
    size_t pos = 0;
    bool forward = false;
    std::string seq;
    std::string qual;
    std::vector<uint64_t> all_hashes;
  };
 
  using HashedKmer = Minimizer;
 
  struct Record
  {
    Record() {}
 
    Record(size_t num,
           std::string id,
           std::string barcode,
           size_t readlen,
           std::vector<Minimizer> minimizers)
      : num(num)
      , id(std::move(id))
      , barcode(std::move(barcode))
      , readlen(readlen)
      , minimizers(std::move(minimizers))
    {
    }
 
    size_t num = 0;
    std::string id;
    std::string barcode;
    size_t readlen = 0;
    std::vector<Minimizer> minimizers;
 
    operator bool() const
    {
      return !id.empty() || !barcode.empty() || !minimizers.empty();
    }
  };
 
  Record read();
 
  Indexlr(std::string seqfile,
          size_t k,
          size_t w,
          unsigned flags = 0,
          unsigned threads = 5,
          bool verbose = false,
          const btllib::BloomFilter& bf1 = Indexlr::dummy_bf(),
          const btllib::BloomFilter& bf2 = Indexlr::dummy_bf());
 
  Indexlr(std::string seqfile,
          size_t k,
          size_t w,
          size_t q,
          unsigned flags = 0,
          unsigned threads = 5,
          bool verbose = false,
          const btllib::BloomFilter& bf1 = Indexlr::dummy_bf(),
          const btllib::BloomFilter& bf2 = Indexlr::dummy_bf());
 
  ~Indexlr();
 
  void close() noexcept;
 
  static const size_t MAX_SIMULTANEOUS_INDEXLRS = 256;
 
  class RecordIterator
  {
  public:
    void operator++() { record = indexlr.read(); }
    bool operator!=(const RecordIterator& i)
    {
      return bool(record) || bool(i.record);
    }
    Record operator*() { return std::move(record); }
    // For wrappers
    Record next()
    {
      auto val = operator*();
      operator++();
      return val;
    }
 
  private:
    friend Indexlr;
 
    RecordIterator(Indexlr& indexlr, bool end)
      : indexlr(indexlr)
    {
      if (!end) {
        operator++();
      }
    }
 
    Indexlr& indexlr;
    Record record;
  };
 
  RecordIterator begin() { return RecordIterator(*this, false); }
  RecordIterator end() { return RecordIterator(*this, true); }
 
private:
  static std::string extract_barcode(const std::string& id,
                                     const std::string& comment);
  static void filter_hashed_kmer(Indexlr::HashedKmer& hk,
                                 bool filter_in,
                                 bool filter_out,
                                 const BloomFilter& filter_in_bf,
                                 const BloomFilter& filter_out_bf);
 
  static void filter_kmer_qual(Indexlr::HashedKmer& hk,
                               const std::string& kmer_qual,
                               size_t q);
  static size_t calc_kmer_quality(const std::string& qual);
 
  static void calc_minimizer(
    const std::vector<Indexlr::HashedKmer>& hashed_kmers_buffer,
    const Indexlr::Minimizer*& min_current,
    size_t idx,
    ssize_t& min_idx_left,
    ssize_t& min_idx_right,
    ssize_t& min_pos_prev,
    size_t w,
    std::vector<Indexlr::Minimizer>& minimizers);
  std::vector<Minimizer> minimize(const std::string& seq,
                                  const std::string& qual) const;
 
  const std::string seqfile;
  const size_t k, w;
  size_t q;
  const unsigned flags;
  const bool verbose;
  const long id;
  std::atomic<bool> closed{ false };
 
  static const BloomFilter& dummy_bf()
  {
    static const BloomFilter var;
    return var;
  }
 
  const std::reference_wrapper<const BloomFilter> filter_in_bf;
  const std::reference_wrapper<const BloomFilter> filter_out_bf;
  bool filter_in_enabled;
  bool filter_out_enabled;
 
  SeqReader reader;
  OrderQueueMPSC<Record> output_queue;
 
  using OutputQueueType = decltype(output_queue);
  static std::unique_ptr<OutputQueueType::Block>* ready_blocks_array()
  {
    thread_local static std::unique_ptr<decltype(output_queue)::Block>
      var[MAX_SIMULTANEOUS_INDEXLRS];
    return var;
  }
 
  static long* ready_blocks_owners()
  {
    thread_local static long var[MAX_SIMULTANEOUS_INDEXLRS] = { 0 };
    return var;
  }
 
  static size_t* ready_blocks_current()
  {
    thread_local static size_t var[MAX_SIMULTANEOUS_INDEXLRS] = { 0 };
    return var;
  }
 
  static std::atomic<long>& last_id()
  {
    static std::atomic<long> var(0);
    return var;
  }
 
  class Worker
  {
  public:
    void start() { t = std::thread(do_work, this); }
    void join() { t.join(); }
    void set_id(const int id) { this->id = id; }
 
    Worker& operator=(const Worker& worker) = delete;
    Worker& operator=(Worker&& worker) = delete;
 
    Worker(Indexlr& indexlr)
      : indexlr(indexlr)
    {
    }
    Worker(const Worker& worker)
      : Worker(worker.indexlr)
    {
    }
    Worker(Worker&& worker) noexcept
      : Worker(worker.indexlr)
    {
    }
 
  private:
    void work();
    static void do_work(Worker* worker) { worker->work(); }
 
    int id = -1;
    Indexlr& indexlr;
    std::thread t;
  };
 
  std::vector<Worker> workers;
  Barrier end_barrier;
  std::mutex last_block_num_mutex;
  uint64_t last_block_num = 0;
  bool last_block_num_valid = false;
};
 
// Constructor for Indexlr class when q is specified
inline Indexlr::Indexlr(std::string seqfile,
                        const size_t k,
                        const size_t w,
                        const size_t q,
                        const unsigned flags,
                        const unsigned threads,
                        const bool verbose,
                        const BloomFilter& bf1,
                        const BloomFilter& bf2)
  : seqfile(std::move(seqfile))
  , k(k)
  , w(w)
  , q(q)
  , flags(flags)
  , verbose(verbose)
  , id(++last_id())
  , filter_in_bf(filter_in() ? bf1 : Indexlr::dummy_bf())
  , filter_out_bf(filter_out() ? filter_in() ? bf2 : bf1 : Indexlr::dummy_bf())
  , filter_in_enabled(filter_in())
  , filter_out_enabled(filter_out())
  , reader(this->seqfile,
           short_mode() ? SeqReader::Flag::SHORT_MODE
                        : SeqReader::Flag::LONG_MODE)
  , output_queue(reader.get_buffer_size(), reader.get_block_size())
  , workers(std::vector<Worker>(threads, Worker(*this)))
  , end_barrier(threads)
{
  check_error(!short_mode() && !long_mode(),
              "Indexlr: no mode selected, either short or long mode flag must "
              "be provided.");
  check_error(short_mode() && long_mode(),
              "Indexlr: short and long mode are mutually exclusive.");
  check_error(threads == 0,
              "Indexlr: Number of processing threads cannot be 0.");
  int id_counter = 0;
  for (auto& worker : workers) {
    worker.set_id(id_counter++);
    worker.start();
  }
}
 
// Constructor for Indexlr class when q is not specified
inline Indexlr::Indexlr(std::string seqfile,
                        const size_t k,
                        const size_t w,
                        const unsigned flags,
                        const unsigned threads,
                        const bool verbose,
                        const BloomFilter& bf1,
                        const BloomFilter& bf2)
  : Indexlr(std::move(seqfile), k, w, 0, flags, threads, verbose, bf1, bf2)
{
}
 
inline Indexlr::~Indexlr()
{
  close();
}
 
inline void
Indexlr::close() noexcept
{
  bool closed_expected = false;
  if (closed.compare_exchange_strong(closed_expected, true)) {
    try {
      reader.close();
      output_queue.close();
      for (auto& worker : workers) {
        worker.join();
      }
    } catch (const std::system_error& e) {
      log_error("Indexlr thread join failure: " + std::string(e.what()));
      std::exit(EXIT_FAILURE); // NOLINT(concurrency-mt-unsafe)
    }
  }
}
 
// Minimerize a sequence: Find the minimizers of a vector of hash values
// representing a sequence.
/* Algorithm
v is a vector of non-negative integers
w is the window size
Invariants
    0 <  w <= v.size() - 1
    0 <= l <= r <= v.size() - 1
Initial conditions
    M    = NIL       Final set of minimizers, empty initially
    min  = -1        Minimum element
    i    = -1        Index of minimum element
    prev = -1        Index of previous minimum element
    l    = 0         Index of left end of window
    r    = l + w - 1 Index of right end of window
Computation
At each window, if the previous minimum is out of scope, find the new,
right-most, minimum or else, check with only the right-most element to determine
if that is the new minimum. A minimizer is added to the final vector only if
it's index has changed. for each window of v bounded by [l, r] if (i < l) i =
index of minimum element in [l, r], furthest from l. else if (v[r] <= v[i]) i =
r min = v[i] if (i != prev) { prev = i M <- M + m
    }
    l = l + 1        Move window's left bound by one element
    r = l + w - 1    Set window's right bound
}*/
 
inline std::string
Indexlr::extract_barcode(const std::string& id, const std::string& comment)
{
  const static std::string barcode_prefix = "BX:Z:";
  if (startswith(comment, barcode_prefix)) {
    const auto space_pos = comment.find(' ');
    if (space_pos != std::string::npos) {
      return comment.substr(barcode_prefix.size(),
                            space_pos - barcode_prefix.size());
    }
    return comment.substr(barcode_prefix.size());
  }
  const auto pound_pos = id.find('#');
  if (pound_pos != std::string::npos) {
    const auto slash_pos = id.find('/');
    if (slash_pos > pound_pos) {
      return id.substr(pound_pos + 1, slash_pos - (pound_pos + 1));
    }
  }
  return "NA";
}
 
inline void
Indexlr::filter_hashed_kmer(Indexlr::HashedKmer& hk,
                            bool filter_in,
                            bool filter_out,
                            const BloomFilter& filter_in_bf,
                            const BloomFilter& filter_out_bf)
{
  if (!filter_in && !filter_out) {
    return;
  }
 
  bool use_all = (filter_in ? filter_in_bf.get_hash_num()
                            : filter_out_bf.get_hash_num()) > 1;
 
  auto check_contains = [&](const BloomFilter& bf) {
    if (use_all) {
      return bf.contains(hk.all_hashes);
    } else {
      return bf.contains({ hk.min_hash });
    }
  };
 
  if (filter_in && filter_out) {
    if (!check_contains(filter_in_bf) || check_contains(filter_out_bf)) {
      hk.min_hash = std::numeric_limits<uint64_t>::max();
    }
  } else if (filter_in) {
    if (!check_contains(filter_in_bf)) {
      hk.min_hash = std::numeric_limits<uint64_t>::max();
    }
  } else if (filter_out) {
    if (check_contains(filter_out_bf)) {
      hk.min_hash = std::numeric_limits<uint64_t>::max();
    }
  }
}
 
inline void
Indexlr::filter_kmer_qual(Indexlr::HashedKmer& hk,
                          const std::string& kmer_qual,
                          size_t q)
{
  if (calc_kmer_quality(kmer_qual) < q) {
    hk.min_hash = std::numeric_limits<uint64_t>::max();
  }
}
 
inline size_t
Indexlr::calc_kmer_quality(const std::string& qual)
{
  // convert the quality scores to integers
  std::vector<int> qual_ints;
  const int thirty_three = 33;
  qual_ints.reserve(qual.size());
  for (auto c : qual) {
    qual_ints.push_back(c - thirty_three);
  }
  // calculate the mean (potential improvement: use other statistics)
  size_t sum = 0;
  for (auto q : qual_ints) {
    sum += q;
  }
  return (sum / qual_ints.size());
}
 
inline void
Indexlr::calc_minimizer(
  const std::vector<Indexlr::HashedKmer>& hashed_kmers_buffer,
  const Indexlr::Minimizer*& min_current,
  const size_t idx,
  ssize_t& min_idx_left,
  ssize_t& min_idx_right,
  ssize_t& min_pos_prev,
  const size_t w,
  std::vector<Indexlr::Minimizer>& minimizers)
{
  min_idx_left = ssize_t(idx + 1 - w);
  min_idx_right = ssize_t(idx + 1);
  const auto& min_left =
    hashed_kmers_buffer[min_idx_left % hashed_kmers_buffer.size()];
  const auto& min_right =
    hashed_kmers_buffer[(min_idx_right - 1) % hashed_kmers_buffer.size()];
 
  if (min_current == nullptr || min_current->pos < min_left.pos) {
    min_current = &min_left;
    // Use of operator '<=' returns the minimum that is furthest from left.
    for (ssize_t i = min_idx_left; i < min_idx_right; i++) {
      const auto& min_i = hashed_kmers_buffer[i % hashed_kmers_buffer.size()];
      if (min_i.min_hash <= min_current->min_hash) {
        min_current = &min_i;
      }
    }
  } else if (min_right.min_hash <= min_current->min_hash) {
    min_current = &min_right;
  }
  if (ssize_t(min_current->pos) > min_pos_prev &&
      min_current->min_hash != std::numeric_limits<uint64_t>::max()) {
    min_pos_prev = ssize_t(min_current->pos);
    minimizers.push_back(*min_current);
  }
}
 
inline std::vector<Indexlr::Minimizer>
Indexlr::minimize(const std::string& seq, const std::string& qual) const
{
  if ((k > seq.size()) || (w > seq.size() - k + 1)) {
    return {};
  }
  std::vector<Minimizer> minimizers;
  minimizers.reserve(2 * (seq.size() - k + 1) / w);
  std::vector<HashedKmer> hashed_kmers_buffer(w + 1);
  ssize_t min_idx_left, min_idx_right, min_pos_prev = -1;
  const Minimizer* min_current = nullptr;
  size_t idx = 0;
  size_t bf_num_hashes = 0;
  if (!filter_in() && !filter_out()) {
    bf_num_hashes = 1;
  } else if (filter_in()) {
    bf_num_hashes = filter_in_bf.get().get_hash_num();
  } else {
    bf_num_hashes = filter_out_bf.get().get_hash_num();
  }
 
  size_t nthash_num_hashes = std::max(static_cast<size_t>(2), bf_num_hashes);
 
  for (NtHash nh(seq, nthash_num_hashes, k); nh.roll(); ++idx) {
    auto& hk = hashed_kmers_buffer[idx % hashed_kmers_buffer.size()];
    if (bf_num_hashes == 1) {
      hk = HashedKmer(nh.hashes()[0],
                      nh.hashes()[1],
                      nh.get_pos(),
                      nh.get_forward_hash() <= nh.get_reverse_hash(),
                      output_seq() ? seq.substr(nh.get_pos(), k) : "",
                      output_qual() ? qual.substr(nh.get_pos(), k) : "");
 
    } else {
      hk = HashedKmer(
        nh.hashes()[0],
        nh.hashes()[1],
        nh.get_pos(),
        nh.get_forward_hash() <= nh.get_reverse_hash(),
        output_seq() ? seq.substr(nh.get_pos(), k) : "",
        output_qual() ? qual.substr(nh.get_pos(), k) : "",
        std::vector<uint64_t>(nh.hashes(), nh.hashes() + bf_num_hashes));
    }
 
    filter_hashed_kmer(
      hk, filter_in(), filter_out(), filter_in_bf.get(), filter_out_bf.get());
 
    if (q > 0) {
      filter_kmer_qual(hk, qual.substr(nh.get_pos(), k), q);
    }
 
    if (idx + 1 >= w) {
      calc_minimizer(hashed_kmers_buffer,
                     min_current,
                     idx,
                     min_idx_left,
                     min_idx_right,
                     min_pos_prev,
                     w,
                     minimizers);
    }
  }
  return minimizers;
}
 
inline Indexlr::Record
Indexlr::read()
{
  if (ready_blocks_owners()[id % MAX_SIMULTANEOUS_INDEXLRS] != id) {
    ready_blocks_array()[id % MAX_SIMULTANEOUS_INDEXLRS] =
      std::unique_ptr< // NOLINT(modernize-make-unique)
        decltype(output_queue)::Block>(
        new decltype(output_queue)::Block(reader.get_block_size()));
    ready_blocks_owners()[id % MAX_SIMULTANEOUS_INDEXLRS] = id;
    ready_blocks_current()[id % MAX_SIMULTANEOUS_INDEXLRS] = 0;
  }
  auto& block = *(ready_blocks_array()[id % MAX_SIMULTANEOUS_INDEXLRS]);
  auto& current = ready_blocks_current()[id % MAX_SIMULTANEOUS_INDEXLRS];
  if (current >= block.count) {
    block.count = 0;
    output_queue.read(block);
    if (block.count == 0) {
      output_queue.close();
      block = decltype(output_queue)::Block(reader.get_block_size());
      return Record();
    }
    current = 0;
  }
  return std::move(block.data[current++]);
}
 
inline void
Indexlr::Worker::work()
{
  decltype(indexlr.output_queue)::Block output_block(
    indexlr.reader.get_block_size());
  uint64_t last_block_num = 0;
  bool last_block_num_valid = false;
  for (;;) {
    auto input_block = indexlr.reader.read_block();
    if (input_block.count == 0) {
      break;
    }
 
    output_block.num = input_block.num;
    for (size_t idx = 0; idx < input_block.count; idx++) {
      Record record;
      auto& reader_record = input_block.data[idx];
      record.num = reader_record.num;
      if (indexlr.output_bx()) {
        record.barcode =
          indexlr.extract_barcode(reader_record.id, reader_record.comment);
      }
      if (indexlr.output_id()) {
        record.id = std::move(reader_record.id);
      }
 
      record.readlen = reader_record.seq.size();
 
      check_info(indexlr.verbose && indexlr.k > record.readlen,
                 "Indexlr: skipped seq " + std::to_string(record.num) +
                   " on line " +
                   std::to_string(record.num * (indexlr.reader.get_format() ==
                                                    SeqReader::Format::FASTA
                                                  ? 2
                                                  : 4) +
                                  2) +
                   "; k (" + std::to_string(indexlr.k) + ") > seq length (" +
                   std::to_string(record.readlen) + ")");
 
      check_info(indexlr.verbose && indexlr.w > record.readlen - indexlr.k + 1,
                 "Indexlr: skipped seq " + std::to_string(record.num) +
                   " on line " +
                   std::to_string(record.num * (indexlr.reader.get_format() ==
                                                    SeqReader::Format::FASTA
                                                  ? 2
                                                  : 4) +
                                  2) +
                   "; w (" + std::to_string(indexlr.w) + ") > # of hashes (" +
                   std::to_string(record.readlen - indexlr.k + 1) + ")");
 
      if (indexlr.k <= record.readlen &&
          indexlr.w <= record.readlen - indexlr.k + 1) {
        record.minimizers =
          indexlr.minimize(reader_record.seq, reader_record.qual);
      } else {
        record.minimizers = {};
      }
 
      output_block.data[output_block.count++] = std::move(record);
    }
    if (output_block.count > 0) {
      last_block_num = output_block.num;
      last_block_num_valid = true;
      indexlr.output_queue.write(output_block);
      output_block.count = 0;
    }
  }
  if (last_block_num_valid) {
    std::unique_lock<std::mutex> lock(indexlr.last_block_num_mutex);
    indexlr.last_block_num = std::max(indexlr.last_block_num, last_block_num);
    indexlr.last_block_num_valid = true;
    lock.unlock();
  }
  indexlr.end_barrier.wait();
  if (last_block_num_valid && indexlr.last_block_num_valid &&
      last_block_num == indexlr.last_block_num) {
    output_block.num = last_block_num + 1;
    indexlr.output_queue.write(output_block);
  } else if (!indexlr.last_block_num_valid && id == 0) {
    output_block.num = 0;
    indexlr.output_queue.write(output_block);
  }
}
 
} // namespace btllib
 
#endif