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}