AhoCorasickAmbiguous.h

Go to the documentation of this file.

// Copyright (c) 2002-present, OpenMS Inc. -- EKU Tuebingen, ETH Zurich, and FU Berlin
// SPDX-License-Identifier: BSD-3-Clause
//
// --------------------------------------------------------------------------
// $Maintainer: Chris Bielow $
// $Authors: Chris Bielow $
// --------------------------------------------------------------------------
 
#pragma once
 
#include <OpenMS/CONCEPT/Macros.h>
#include <bitset>
#include <cassert>
#include <cstring>
#include <functional> // for std::hash
#include <limits>
#include <queue>
#include <string>
#include <cstdint>
#include <unordered_map>
#include <vector>
#include <bit> // for std::popcount
 
namespace OpenMS
{
  constexpr char const AAtoChar[28] = {
    'A', // 00 Ala Alanine
    'Y', // 01 Tyr Tyrosine
    'C', // 02 Cys Cysteine
    'D', // 03 Asp Aspartic Acid   // B
    'N', // 04 Asn Asparagine      // B
    'F', // 05 Phe Phenylalanine
    'G', // 06 Gly Glycine
    'H', // 07 His Histidine
    'I', // 08 Ile Isoleucine      // J
    'L', // 09 Leu Leucine         // J
    'K', // 10 Lys Lysine
    'W', // 11 Trp Tryptophan
    'M', // 12 Met Methionine
    'O', // 13 Pyl Pyrrolysine
    'P', // 14 Pro Proline
    'E', // 15 Glu Glutamic Acid   // Z
    'Q', // 16 Gln Glutamine       // Z
    'R', // 17 Arg Arginine
    'S', // 18 Ser Serine
    'T', // 19 Thr Threonine
    'U', // 20 Selenocysteine
    'V', // 21 Val Valine
    // ambiguous AAs start here (index: 22...25)
    'B', // 22 Aspartic Acid, Asparagine  $   // the ambAA's need to be consecutive (B,J,Z,X,$)
    'J', // 23 Leucine, Isoleucine        $
    'Z', // 24 Glutamic Acid, Glutamine   $
    'X', // 25 Unknown 
    // non-regular AA's, which are special
    '$', // 26 superAA, i.e. it models a mismatch, which can be anything, including AAAs
    '?', // 27 invalid AA (will usually be skipped) -- must be the last AA (AA::operator++ and others rely on it)
  };
 
  constexpr char const CharToAA[256] = {
    // ASCII char (7-bit Int with values from 0..127) --> amino acid 
    27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, // 0
    27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, // 1
  //                 $
    27, 27, 27, 27, 26, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, // 2
    27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, // 3
 
  //  ,  A,  B,  C,  D,  E,  F,  G,  H,  I,  J,  K,  L,  M,  N,  O,
    27, 00, 22, 02, 03, 15, 05, 06, 07,  8, 23, 10,  9, 12, 04, 13, // 4
 
  // P,  Q,  R,  S,  T,  U,  V,  W,  X,  Y,  Z,   ,   ,   ,   ,   ,
    14, 16, 17, 18, 19, 20, 21, 11, 25, 01, 24, 27, 27, 27, 27, 27, // 5
 
  //  ,  a,  b,  c,  d,  e,  f,  g,  h,  i,  j,  k,  l,  m,  n,  o,
    27, 00, 22, 02, 03, 15, 05, 06, 07,  8, 23, 10,  9, 12, 04, 13,   // 6
 
  // p,  q,  r,  s,  t,  u,  v,  w,  x,  y,  z,   ,   ,   ,   ,   ,
    14, 16, 17, 18, 19, 20, 21, 11, 25, 01, 24, 27, 27, 27, 27, 27, // 7
 
    // bytes 128..255 (high bit set: extended ASCII / UTF-8 continuation bytes) --> invalid AA ('?')
    27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, // 8
    27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, // 9
    27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, // 10
    27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, // 11
    27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, // 12
    27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, // 13
    27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, // 14
    27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, 27, // 15
  };
 
  struct OPENMS_DLLAPI Hit {
    using T = uint32_t;
    Hit() = default;
    Hit(T needle_index, T needle_length, T query_pos) : needle_index(needle_index), needle_length(needle_length), query_pos(query_pos) {};
    T needle_index;
    T needle_length;
    T query_pos;
    bool operator==(const Hit& rhs) const
    {
      return needle_index == rhs.needle_index && needle_length == rhs.needle_length && query_pos == rhs.query_pos;
    }
  };
 
  struct OPENMS_DLLAPI AA
  {
    constexpr AA() : aa_(AA('?').aa_)
    {
    }
    
    constexpr explicit AA(const char c) : aa_(CharToAA[(unsigned char)c])
    {
    }
 
    constexpr uint8_t operator()() const
    {
      return aa_;
    }
 
    constexpr bool operator==(const AA other) const
    {
      return aa_ == other.aa_;
    }
 
    constexpr bool operator<=(const AA other) const
    {
      return aa_ <= other.aa_;
    }
 
    constexpr bool isAmbiguous() const
    {
      return aa_ >= AA('B').aa_;
    }
    
    constexpr bool isValid() const
    {
      return aa_ != AA('?').aa_;
    }
 
    constexpr bool isValidForPeptide() const
    {
      return aa_ <= AA('X').aa_; // '$' or '?'
    }
 
    constexpr AA& operator++()
    {
      ++aa_;
      assert(aa_ <= AA('?').aa_); // make sure we don't overflow
      return *this;
    }
 
    constexpr AA operator++(int)
    {
      AA r(*this);
      ++aa_;
      assert(aa_ <= AA('?').aa_); // make sure we don't overflow
      return r;
    }
 
    constexpr AA operator-(const AA rhs) const
    {
      AA r(*this);
      r.aa_ -= rhs.aa_;
      return r;
    }
 
    constexpr char toChar() const
    {
      return AAtoChar[aa_];
    }
 
    static AA fromIndex(size_t index) noexcept
    {
      OPENMS_PRECONDITION(index < unambiguousAACount(), "AA::fromIndex(): index must be in [0, unambiguousAACount())");
      AA r;
      r.aa_ = static_cast<uint8_t>(index);
      return r;
    }
 
    static size_t unambiguousAACount()
    {
      static_assert(AA('V').aa_ - AA('A').aa_ + 1 == 22, "Unambiguous AA count must be 22 (A-Z, excluding B, J, Z, X)");
      static_assert(AA('A').aa_ == 0, "A must be the first AA in the enumeration");
      static_assert(AA('V').aa_ == 21, "V must be the last unambiguous AA in the enumeration");
      return AA('V').aa_ - AA('A').aa_ + 1;
    }
 
  private:
    uint8_t aa_; 
  };
 
  class OPENMS_DLLAPI Index
  {
  public:
    using T = uint32_t;
    constexpr Index() = default;
    
    constexpr Index(T val) : i_(val) {};
 
    constexpr bool isInvalid() const
    {
      return i_ == std::numeric_limits<T>::max();
    }
 
    constexpr bool isValid() const
    {
      return i_ != std::numeric_limits<T>::max();
    }
 
    constexpr T operator()() const
    {
      return i_;
    }
 
    constexpr bool operator==(const Index other) const
    {
      return i_ == other.i_;
    }
 
    constexpr T& pos()
    {
      return i_;
    }
 
    constexpr T pos() const
    {
      return i_;
    }
  private:
    T i_ = std::numeric_limits<T>::max(); 
  };
 } // namespace OpenMS
 
// this needs to go into namespace std
template<> struct std::hash<OpenMS::Index>
{
  std::size_t operator()(OpenMS::Index const& s) const noexcept
  {
    return std::hash<OpenMS::Index::T> {}(s());
  }
};
 
namespace OpenMS
{
  struct Bitset {
    uint32_t bits = 0;
 
    // Set bit at position `pos` (0-based)
    void set(int pos) {
      bits |= (1u << pos);
    }
 
    // Clear bit at position `pos`
    void clear(int pos) {
      bits &= ~(1u << pos);
    }
 
    // Check if bit at position `pos` is set
    bool test(int pos) const {
      return (bits >> pos) & 1u;
    }
 
    // Left shift the bitset by `n` positions
    Bitset operator<<(int n) const {
      Bitset o = *this;
      return o <<= n;
    }
 
    // Left shift the bitset by `n` positions
    Bitset& operator<<=(int n) {
      if (n > 31)
      { // if we shift more than 31 bits, behaviour is undefined (e.g. GCC/Clang/MSVC will not modify 'bits' if n==32)
        bits = 0; // shift out all bits
      }
      else
      {
        bits <<= n;
      }
      return *this;
    }
 
    // Count number of set bits (population count)
    int pop_count() const {
      return std::popcount(bits);
    }
  }; // end Bitset
 
 
  struct OPENMS_DLLAPI ACNode
  {
    ACNode() {};
 
    ACNode(const AA label, const uint8_t depth) : edge(label)
    {
      depth_and_hits.depth = depth;
    }
 
    struct DepthHits {
      DepthHits()
      {
        memset(this, 0, sizeof *this); // make sure bitfields are 0; C++20 allows {} initialization ...
      };
      uint8_t has_hit: 1; 
      // we could add another bit here to distinguish between a local hit and suffix hit, but on Windows, this slows it down
      uint8_t depth : 7; 
    };
 
    using ChildCountType = uint8_t; 
 
    // do not use std::bitset, since GCC/clang wastes 8 bytes here.
    //std::bitset<26> children_bitset; ///< bitfield of children (if tree is in BFS order); 26 bits are enough to cover all AA incl. (B,J,X,Z)
    Bitset children_bitset;  
    Index suffix {0};        
    Index first_child {0};   
    AA edge {0};             
    ChildCountType nr_children = 0; 
    DepthHits depth_and_hits; 
  };
  
  // forward declaration
  struct ACTrieState;
 
  struct OPENMS_DLLAPI ACScout
  {
    ACScout() = delete;
 
    ACScout(const char* query_pos, Index tree_pos, uint8_t max_aa, uint8_t max_mm, uint8_t max_prefix_loss);
 
    size_t textPos(const ACTrieState& state) const;
 
    AA nextValidAA();
 
    const char* it_query = 0; 
    Index tree_pos;                       
    uint8_t max_aaa_leftover {0};         
    uint8_t max_mm_leftover {0};          
    uint8_t max_prefix_loss_leftover {0}; 
  };
 
  AA nextValidAA(const char*& it_q);
 
  struct OPENMS_DLLAPI ACTrieState
  {
    friend ACScout;
    void setQuery(const std::string& haystack);
 
    size_t textPos() const;
 
    const char* textPosIt() const;
 
    const std::string& getQuery() const;
 
    AA nextValidAA();
 
    std::vector<Hit> hits;             
    std::queue<ACScout> scouts;
    Index tree_pos;                    
  private:
    const char* it_q_ = "";
    std::string query_;                
  };
 
  class OPENMS_DLLAPI ACTrie
  {
  public:
    ACTrie(uint32_t max_aaa = 0, uint32_t max_mm = 0);
 
    ~ACTrie();
 
    void addNeedle(const std::string& needle);
 
    void addNeedles(const std::vector<std::string>& needles);
 
    void addNeedlesAndCompress(const std::vector<std::string>& needles);
 
    size_t size() const
    {
      return trie_.size();
    }
 
    void compressTrie();
 
    size_t getNeedleCount() const;
 
    void setMaxAAACount(const uint32_t max_aaa);
 
    uint32_t getMaxAAACount() const;
    
    void setMaxMMCount(const uint32_t max_mm);
 
    uint32_t getMaxMMCount() const;
 
    bool nextHits(ACTrieState& state) const;
 
    void getAllHits(ACTrieState& state) const;
 
  private:
    bool nextHitsNoClear_(ACTrieState& state) const;
 
    Index add_(const Index from, const AA edge);
 
    bool addHits_(Index i, const size_t text_pos, std::vector<Hit>& hits) const;
 
    bool addHitsScout_(Index i, const ACScout& scout, const size_t text_pos, std::vector<Hit>& hits, const int current_scout_depths) const;
 
    Index follow_(const Index i, const AA edge) const;
 
    bool followScout_(ACScout& scout, const AA edge, ACTrieState& state) const;
 
    Index stepPrimary_(const Index i, const AA edge, ACTrieState& state) const;
 
    bool stepScout_(ACScout& scout, ACTrieState& state) const;
 
    void createScouts_(const Index i, const AA fromAA, const AA toAA, ACTrieState& state, const uint32_t aaa_left, const uint32_t mm_left) const;
 
    void createSubScouts_(const ACScout& prototype, const AA fromAA, const AA toAA, ACTrieState& state) const;
 
    void createMMScouts_(const Index i, const AA except_fromAA, const AA except_toAA, const AA except_edge, ACTrieState& state, const uint32_t aaa_left, const uint32_t mm_left) const;
 
    void createMMSubScouts_(const ACScout& prototype, const AA except_fromAA, const AA except_toAA, const AA except_edge, ACTrieState& state) const;
 
    Index findChildNaive_(Index parent, AA child_label);
 
    Index findChildBFS_(const Index parent, const AA child_label) const;
 
    std::vector<ACNode> trie_;  
    uint32_t needle_count_ {0}; 
    uint32_t max_aaa_ {0};      
    uint32_t max_mm_ {0};       
 
    std::vector<std::vector<uint32_t>> vec_index2needles_ = { {} }; // one empty element to allow capacity doubling; note: sparse vector but still much faster and less memory than unordered_map
    std::vector<std::vector<Index>> vec_index2children_naive_ = { {} }; // one empty element to allow capacity doubling
  };
 
} // namespace OpenMS