ktrs_syntax/tree/
kind_scan.rs

1//! [`KindScan`]: the iterator behind [`Tree::find_kinds`](super::Tree::find_kinds).
2
3use super::ElementId;
4
5const CHUNK: usize = 64;
6/// Raw kinds a scan can look for; unused slots hold [`NEVER`].
7pub(super) const MAX_WANTED: usize = 12;
8/// No raw kind: kinds are under 512, with or without the token bit.
9pub(super) const NEVER: u16 = u16::MAX;
10
11/// Matches 64 raw kinds at a time into a bitmask (the compares vectorize), then yields its set bits.
12pub struct KindScan<'t> {
13    kinds: &'t [u16],
14    /// Element id of `kinds[0]`.
15    base: usize,
16    /// Start of the chunk `mask` describes.
17    chunk: usize,
18    mask: u64,
19    wanted: [u16; MAX_WANTED],
20}
21
22impl<'t> KindScan<'t> {
23    pub(super) fn new(kinds: &'t [u16], base: usize, wanted: [u16; MAX_WANTED]) -> Self {
24        let mut scan = KindScan { kinds, base, chunk: 0, mask: 0, wanted };
25        scan.mask = scan.chunk_mask(0);
26        scan
27    }
28
29    fn chunk_mask(&self, start: usize) -> u64 {
30        let wanted = self.wanted;
31        let is_wanted = |k: u16| wanted.iter().fold(false, |a, &w| a | (w == k)) as u8;
32        let mut hits = [0u8; CHUNK];
33        match self.kinds.get(start..start + CHUNK) {
34            Some(chunk) => {
35                let chunk: &[u16; CHUNK] = chunk.try_into().unwrap();
36                for (hit, &k) in hits.iter_mut().zip(chunk) {
37                    *hit = is_wanted(k);
38                }
39            }
40            None => {
41                for (hit, &k) in hits.iter_mut().zip(self.kinds.get(start..).unwrap_or(&[])) {
42                    *hit = is_wanted(k);
43                }
44            }
45        }
46        let mut mask = 0;
47        for (g, bytes) in hits.chunks_exact(8).enumerate() {
48            // Gathers the 0/1 bytes into the top byte: byte i lands in bit 56 + i.
49            let packed = u64::from_le_bytes(bytes.try_into().unwrap()).wrapping_mul(0x0102_0408_1020_4080) >> 56;
50            mask |= packed << (g * 8);
51        }
52        mask
53    }
54}
55
56impl Iterator for KindScan<'_> {
57    type Item = ElementId;
58
59    fn next(&mut self) -> Option<ElementId> {
60        while self.mask == 0 {
61            self.chunk += CHUNK;
62            if self.chunk >= self.kinds.len() {
63                return None;
64            }
65            self.mask = self.chunk_mask(self.chunk);
66        }
67        let bit = self.mask.trailing_zeros() as usize;
68        self.mask &= self.mask - 1;
69        Some((self.base + self.chunk + bit) as ElementId)
70    }
71}
ktrs_syntax/tree/kind_scan.rs

ktrs_syntax/tree/
kind_scan.rs
