musli_zerocopy/trie/mod.rs
1//! A serialized prefix-trie.
2
3#[cfg(test)]
4mod tests;
5
6#[cfg(feature = "alloc")]
7pub use self::factory::{store, Builder};
8#[cfg(feature = "alloc")]
9mod factory;
10
11use self::walk::Walk;
12mod walk;
13
14use core::fmt;
15use core::marker::PhantomData;
16
17#[cfg(feature = "alloc")]
18use alloc::vec::Vec;
19
20use crate::endian::Native;
21use crate::lossy_str::LossyStr;
22use crate::slice::{binary_search_by, BinarySearch, Slice};
23use crate::stack::ArrayStack;
24use crate::{Buf, ByteOrder, DefaultSize, Error, Ref, Size, ZeroCopy};
25
26type StackEntry<'buf, T, F> = (LinksRef<T, F>, usize, &'buf [u8]);
27
28/// The flavor of a trie. Allows for customization of implementation details to
29/// for example use a more compact representation than the one provided by
30/// default in case there is a known upper bound on the number of elements or
31/// values.
32///
33/// # Examples
34///
35/// ```
36/// use musli_zerocopy::{trie, Error, OwnedBuf, ZeroCopy};
37/// use musli_zerocopy::slice::Packed;
38///
39/// struct PackedTrie;
40///
41/// impl trie::Flavor for PackedTrie {
42/// // The maximum length of a string slice stored in the trie is `u8::MAX`.
43/// type String = Packed<[u8], u32, u8>;
44///
45/// // The max number of values stored in a single node is `u16::MAX`.
46/// type Values<T> = Packed<[T], u32, u16>
47/// where
48/// T: ZeroCopy;
49///
50/// // The maximum number of children for a single node is `u8::MAX`.
51/// type Children<T> = Packed<[T], u32, u8>
52/// where
53/// T: ZeroCopy;
54/// }
55///
56/// fn populate<F>(buf: &mut OwnedBuf, mut trie: trie::Builder<u32, F>) -> Result<trie::TrieRef<u32, F>, Error>
57/// where
58/// F: trie::Flavor
59/// {
60/// let a = buf.store_unsized("hello");
61/// let b = buf.store_unsized("hello world");
62/// trie.insert(buf, a, 1)?;
63/// trie.insert(buf, b, 2)?;
64/// trie.build(buf)
65/// }
66///
67/// let mut b1 = OwnedBuf::new();
68///
69/// let mut trie = trie::Builder::<u32, PackedTrie>::with_flavor();
70/// let trie = populate(&mut b1, trie)?;
71///
72/// assert_eq!(trie.get(&b1, "hello")?, Some(&[1][..]));
73/// assert_eq!(trie.get(&b1, "hello world")?, Some(&[2][..]));
74///
75/// let mut b2 = OwnedBuf::new();
76///
77/// let mut trie = trie::Builder::new();
78/// let trie = populate(&mut b2, trie)?;
79///
80/// assert_eq!(trie.get(&b2, "hello")?, Some(&[1][..]));
81/// assert_eq!(trie.get(&b2, "hello world")?, Some(&[2][..]));
82///
83/// assert!(b1.len() < b2.len());
84/// # Ok::<_, musli_zerocopy::Error>(())
85/// ```
86pub trait Flavor {
87 /// The type representing a string in the trie.
88 type String: Slice<Item = u8>;
89
90 /// The type representing a collection of values in the trie.
91 type Values<T>: Slice<Item = T>
92 where
93 T: ZeroCopy;
94
95 /// The type representing a collection of children in the trie.
96 type Children<T>: Slice<Item = T>
97 where
98 T: ZeroCopy;
99}
100
101/// Marker type indicating the default trie [`Flavor`] to use for a given
102/// [`ByteOrder`] and [`Size`].
103pub struct DefaultFlavor<E = Native, O = DefaultSize>(PhantomData<(E, O)>)
104where
105 E: ByteOrder,
106 O: Size;
107
108impl<E, O> Flavor for DefaultFlavor<E, O>
109where
110 E: ByteOrder,
111 O: Size,
112{
113 type String = Ref<[u8], E, O>;
114 type Values<T>
115 = Ref<[T], E, O>
116 where
117 T: ZeroCopy;
118 type Children<T>
119 = Ref<[T], E, O>
120 where
121 T: ZeroCopy;
122}
123
124/// A stored reference to a trie.
125#[derive(ZeroCopy)]
126#[zero_copy(crate)]
127#[repr(C)]
128pub struct TrieRef<T, F = DefaultFlavor>
129where
130 T: ZeroCopy,
131 F: Flavor,
132{
133 links: LinksRef<T, F>,
134}
135
136impl<T, F> TrieRef<T, F>
137where
138 T: ZeroCopy,
139 F: Flavor,
140{
141 /// Debug print the current trie.
142 ///
143 /// This treats the keys as strings, with illegal unicode sequences being
144 /// replaced with the `U+FFFD` escape sequence.
145 ///
146 /// # Examples
147 ///
148 /// ```
149 /// use musli_zerocopy::{trie, OwnedBuf};
150 ///
151 /// let mut buf = OwnedBuf::new();
152 /// let a = buf.store(b"\xe2\x28\xa1").array_into_slice();
153 /// let b = buf.store_unsized("食べない");
154 ///
155 /// let mut trie = trie::Builder::new();
156 ///
157 /// trie.insert(&buf, b, 2)?;
158 /// trie.insert(&buf, a, 1)?;
159 ///
160 /// let trie = trie.build(&mut buf)?;
161 /// assert_eq!(format!("{:?}", trie.debug(&buf)), "{\"�(�\": 1, \"食べない\": 2}");
162 /// # Ok::<_, musli_zerocopy::Error>(())
163 /// ```
164 #[cfg(feature = "alloc")]
165 pub fn debug<'a, 'buf>(&'a self, buf: &'buf Buf) -> Debug<'a, 'buf, T, F>
166 where
167 T: fmt::Debug,
168 {
169 Debug { trie: self, buf }
170 }
171
172 /// Debug print the current trie, with a fixed iteration depth of `N`.
173 ///
174 /// This treats the keys as strings, with illegal unicode sequences being
175 /// replaced with the `U+FFFD` escape sequence.
176 ///
177 /// # Examples
178 ///
179 /// ```
180 /// use musli_zerocopy::{trie, OwnedBuf};
181 ///
182 /// let mut buf = OwnedBuf::new();
183 /// let a = buf.store(b"\xe2\x28\xa1").array_into_slice();
184 /// let b = buf.store_unsized("食べない");
185 ///
186 /// let mut trie = trie::Builder::new();
187 ///
188 /// trie.insert(&buf, b, 2)?;
189 /// trie.insert(&buf, a, 1)?;
190 ///
191 /// let trie = trie.build(&mut buf)?;
192 /// assert_eq!(format!("{:?}", trie.debug_fixed::<16>(&buf)), "{\"�(�\": 1, \"食べない\": 2}");
193 /// # Ok::<_, musli_zerocopy::Error>(())
194 /// ```
195 pub fn debug_fixed<'a, 'buf, const N: usize>(
196 &'a self,
197 buf: &'buf Buf,
198 ) -> DebugFixed<'a, 'buf, N, T, F>
199 where
200 T: fmt::Debug,
201 {
202 DebugFixed { trie: self, buf }
203 }
204
205 /// Get all values associated with the given string.
206 ///
207 /// # Examples
208 ///
209 /// ```
210 /// use musli_zerocopy::{trie, OwnedBuf};
211 ///
212 /// let mut buf = OwnedBuf::new();
213 ///
214 /// let values = [
215 /// (buf.store_unsized("work"), 1),
216 /// (buf.store_unsized("worker"), 2),
217 /// (buf.store_unsized("workers"), 3),
218 /// (buf.store_unsized("working"), 4),
219 /// (buf.store_unsized("working"), 5),
220 /// (buf.store_unsized("working man"), 6),
221 /// ];
222 ///
223 /// let trie = trie::store(&mut buf, values)?;
224 ///
225 /// assert_eq!(trie.get(&buf, "aard")?, None);
226 /// assert_eq!(trie.get(&buf, "worker")?, Some(&[2][..]));
227 /// assert_eq!(trie.get(&buf, "working")?, Some(&[4, 5][..]));
228 /// # Ok::<_, musli_zerocopy::Error>(())
229 /// ```
230 pub fn get<'buf, S>(&self, buf: &'buf Buf, string: &S) -> Result<Option<&'buf [T]>, Error>
231 where
232 S: ?Sized + AsRef<[u8]>,
233 {
234 let mut this = self.links;
235 let mut string = string.as_ref();
236
237 loop {
238 let search =
239 binary_search_by(buf, this.children, |c| Ok(buf.load(c.string)?.cmp(string)))?;
240
241 match search {
242 BinarySearch::Found(n) => {
243 let child = this.children.get_unchecked(n);
244 let child = buf.load(child)?;
245 let values = buf.load(child.links.values)?;
246 return Ok(Some(values));
247 }
248 BinarySearch::Missing(0) => {
249 return Ok(None);
250 }
251 BinarySearch::Missing(n) => {
252 let child = this.children.get_unchecked(n - 1);
253 let child = buf.load(child)?;
254
255 // Find common prefix and split nodes if necessary.
256 let prefix = prefix(buf.load(child.string)?, string);
257
258 if prefix == 0 {
259 return Ok(None);
260 }
261
262 string = &string[prefix..];
263 this = child.links;
264 }
265 };
266 }
267 }
268
269 /// Construct an iterator over all values in the trie.
270 ///
271 /// Note that the iteration order is unspecified and might change in future
272 /// versions.
273 ///
274 /// # Errors
275 ///
276 /// This errors in case the trie being iterated over is structurally
277 /// invalid.
278 ///
279 /// # Examples
280 ///
281 /// ```
282 /// use musli_zerocopy::{trie, OwnedBuf};
283 ///
284 /// let mut buf = OwnedBuf::new();
285 ///
286 /// let values = [
287 /// (buf.store_unsized("work"), 1),
288 /// (buf.store_unsized("worker"), 2),
289 /// (buf.store_unsized("workers"), 3),
290 /// (buf.store_unsized("working"), 4),
291 /// (buf.store_unsized("working"), 5),
292 /// (buf.store_unsized("working man"), 6),
293 /// (buf.store_unsized("run"), 7),
294 /// (buf.store_unsized("running"), 8),
295 /// ];
296 ///
297 /// let trie = trie::store(&mut buf, values)?;
298 ///
299 /// let mut values = trie.values(&buf).collect::<Result<Vec<_>, _>>()?;
300 /// values.sort();
301 /// assert!(values.into_iter().copied().eq([1, 2, 3, 4, 5, 6, 7, 8]));
302 /// # Ok::<_, musli_zerocopy::Error>(())
303 /// ```
304 #[cfg(feature = "alloc")]
305 pub fn values<'buf>(&self, buf: &'buf Buf) -> Values<'buf, T, F> {
306 Values {
307 iter: Walk::find(buf, self.links, &[]),
308 }
309 }
310
311 /// Construct an iterator over all values in the trie using a fixed max
312 /// iteration depth of `N`.
313 ///
314 /// Note that the iteration order is unspecified and might change in future
315 /// versions.
316 ///
317 /// # Errors
318 ///
319 /// This errors in case the trie being iterated over is structurally invalid
320 /// or if the iteration depth exceeds `N`.
321 ///
322 /// # Examples
323 ///
324 /// ```
325 /// use musli_zerocopy::{trie, OwnedBuf};
326 ///
327 /// let mut buf = OwnedBuf::new();
328 ///
329 /// let values = [
330 /// (buf.store_unsized("work"), 1),
331 /// (buf.store_unsized("worker"), 2),
332 /// (buf.store_unsized("workers"), 3),
333 /// (buf.store_unsized("working"), 4),
334 /// (buf.store_unsized("working"), 5),
335 /// (buf.store_unsized("working man"), 6),
336 /// (buf.store_unsized("run"), 7),
337 /// (buf.store_unsized("running"), 8),
338 /// ];
339 ///
340 /// let trie = trie::store(&mut buf, values)?;
341 ///
342 /// let mut values = trie.values_fixed::<16>(&buf).collect::<Result<Vec<_>, _>>()?;
343 /// values.sort();
344 /// assert!(values.into_iter().copied().eq([1, 2, 3, 4, 5, 6, 7, 8]));
345 /// # Ok::<_, musli_zerocopy::Error>(())
346 /// ```
347 pub fn values_fixed<'buf, const N: usize>(&self, buf: &'buf Buf) -> ValuesFixed<'buf, N, T, F> {
348 ValuesFixed {
349 iter: Walk::find(buf, self.links, &[]),
350 }
351 }
352
353 /// Construct an iterator over values that are inside of the specified
354 /// `prefix` in the trie.
355 ///
356 /// Note that the iteration order is unspecified and might change in future
357 /// versions.
358 ///
359 /// # Errors
360 ///
361 /// This errors in case the trie being iterated over is structurally
362 /// invalid.
363 ///
364 /// # Examples
365 ///
366 /// ```
367 /// use musli_zerocopy::{trie, OwnedBuf};
368 ///
369 /// let mut buf = OwnedBuf::new();
370 ///
371 /// let values = [
372 /// (buf.store_unsized("work"), 1),
373 /// (buf.store_unsized("worker"), 2),
374 /// (buf.store_unsized("workers"), 3),
375 /// (buf.store_unsized("working"), 4),
376 /// (buf.store_unsized("working"), 5),
377 /// (buf.store_unsized("working man"), 6),
378 /// (buf.store_unsized("run"), 7),
379 /// (buf.store_unsized("running"), 8),
380 /// ];
381 ///
382 /// let trie = trie::store(&mut buf, values)?;
383 ///
384 /// let values = trie.values_in(&buf, "workin").collect::<Result<Vec<_>, _>>()?;
385 /// assert!(values.into_iter().copied().eq([4, 5, 6]));
386 ///
387 /// let values = trie.values_in(&buf, "wor").collect::<Result<Vec<_>, _>>()?;
388 /// assert!(values.into_iter().copied().eq([1, 2, 3, 4, 5, 6]));
389 ///
390 /// let values = trie.values_in(&buf, "runn").collect::<Result<Vec<_>, _>>()?;
391 /// assert!(values.into_iter().copied().eq([8]));
392 /// # Ok::<_, musli_zerocopy::Error>(())
393 /// ```
394 #[cfg(feature = "alloc")]
395 pub fn values_in<'a, 'buf, S>(&self, buf: &'buf Buf, prefix: &'a S) -> ValuesIn<'a, 'buf, T, F>
396 where
397 S: ?Sized + AsRef<[u8]>,
398 {
399 ValuesIn {
400 iter: Walk::find(buf, self.links, prefix.as_ref()),
401 }
402 }
403
404 /// Construct an iterator over values that are inside of the specified
405 /// `prefix` in the trie using a fixed max iteration depth of `N`.
406 ///
407 /// Note that the iteration order is unspecified and might change in future
408 /// versions.
409 ///
410 /// # Errors
411 ///
412 /// This errors in case the trie being iterated over is structurally
413 /// invalid or if the iteration depth exceeds `N`.
414 ///
415 /// # Examples
416 ///
417 /// ```
418 /// use musli_zerocopy::{trie, OwnedBuf};
419 ///
420 /// let mut buf = OwnedBuf::new();
421 ///
422 /// let values = [
423 /// (buf.store_unsized("work"), 1),
424 /// (buf.store_unsized("worker"), 2),
425 /// (buf.store_unsized("workers"), 3),
426 /// (buf.store_unsized("working"), 4),
427 /// (buf.store_unsized("working"), 5),
428 /// (buf.store_unsized("working man"), 6),
429 /// (buf.store_unsized("run"), 7),
430 /// (buf.store_unsized("running"), 8),
431 /// ];
432 ///
433 /// let trie = trie::store(&mut buf, values)?;
434 ///
435 /// let values = trie.values_in_fixed::<16, _>(&buf, "workin").collect::<Result<Vec<_>, _>>()?;
436 /// assert!(values.into_iter().copied().eq([4, 5, 6]));
437 ///
438 /// let values = trie.values_in_fixed::<16, _>(&buf, "wor").collect::<Result<Vec<_>, _>>()?;
439 /// assert!(values.into_iter().copied().eq([1, 2, 3, 4, 5, 6]));
440 ///
441 /// let values = trie.values_in_fixed::<16, _>(&buf, "runn").collect::<Result<Vec<_>, _>>()?;
442 /// assert!(values.into_iter().copied().eq([8]));
443 /// # Ok::<_, musli_zerocopy::Error>(())
444 /// ```
445 pub fn values_in_fixed<'a, 'buf, const N: usize, S>(
446 &self,
447 buf: &'buf Buf,
448 prefix: &'a S,
449 ) -> ValuesInFixed<'a, 'buf, N, T, F>
450 where
451 S: ?Sized + AsRef<[u8]>,
452 {
453 ValuesInFixed {
454 iter: Walk::find(buf, self.links, prefix.as_ref()),
455 }
456 }
457
458 /// Construct an iterator over all entries in the trie.
459 ///
460 /// Note that the iteration order is unspecified and might change in future
461 /// versions.
462 ///
463 /// # Errors
464 ///
465 /// This errors in case the trie being iterated over is structurally
466 /// invalid.
467 ///
468 /// # Examples
469 ///
470 /// ```
471 /// use std::str::from_utf8;
472 ///
473 /// use anyhow::Result;
474 /// use musli_zerocopy::{trie, OwnedBuf};
475 ///
476 /// // Helper to convert output to utf-8.
477 /// fn to_utf8<'buf, E>(result: Result<(&'buf [u8], &'buf i32), E>) -> Result<(&'buf str, i32)>
478 /// where
479 /// anyhow::Error: From<E>
480 /// {
481 /// let (k, v) = result?;
482 /// Ok((from_utf8(k)?, *v))
483 /// }
484 ///
485 /// let mut buf = OwnedBuf::new();
486 ///
487 /// let values = [
488 /// (buf.store_unsized("work"), 1),
489 /// (buf.store_unsized("worker"), 2),
490 /// (buf.store_unsized("workers"), 3),
491 /// (buf.store_unsized("working"), 4),
492 /// (buf.store_unsized("working"), 5),
493 /// (buf.store_unsized("working man"), 6),
494 /// (buf.store_unsized("run"), 7),
495 /// (buf.store_unsized("running"), 8),
496 /// ];
497 ///
498 /// let trie = trie::store(&mut buf, values)?;
499 ///
500 /// let mut values = trie.iter(&buf)
501 /// .map(to_utf8)
502 /// .collect::<Result<Vec<_>>>()?;
503 /// values.sort();
504 ///
505 /// assert_eq! {
506 /// values,
507 /// [
508 /// ("run", 7),
509 /// ("running", 8),
510 /// ("work", 1),
511 /// ("worker", 2),
512 /// ("workers", 3),
513 /// ("working", 4),
514 /// ("working", 5),
515 /// ("working man", 6),
516 /// ]
517 /// };
518 /// # Ok::<_, anyhow::Error>(())
519 /// ```
520 #[cfg(feature = "alloc")]
521 pub fn iter<'buf>(&self, buf: &'buf Buf) -> Iter<'buf, T, F> {
522 Iter {
523 iter: Walk::find(buf, self.links, &[]),
524 }
525 }
526
527 /// Construct an iterator over all entries in the trie using a fixed max
528 /// iteration depth of `N`
529 ///
530 /// Note that the iteration order is unspecified and might change in future
531 /// versions.
532 ///
533 /// # Errors
534 ///
535 /// This errors in case the trie being iterated over is structurally
536 /// invalid or if the iteration depth exceeds `N`.
537 ///
538 /// # Examples
539 ///
540 /// ```
541 /// use std::str::from_utf8;
542 ///
543 /// use anyhow::Result;
544 /// use musli_zerocopy::{trie, OwnedBuf};
545 ///
546 /// // Helper to convert output to utf-8.
547 /// fn to_utf8<'buf, E>(result: Result<(&'buf [u8], &'buf i32), E>) -> Result<(&'buf str, i32)>
548 /// where
549 /// anyhow::Error: From<E>
550 /// {
551 /// let (k, v) = result?;
552 /// Ok((from_utf8(k)?, *v))
553 /// }
554 ///
555 /// let mut buf = OwnedBuf::new();
556 ///
557 /// let values = [
558 /// (buf.store_unsized("work"), 1),
559 /// (buf.store_unsized("worker"), 2),
560 /// (buf.store_unsized("workers"), 3),
561 /// (buf.store_unsized("working"), 4),
562 /// (buf.store_unsized("working"), 5),
563 /// (buf.store_unsized("working man"), 6),
564 /// (buf.store_unsized("run"), 7),
565 /// (buf.store_unsized("running"), 8),
566 /// ];
567 ///
568 /// let trie = trie::store(&mut buf, values)?;
569 ///
570 /// let mut values = trie.iter_fixed::<16>(&buf)
571 /// .map(to_utf8)
572 /// .collect::<Result<Vec<_>>>()?;
573 /// values.sort();
574 ///
575 /// assert_eq! {
576 /// values,
577 /// [
578 /// ("run", 7),
579 /// ("running", 8),
580 /// ("work", 1),
581 /// ("worker", 2),
582 /// ("workers", 3),
583 /// ("working", 4),
584 /// ("working", 5),
585 /// ("working man", 6),
586 /// ]
587 /// };
588 /// # Ok::<_, anyhow::Error>(())
589 /// ```
590 pub fn iter_fixed<'buf, const N: usize>(&self, buf: &'buf Buf) -> IterFixed<'buf, N, T, F> {
591 IterFixed {
592 iter: Walk::find(buf, self.links, &[]),
593 }
594 }
595
596 /// Construct an iterator over all matching string prefixes in the trie
597 /// which also returns the string of the entries.
598 ///
599 /// Note that the iteration order is unspecified and might change in future
600 /// versions.
601 ///
602 /// # Errors
603 ///
604 /// This errors in case the trie being iterated over is structurally
605 /// invalid.
606 ///
607 /// # Examples
608 ///
609 /// ```
610 /// use std::str::from_utf8;
611 ///
612 /// use anyhow::Result;
613 /// use musli_zerocopy::{trie, OwnedBuf};
614 ///
615 /// // Helper to convert output to utf-8.
616 /// fn to_utf8<'buf, E>(result: Result<(&'buf [u8], &'buf i32), E>) -> Result<(&'buf str, i32)>
617 /// where
618 /// anyhow::Error: From<E>
619 /// {
620 /// let (k, v) = result?;
621 /// Ok((from_utf8(k)?, *v))
622 /// }
623 ///
624 /// let mut buf = OwnedBuf::new();
625 ///
626 /// let values = [
627 /// (buf.store_unsized("work"), 1),
628 /// (buf.store_unsized("worker"), 2),
629 /// (buf.store_unsized("workers"), 3),
630 /// (buf.store_unsized("working"), 4),
631 /// (buf.store_unsized("working"), 5),
632 /// (buf.store_unsized("working man"), 6),
633 /// (buf.store_unsized("run"), 7),
634 /// (buf.store_unsized("running"), 8),
635 /// ];
636 ///
637 /// let trie = trie::store(&mut buf, values)?;
638 ///
639 /// let mut values = trie
640 /// .iter_in(&buf, "workin")
641 /// .map(to_utf8)
642 /// .collect::<Result<Vec<_>>>()?;
643 /// values.sort();
644 ///
645 /// assert_eq!(values, [("working", 4), ("working", 5), ("working man", 6)]);
646 ///
647 /// let mut values = trie
648 /// .iter_in(&buf, "wor")
649 /// .map(to_utf8)
650 /// .collect::<Result<Vec<_>>>()?;
651 /// values.sort();
652 ///
653 /// assert_eq! {
654 /// values,
655 /// [
656 /// ("work", 1),
657 /// ("worker", 2),
658 /// ("workers", 3),
659 /// ("working", 4),
660 /// ("working", 5),
661 /// ("working man", 6),
662 /// ]
663 /// };
664 ///
665 /// let mut values = trie
666 /// .iter_in(&buf, "runn")
667 /// .map(to_utf8)
668 /// .collect::<Result<Vec<_>>>()?;
669 /// values.sort();
670 ///
671 /// assert_eq!(values, [("running", 8)]);
672 /// # Ok::<_, anyhow::Error>(())
673 /// ```
674 #[cfg(feature = "alloc")]
675 pub fn iter_in<'a, 'buf, S>(&self, buf: &'buf Buf, prefix: &'a S) -> IterIn<'a, 'buf, T, F>
676 where
677 S: ?Sized + AsRef<[u8]>,
678 {
679 IterIn {
680 iter: Walk::find(buf, self.links, prefix.as_ref()),
681 }
682 }
683
684 /// Construct an iterator over all matching string prefixes in the trie
685 /// which also returns the string of the entries using a fixed max iteration
686 /// depth of `N`.
687 ///
688 /// Note that the iteration order is unspecified and might change in future
689 /// versions.
690 ///
691 /// # Errors
692 ///
693 /// This errors in case the trie being iterated over is structurally
694 /// invalid or if the iteration depth exceeds `N`.
695 ///
696 /// # Examples
697 ///
698 /// ```
699 /// use std::str::from_utf8;
700 ///
701 /// use anyhow::Result;
702 /// use musli_zerocopy::{trie, OwnedBuf};
703 ///
704 /// // Helper to convert output to utf-8.
705 /// fn to_utf8<'buf, E>(result: Result<(&'buf [u8], &'buf i32), E>) -> Result<(&'buf str, i32)>
706 /// where
707 /// anyhow::Error: From<E>
708 /// {
709 /// let (k, v) = result?;
710 /// Ok((from_utf8(k)?, *v))
711 /// }
712 ///
713 /// let mut buf = OwnedBuf::new();
714 ///
715 /// let values = [
716 /// (buf.store_unsized("work"), 1),
717 /// (buf.store_unsized("worker"), 2),
718 /// (buf.store_unsized("workers"), 3),
719 /// (buf.store_unsized("working"), 4),
720 /// (buf.store_unsized("working"), 5),
721 /// (buf.store_unsized("working man"), 6),
722 /// (buf.store_unsized("run"), 7),
723 /// (buf.store_unsized("running"), 8),
724 /// ];
725 ///
726 /// let trie = trie::store(&mut buf, values)?;
727 ///
728 /// let mut values = trie
729 /// .iter_in_fixed::<16, _>(&buf, "workin")
730 /// .map(to_utf8)
731 /// .collect::<Result<Vec<_>>>()?;
732 /// values.sort();
733 ///
734 /// assert_eq!(values, [("working", 4), ("working", 5), ("working man", 6)]);
735 ///
736 /// let mut values = trie
737 /// .iter_in_fixed::<16, _>(&buf, "wor")
738 /// .map(to_utf8)
739 /// .collect::<Result<Vec<_>>>()?;
740 /// values.sort();
741 ///
742 /// assert_eq! {
743 /// values,
744 /// [
745 /// ("work", 1),
746 /// ("worker", 2),
747 /// ("workers", 3),
748 /// ("working", 4),
749 /// ("working", 5),
750 /// ("working man", 6),
751 /// ]
752 /// };
753 ///
754 /// let mut values = trie
755 /// .iter_in_fixed::<16, _>(&buf, "runn")
756 /// .map(to_utf8)
757 /// .collect::<Result<Vec<_>>>()?;
758 /// values.sort();
759 ///
760 /// assert_eq!(values, [("running", 8)]);
761 /// # Ok::<_, anyhow::Error>(())
762 /// ```
763 pub fn iter_in_fixed<'a, 'buf, const N: usize, S>(
764 &self,
765 buf: &'buf Buf,
766 prefix: &'a S,
767 ) -> IterInFixed<'a, 'buf, N, T, F>
768 where
769 S: ?Sized + AsRef<[u8]>,
770 {
771 IterInFixed {
772 iter: Walk::find(buf, self.links, prefix.as_ref()),
773 }
774 }
775}
776
777/// An iterator over values matching a `prefix` in a [`TrieRef`].
778///
779/// See [`TrieRef::values_in()`].
780#[cfg(feature = "alloc")]
781pub struct ValuesIn<'a, 'buf, T, F>
782where
783 T: ZeroCopy,
784 F: Flavor,
785{
786 iter: Walk<'a, 'buf, T, F, Vec<StackEntry<'buf, T, F>>>,
787}
788
789#[cfg(feature = "alloc")]
790impl<'buf, T, F> Iterator for ValuesIn<'_, 'buf, T, F>
791where
792 T: ZeroCopy,
793 F: Flavor,
794{
795 type Item = Result<&'buf T, Error>;
796
797 #[inline]
798 fn next(&mut self) -> Option<Self::Item> {
799 let (_, value) = match self.iter.poll() {
800 Ok(entry) => entry?,
801 Err(error) => return Some(Err(error)),
802 };
803
804 Some(Ok(value))
805 }
806}
807
808/// An iterator over values matching a `prefix` in a [`TrieRef`] using a fixed
809/// max iteration depth of `N`
810///
811/// See [`TrieRef::values_in_fixed()`].
812pub struct ValuesInFixed<'a, 'buf, const N: usize, T, F>
813where
814 T: ZeroCopy,
815 F: Flavor,
816{
817 iter: Walk<'a, 'buf, T, F, ArrayStack<StackEntry<'buf, T, F>, N>>,
818}
819
820impl<'buf, const N: usize, T, F> Iterator for ValuesInFixed<'_, 'buf, N, T, F>
821where
822 T: ZeroCopy,
823 F: Flavor,
824{
825 type Item = Result<&'buf T, Error>;
826
827 #[inline]
828 fn next(&mut self) -> Option<Self::Item> {
829 let (_, value) = match self.iter.poll() {
830 Ok(entry) => entry?,
831 Err(error) => return Some(Err(error)),
832 };
833
834 Some(Ok(value))
835 }
836}
837
838/// An iterator over all values in a [`TrieRef`].
839///
840/// See [`TrieRef::values()`].
841#[cfg(feature = "alloc")]
842pub struct Values<'buf, T, F>
843where
844 T: ZeroCopy,
845 F: Flavor,
846{
847 iter: Walk<'static, 'buf, T, F, Vec<StackEntry<'buf, T, F>>>,
848}
849
850#[cfg(feature = "alloc")]
851impl<'buf, T, F> Iterator for Values<'buf, T, F>
852where
853 T: ZeroCopy,
854 F: Flavor,
855{
856 type Item = Result<&'buf T, Error>;
857
858 #[inline]
859 fn next(&mut self) -> Option<Self::Item> {
860 let (_, value) = match self.iter.poll() {
861 Ok(entry) => entry?,
862 Err(error) => return Some(Err(error)),
863 };
864
865 Some(Ok(value))
866 }
867}
868
869/// An iterator over all values in a [`TrieRef`] using a fixed max iteration
870/// depth of `N`
871///
872/// See [`TrieRef::values_fixed()`].
873pub struct ValuesFixed<'buf, const N: usize, T, F>
874where
875 T: ZeroCopy,
876 F: Flavor,
877{
878 iter: Walk<'static, 'buf, T, F, ArrayStack<StackEntry<'buf, T, F>, N>>,
879}
880
881impl<'buf, const N: usize, T, F> Iterator for ValuesFixed<'buf, N, T, F>
882where
883 T: ZeroCopy,
884 F: Flavor,
885{
886 type Item = Result<&'buf T, Error>;
887
888 #[inline]
889 fn next(&mut self) -> Option<Self::Item> {
890 let (_, value) = match self.iter.poll() {
891 Ok(entry) => entry?,
892 Err(error) => return Some(Err(error)),
893 };
894
895 Some(Ok(value))
896 }
897}
898
899/// An iterator over all entries in a [`TrieRef`].
900///
901/// See [`TrieRef::iter()`].
902#[cfg(feature = "alloc")]
903pub struct Iter<'buf, T, F>
904where
905 T: ZeroCopy,
906 F: Flavor,
907{
908 iter: Walk<'static, 'buf, T, F, Vec<StackEntry<'buf, T, F>>>,
909}
910
911#[cfg(feature = "alloc")]
912impl<'buf, T, F> Iterator for Iter<'buf, T, F>
913where
914 T: ZeroCopy,
915 F: Flavor,
916{
917 type Item = Result<(&'buf [u8], &'buf T), Error>;
918
919 #[inline]
920 fn next(&mut self) -> Option<Self::Item> {
921 self.iter.poll().transpose()
922 }
923}
924
925/// An iterator over all entries in a [`TrieRef`] using a fixed max iteration
926/// depth of `N`
927///
928/// See [`TrieRef::iter_fixed()`].
929pub struct IterFixed<'buf, const N: usize, T, F>
930where
931 T: ZeroCopy,
932 F: Flavor,
933{
934 iter: Walk<'static, 'buf, T, F, ArrayStack<StackEntry<'buf, T, F>, N>>,
935}
936
937impl<'buf, const N: usize, T, F> Iterator for IterFixed<'buf, N, T, F>
938where
939 T: ZeroCopy,
940 F: Flavor,
941{
942 type Item = Result<(&'buf [u8], &'buf T), Error>;
943
944 #[inline]
945 fn next(&mut self) -> Option<Self::Item> {
946 self.iter.poll().transpose()
947 }
948}
949
950/// An iterator over all entries inside of a `prefix` in a [`TrieRef`].
951///
952/// See [`TrieRef::iter_in()`].
953#[cfg(feature = "alloc")]
954pub struct IterIn<'a, 'buf, T, F>
955where
956 T: ZeroCopy,
957 F: Flavor,
958{
959 iter: Walk<'a, 'buf, T, F, Vec<StackEntry<'buf, T, F>>>,
960}
961
962#[cfg(feature = "alloc")]
963impl<'buf, T, F> Iterator for IterIn<'_, 'buf, T, F>
964where
965 T: ZeroCopy,
966 F: Flavor,
967{
968 type Item = Result<(&'buf [u8], &'buf T), Error>;
969
970 #[inline]
971 fn next(&mut self) -> Option<Self::Item> {
972 self.iter.poll().transpose()
973 }
974}
975
976/// An iterator over all entries inside of a `prefix` in a [`TrieRef`] using a
977/// fixed max iteration depth of `N`
978///
979/// See [`TrieRef::iter_in_fixed()`].
980pub struct IterInFixed<'a, 'buf, const N: usize, T, F>
981where
982 T: ZeroCopy,
983 F: Flavor,
984{
985 iter: Walk<'a, 'buf, T, F, ArrayStack<StackEntry<'buf, T, F>, N>>,
986}
987
988impl<'buf, const N: usize, T, F> Iterator for IterInFixed<'_, 'buf, N, T, F>
989where
990 T: ZeroCopy,
991 F: Flavor,
992{
993 type Item = Result<(&'buf [u8], &'buf T), Error>;
994
995 #[inline]
996 fn next(&mut self) -> Option<Self::Item> {
997 self.iter.poll().transpose()
998 }
999}
1000
1001/// Debug printing of a trie.
1002///
1003/// See [`TrieRef::debug()`].
1004#[cfg(feature = "alloc")]
1005pub struct Debug<'a, 'buf, T, F>
1006where
1007 T: ZeroCopy,
1008 F: Flavor,
1009{
1010 trie: &'a TrieRef<T, F>,
1011 buf: &'buf Buf,
1012}
1013
1014#[cfg(feature = "alloc")]
1015impl<T, F> fmt::Debug for Debug<'_, '_, T, F>
1016where
1017 T: fmt::Debug + ZeroCopy,
1018 F: Flavor,
1019{
1020 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1021 let mut f = f.debug_map();
1022
1023 for result in self.trie.iter(self.buf) {
1024 let (key, value) = result.map_err(|_| fmt::Error)?;
1025 f.entry(&LossyStr::new(key), value);
1026 }
1027
1028 f.finish()
1029 }
1030}
1031
1032/// Debug printing of a trie with a fixed iteration depth of `N`.
1033///
1034/// See [`TrieRef::debug_fixed()`].
1035pub struct DebugFixed<'a, 'buf, const N: usize, T, F>
1036where
1037 T: ZeroCopy,
1038 F: Flavor,
1039{
1040 trie: &'a TrieRef<T, F>,
1041 buf: &'buf Buf,
1042}
1043
1044impl<const N: usize, T, F> fmt::Debug for DebugFixed<'_, '_, N, T, F>
1045where
1046 T: fmt::Debug + ZeroCopy,
1047 F: Flavor,
1048{
1049 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1050 let mut f = f.debug_map();
1051
1052 for result in self.trie.iter_fixed::<N>(self.buf) {
1053 let (key, value) = result.map_err(|_| fmt::Error)?;
1054 f.entry(&LossyStr::new(key), value);
1055 }
1056
1057 f.finish()
1058 }
1059}
1060
1061impl<T, F> Clone for TrieRef<T, F>
1062where
1063 T: ZeroCopy,
1064 F: Flavor,
1065 F::Values<T>: Clone,
1066 F::Children<NodeRef<T, F>>: Clone,
1067{
1068 #[inline]
1069 fn clone(&self) -> Self {
1070 *self
1071 }
1072}
1073
1074impl<T, F> Copy for TrieRef<T, F>
1075where
1076 T: ZeroCopy,
1077 F: Flavor,
1078 F::Values<T>: Copy,
1079 F::Children<NodeRef<T, F>>: Copy,
1080{
1081}
1082
1083#[derive(ZeroCopy)]
1084#[zero_copy(crate)]
1085#[repr(C)]
1086struct LinksRef<T, F>
1087where
1088 T: ZeroCopy,
1089 F: Flavor,
1090{
1091 values: F::Values<T>,
1092 children: F::Children<NodeRef<T, F>>,
1093}
1094
1095impl<T, F> Clone for LinksRef<T, F>
1096where
1097 T: ZeroCopy,
1098 F: Flavor,
1099 F::Values<T>: Copy,
1100 F::Children<NodeRef<T, F>>: Copy,
1101{
1102 #[inline]
1103 fn clone(&self) -> Self {
1104 *self
1105 }
1106}
1107
1108impl<T, F> Copy for LinksRef<T, F>
1109where
1110 T: ZeroCopy,
1111 F: Flavor,
1112 F::Values<T>: Copy,
1113 F::Children<NodeRef<T, F>>: Copy,
1114{
1115}
1116
1117#[derive(ZeroCopy)]
1118#[zero_copy(crate)]
1119#[repr(C)]
1120struct NodeRef<T, F>
1121where
1122 T: ZeroCopy,
1123 F: Flavor,
1124{
1125 string: F::String,
1126 links: LinksRef<T, F>,
1127}
1128
1129/// Calculate the common prefix between two strings.
1130fn prefix(a: &[u8], b: &[u8]) -> usize {
1131 a.iter().zip(b.iter()).take_while(|(a, b)| a == b).count()
1132}