Skip to main content

seq_hash/
lib.rs

1//! A crate for streaming hashing of k-mers via `KmerHasher`.
2//!
3//! This builds on [`packed_seq`] and is used by e.g. [`simd_minimizers`].
4//!
5//! The default [`NtHasher`] is canonical.
6//! If that's not needed, [`NtHasher<false>`] will be slightly faster.
7//! For non-DNA sequences with >2-bit alphabets, use [`MulHasher`] instead.
8//!
9//! Note that [`KmerHasher`] objects need `k` on their construction, so that they can precompute required constants.
10//! Prefer reusing the same [`KmerHasher`].
11//!
12//! This crate also includes [`AntiLexHasher`], see [this blogpost](https://curiouscoding.nl/posts/practical-minimizers/).
13//!
14//! ## Typical usage
15//!
16//! Construct a default [`NtHasher`] via `let hasher = <NtHasher>::new(k)`.
17//! Then call either `hasher.hash_kmers_simd(seq, context)`,
18//! or use the underlying 'mapper' via `hasher.in_out_mapper_simd(seq)`.
19//! ```
20//! use seq_hash::packed_seq;
21//! use packed_seq::{AsciiSeqVec, PackedSeqVec, SeqVec};
22//! use seq_hash::{KmerHasher, NtHasher};
23//! let k = 3;
24//!
25//! // Default `NtHasher` is canonical.
26//! let hasher = <NtHasher>::new(k);
27//! let kmer = PackedSeqVec::from_ascii(b"ACG");
28//! let kmer_rc = PackedSeqVec::from_ascii(b"CGT");
29//! // Normally, prefer `hash_kmers_simd` over `hash_seq`.
30//! assert_eq!(
31//!     hasher.hash_seq(kmer.as_slice()),
32//!     hasher.hash_seq(kmer_rc.as_slice())
33//! );
34//!
35//! let fwd_hasher = NtHasher::<false>::new(k);
36//! assert_ne!(
37//!     fwd_hasher.hash_seq(kmer.as_slice()),
38//!     fwd_hasher.hash_seq(kmer_rc.as_slice())
39//! );
40//!
41//! let seq = b"ACGGCAGCGCATATGTAGT";
42//! let ascii_seq = AsciiSeqVec::from_ascii(seq);
43//! let packed_seq = PackedSeqVec::from_ascii(seq);
44//!
45//! // hasher.hash_kmers_scalar(seq.as_slice()); // Panics since `NtHasher` does not support ASCII.
46//! let hashes_1: Vec<_> = hasher.hash_kmers_scalar(ascii_seq.as_slice()).collect();
47//! let hashes_2: Vec<_> = hasher.hash_kmers_scalar(packed_seq.as_slice()).collect();
48//! // Hashes are equal for [`packed_seq::AsciiSeq`] and [`packed_seq::PackedSeq`].
49//! assert_eq!(hashes_1, hashes_2);
50//! assert_eq!(hashes_1.len(), seq.len() - (k-1));
51//!
52//! // Consider a 'context' of a single kmer.
53//! let hashes_3: Vec<_> = hasher.hash_kmers_simd(ascii_seq.as_slice(), 1).collect();
54//! let hashes_4: Vec<_> = hasher.hash_kmers_simd(packed_seq.as_slice(), 1).collect();
55//! assert_eq!(hashes_1, hashes_3);
56//! assert_eq!(hashes_1, hashes_4);
57//! ```
58
59mod anti_lex;
60mod intrinsics;
61mod nthash;
62#[cfg(test)]
63mod test;
64
65pub use anti_lex::AntiLexHasher;
66pub use nthash::{MulHasher, NtHasher};
67
68/// Re-export of the `packed-seq` crate.
69pub use packed_seq;
70
71use packed_seq::{ChunkIt, Delay, PackedNSeq, PaddedIt, Seq};
72use std::iter::{repeat, zip};
73
74type S = packed_seq::u32x8;
75
76/// A hasher that can hash all k-mers in a string.
77///
78/// Note that a `KmerHasher` must be initialized with a specific `k`,
79/// so that it can precompute associated constants.
80pub trait KmerHasher {
81    /// True when the hash function is invariant under reverse-complement.
82    const CANONICAL: bool;
83
84    fn new(k: usize) -> Self;
85
86    /// Helper function returning [`Self::CANONICAL`].
87    #[inline(always)]
88    fn is_canonical(&self) -> bool {
89        Self::CANONICAL
90    }
91
92    /// The value of `k` for this hasher.
93    fn k(&self) -> usize;
94
95    /// The delay of the 'out' character passed to the `in_out_mapper` functions.
96    /// Defaults to `k-1`.
97    #[inline(always)]
98    fn delay(&self) -> Delay {
99        Delay(self.k() - 1)
100    }
101
102    /// A scalar mapper function that should be called with each `(in, out)` base.
103    ///
104    /// The delay should be [`Self::delay()`]. The first `delay` calls should have `out=0`.
105    /// `seq` is only used to ensure that the hasher can handle the underlying alphabet.
106    fn in_out_mapper_scalar<'s>(&self, seq: impl Seq<'s>) -> impl FnMut((u8, u8)) -> u32;
107    /// A SIMD mapper function that should be called with a `(in, out)` base per lane.
108    ///
109    /// The delay should be [`Self::delay()`]. The first `delay` calls should have `out=u32x8::splat(0)`.
110    /// `seq` is only used to ensure that the hasher can handle the underlying alphabet.
111    fn in_out_mapper_simd<'s>(&self, seq: impl Seq<'s>) -> impl FnMut((S, S)) -> S;
112
113    fn in_out_mapper_ambiguous_scalar<'s>(
114        &self,
115        nseq: PackedNSeq<'s>,
116    ) -> impl FnMut((u8, u8)) -> u32 {
117        let mut mapper = self.in_out_mapper_scalar(nseq.seq);
118        let mut ambiguous = nseq.ambiguous.iter_kmer_ambiguity(self.k());
119        let k = self.k();
120        let mut i = 0;
121        move |(a, r)| {
122            let hash = mapper((a, r));
123            let ambiguous = if i > k - 1 {
124                ambiguous.next().unwrap()
125            } else {
126                false
127            };
128            i += 1;
129            if ambiguous { u32::MAX } else { hash }
130        }
131    }
132
133    #[inline(always)]
134    fn in_out_mapper_ambiguous_simd<'s>(
135        &self,
136        nseq: PackedNSeq<'s>,
137        context: usize,
138    ) -> impl FnMut((S, S)) -> S {
139        let mut mapper = self.in_out_mapper_simd(nseq.seq);
140        let mut ambiguous = nseq.ambiguous.par_iter_kmer_ambiguity(self.k(), context, 0);
141        move |(a, r)| {
142            let hash = mapper((a, r));
143            let ambiguous = ambiguous.it.next().unwrap();
144            ambiguous.select(S::MAX, hash)
145        }
146    }
147
148    /// A scalar iterator over all k-mer hashes in `seq`.
149    #[inline(always)]
150    fn hash_kmers_scalar<'s>(&self, seq: impl Seq<'s>) -> impl ExactSizeIterator<Item = u32> {
151        let k = self.k();
152        let delay = self.delay();
153        let mut add = seq.iter_bp();
154        let mut remove = seq.iter_bp();
155        let mut mapper = self.in_out_mapper_scalar(seq);
156        zip(add.by_ref().take(delay.0), repeat(0)).for_each(|a| {
157            mapper(a);
158        });
159        zip(add.by_ref(), remove.by_ref())
160            .take(k - 1 - delay.0)
161            .for_each(|a| {
162                mapper(a);
163            });
164        zip(add, remove).map(mapper)
165    }
166
167    /// A SIMD-parallel iterator over all k-mer hashes in `seq`.
168    #[inline(always)]
169    fn hash_kmers_simd<'s>(&self, seq: impl Seq<'s>, context: usize) -> PaddedIt<impl ChunkIt<S>> {
170        let k = self.k();
171        let delay = self.delay();
172        seq.par_iter_bp_delayed(context + k - 1, delay)
173            .map(self.in_out_mapper_simd(seq))
174            .advance(k - 1)
175    }
176
177    /// An iterator over all k-mer hashes in `seq`.
178    /// Ambiguous kmers get hash `u32::MAX`.
179    #[inline(always)]
180    fn hash_valid_kmers_scalar<'s>(
181        &self,
182        nseq: PackedNSeq<'s>,
183    ) -> impl ExactSizeIterator<Item = u32> {
184        let k = self.k();
185        let delay = self.delay();
186        assert!(delay.0 < k);
187
188        let mut mapper = self.in_out_mapper_scalar(nseq.seq);
189
190        let mut a = nseq.seq.iter_bp();
191        let mut r = nseq.seq.iter_bp();
192
193        a.by_ref().take(delay.0).for_each(
194            #[inline(always)]
195            |a| {
196                mapper((a, 0));
197            },
198        );
199
200        zip(a.by_ref(), r.by_ref())
201            .take((k - 1) - delay.0)
202            .for_each(
203                #[inline(always)]
204                |(a, r)| {
205                    mapper((a, r));
206                },
207            );
208
209        zip(zip(a, r), nseq.ambiguous.iter_kmer_ambiguity(k)).map(
210            #[inline(always)]
211            move |(ar, ambiguous)| {
212                let hash = mapper(ar);
213                if ambiguous { u32::MAX } else { hash }
214            },
215        )
216    }
217
218    /// A SIMD-parallel iterator over all k-mer hashes in `seq`.
219    /// Ambiguous kmers get hash `u32::MAX`.
220    #[inline(always)]
221    fn hash_valid_kmers_simd<'s, 't>(
222        &'t self,
223        nseq: PackedNSeq<'s>,
224        context: usize,
225    ) -> PaddedIt<impl ChunkIt<S> + use<'s, 't, Self>> {
226        let k = self.k();
227        let delay = self.delay();
228        let mut hash_mapper = self.in_out_mapper_simd(nseq.seq);
229        let mut ambiguity_it = nseq
230            .ambiguous
231            .par_iter_kmer_ambiguity(k, context + k - 1, 0);
232        nseq.seq
233            .par_iter_bp_delayed_with_factor(context + k - 1, delay, 2)
234            .map(
235                #[inline(always)]
236                move |(a, r)| {
237                    // SAFETY: these iterators have the same length.
238                    let is_ambiguous = unsafe { ambiguity_it.it.next().unwrap_unchecked() };
239                    let hash = hash_mapper((a, r));
240                    is_ambiguous.select(S::MAX, hash)
241                },
242            )
243            .advance(k - 1)
244    }
245
246    /// Hash a sequence one character at a time. Ignores `k`.
247    ///
248    /// `seq` is only used to ensure that the hasher can handle the underlying alphabet.
249    fn mapper<'s>(&self, seq: impl Seq<'s>) -> impl FnMut(u8) -> u32;
250
251    /// Hash the given sequence. Ignores `k`.
252    ///
253    /// This is slightly inefficient because it recomputes the constants based on the sequence length.
254    #[inline(always)]
255    fn hash_seq<'s>(&self, seq: impl Seq<'s>) -> u32 {
256        seq.iter_bp().map(self.mapper(seq)).last().unwrap_or(0)
257    }
258    /// Hash all non-empty prefixes of the given sequence. Ignores `k`.
259    #[inline(always)]
260    fn hash_prefixes<'s>(&self, seq: impl Seq<'s>) -> impl ExactSizeIterator<Item = u32> {
261        seq.iter_bp().map(self.mapper(seq))
262    }
263}