Skip to main content

seq_hash/
anti_lex.rs

1//! 'Anti lexicographic' hashing:
2//!
3//! A kmer's hash found by reading it's characters right to left, and by inverting the last (most significant) character.
4//! When k > 16, only the last 16 characters are used.
5
6use std::cmp::min;
7
8use crate::{KmerHasher, S};
9use packed_seq::{Delay, Seq};
10
11/// A hash function that compares strings reverse-lexicographically,
12/// with the last (most significant) character inverted.
13///
14/// Only supports 2-bit DNA sequences ([`packed_seq::AsciiSeq`] and [`packed_seq::PackedSeq`]).
15///
16/// The canonical version (with `CANONICAL=true`) returns the minimum of the forward and reverse-complement hashes.
17/// TODO: Test minimum vs maximum.
18pub struct AntiLexHasher<const CANONICAL: bool> {
19    k: usize,
20    /// Number of bits of each character.
21    b: usize,
22    /// Number of bits to shift each new character up to make it the most significant one.
23    shift: u32,
24    /// Shift to get everything into the most significant positions.
25    left_shift: u32,
26    /// Mask to flip the bits of the most significant character.
27    anti: u32,
28    /// Mask to keep only the lowest k*b bits.
29    mask: u32,
30}
31
32impl<const CANONICAL: bool> AntiLexHasher<CANONICAL> {
33    /// Create a new [`AntiLexHasher`] for kmers of length `k`.
34    #[inline(always)]
35    pub const fn new(k: usize) -> Self {
36        let b = 2;
37        let shift = if b * k <= 32 { b * (k - 1) } else { 32 - b } as u32;
38        let anti = ((1 << b) - 1) << shift;
39        let mask = if b * k < 32 {
40            (1 << (b * k)) - 1
41        } else {
42            u32::MAX
43        };
44        Self {
45            k,
46            b,
47            shift,
48            left_shift: 32u32.saturating_sub((b * k) as u32),
49            anti,
50            mask,
51        }
52    }
53}
54
55impl KmerHasher for AntiLexHasher<false> {
56    const CANONICAL: bool = false;
57
58    fn new(k: usize) -> Self {
59        Self::new(k)
60    }
61
62    #[inline(always)]
63    fn k(&self) -> usize {
64        self.k
65    }
66
67    #[inline(always)]
68    fn in_out_mapper_scalar<'s>(&self, seq: impl Seq<'s>) -> impl FnMut((u8, u8)) -> u32 {
69        assert!(seq.bits_per_char() <= self.b);
70
71        let mut fw: u32 = 0;
72        move |(a, _r)| {
73            fw = (fw >> self.b) ^ ((a as u32) << self.shift);
74            (fw ^ self.anti) << self.left_shift
75        }
76    }
77
78    #[inline(always)]
79    fn in_out_mapper_simd<'s>(&self, seq: impl Seq<'s>) -> impl FnMut((S, S)) -> S {
80        assert!(seq.bits_per_char() <= self.b);
81
82        let mut fw: S = S::splat(0);
83        move |(a, _r)| {
84            fw = (fw >> self.b as u32) ^ (a << self.shift);
85            (fw ^ S::splat(self.anti)) << self.left_shift
86        }
87    }
88
89    #[inline(always)]
90    fn mapper<'s>(&self, seq: impl Seq<'s>) -> impl FnMut(u8) -> u32 {
91        assert!(seq.bits_per_char() <= self.b);
92        let k = seq.len();
93        let shift = if self.b * k <= 32 {
94            self.b * (k - 1)
95        } else {
96            32 - self.b
97        } as u32;
98        let anti = ((1 << self.b) - 1) << shift;
99
100        let mut fw: u32 = 0;
101        move |a| {
102            fw = (fw >> self.b) ^ ((a as u32) << shift);
103            (fw ^ anti) << self.left_shift
104        }
105    }
106}
107
108impl KmerHasher for AntiLexHasher<true> {
109    const CANONICAL: bool = true;
110
111    fn new(k: usize) -> Self {
112        Self::new(k)
113    }
114
115    #[inline(always)]
116    fn k(&self) -> usize {
117        self.k
118    }
119
120    #[inline(always)]
121    fn delay(&self) -> Delay {
122        Delay(self.k.saturating_sub(32 / self.b))
123    }
124
125    #[inline(always)]
126    fn mapper<'s>(&self, seq: impl Seq<'s>) -> impl FnMut(u8) -> u32 {
127        assert!(seq.bits_per_char() <= self.b);
128        let mut shift = 0;
129        let mut anti = (1 << self.b) - 1;
130        let mut mask = anti;
131
132        let mut fw: u32 = 0;
133        let mut rc: u32 = 0;
134        let mut i = 0;
135        move |a| {
136            if i * self.b >= 32 {
137                fw >>= self.b;
138            }
139            fw ^= (a as u32) << shift;
140            if i * self.b < 32 {
141                // ^2 for complement.
142                rc = ((rc << self.b) & mask) ^ (a as u32 ^ 2);
143            }
144            let out = min(fw ^ anti, rc ^ anti);
145
146            if (i + 1) * self.b < 32 {
147                shift += self.b as u32;
148                anti <<= self.b;
149                mask = (mask << self.b) | ((1 << self.b) - 1);
150            }
151            i += 1;
152
153            out << self.left_shift
154        }
155    }
156
157    #[inline(always)]
158    fn in_out_mapper_scalar<'s>(&self, seq: impl Seq<'s>) -> impl FnMut((u8, u8)) -> u32 {
159        assert!(seq.bits_per_char() <= self.b);
160
161        let mut fw: u32 = 0;
162        let mut rc: u32 = 0;
163        move |(a, r)| {
164            fw = (fw >> self.b) ^ ((a as u32) << self.shift);
165            // ^2 for complement.
166            rc = ((rc << self.b) & self.mask) ^ (r as u32 ^ 2);
167            min(fw ^ self.anti, rc ^ self.anti) << self.left_shift
168        }
169    }
170
171    #[inline(always)]
172    fn in_out_mapper_simd<'s>(&self, seq: impl Seq<'s>) -> impl FnMut((S, S)) -> S {
173        assert!(seq.bits_per_char() <= self.b);
174
175        let mut fw: S = S::splat(0);
176        let mut rc: S = S::splat(0);
177        move |(a, r)| {
178            fw = (fw >> self.b as u32) ^ (a << self.shift);
179            rc = ((rc << self.b as u32) & S::splat(self.mask)) ^ (r ^ S::splat(2));
180            (fw ^ S::splat(self.anti)).min(rc ^ S::splat(self.anti)) << self.left_shift
181        }
182    }
183}