sudachi/dic/lexicon/
trie.rs

1/*
2 * Copyright (c) 2021-2026 Works Applications Co., Ltd.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *     http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17use crate::util::cow_array::CowArray;
18use crate::util::prefetch::prefetch_l1;
19use std::iter::FusedIterator;
20
21#[derive(Debug, Eq, PartialEq, Clone)]
22pub struct TrieEntry {
23    /// Value of Trie, this is not the pointer to WordId, but the offset in WordId table
24    pub value: u32,
25    /// Offset of word end
26    pub end: usize,
27}
28
29impl TrieEntry {
30    #[inline]
31    pub fn new(value: u32, offset: usize) -> TrieEntry {
32        TrieEntry { value, end: offset }
33    }
34}
35
36pub struct Trie<'a> {
37    array: CowArray<'a, u32>,
38}
39
40pub struct TrieEntryIter<'a> {
41    trie: &'a [u32],
42    node_pos: usize,
43    data: &'a [u8],
44    offset: usize,
45}
46
47impl<'a> Iterator for TrieEntryIter<'a> {
48    type Item = TrieEntry;
49
50    #[inline]
51    fn next(&mut self) -> Option<Self::Item> {
52        let mut node_pos = self.node_pos;
53
54        for i in self.offset..self.data.len() {
55            // Unwrap is safe: access is always in bounds
56            // It is optimized away: https://rust.godbolt.org/z/va9K3az4n
57            let k = *self.data.get(i).unwrap();
58            match step_once(self.trie, k, &mut node_pos) {
59                Step::Dead => return None,
60                Step::Continue => {}
61                Step::Match { value } => {
62                    let r = TrieEntry::new(value, i + 1);
63                    self.offset = r.end;
64                    self.node_pos = node_pos;
65                    return Some(r);
66                }
67            }
68        }
69        None
70    }
71}
72
73impl FusedIterator for TrieEntryIter<'_> {}
74
75/// Outcome of consuming one input byte during a double-array trie walk.
76enum Step {
77    Dead,
78    Continue,
79    Match { value: u32 },
80}
81
82/// Consume one input byte `k` from the node at `node_pos`, advancing it. Shared
83/// by [`TrieEntryIter`] and [`Trie::common_prefix_batch`] so the scalar and
84/// pipelined walks are identical by construction.
85#[inline(always)]
86fn step_once(trie: &[u32], k: u8, node_pos: &mut usize) -> Step {
87    let k = k as usize;
88    *node_pos ^= k;
89    // Safe: the trie is built so every reachable index is in bounds.
90    let unit = *unsafe { trie.get_unchecked(*node_pos) } as usize;
91    if Trie::label(unit) != k {
92        return Step::Dead;
93    }
94    *node_pos ^= Trie::offset(unit);
95    if Trie::has_leaf(unit) {
96        Step::Match {
97            value: Trie::value(*unsafe { trie.get_unchecked(*node_pos) }),
98        }
99    } else {
100        Step::Continue
101    }
102}
103
104/// Speculatively prefetch the cache line a walk will load next: `node_pos ^
105/// input[pos]` is the index `step_once` will dereference on this lane's next
106/// visit. The index is clamped; the prefetch is only a hint, so a wrong address
107/// is harmless.
108#[inline(always)]
109fn prefetch_lane(trie: &[u32], input: &[u8], node_pos: usize, pos: usize) {
110    if pos < input.len() {
111        let k = *unsafe { input.get_unchecked(pos) } as usize;
112        let idx = (node_pos ^ k).min(trie.len().saturating_sub(1));
113        prefetch_l1(unsafe { trie.as_ptr().add(idx) });
114    }
115}
116
117impl<'a> Trie<'a> {
118    /// Number of independent walks [`Trie::common_prefix_batch`] keeps in flight.
119    pub const DEFAULT_PREFETCH_LANES: usize = 4;
120    /// Upper bound on lanes so the scheduler can keep lane state on the stack.
121    pub const MAX_PREFETCH_LANES: usize = 16;
122
123    pub fn from_bytes(data: &'a [u8]) -> Trie<'a> {
124        Trie {
125            array: CowArray::from_bytes(data, 0, data.len() / std::mem::size_of::<u32>()),
126        }
127    }
128
129    pub fn new(data: &'a [u8], size: usize) -> Trie<'a> {
130        Trie {
131            array: CowArray::from_bytes(data, 0, size),
132        }
133    }
134
135    pub fn new_owned(data: Vec<u32>) -> Trie<'a> {
136        Trie {
137            array: CowArray::from_owned(data),
138        }
139    }
140
141    pub fn total_size(&self) -> usize {
142        4 * self.array.len()
143    }
144
145    #[inline]
146    pub fn common_prefix_iterator<'b>(&'a self, input: &'b [u8], offset: usize) -> TrieEntryIter<'b>
147    where
148        'a: 'b,
149    {
150        let unit: usize = self.get(0) as usize;
151
152        TrieEntryIter {
153            node_pos: Trie::offset(unit),
154            data: input,
155            trie: &self.array,
156            offset,
157        }
158    }
159
160    /// Run common-prefix search from many start positions, overlapping their
161    /// memory latency.
162    ///
163    /// `emit(bucket, value, end)` fires once per match; `bucket` indexes
164    /// `starts`. Within a bucket, matches keep [`Trie::common_prefix_iterator`]
165    /// order, so grouping by `bucket` reproduces the scalar result.
166    #[inline]
167    pub fn common_prefix_batch<F: FnMut(usize, u32, usize)>(
168        &self,
169        input: &[u8],
170        starts: &[usize],
171        emit: F,
172    ) {
173        self.common_prefix_batch_cfg(input, starts, Self::DEFAULT_PREFETCH_LANES, true, emit)
174    }
175
176    /// [`Trie::common_prefix_batch`] with an explicit lane count and prefetch
177    /// toggle. `lanes` is rounded to the nearest compile-time-specialized count;
178    /// values above [`Trie::MAX_PREFETCH_LANES`] are clamped.
179    #[inline]
180    pub fn common_prefix_batch_cfg<F: FnMut(usize, u32, usize)>(
181        &self,
182        input: &[u8],
183        starts: &[usize],
184        lanes: usize,
185        prefetch: bool,
186        emit: F,
187    ) {
188        // Specialize on lanes and prefetch so both are compile-time constants.
189        macro_rules! dispatch {
190            ($k:literal) => {
191                if prefetch {
192                    self.batch_impl::<$k, true, F>(input, starts, emit)
193                } else {
194                    self.batch_impl::<$k, false, F>(input, starts, emit)
195                }
196            };
197        }
198        match lanes.clamp(1, Self::MAX_PREFETCH_LANES) {
199            1 => dispatch!(1),
200            2 => dispatch!(2),
201            3..=4 => dispatch!(4),
202            5..=6 => dispatch!(6),
203            7..=8 => dispatch!(8),
204            9..=12 => dispatch!(12),
205            _ => dispatch!(16),
206        }
207    }
208
209    fn batch_impl<const K: usize, const PF: bool, F: FnMut(usize, u32, usize)>(
210        &self,
211        input: &[u8],
212        starts: &[usize],
213        mut emit: F,
214    ) {
215        if starts.is_empty() {
216            return;
217        }
218        let trie: &[u32] = &self.array;
219        let root = Trie::offset(self.get(0) as usize);
220        let n_in = input.len();
221
222        // Per-lane state, structure-of-arrays for constant `K`.
223        let mut node = [0usize; K];
224        let mut pos = [0usize; K];
225        let mut bucket = [0usize; K];
226        let mut active = [false; K];
227        let mut next_start = 0usize;
228        let mut live = 0usize;
229
230        for i in 0..K {
231            if next_start < starts.len() {
232                node[i] = root;
233                pos[i] = starts[next_start];
234                bucket[i] = next_start;
235                active[i] = true;
236                next_start += 1;
237                live += 1;
238                if PF {
239                    prefetch_lane(trie, input, root, pos[i]);
240                }
241            }
242        }
243
244        while live > 0 {
245            for i in 0..K {
246                if !active[i] {
247                    continue;
248                }
249                let advanced = if pos[i] < n_in {
250                    let k = *unsafe { input.get_unchecked(pos[i]) };
251                    let mut np = node[i];
252                    match step_once(trie, k, &mut np) {
253                        Step::Dead => false,
254                        Step::Continue => {
255                            node[i] = np;
256                            pos[i] += 1;
257                            if PF {
258                                prefetch_lane(trie, input, np, pos[i]);
259                            }
260                            true
261                        }
262                        Step::Match { value } => {
263                            let end = pos[i] + 1;
264                            emit(bucket[i], value, end);
265                            node[i] = np;
266                            pos[i] = end;
267                            if PF {
268                                prefetch_lane(trie, input, np, end);
269                            }
270                            true
271                        }
272                    }
273                } else {
274                    false
275                };
276
277                if !advanced {
278                    if next_start < starts.len() {
279                        node[i] = root;
280                        pos[i] = starts[next_start];
281                        bucket[i] = next_start;
282                        next_start += 1;
283                        if PF {
284                            prefetch_lane(trie, input, root, pos[i]);
285                        }
286                    } else {
287                        active[i] = false;
288                        live -= 1;
289                    }
290                }
291            }
292        }
293    }
294
295    #[inline(always)]
296    fn get(&self, index: usize) -> u32 {
297        debug_assert!(index < self.array.len());
298        // UB if out of bounds
299        // Should we panic in release builds here instead?
300        // Safe version is not optimized away
301        *unsafe { self.array.get_unchecked(index) }
302    }
303
304    #[inline(always)]
305    fn has_leaf(unit: usize) -> bool {
306        ((unit >> 8) & 1) == 1
307    }
308
309    #[inline(always)]
310    fn value(unit: u32) -> u32 {
311        unit & ((1 << 31) - 1)
312    }
313
314    #[inline(always)]
315    fn label(unit: usize) -> usize {
316        unit & ((1 << 31) | 0xFF)
317    }
318
319    #[inline(always)]
320    fn offset(unit: usize) -> usize {
321        (unit >> 10) << ((unit & (1 << 9)) >> 6)
322    }
323}
324
325#[cfg(test)]
326mod tests {
327    use super::Trie;
328
329    /// Build a double array from `(key, value)` pairs.
330    fn build_trie(mut keys: Vec<(&str, u32)>) -> Vec<u8> {
331        keys.sort_by(|a, b| a.0.cmp(b.0));
332        yada::builder::DoubleArrayBuilder::build(&keys).expect("trie build failed")
333    }
334
335    /// Per-start reference output of the scalar iterator.
336    fn scalar_reference(trie: &Trie, input: &[u8], starts: &[usize]) -> Vec<Vec<(u32, usize)>> {
337        starts
338            .iter()
339            .map(|&s| {
340                trie.common_prefix_iterator(input, s)
341                    .map(|e| (e.value, e.end))
342                    .collect()
343            })
344            .collect()
345    }
346
347    fn batch_grouped(
348        trie: &Trie,
349        input: &[u8],
350        starts: &[usize],
351        lanes: usize,
352        prefetch: bool,
353    ) -> Vec<Vec<(u32, usize)>> {
354        let mut got: Vec<Vec<(u32, usize)>> = vec![Vec::new(); starts.len()];
355        trie.common_prefix_batch_cfg(input, starts, lanes, prefetch, |bucket, value, end| {
356            got[bucket].push((value, end));
357        });
358        got
359    }
360
361    #[test]
362    fn batch_matches_scalar_across_lanes_and_prefetch() {
363        // Shared prefixes, a prefix-of-another key, multibyte UTF-8, and misses.
364        let bytes = build_trie(vec![
365            ("a", 1),
366            ("ab", 2),
367            ("abc", 3),
368            ("abd", 4),
369            ("b", 5),
370            ("bc", 6),
371            ("東", 7),
372            ("東京", 8),
373            ("東京都", 9),
374        ]);
375        let trie = Trie::from_bytes(&bytes);
376
377        for text in ["abcabd b 東京都 abZ", "東京", "", "zzz", "ab東京都bca"] {
378            let input = text.as_bytes();
379            let starts: Vec<usize> = (0..=input.len()).collect();
380            let reference = scalar_reference(&trie, input, &starts);
381            for &lanes in &[1usize, 2, 4, 6, 8, 12, 16] {
382                for &prefetch in &[false, true] {
383                    let got = batch_grouped(&trie, input, &starts, lanes, prefetch);
384                    assert_eq!(
385                        got, reference,
386                        "mismatch for text {text:?} lanes={lanes} prefetch={prefetch}"
387                    );
388                }
389            }
390        }
391    }
392
393    #[test]
394    fn batch_handles_sparse_and_unordered_starts() {
395        let bytes = build_trie(vec![("ab", 2), ("abc", 3), ("xy", 10)]);
396        let trie = Trie::from_bytes(&bytes);
397        let input = b"abcxyab";
398        // Non-contiguous, repeated, and out-of-order starts.
399        let starts = [6usize, 0, 3, 0, 3];
400        let reference = scalar_reference(&trie, input, &starts);
401        for &lanes in &[1usize, 2, 4, 16] {
402            for &prefetch in &[false, true] {
403                let got = batch_grouped(&trie, input, &starts, lanes, prefetch);
404                assert_eq!(got, reference, "lanes={lanes} prefetch={prefetch}");
405            }
406        }
407    }
408}
sudachi/dic/lexicon/trie.rs

sudachi/dic/lexicon/
trie.rs
