nthash_kmer.hpp

#pragma once
 
#include <array>
#include <cstdint>
#include <cstring>
#include <deque>
#include <memory>
#include <string>
#include <vector>
 
#include <btllib/hashing_internals.hpp>
#include <btllib/status.hpp>
 
namespace btllib::hashing_internals {
 
inline uint64_t
base_forward_hash(const char* seq,
                  unsigned k,
                  const uint64_t* primitive_tab = SEED_TAB,
                  const uint64_t* dimer_tab = DIMER_TAB,
                  const uint64_t* trimer_tab = TRIMER_TAB,
                  const uint64_t* tetramer_tab = TETRAMER_TAB,
                  const uint8_t* convert_tab = CONVERT_TAB)
{
  uint64_t h_val = 0;
  for (unsigned i = 0; i < k - 3; i += 4) {
    h_val = srol(h_val, 4);
    uint8_t loc = 0;
    loc += 64 * convert_tab[(unsigned char)seq[i]];     // NOLINT
    loc += 16 * convert_tab[(unsigned char)seq[i + 1]]; // NOLINT
    loc += 4 * convert_tab[(unsigned char)seq[i + 2]];
    loc += convert_tab[(unsigned char)seq[i + 3]];
    h_val ^= tetramer_tab[loc];
  }
  const unsigned remainder = k % 4;
  if (remainder > 0) {
    h_val = srol(h_val, remainder);
  }
  if (remainder == 3) {
    uint8_t trimer_loc = 0;
    trimer_loc += 16 * convert_tab[(unsigned char)seq[k - 3]]; // NOLINT
    trimer_loc += 4 * convert_tab[(unsigned char)seq[k - 2]];
    trimer_loc += convert_tab[(unsigned char)seq[k - 1]];
    h_val ^= trimer_tab[trimer_loc];
  } else if (remainder == 2) {
    uint8_t dimer_loc = 0;
    dimer_loc += 4 * convert_tab[(unsigned char)seq[k - 2]];
    dimer_loc += convert_tab[(unsigned char)seq[k - 1]];
    h_val ^= dimer_tab[dimer_loc];
  } else if (remainder == 1) {
    h_val ^= primitive_tab[(unsigned char)seq[k - 1]];
  }
  return h_val;
}
 
inline uint64_t
next_forward_hash(uint64_t fh_val,
                  unsigned k,
                  unsigned char char_out,
                  unsigned char char_in,
                  const uint64_t* primitive_tab = SEED_TAB,
                  const uint64_t* const* right_table = MS_TAB_33R,
                  const uint64_t* const* left_table = MS_TAB_31L)
{
  uint64_t h_val = srol(fh_val);
  h_val ^= primitive_tab[char_in];
  h_val ^= srol_table(char_out, k, right_table, left_table);
  return h_val;
}
 
inline uint64_t
prev_forward_hash(uint64_t fh_val,
                  unsigned k,
                  unsigned char char_out,
                  unsigned char char_in,
                  const uint64_t* primitive_tab = SEED_TAB,
                  const uint64_t* const* right_table = MS_TAB_33R,
                  const uint64_t* const* left_table = MS_TAB_31L)
{
  uint64_t h_val = fh_val ^ srol_table(char_in, k, right_table, left_table);
  h_val ^= primitive_tab[char_out];
  h_val = sror(h_val);
  return h_val;
}
 
inline uint64_t
base_reverse_hash(const char* seq,
                  unsigned k,
                  const uint64_t* primitive_tab = SEED_TAB,
                  const uint64_t* dimer_tab = DIMER_TAB,
                  const uint64_t* trimer_tab = TRIMER_TAB,
                  const uint64_t* tetramer_tab = TETRAMER_TAB,
                  const uint8_t* rc_convert_tab = RC_CONVERT_TAB)
{
  uint64_t h_val = 0;
  const unsigned remainder = k % 4;
  if (remainder == 3) {
    uint8_t trimer_loc = 0;
    trimer_loc += 16 * rc_convert_tab[(unsigned char)seq[k - 1]]; // NOLINT
    trimer_loc += 4 * rc_convert_tab[(unsigned char)seq[k - 2]];
    trimer_loc += rc_convert_tab[(unsigned char)seq[k - 3]];
    h_val ^= trimer_tab[trimer_loc];
  } else if (remainder == 2) {
    uint8_t dimer_loc = 0;
    dimer_loc += 4 * rc_convert_tab[(unsigned char)seq[k - 1]];
    dimer_loc += rc_convert_tab[(unsigned char)seq[k - 2]];
    h_val ^= dimer_tab[dimer_loc];
  } else if (remainder == 1) {
    h_val ^= primitive_tab[(unsigned char)seq[k - 1] & CP_OFF];
  }
  for (int i = (int)(k - remainder) - 1; i >= 3; i -= 4) {
    h_val = srol(h_val, 4);
    uint8_t loc = 0;
    loc += 64 * rc_convert_tab[(unsigned char)seq[i]];     // NOLINT
    loc += 16 * rc_convert_tab[(unsigned char)seq[i - 1]]; // NOLINT
    loc += 4 * rc_convert_tab[(unsigned char)seq[i - 2]];
    loc += rc_convert_tab[(unsigned char)seq[i - 3]];
    h_val ^= tetramer_tab[loc];
  }
  return h_val;
}
 
inline uint64_t
next_reverse_hash(uint64_t rh_val,
                  unsigned k,
                  unsigned char char_out,
                  unsigned char char_in,
                  const uint64_t* primitive_tab = SEED_TAB,
                  const uint64_t* const* right_table = MS_TAB_33R,
                  const uint64_t* const* left_table = MS_TAB_31L)
{
  uint64_t h_val =
    rh_val ^ srol_table(char_in & CP_OFF, k, right_table, left_table);
  h_val ^= primitive_tab[char_out & CP_OFF];
  h_val = sror(h_val);
  return h_val;
}
 
inline uint64_t
prev_reverse_hash(uint64_t rh_val,
                  unsigned k,
                  unsigned char char_out,
                  unsigned char char_in,
                  const uint64_t* primitive_tab = SEED_TAB,
                  const uint64_t* const* right_table = MS_TAB_33R,
                  const uint64_t* const* left_table = MS_TAB_31L)
{
  uint64_t h_val = srol(rh_val);
  h_val ^= primitive_tab[char_in & CP_OFF];
  h_val ^= srol_table(char_out & CP_OFF, k, right_table, left_table);
  return h_val;
}
 
inline void
sub_hash(uint64_t fh_val,
         uint64_t rh_val,
         const char* kmer_seq,
         const std::vector<unsigned>& positions,
         const std::vector<unsigned char>& new_bases,
         unsigned k,
         unsigned m,
         uint64_t* h_val,
         const uint64_t* const* right_table = MS_TAB_33R,
         const uint64_t* const* left_table = MS_TAB_31L)
{
  uint64_t b_val = 0;
 
  for (size_t i = 0; i < positions.size(); i++) {
    const auto pos = positions[i];
    const auto new_base = new_bases[i];
 
    fh_val ^= srol_table(
      (unsigned char)kmer_seq[pos], k - 1 - pos, right_table, left_table);
    fh_val ^= srol_table(new_base, k - 1 - pos, right_table, left_table);
 
    rh_val ^= srol_table(
      (unsigned char)kmer_seq[pos] & CP_OFF, pos, right_table, left_table);
    rh_val ^= srol_table(new_base & CP_OFF, pos, right_table, left_table);
  }
 
  b_val = canonical(fh_val, rh_val);
  extend_hashes(b_val, k, m, h_val);
}
 
} // namespace btllib::hashing_internals
 
namespace btllib {
 
using hashing_internals::base_forward_hash;
using hashing_internals::base_reverse_hash;
using hashing_internals::extend_hashes;
using hashing_internals::next_forward_hash;
using hashing_internals::next_reverse_hash;
using hashing_internals::prev_forward_hash;
using hashing_internals::prev_reverse_hash;
using hashing_internals::SEED_N;
using hashing_internals::SEED_TAB;
using hashing_internals::sub_hash;
 
class NtHash
{
 
public:
  NtHash(const char* seq,
         size_t seq_len,
         hashing_internals::NUM_HASHES_TYPE num_hashes,
         hashing_internals::K_TYPE k,
         size_t pos = 0)
    : seq(seq)
    , seq_len(seq_len)
    , num_hashes(num_hashes)
    , k(k)
    , pos(pos)
    , initialized(false)
    , hash_arr(new uint64_t[num_hashes])
  {
    check_error(k == 0, "NtHash: k must be greater than 0");
    check_error(this->seq_len < k,
                "NtHash: sequence length (" + std::to_string(this->seq_len) +
                  ") is smaller than k (" + std::to_string(k) + ")");
    check_error(pos > this->seq_len - k,
                "NtHash: passed position (" + std::to_string(pos) +
                  ") is larger than sequence length (" +
                  std::to_string(this->seq_len) + ")");
  }
 
  NtHash(const std::string& seq,
         hashing_internals::NUM_HASHES_TYPE num_hashes,
         hashing_internals::K_TYPE k,
         size_t pos = 0)
    : NtHash(seq.data(), seq.size(), num_hashes, k, pos)
  {
  }
 
  NtHash(const NtHash& obj)
    : seq(obj.seq)
    , seq_len(obj.seq_len)
    , num_hashes(obj.num_hashes)
    , k(obj.k)
    , pos(obj.pos)
    , initialized(obj.initialized)
    , fwd_hash(obj.fwd_hash)
    , rev_hash(obj.rev_hash)
    , hash_arr(new uint64_t[obj.num_hashes])
  {
    std::memcpy(
      hash_arr.get(), obj.hash_arr.get(), num_hashes * sizeof(uint64_t));
  }
 
  NtHash(NtHash&&) = default;
 
  bool roll()
  {
    if (!initialized) {
      return init();
    }
    if (pos >= seq_len - k) {
      return false;
    }
    if (hashing_internals::SEED_TAB[(unsigned char)seq[pos + k]] ==
        hashing_internals::SEED_N) {
      pos += k;
      return init();
    }
    fwd_hash = next_forward_hash(fwd_hash, k, seq[pos], seq[pos + k]);
    rev_hash = next_reverse_hash(rev_hash, k, seq[pos], seq[pos + k]);
    extend_hashes(fwd_hash, rev_hash, k, num_hashes, hash_arr.get());
    ++pos;
    return true;
  }
 
  bool roll_back()
  {
    if (!initialized) {
      return init();
    }
    if (pos == 0) {
      return false;
    }
    if (SEED_TAB[(unsigned char)seq[pos - 1]] == SEED_N && pos >= k) {
      pos -= k;
      return init();
    }
    if (SEED_TAB[(unsigned char)seq[pos - 1]] == SEED_N) {
      return false;
    }
    fwd_hash = prev_forward_hash(fwd_hash, k, seq[pos + k - 1], seq[pos - 1]);
    rev_hash = prev_reverse_hash(rev_hash, k, seq[pos + k - 1], seq[pos - 1]);
    extend_hashes(fwd_hash, rev_hash, k, num_hashes, hash_arr.get());
    --pos;
    return true;
  }
 
  bool peek()
  {
    if (pos >= seq_len - k) {
      return false;
    }
    return peek(seq[pos + k]);
  }
 
  bool peek_back()
  {
    if (pos == 0) {
      return false;
    }
    return peek_back(seq[pos - 1]);
  }
 
  bool peek(char char_in)
  {
    if (!initialized) {
      return init();
    }
    if (SEED_TAB[(unsigned char)char_in] == SEED_N) {
      return false;
    }
    const uint64_t fwd = next_forward_hash(fwd_hash, k, seq[pos], char_in);
    const uint64_t rev = next_reverse_hash(rev_hash, k, seq[pos], char_in);
    extend_hashes(fwd, rev, k, num_hashes, hash_arr.get());
    return true;
  }
 
  bool peek_back(char char_in)
  {
    if (!initialized) {
      return init();
    }
    if (SEED_TAB[(unsigned char)char_in] == SEED_N) {
      return false;
    }
    const unsigned char char_out = seq[pos + k - 1];
    const uint64_t fwd = prev_forward_hash(fwd_hash, k, char_out, char_in);
    const uint64_t rev = prev_reverse_hash(rev_hash, k, char_out, char_in);
    extend_hashes(fwd, rev, k, num_hashes, hash_arr.get());
    return true;
  }
 
  void sub(const std::vector<unsigned>& positions,
           const std::vector<unsigned char>& new_bases)
  {
    sub_hash(fwd_hash,
             rev_hash,
             seq + pos,
             positions,
             new_bases,
             get_k(),
             get_hash_num(),
             hash_arr.get());
  }
 
  const uint64_t* hashes() const { return hash_arr.get(); }
 
  size_t get_pos() const { return pos; }
 
  hashing_internals::NUM_HASHES_TYPE get_hash_num() const { return num_hashes; }
 
  hashing_internals::K_TYPE get_k() const { return k; }
 
  uint64_t get_forward_hash() const { return fwd_hash; }
 
  uint64_t get_reverse_hash() const { return rev_hash; }
 
private:
  const char* seq;
  const size_t seq_len;
  hashing_internals::NUM_HASHES_TYPE num_hashes;
  hashing_internals::K_TYPE k;
  size_t pos;
  bool initialized;
  uint64_t fwd_hash = 0;
  uint64_t rev_hash = 0;
  std::unique_ptr<uint64_t[]> hash_arr;
 
  bool init()
  {
    bool has_n = true;
    while (pos <= seq_len - k + 1 && has_n) {
      has_n = false;
      for (unsigned i = 0; i < k && pos <= seq_len - k + 1; i++) {
        if (SEED_TAB[(unsigned char)seq[pos + k - i - 1]] == SEED_N) {
          pos += k - i;
          has_n = true;
        }
      }
    }
    if (pos > seq_len - k) {
      return false;
    }
    fwd_hash = base_forward_hash(seq + pos, k);
    rev_hash = base_reverse_hash(seq + pos, k);
    extend_hashes(fwd_hash, rev_hash, k, num_hashes, hash_arr.get());
    initialized = true;
    return true;
  }
};
 
class BlindNtHash
{
 
public:
  BlindNtHash(const std::string& seq,
              hashing_internals::NUM_HASHES_TYPE num_hashes,
              hashing_internals::K_TYPE k,
              long pos = 0)
    : seq(seq.data() + pos, seq.data() + pos + k)
    , num_hashes(num_hashes)
    , pos(pos)
    , fwd_hash(base_forward_hash(seq.data(), k))
    , rev_hash(base_reverse_hash(seq.data(), k))
    , hash_arr(new uint64_t[num_hashes])
  {
    check_error(k == 0, "BlindNtHash: k must be greater than 0");
    extend_hashes(fwd_hash, rev_hash, k, num_hashes, hash_arr.get());
  }
 
  BlindNtHash(const BlindNtHash& obj)
    : seq(obj.seq)
    , num_hashes(obj.num_hashes)
    , pos(obj.pos)
    , fwd_hash(obj.fwd_hash)
    , rev_hash(obj.rev_hash)
    , hash_arr(new uint64_t[obj.num_hashes])
  {
    std::memcpy(
      hash_arr.get(), obj.hash_arr.get(), num_hashes * sizeof(uint64_t));
  }
 
  BlindNtHash(BlindNtHash&&) = default;
 
  void roll(char char_in)
  {
    fwd_hash = next_forward_hash(fwd_hash, seq.size(), seq.front(), char_in);
    rev_hash = next_reverse_hash(rev_hash, seq.size(), seq.front(), char_in);
    extend_hashes(fwd_hash, rev_hash, seq.size(), num_hashes, hash_arr.get());
    seq.pop_front();
    seq.push_back(char_in);
    ++pos;
  }
 
  void roll_back(char char_in)
  {
    fwd_hash = prev_forward_hash(fwd_hash, seq.size(), seq.back(), char_in);
    rev_hash = prev_reverse_hash(rev_hash, seq.size(), seq.back(), char_in);
    extend_hashes(fwd_hash, rev_hash, seq.size(), num_hashes, hash_arr.get());
    seq.pop_back();
    seq.push_front(char_in);
    --pos;
  }
 
  void peek(char char_in)
  {
    const hashing_internals::K_TYPE k = seq.size();
    const uint64_t fwd = next_forward_hash(fwd_hash, k, seq.front(), char_in);
    const uint64_t rev = next_reverse_hash(rev_hash, k, seq.front(), char_in);
    extend_hashes(fwd, rev, seq.size(), num_hashes, hash_arr.get());
  }
 
  void peek_back(char char_in)
  {
    const hashing_internals::K_TYPE k = seq.size();
    const uint64_t fwd = prev_forward_hash(fwd_hash, k, seq.back(), char_in);
    const uint64_t rev = prev_reverse_hash(rev_hash, k, seq.back(), char_in);
    extend_hashes(fwd, rev, seq.size(), num_hashes, hash_arr.get());
  }
 
  const uint64_t* hashes() const { return hash_arr.get(); }
 
  long get_pos() const { return pos; }
 
  hashing_internals::NUM_HASHES_TYPE get_hash_num() const { return num_hashes; }
 
  hashing_internals::K_TYPE get_k() const { return seq.size(); }
 
  uint64_t get_forward_hash() const { return fwd_hash; }
 
  uint64_t get_reverse_hash() const { return rev_hash; }
 
private:
  std::deque<char> seq;
  hashing_internals::NUM_HASHES_TYPE num_hashes;
  long pos;
  uint64_t fwd_hash;
  uint64_t rev_hash;
  std::unique_ptr<uint64_t[]> hash_arr;
};
 
} // namespace btllib