sassy 0.2.1

Approximate string matching using SIMD
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152

use crate::pattern_tiling::backend::SimdBackend;
use crate::profiles::{Profile, iupac::reverse_complement};
use std::marker::PhantomData;

pub(crate) const IUPAC_MASKS: usize = 16;

#[derive(Debug, Clone, Default)]
pub struct TQueries<B: SimdBackend, P: Profile> {
    /// Holds [position][transposition_block_index]
    pub pattern_length: usize,
    pub n_queries: usize,
    pub n_original_queries: usize,
    /// Number of SIMD vectors needed to hold all queries (for Peqs)
    pub n_simd_blocks: usize,
    pub peq_masks: [Vec<u64>; IUPAC_MASKS],
    pub peqs: Vec<B::Simd>,
    /// Pre-computed offsets into peqs for each character
    /// peq_offsets[encoded_char] is the starting index for that character's peq blocks
    pub peq_offsets: [usize; IUPAC_MASKS],
    // Not sure if we should keep it here, but might be handy with
    // get_pattern_seq(i) to make sure it aligns with the indices of `scan` result bools
    pub queries: Vec<Vec<u8>>,
    _marker: PhantomData<(B, P)>,
}

#[inline(always)]
fn build_flat_peqs<B: SimdBackend>(
    peq_masks: &[Vec<u64>; IUPAC_MASKS],
    nq: usize,
    n_simd_blocks: usize,
) -> Vec<B::Simd> {
    let mut peqs = Vec::with_capacity(IUPAC_MASKS * n_simd_blocks);

    for mask_vec in peq_masks.iter() {
        for block_idx in 0..n_simd_blocks {
            let base = block_idx * B::LANES;
            let limit = (nq - base).min(B::LANES);

            let mut lane = B::LaneArray::default();
            let lane_slice = lane.as_mut();

            for lane_idx in 0..limit {
                lane_slice[lane_idx] = B::mask_word_to_scalar(mask_vec[base + lane_idx]);
            }
            peqs.push(B::from_array(lane));
        }
    }
    peqs
}

#[allow(clippy::needless_range_loop)]
impl<B: SimdBackend, P: Profile> TQueries<B, P> {
    pub fn new(queries: &[Vec<u8>], include_rc: bool) -> Self {
        if P::N_CHARS > IUPAC_MASKS {
            panic!("Please raise issue on GitHub if you are using this for ASCII profile");
        }

        assert!(!queries.is_empty(), "No queries provided");
        let pattern_length = queries[0].len();
        assert!(
            pattern_length > 0 && pattern_length <= B::LIMB_BITS,
            "Invalid pattern length {} (must be <= {})",
            pattern_length,
            B::LIMB_BITS
        );
        assert!(
            queries.iter().all(|q| q.len() == pattern_length),
            "All queries must have the same length"
        );

        let n_original_queries = queries.len();
        let est_capacity = n_original_queries * (1 + include_rc as usize);
        let mut all_queries = Vec::with_capacity(est_capacity);
        all_queries.extend_from_slice(queries);
        if include_rc {
            for q in queries {
                let q_rc = reverse_complement(q);
                all_queries.push(q_rc);
            }
        }

        let n_queries = all_queries.len();
        let n_transposition_blocks = n_queries.div_ceil(B::LIMB_BITS);
        let n_simd_blocks = n_queries.div_ceil(B::LANES);

        let mut peq_masks: [Vec<u64>; IUPAC_MASKS] = std::array::from_fn(|_| vec![0u64; n_queries]);

        let mut temp_block_buffer = vec![0u8; B::LIMB_BITS];

        for pos in 0..pattern_length {
            for block_idx in 0..n_transposition_blocks {
                let start_q = block_idx * B::LIMB_BITS;
                let end_q = (start_q + B::LIMB_BITS).min(n_queries);
                let count = end_q - start_q;

                temp_block_buffer.fill(0);

                // iter queries
                for i in 0..count {
                    let qi = start_q + i;
                    let encoded = P::encode_char(all_queries[qi][pos]);
                    temp_block_buffer[i] = encoded;
                    if encoded != 0 {
                        // Skipping X
                        let bit = 1u64 << pos;
                        for (mask_idx, mask_vec) in peq_masks.iter_mut().enumerate().skip(1) {
                            if (encoded & (mask_idx as u8)) != 0 {
                                mask_vec[qi] |= bit;
                            }
                        }
                    }
                }
            }
        }

        let peqs = build_flat_peqs::<B>(&peq_masks, n_queries, n_simd_blocks);

        let mut peq_offsets = [0usize; IUPAC_MASKS];
        for i in 0..IUPAC_MASKS {
            peq_offsets[i] = i * n_simd_blocks;
        }

        Self {
            pattern_length,
            n_queries,
            n_original_queries,
            n_simd_blocks,
            peq_masks,
            peqs,
            peq_offsets,
            queries: all_queries,
            _marker: PhantomData,
        }
    }

    pub fn get_pattern_seq(&self, pattern_idx: usize) -> &[u8] {
        &self.queries[pattern_idx]
    }

    pub fn reduce_to_suffix<NB: SimdBackend, NP: Profile>(&self) -> TQueries<NB, NP> {
        let suffix_len = NB::LIMB_BITS;
        let suffix_queries: Vec<Vec<u8>> = self
            .queries
            .iter()
            .map(|q| {
                let start = q.len().saturating_sub(suffix_len);
                q[start..].to_vec()
            })
            .collect();
        TQueries::<NB, NP>::new(&suffix_queries, false)
    }
}