keyhammer/trie.rs
1// SPDX-License-Identifier: AGPL-3.0-or-later
2// Copyright (C) 2026 Robson Trasel
3
4//! A flat, breadth-first trie stored as parallel arrays.
5//!
6//! Children of a node are contiguous and always have larger indices than the
7//! node, so the structure has no cycles by construction.
8//!
9//! Edges are labelled with symbols: the Unicode scalar values of the terms,
10//! one per `char` (see `docs/design/unicode.md`). For ASCII a symbol is the
11//! byte.
12
13use alloc::collections::VecDeque;
14use alloc::string::String;
15use alloc::vec;
16use alloc::vec::Vec;
17use core::fmt;
18use core::ops::Range;
19
20use crate::cost::symbol_class;
21use crate::text::Normalizer;
22
23/// Value of [`Trie::term_id`] for nodes where no term ends.
24pub const NO_TERM: u32 = u32::MAX;
25
26/// Why a trie could not be built.
27#[derive(Clone, Copy, Debug, PartialEq, Eq)]
28#[non_exhaustive]
29pub enum BuildError {
30 /// No terms were given.
31 Empty,
32 /// A term was the empty string (for [`Trie::build_normalized`]: after
33 /// normalisation, as for a term made only of combining marks).
34 EmptyTerm,
35 /// A term is longer than 65535 bytes (for [`Trie::build_normalized`]:
36 /// after normalisation).
37 TermTooLong,
38 /// More than `u32::MAX` terms were given: term ids and input indices are
39 /// `u32`, and `u32::MAX` itself is reserved for [`NO_TERM`].
40 TooManyTerms,
41}
42
43impl fmt::Display for BuildError {
44 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
45 match self {
46 BuildError::Empty => f.write_str("cannot build an index from an empty term list"),
47 BuildError::EmptyTerm => f.write_str("terms must not be empty"),
48 BuildError::TermTooLong => f.write_str("terms must be at most 65535 bytes"),
49 BuildError::TooManyTerms => f.write_str("an index holds at most 4294967295 terms"),
50 }
51 }
52}
53
54/// Converts a term position or input index to `u32`, refusing values that do
55/// not fit or that would equal [`NO_TERM`].
56fn index_u32(i: usize) -> Result<u32, BuildError> {
57 match u32::try_from(i) {
58 Ok(v) if v != NO_TERM => Ok(v),
59 _ => Err(BuildError::TooManyTerms),
60 }
61}
62
63/// Edge labels: one byte per node while every label is at most U+00FF (every
64/// ASCII or Latin-1 dictionary), four bytes otherwise.
65#[derive(Clone, Debug)]
66enum Labels {
67 Narrow(Vec<u8>),
68 Wide(Vec<u32>),
69}
70
71impl Labels {
72 #[inline]
73 fn get(&self, v: usize) -> u32 {
74 match self {
75 Labels::Narrow(l) => u32::from(l[v]),
76 Labels::Wide(l) => l[v],
77 }
78 }
79
80 fn len(&self) -> usize {
81 match self {
82 Labels::Narrow(l) => l.len(),
83 Labels::Wide(l) => l.len(),
84 }
85 }
86}
87
88/// The structure of a trie before its aggregates are computed.
89struct Shape {
90 child_start: Vec<u32>,
91 child_count: Vec<u16>,
92 term_id: Vec<u32>,
93 depth_of: Vec<u16>,
94}
95
96/// Lays out the trie of `terms` (sorted, distinct, non-empty) breadth first;
97/// returns the edge labels and the structure.
98fn shape<T: Copy + Default + PartialEq>(terms: &[&[T]]) -> Result<(Vec<T>, Shape), BuildError> {
99 let mut label = vec![T::default()];
100 let mut child_start = vec![0u32];
101 let mut child_count = vec![0u16];
102 let mut term_id = vec![NO_TERM];
103 let mut depth_of = vec![0u16];
104
105 let mut queue: VecDeque<(usize, usize, usize, usize)> = VecDeque::new();
106 queue.push_back((0, 0, terms.len(), 0));
107 while let Some((node, mut lo, hi, depth)) = queue.pop_front() {
108 if terms[lo].len() == depth {
109 term_id[node] = index_u32(lo)?;
110 lo += 1;
111 }
112 let first_child = label.len() as u32;
113 let mut count = 0u16;
114 let mut i = lo;
115 while i < hi {
116 let b = terms[i][depth];
117 let mut j = i + 1;
118 while j < hi && terms[j][depth] == b {
119 j += 1;
120 }
121 let child = label.len();
122 label.push(b);
123 child_start.push(0);
124 child_count.push(0);
125 term_id.push(NO_TERM);
126 depth_of.push((depth + 1) as u16);
127 queue.push_back((child, i, j, depth + 1));
128 count += 1;
129 i = j;
130 }
131 child_start[node] = first_child;
132 child_count[node] = count;
133 }
134 Ok((
135 label,
136 Shape {
137 child_start,
138 child_count,
139 term_id,
140 depth_of,
141 },
142 ))
143}
144
145/// A trie over strings, split into Unicode scalar values, with a weight per
146/// term.
147#[derive(Clone, Debug)]
148pub struct Trie {
149 terms: Vec<String>,
150 weights: Vec<u16>,
151 input_index: Vec<u32>,
152 label: Labels,
153 child_start: Vec<u32>,
154 child_count: Vec<u16>,
155 term_id: Vec<u32>,
156 max_weight: Vec<u16>,
157 len_min: Vec<u16>,
158 len_max: Vec<u16>,
159 below_mask: Vec<u64>,
160 normalizer: Option<Normalizer>,
161}
162
163impl Trie {
164 /// Builds a trie. Duplicate terms are merged, keeping the entry with the
165 /// highest weight (the first one among equal weights).
166 ///
167 /// Term ids are positions in the byte-sorted, deduplicated list, not
168 /// positions in `items`; [`Trie::input_index`] maps an id back to the
169 /// index in `items` of the entry that was kept.
170 ///
171 /// Terms are stored as given, one symbol per code point, without any
172 /// folding: for ASCII this is the behaviour of earlier versions exactly.
173 /// [`Trie::build_normalized`] folds case and diacritics.
174 pub fn build(items: &[(&str, u16)]) -> Result<Trie, BuildError> {
175 if items.is_empty() {
176 return Err(BuildError::Empty);
177 }
178 for (t, _) in items {
179 if t.is_empty() {
180 return Err(BuildError::EmptyTerm);
181 }
182 if t.len() > usize::from(u16::MAX) {
183 return Err(BuildError::TermTooLong);
184 }
185 }
186 // (term, weight, input index)
187 let mut pairs: Vec<(&str, u16, u32)> = items
188 .iter()
189 .enumerate()
190 .map(|(i, &(t, w))| Ok((t, w, index_u32(i)?)))
191 .collect::<Result<_, BuildError>>()?;
192 pairs.sort_by(|a, b| {
193 a.0.as_bytes()
194 .cmp(b.0.as_bytes())
195 .then(b.1.cmp(&a.1))
196 .then(a.2.cmp(&b.2))
197 });
198 pairs.dedup_by(|cur, prev| cur.0 == prev.0);
199
200 // ASCII terms are walked as bytes (the symbols are the bytes); any
201 // other dictionary as code points, in one buffer.
202 let (
203 label,
204 Shape {
205 child_start,
206 child_count,
207 term_id,
208 depth_of,
209 },
210 ) = if pairs.iter().all(|p| p.0.is_ascii()) {
211 let terms: Vec<&[u8]> = pairs.iter().map(|p| p.0.as_bytes()).collect();
212 let (label, shape) = shape(&terms)?;
213 (Labels::Narrow(label), shape)
214 } else {
215 let mut syms: Vec<u32> = Vec::new();
216 let mut start: Vec<usize> = Vec::with_capacity(pairs.len() + 1);
217 for p in &pairs {
218 start.push(syms.len());
219 syms.extend(p.0.chars().map(u32::from));
220 }
221 start.push(syms.len());
222 let terms: Vec<&[u32]> = start.windows(2).map(|w| &syms[w[0]..w[1]]).collect();
223 let (label, shape) = shape(&terms)?;
224 let label = if label.iter().any(|&c| c > 0xFF) {
225 Labels::Wide(label)
226 } else {
227 Labels::Narrow(label.into_iter().map(|c| c as u8).collect())
228 };
229 (label, shape)
230 };
231
232 let weights: Vec<u16> = pairs.iter().map(|p| p.1).collect();
233 let n = label.len();
234 let mut max_weight = vec![0u16; n];
235 let mut len_min = vec![u16::MAX; n];
236 let mut len_max = vec![0u16; n];
237 let mut below_mask = vec![0u64; n];
238 for v in (0..n).rev() {
239 if term_id[v] != NO_TERM {
240 max_weight[v] = weights[term_id[v] as usize];
241 len_min[v] = depth_of[v];
242 len_max[v] = depth_of[v];
243 }
244 let start = child_start[v] as usize;
245 for c in start..start + usize::from(child_count[v]) {
246 max_weight[v] = max_weight[v].max(max_weight[c]);
247 len_min[v] = len_min[v].min(len_min[c]);
248 len_max[v] = len_max[v].max(len_max[c]);
249 below_mask[v] |= symbol_class(label.get(c)) | below_mask[c];
250 }
251 }
252
253 let terms = pairs.iter().map(|p| String::from(p.0)).collect();
254 let input_index = pairs.iter().map(|p| p.2).collect();
255 Ok(Trie {
256 terms,
257 weights,
258 input_index,
259 label,
260 child_start,
261 child_count,
262 term_id,
263 max_weight,
264 len_min,
265 len_max,
266 below_mask,
267 normalizer: None,
268 })
269 }
270
271 /// Builds a trie of the terms normalised by `normalizer`, and remembers
272 /// it so that [`Searcher::search_text`](crate::search::Searcher::search_text)
273 /// normalises queries the same way.
274 ///
275 /// Terms that are equal after normalisation are merged as by
276 /// [`Trie::build`] (highest weight, then first); [`Trie::term`] returns the
277 /// normalised text and [`Trie::input_index`] the index in `items` of the
278 /// entry kept, from which the caller recovers its original spelling.
279 ///
280 /// ```
281 /// use keyhammer::text::Normalizer;
282 /// use keyhammer::trie::Trie;
283 ///
284 /// let items = [("Café", 5), ("cafe", 9), ("São Paulo", 1)];
285 /// let trie = Trie::build_normalized(&items, &Normalizer::new()).unwrap();
286 /// assert_eq!(trie.len(), 2);
287 /// assert_eq!(trie.term(0), "cafe");
288 /// assert_eq!(trie.input_index(0), 1); // "cafe" (weight 9) was kept
289 /// assert_eq!(trie.term(1), "sao paulo");
290 /// assert_eq!(trie.normalizer(), Some(Normalizer::new()));
291 /// ```
292 pub fn build_normalized(
293 items: &[(&str, u16)],
294 normalizer: &Normalizer,
295 ) -> Result<Trie, BuildError> {
296 let texts: Vec<String> = items.iter().map(|(t, _)| normalizer.normalize(t)).collect();
297 let normalized: Vec<(&str, u16)> = texts
298 .iter()
299 .zip(items)
300 .map(|(t, &(_, w))| (t.as_str(), w))
301 .collect();
302 let mut trie = Trie::build(&normalized)?;
303 trie.normalizer = Some(*normalizer);
304 Ok(trie)
305 }
306
307 /// The normaliser the trie was built with ([`Trie::build_normalized`]),
308 /// or `None` ([`Trie::build`]).
309 pub fn normalizer(&self) -> Option<Normalizer> {
310 self.normalizer
311 }
312
313 /// Number of distinct terms.
314 pub fn len(&self) -> usize {
315 self.terms.len()
316 }
317
318 /// Whether the trie holds no terms (never true for a built trie).
319 pub fn is_empty(&self) -> bool {
320 self.terms.is_empty()
321 }
322
323 /// Number of trie nodes, root included.
324 pub fn node_count(&self) -> usize {
325 self.label.len()
326 }
327
328 /// The term with the given id.
329 pub fn term(&self, id: u32) -> &str {
330 self.terms.get(id as usize).map_or("", String::as_str)
331 }
332
333 /// The weight of the term with the given id.
334 pub fn weight(&self, id: u32) -> u16 {
335 self.weights.get(id as usize).copied().unwrap_or(0)
336 }
337
338 /// The index, in the slice given to [`Trie::build`], of the entry kept
339 /// for the term with the given id: for duplicates, the entry with the
340 /// highest weight, and the first one among equal weights. Returns
341 /// `u32::MAX` for an out-of-range id.
342 pub fn input_index(&self, id: u32) -> u32 {
343 self.input_index
344 .get(id as usize)
345 .copied()
346 .unwrap_or(u32::MAX)
347 }
348
349 /// The symbol on the edge leading into node `v`, as a `char` (`'\0'` for
350 /// the root).
351 ///
352 /// # Panics
353 ///
354 /// Panics if `v` is not a node index, that is, if `v >= self.node_count()`.
355 /// Unlike [`Trie::term`], [`Trie::weight`] and [`Trie::input_index`], which
356 /// answer for an unknown id, the per-node accessors index directly.
357 ///
358 /// # Examples
359 ///
360 /// ```
361 /// use keyhammer::trie::Trie;
362 /// let trie = Trie::build(&[("car", 9), ("cat", 4)]).unwrap();
363 /// // Nodes in breadth-first order: root, c, a, r, t.
364 /// assert_eq!(trie.node_count(), 5);
365 /// assert_eq!(trie.label(0), '\0');
366 /// assert_eq!(trie.label(1), 'c');
367 /// assert_eq!(trie.label(4), 't');
368 /// // One node per code point, not per UTF-8 byte.
369 /// let trie = Trie::build(&[("né", 1)]).unwrap();
370 /// assert_eq!(trie.node_count(), 3);
371 /// assert_eq!(trie.label(2), 'é');
372 /// ```
373 ///
374 /// An index past the last node panics:
375 ///
376 /// ```should_panic
377 /// # use keyhammer::trie::Trie;
378 /// let trie = Trie::build(&[("a", 1)]).unwrap();
379 /// trie.label(trie.node_count());
380 /// ```
381 pub fn label(&self, v: usize) -> char {
382 char::from_u32(self.label.get(v)).unwrap_or('\0')
383 }
384
385 /// The symbol on the edge into `v` as a `u32` (what the search compares).
386 #[inline]
387 pub(crate) fn symbol(&self, v: usize) -> u32 {
388 self.label.get(v)
389 }
390
391 /// The node indices of the children of `v`.
392 ///
393 /// # Panics
394 ///
395 /// Panics if `v` is not a node index, that is, if `v >= self.node_count()`.
396 /// Unlike [`Trie::term`], [`Trie::weight`] and [`Trie::input_index`], which
397 /// answer for an unknown id, the per-node accessors index directly.
398 ///
399 /// # Examples
400 ///
401 /// ```
402 /// use keyhammer::trie::Trie;
403 /// let trie = Trie::build(&[("car", 9), ("cat", 4)]).unwrap();
404 /// // Nodes in breadth-first order: root, c, a, r, t.
405 /// assert_eq!(trie.node_count(), 5);
406 /// assert_eq!(trie.children(0), 1..2);
407 /// assert_eq!(trie.children(2), 3..5); // a -> r, t
408 /// assert!(trie.children(3).is_empty());
409 /// ```
410 pub fn children(&self, v: usize) -> Range<usize> {
411 let start = self.child_start[v] as usize;
412 start..start + usize::from(self.child_count[v])
413 }
414
415 /// The id of the term that ends at `v`, or [`NO_TERM`].
416 ///
417 /// # Panics
418 ///
419 /// Panics if `v` is not a node index, that is, if `v >= self.node_count()`.
420 /// Unlike [`Trie::term`], [`Trie::weight`] and [`Trie::input_index`], which
421 /// answer for an unknown id, the per-node accessors index directly.
422 ///
423 /// # Examples
424 ///
425 /// ```
426 /// use keyhammer::trie::{NO_TERM, Trie};
427 /// let trie = Trie::build(&[("car", 9), ("cat", 4)]).unwrap();
428 /// // Nodes in breadth-first order: root, c, a, r, t.
429 /// assert_eq!(trie.node_count(), 5);
430 /// assert_eq!(trie.term_id(0), NO_TERM);
431 /// assert_eq!(trie.term_id(3), 0); // "car"
432 /// assert_eq!(trie.term_id(4), 1); // "cat"
433 /// ```
434 pub fn term_id(&self, v: usize) -> u32 {
435 self.term_id[v]
436 }
437
438 /// The largest weight among the terms at or below `v`.
439 ///
440 /// # Panics
441 ///
442 /// Panics if `v` is not a node index, that is, if `v >= self.node_count()`.
443 /// Unlike [`Trie::term`], [`Trie::weight`] and [`Trie::input_index`], which
444 /// answer for an unknown id, the per-node accessors index directly.
445 ///
446 /// # Examples
447 ///
448 /// ```
449 /// use keyhammer::trie::Trie;
450 /// let trie = Trie::build(&[("car", 9), ("cat", 4)]).unwrap();
451 /// // Nodes in breadth-first order: root, c, a, r, t.
452 /// assert_eq!(trie.node_count(), 5);
453 /// assert_eq!(trie.max_weight(0), 9);
454 /// assert_eq!(trie.max_weight(4), 4);
455 /// ```
456 pub fn max_weight(&self, v: usize) -> u16 {
457 self.max_weight[v]
458 }
459
460 /// The length, in symbols, of the shortest term at or below `v`.
461 ///
462 /// # Panics
463 ///
464 /// Panics if `v` is not a node index, that is, if `v >= self.node_count()`.
465 /// Unlike [`Trie::term`], [`Trie::weight`] and [`Trie::input_index`], which
466 /// answer for an unknown id, the per-node accessors index directly.
467 ///
468 /// # Examples
469 ///
470 /// ```
471 /// use keyhammer::trie::Trie;
472 /// let trie = Trie::build(&[("a", 1), ("abc", 1)]).unwrap();
473 /// assert_eq!(trie.len_min(0), 1);
474 /// assert_eq!(trie.len_min(trie.children(0).start), 1);
475 /// ```
476 pub fn len_min(&self, v: usize) -> u16 {
477 self.len_min[v]
478 }
479
480 /// The length, in symbols, of the longest term at or below `v`.
481 ///
482 /// # Panics
483 ///
484 /// Panics if `v` is not a node index, that is, if `v >= self.node_count()`.
485 /// Unlike [`Trie::term`], [`Trie::weight`] and [`Trie::input_index`], which
486 /// answer for an unknown id, the per-node accessors index directly.
487 ///
488 /// # Examples
489 ///
490 /// ```
491 /// use keyhammer::trie::Trie;
492 /// let trie = Trie::build(&[("a", 1), ("abc", 1)]).unwrap();
493 /// assert_eq!(trie.len_max(0), 3);
494 /// let a = trie.children(0).start;
495 /// assert_eq!(trie.len_max(a), 3);
496 /// ```
497 pub fn len_max(&self, v: usize) -> u16 {
498 self.len_max[v]
499 }
500
501 /// The symbol classes ([`symbol_class`]) that appear on edges strictly
502 /// below `v`.
503 ///
504 /// # Panics
505 ///
506 /// Panics if `v` is not a node index, that is, if `v >= self.node_count()`.
507 /// Unlike [`Trie::term`], [`Trie::weight`] and [`Trie::input_index`], which
508 /// answer for an unknown id, the per-node accessors index directly.
509 ///
510 /// # Examples
511 ///
512 /// ```
513 /// use keyhammer::cost::symbol_class;
514 /// use keyhammer::trie::Trie;
515 /// let trie = Trie::build(&[("ab", 1)]).unwrap();
516 /// let a = trie.children(0).start;
517 /// let class = |c: char| symbol_class(c.into());
518 /// assert_eq!(trie.below_mask(0), class('a') | class('b'));
519 /// assert_eq!(trie.below_mask(a), class('b')); // not its own edge
520 /// let b = trie.children(a).start;
521 /// assert_eq!(trie.below_mask(b), 0);
522 /// ```
523 pub fn below_mask(&self, v: usize) -> u64 {
524 self.below_mask[v]
525 }
526}
527
528#[cfg(test)]
529mod tests {
530 use super::*;
531
532 #[test]
533 fn index_conversion_refuses_what_does_not_fit_or_equals_no_term() {
534 assert_eq!(index_u32(0), Ok(0));
535 assert_eq!(index_u32(NO_TERM as usize - 1), Ok(NO_TERM - 1));
536 assert_eq!(index_u32(NO_TERM as usize), Err(BuildError::TooManyTerms));
537 if let Some(big) = (NO_TERM as usize).checked_add(1) {
538 assert_eq!(index_u32(big), Err(BuildError::TooManyTerms));
539 }
540 }
541}