zerotrie/
zerotrie.rs

1// This file is part of ICU4X. For terms of use, please see the file
2// called LICENSE at the top level of the ICU4X source tree
3// (online at: https://github.com/unicode-org/icu4x/blob/main/LICENSE ).
4
5use crate::reader;
6
7use core::borrow::Borrow;
8
9#[cfg(feature = "alloc")]
10use crate::{
11    builder::nonconst::ZeroTrieBuilder, builder::slice_indices::ByteSliceWithIndices,
12    error::ZeroTrieBuildError,
13};
14#[cfg(feature = "alloc")]
15use alloc::{boxed::Box, collections::BTreeMap, collections::VecDeque, string::String, vec::Vec};
16#[cfg(feature = "litemap")]
17use litemap::LiteMap;
18
19/// A data structure that compactly maps from byte sequences to integers.
20///
21/// There are several variants of [`ZeroTrie`] which are very similar but are optimized
22/// for different use cases:
23///
24/// - [`ZeroTrieSimpleAscii`] is the most compact structure. Very fast for small data.
25///   Only stores ASCII-encoded strings. Can be const-constructed!
26/// - [`ZeroTriePerfectHash`] is also compact, but it also supports arbitrary binary
27///   strings. It also scales better to large data. Cannot be const-constructed.
28/// - [`ZeroTrieExtendedCapacity`] can be used if more than 2^32 bytes are required.
29///
30/// You can create a [`ZeroTrie`] directly, in which case the most appropriate
31/// backing implementation will be chosen.
32///
33/// # Backing Store
34///
35/// The data structure has a flexible backing data store. The only requirement for most
36/// functionality is that it implement `AsRef<[u8]>`. All of the following are valid
37/// [`ZeroTrie`] types:
38///
39/// - `ZeroTrie<[u8]>` (dynamically sized type: must be stored in a reference or Box)
40/// - `ZeroTrie<&[u8]>` (borrows its data from a u8 buffer)
41/// - `ZeroTrie<Vec<u8>>` (fully owned data)
42/// - `ZeroTrie<ZeroVec<u8>>` (the recommended borrowed-or-owned signature)
43/// - `Cow<ZeroTrie<[u8]>>` (another borrowed-or-owned signature)
44/// - `ZeroTrie<Cow<[u8]>>` (another borrowed-or-owned signature)
45///
46/// # Examples
47///
48/// ```
49/// use litemap::LiteMap;
50/// use zerotrie::ZeroTrie;
51///
52/// let mut map = LiteMap::<&[u8], usize>::new_vec();
53/// map.insert("foo".as_bytes(), 1);
54/// map.insert("bar".as_bytes(), 2);
55/// map.insert("bazzoo".as_bytes(), 3);
56///
57/// let trie = ZeroTrie::try_from(&map)?;
58///
59/// assert_eq!(trie.get("foo"), Some(1));
60/// assert_eq!(trie.get("bar"), Some(2));
61/// assert_eq!(trie.get("bazzoo"), Some(3));
62/// assert_eq!(trie.get("unknown"), None);
63///
64/// # Ok::<_, zerotrie::ZeroTrieBuildError>(())
65/// ```
66#[derive(Debug, Clone, Copy, PartialEq, Eq)]
67// Note: The absence of the following derive does not cause any test failures in this crate
68#[cfg_attr(feature = "yoke", derive(yoke::Yokeable))]
69pub struct ZeroTrie<Store>(pub(crate) ZeroTrieFlavor<Store>);
70
71#[derive(Debug, Clone, Copy, PartialEq, Eq)]
72pub(crate) enum ZeroTrieFlavor<Store> {
73    SimpleAscii(ZeroTrieSimpleAscii<Store>),
74    PerfectHash(ZeroTriePerfectHash<Store>),
75    ExtendedCapacity(ZeroTrieExtendedCapacity<Store>),
76}
77
78/// A data structure that compactly maps from ASCII strings to integers.
79///
80/// For more information, see [`ZeroTrie`].
81///
82/// # Examples
83///
84/// ```
85/// use litemap::LiteMap;
86/// use zerotrie::ZeroTrieSimpleAscii;
87///
88/// let mut map = LiteMap::new_vec();
89/// map.insert(&b"foo"[..], 1);
90/// map.insert(b"bar", 2);
91/// map.insert(b"bazzoo", 3);
92///
93/// let trie = ZeroTrieSimpleAscii::try_from(&map)?;
94///
95/// assert_eq!(trie.get(b"foo"), Some(1));
96/// assert_eq!(trie.get(b"bar"), Some(2));
97/// assert_eq!(trie.get(b"bazzoo"), Some(3));
98/// assert_eq!(trie.get(b"unknown"), None);
99///
100/// # Ok::<_, zerotrie::ZeroTrieBuildError>(())
101/// ```
102///
103/// The trie can only store ASCII bytes; a string with non-ASCII always returns None:
104///
105/// ```
106/// use zerotrie::ZeroTrieSimpleAscii;
107///
108/// // A trie with two values: "abc" and "abcdef"
109/// let trie = ZeroTrieSimpleAscii::from_bytes(b"abc\x80def\x81");
110///
111/// assert!(trie.get(b"ab\xFF").is_none());
112/// ```
113#[repr(transparent)]
114#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
115#[cfg_attr(feature = "databake", derive(databake::Bake))]
116#[cfg_attr(feature = "databake", databake(path = zerotrie))]
117#[allow(clippy::exhaustive_structs)] // databake hidden fields
118pub struct ZeroTrieSimpleAscii<Store: ?Sized> {
119    #[doc(hidden)] // for databake, but there are no invariants
120    pub store: Store,
121}
122
123impl<Store: ?Sized> ZeroTrieSimpleAscii<Store> {
124    #[allow(unsafe_code)] // transparent newtype casts are documented
125    fn transparent_ref_from_store(s: &Store) -> &Self {
126        unsafe {
127            // Safety: Self is transparent over Store
128            &*(s as *const Store as *const Self)
129        }
130    }
131}
132
133impl<Store> ZeroTrieSimpleAscii<Store> {
134    /// Wrap this specific [`ZeroTrie`] variant into a [`ZeroTrie`].
135    #[inline]
136    pub const fn into_zerotrie(self) -> ZeroTrie<Store> {
137        ZeroTrie(ZeroTrieFlavor::SimpleAscii(self))
138    }
139}
140
141/// A data structure that compactly maps from ASCII strings to integers
142/// in a case-insensitive way.
143///
144/// # Examples
145///
146/// ```
147/// use litemap::LiteMap;
148/// use zerotrie::ZeroAsciiIgnoreCaseTrie;
149///
150/// let mut map = LiteMap::new_vec();
151/// map.insert(&b"foo"[..], 1);
152/// map.insert(b"Bar", 2);
153/// map.insert(b"Bazzoo", 3);
154///
155/// let trie = ZeroAsciiIgnoreCaseTrie::try_from(&map)?;
156///
157/// assert_eq!(trie.get(b"foo"), Some(1));
158/// assert_eq!(trie.get(b"bar"), Some(2));
159/// assert_eq!(trie.get(b"BAR"), Some(2));
160/// assert_eq!(trie.get(b"bazzoo"), Some(3));
161/// assert_eq!(trie.get(b"unknown"), None);
162///
163/// # Ok::<_, zerotrie::ZeroTrieBuildError>(())
164/// ```
165///
166/// Strings with different cases of the same character at the same offset are not allowed:
167///
168/// ```
169/// use litemap::LiteMap;
170/// use zerotrie::ZeroAsciiIgnoreCaseTrie;
171///
172/// let mut map = LiteMap::new_vec();
173/// map.insert(&b"bar"[..], 1);
174/// // OK: 'r' and 'Z' are different letters
175/// map.insert(b"baZ", 2);
176/// // Bad: we already inserted 'r' so we cannot also insert 'R' at the same position
177/// map.insert(b"baR", 2);
178///
179/// ZeroAsciiIgnoreCaseTrie::try_from(&map).expect_err("mixed-case strings!");
180/// ```
181#[repr(transparent)]
182#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
183#[cfg_attr(feature = "databake", derive(databake::Bake))]
184#[cfg_attr(feature = "databake", databake(path = zerotrie))]
185#[allow(clippy::exhaustive_structs)] // databake hidden fields
186pub struct ZeroAsciiIgnoreCaseTrie<Store: ?Sized> {
187    #[doc(hidden)] // for databake, but there are no invariants
188    pub store: Store,
189}
190
191impl<Store: ?Sized> ZeroAsciiIgnoreCaseTrie<Store> {
192    #[allow(unsafe_code)] // transparent newtype casts are documented
193    fn transparent_ref_from_store(s: &Store) -> &Self {
194        unsafe {
195            // Safety: Self is transparent over Store
196            &*(s as *const Store as *const Self)
197        }
198    }
199}
200
201// Note: ZeroAsciiIgnoreCaseTrie is not a variant of ZeroTrie so there is no `into_zerotrie`
202
203/// A data structure that compactly maps from byte strings to integers.
204///
205/// For more information, see [`ZeroTrie`].
206///
207/// # Examples
208///
209/// ```
210/// use litemap::LiteMap;
211/// use zerotrie::ZeroTriePerfectHash;
212///
213/// let mut map = LiteMap::<&[u8], usize>::new_vec();
214/// map.insert("foo".as_bytes(), 1);
215/// map.insert("bår".as_bytes(), 2);
216/// map.insert("båzzøø".as_bytes(), 3);
217///
218/// let trie = ZeroTriePerfectHash::try_from(&map)?;
219///
220/// assert_eq!(trie.get("foo".as_bytes()), Some(1));
221/// assert_eq!(trie.get("bår".as_bytes()), Some(2));
222/// assert_eq!(trie.get("båzzøø".as_bytes()), Some(3));
223/// assert_eq!(trie.get("bazzoo".as_bytes()), None);
224///
225/// # Ok::<_, zerotrie::ZeroTrieBuildError>(())
226/// ```
227#[repr(transparent)]
228#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
229#[cfg_attr(feature = "databake", derive(databake::Bake))]
230#[cfg_attr(feature = "databake", databake(path = zerotrie))]
231#[allow(clippy::exhaustive_structs)] // databake hidden fields
232pub struct ZeroTriePerfectHash<Store: ?Sized> {
233    #[doc(hidden)] // for databake, but there are no invariants
234    pub store: Store,
235}
236
237impl<Store: ?Sized> ZeroTriePerfectHash<Store> {
238    #[allow(unsafe_code)] // transparent newtype casts are documented
239    fn transparent_ref_from_store(s: &Store) -> &Self {
240        unsafe {
241            // Safety: Self is transparent over Store
242            &*(s as *const Store as *const Self)
243        }
244    }
245}
246
247impl<Store> ZeroTriePerfectHash<Store> {
248    /// Wrap this specific [`ZeroTrie`] variant into a [`ZeroTrie`].
249    #[inline]
250    pub const fn into_zerotrie(self) -> ZeroTrie<Store> {
251        ZeroTrie(ZeroTrieFlavor::PerfectHash(self))
252    }
253}
254
255/// A data structure that maps from a large number of byte strings to integers.
256///
257/// For more information, see [`ZeroTrie`].
258#[repr(transparent)]
259#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
260#[cfg_attr(feature = "databake", derive(databake::Bake))]
261#[cfg_attr(feature = "databake", databake(path = zerotrie))]
262#[allow(clippy::exhaustive_structs)] // databake hidden fields
263pub struct ZeroTrieExtendedCapacity<Store: ?Sized> {
264    #[doc(hidden)] // for databake, but there are no invariants
265    pub store: Store,
266}
267
268impl<Store: ?Sized> ZeroTrieExtendedCapacity<Store> {
269    #[allow(unsafe_code)] // transparent newtype casts are documented
270    fn transparent_ref_from_store(s: &Store) -> &Self {
271        unsafe {
272            // Safety: Self is transparent over Store
273            &*(s as *const Store as *const Self)
274        }
275    }
276}
277
278impl<Store> ZeroTrieExtendedCapacity<Store> {
279    /// Wrap this specific [`ZeroTrie`] variant into a [`ZeroTrie`].
280    #[inline]
281    pub const fn into_zerotrie(self) -> ZeroTrie<Store> {
282        ZeroTrie(ZeroTrieFlavor::ExtendedCapacity(self))
283    }
284}
285
286macro_rules! impl_zerotrie_subtype {
287    ($name:ident, $iter_element:ty, $iter_fn:path, $iter_ty:ty, $cnv_fn:path) => {
288        impl<Store> $name<Store> {
289            /// Create a trie directly from a store.
290            ///
291            /// If the store does not contain valid bytes, unexpected behavior may occur.
292            #[inline]
293            pub const fn from_store(store: Store) -> Self {
294                Self { store }
295            }
296            /// Takes the byte store from this trie.
297            #[inline]
298            pub fn into_store(self) -> Store {
299                self.store
300            }
301            /// Converts this trie's store to a different store implementing the `From` trait.
302            ///
303            #[doc = concat!("For example, use this to change `", stringify!($name), "<Vec<u8>>` to `", stringify!($name), "<Cow<[u8]>>`.")]
304            ///
305            /// # Examples
306            ///
307            /// ```
308            /// use std::borrow::Cow;
309            #[doc = concat!("use zerotrie::", stringify!($name), ";")]
310            ///
311            #[doc = concat!("let trie: ", stringify!($name), "<Vec<u8>> = ", stringify!($name), "::from_bytes(b\"abc\\x85\").to_owned();")]
312            #[doc = concat!("let cow: ", stringify!($name), "<Cow<[u8]>> = trie.convert_store();")]
313            ///
314            /// assert_eq!(cow.get(b"abc"), Some(5));
315            /// ```
316            pub fn convert_store<X: From<Store>>(self) -> $name<X> {
317                $name::<X>::from_store(X::from(self.store))
318            }
319        }
320        impl<Store> $name<Store>
321        where
322        Store: AsRef<[u8]> + ?Sized,
323        {
324            /// Queries the trie for a string.
325            // Note: We do not need the Borrow trait's guarantees, so we use
326            // the more general AsRef trait.
327            pub fn get<K>(&self, key: K) -> Option<usize> where K: AsRef<[u8]> {
328                reader::get_parameterized::<Self>(self.store.as_ref(), key.as_ref())
329            }
330            /// Returns `true` if the trie is empty.
331            #[inline]
332            pub fn is_empty(&self) -> bool {
333                self.store.as_ref().is_empty()
334            }
335            /// Returns the size of the trie in number of bytes.
336            ///
337            /// To get the number of keys in the trie, use `.iter().count()`:
338            ///
339            /// ```
340            #[doc = concat!("use zerotrie::", stringify!($name), ";")]
341            ///
342            /// // A trie with two values: "abc" and "abcdef"
343            #[doc = concat!("let trie: &", stringify!($name), "<[u8]> = ", stringify!($name), "::from_bytes(b\"abc\\x80def\\x81\");")]
344            ///
345            /// assert_eq!(8, trie.byte_len());
346            /// assert_eq!(2, trie.iter().count());
347            /// ```
348            #[inline]
349            pub fn byte_len(&self) -> usize {
350                self.store.as_ref().len()
351            }
352            /// Returns the bytes contained in the underlying store.
353            #[inline]
354            pub fn as_bytes(&self) -> &[u8] {
355                self.store.as_ref()
356            }
357            /// Returns this trie as a reference transparent over a byte slice.
358            #[inline]
359            pub fn as_borrowed(&self) -> &$name<[u8]> {
360                $name::from_bytes(self.store.as_ref())
361            }
362            /// Returns a trie with a store borrowing from this trie.
363            #[inline]
364            pub fn as_borrowed_slice(&self) -> $name<&[u8]> {
365                $name::from_store(self.store.as_ref())
366            }
367        }
368        impl<Store> AsRef<$name<[u8]>> for $name<Store>
369        where
370        Store: AsRef<[u8]> + ?Sized,
371        {
372            #[inline]
373            fn as_ref(&self) -> &$name<[u8]> {
374                self.as_borrowed()
375            }
376        }
377        #[cfg(feature = "alloc")]
378        impl<Store> $name<Store>
379        where
380        Store: AsRef<[u8]> + ?Sized,
381        {
382            /// Converts a possibly-borrowed $name to an owned one.
383            ///
384            /// ✨ *Enabled with the `alloc` Cargo feature.*
385            ///
386            /// # Examples
387            ///
388            /// ```
389            #[doc = concat!("use zerotrie::", stringify!($name), ";")]
390            ///
391            #[doc = concat!("let trie: &", stringify!($name), "<[u8]> = ", stringify!($name), "::from_bytes(b\"abc\\x85\");")]
392            #[doc = concat!("let owned: ", stringify!($name), "<Vec<u8>> = trie.to_owned();")]
393            ///
394            /// assert_eq!(trie.get(b"abc"), Some(5));
395            /// assert_eq!(owned.get(b"abc"), Some(5));
396            /// ```
397            #[inline]
398            pub fn to_owned(&self) -> $name<Vec<u8>> {
399                $name::from_store(
400                    Vec::from(self.store.as_ref()),
401                )
402            }
403            /// Returns an iterator over the key/value pairs in this trie.
404            ///
405            /// ✨ *Enabled with the `alloc` Cargo feature.*
406            ///
407            /// # Examples
408            ///
409            /// ```
410            #[doc = concat!("use zerotrie::", stringify!($name), ";")]
411            ///
412            /// // A trie with two values: "abc" and "abcdef"
413            #[doc = concat!("let trie: &", stringify!($name), "<[u8]> = ", stringify!($name), "::from_bytes(b\"abc\\x80def\\x81\");")]
414            ///
415            /// let mut it = trie.iter();
416            /// assert_eq!(it.next(), Some(("abc".into(), 0)));
417            /// assert_eq!(it.next(), Some(("abcdef".into(), 1)));
418            /// assert_eq!(it.next(), None);
419            /// ```
420            #[inline]
421            pub fn iter(&self) -> $iter_ty {
422                 $iter_fn(self.as_bytes())
423            }
424        }
425        impl $name<[u8]> {
426            /// Casts from a byte slice to a reference to a trie with the same lifetime.
427            ///
428            /// If the bytes are not a valid trie, unexpected behavior may occur.
429            #[inline]
430            pub fn from_bytes(trie: &[u8]) -> &Self {
431                Self::transparent_ref_from_store(trie)
432            }
433        }
434        #[cfg(feature = "alloc")]
435        impl $name<Vec<u8>> {
436            pub(crate) fn try_from_tuple_slice(items: ByteSliceWithIndices) -> Result<Self, ZeroTrieBuildError> {
437                use crate::options::ZeroTrieWithOptions;
438                ZeroTrieBuilder::<VecDeque<u8>>::from_sorted_tuple_slice(
439                    items,
440                    Self::OPTIONS,
441                )
442                .map(|s| Self {
443                    store: s.to_bytes(),
444                })
445            }
446            /// Creates a trie from a [`BTreeMap`] of string keys.
447            ///
448            /// See also the [`TryFrom`] and [`FromIterator`] impls.
449            pub fn try_from_btree_map_str<K>(items: &BTreeMap<K, usize>) -> Result<Self, ZeroTrieBuildError>
450            where
451                K: Borrow<str>
452            {
453                let tuples: Vec<(&[u8], usize)> = items
454                    .iter()
455                    .map(|(k, v)| (k.borrow().as_bytes(), *v))
456                    .collect();
457                let byte_str_slice = ByteSliceWithIndices::from_byte_slice(&tuples);
458                Self::try_from_tuple_slice(byte_str_slice)
459            }
460        }
461        #[cfg(feature = "alloc")]
462        impl<K> FromIterator<(K, usize)> for $name<Vec<u8>>
463        where
464            K: AsRef<[u8]>
465        {
466            fn from_iter<T: IntoIterator<Item = (K, usize)>>(iter: T) -> Self {
467                use crate::options::ZeroTrieWithOptions;
468                use crate::builder::nonconst::ZeroTrieBuilder;
469                ZeroTrieBuilder::<VecDeque<u8>>::from_bytes_iter(
470                    iter,
471                    Self::OPTIONS
472                )
473                .map(|s| Self {
474                    store: s.to_bytes(),
475                })
476                .unwrap()
477            }
478        }
479        #[cfg(feature = "alloc")]
480        impl<'a, K> TryFrom<&'a BTreeMap<K, usize>> for $name<Vec<u8>>
481        where
482            K: Borrow<[u8]>
483        {
484            type Error = crate::error::ZeroTrieBuildError;
485            fn try_from(map: &'a BTreeMap<K, usize>) -> Result<Self, Self::Error> {
486                let tuples: Vec<(&[u8], usize)> = map
487                    .iter()
488                    .map(|(k, v)| (k.borrow(), *v))
489                    .collect();
490                let byte_str_slice = ByteSliceWithIndices::from_byte_slice(&tuples);
491                Self::try_from_tuple_slice(byte_str_slice)
492            }
493        }
494        #[cfg(feature = "alloc")]
495        impl<Store> $name<Store>
496        where
497            Store: AsRef<[u8]> + ?Sized
498        {
499            /// Exports the data from this [`ZeroTrie`] type into a [`BTreeMap`].
500            ///
501            /// ✨ *Enabled with the `alloc` Cargo feature.*
502            ///
503            /// # Examples
504            ///
505            /// ```
506            #[doc = concat!("use zerotrie::", stringify!($name), ";")]
507            /// use std::collections::BTreeMap;
508            ///
509            #[doc = concat!("let trie = ", stringify!($name), "::from_bytes(b\"abc\\x81def\\x82\");")]
510            /// let items = trie.to_btreemap();
511            ///
512            /// assert_eq!(items.len(), 2);
513            ///
514            #[doc = concat!("let recovered_trie: ", stringify!($name), "<Vec<u8>> = items")]
515            ///     .into_iter()
516            ///     .collect();
517            /// assert_eq!(trie.as_bytes(), recovered_trie.as_bytes());
518            /// ```
519            pub fn to_btreemap(&self) -> BTreeMap<$iter_element, usize> {
520                self.iter().collect()
521            }
522            #[allow(dead_code)] // not needed for ZeroAsciiIgnoreCaseTrie
523            pub(crate) fn to_btreemap_bytes(&self) -> BTreeMap<Box<[u8]>, usize> {
524                self.iter().map(|(k, v)| ($cnv_fn(k), v)).collect()
525            }
526        }
527        #[cfg(feature = "alloc")]
528        impl<Store> From<&$name<Store>> for BTreeMap<$iter_element, usize>
529        where
530            Store: AsRef<[u8]> + ?Sized,
531        {
532            #[inline]
533            fn from(other: &$name<Store>) -> Self {
534                other.to_btreemap()
535            }
536        }
537        #[cfg(feature = "litemap")]
538        impl<'a, K, S> TryFrom<&'a LiteMap<K, usize, S>> for $name<Vec<u8>>
539        where
540            K: Borrow<[u8]>,
541            S: litemap::store::StoreIterable<'a, K, usize>,
542        {
543            type Error = crate::error::ZeroTrieBuildError;
544            fn try_from(map: &'a LiteMap<K, usize, S>) -> Result<Self, Self::Error> {
545                let tuples: Vec<(&[u8], usize)> = map
546                    .iter()
547                    .map(|(k, v)| (k.borrow(), *v))
548                    .collect();
549                let byte_str_slice = ByteSliceWithIndices::from_byte_slice(&tuples);
550                Self::try_from_tuple_slice(byte_str_slice)
551            }
552        }
553        #[cfg(feature = "litemap")]
554        impl<Store> $name<Store>
555        where
556            Store: AsRef<[u8]> + ?Sized,
557        {
558            /// Exports the data from this [`ZeroTrie`] type into a [`LiteMap`].
559            ///
560            /// ✨ *Enabled with the `litemap` Cargo feature.*
561            ///
562            /// # Examples
563            ///
564            /// ```
565            #[doc = concat!("use zerotrie::", stringify!($name), ";")]
566            /// use litemap::LiteMap;
567            ///
568            #[doc = concat!("let trie = ", stringify!($name), "::from_bytes(b\"abc\\x81def\\x82\");")]
569            ///
570            /// let items = trie.to_litemap();
571            /// assert_eq!(items.len(), 2);
572            ///
573            #[doc = concat!("let recovered_trie: ", stringify!($name), "<Vec<u8>> = items")]
574            ///     .iter()
575            ///     .map(|(k, v)| (k, *v))
576            ///     .collect();
577            /// assert_eq!(trie.as_bytes(), recovered_trie.as_bytes());
578            /// ```
579            pub fn to_litemap(&self) -> LiteMap<$iter_element, usize> {
580                self.iter().collect()
581            }
582            #[allow(dead_code)] // not needed for ZeroAsciiIgnoreCaseTrie
583            pub(crate) fn to_litemap_bytes(&self) -> LiteMap<Box<[u8]>, usize> {
584                self.iter().map(|(k, v)| ($cnv_fn(k), v)).collect()
585            }
586            #[cfg(feature = "serde")]
587            pub(crate) fn to_litemap_serde(&self) -> LiteMap<crate::serde::SerdeByteStrOwned, usize> {
588                self.iter().map(|(k, v)| (crate::serde::SerdeByteStrOwned($cnv_fn(k)), v)).collect()
589            }
590        }
591        #[cfg(feature = "litemap")]
592        impl<Store> From<&$name<Store>> for LiteMap<$iter_element, usize>
593        where
594            Store: AsRef<[u8]> + ?Sized,
595        {
596            #[inline]
597            fn from(other: &$name<Store>) -> Self {
598                other.to_litemap()
599            }
600        }
601        #[cfg(feature = "litemap")]
602        impl $name<Vec<u8>>
603        {
604            #[cfg(feature = "serde")]
605            pub(crate) fn try_from_serde_litemap(items: &LiteMap<crate::serde::SerdeByteStrOwned, usize>) -> Result<Self, ZeroTrieBuildError> {
606                let tuples: Vec<(&[u8], usize)> = items.iter().map(|(k, v)| (k.as_bytes(), *v)).collect();
607                let byte_str_slice = ByteSliceWithIndices::from_byte_slice(&tuples);
608                Self::try_from_tuple_slice(byte_str_slice)
609            }
610        }
611        // Note: Can't generalize this impl due to the `core::borrow::Borrow` blanket impl.
612        impl Borrow<$name<[u8]>> for $name<&[u8]> {
613            #[inline]
614            fn borrow(&self) -> &$name<[u8]> {
615                self.as_borrowed()
616            }
617        }
618        // Note: Can't generalize this impl due to the `core::borrow::Borrow` blanket impl.
619        #[cfg(feature = "alloc")]
620        impl Borrow<$name<[u8]>> for $name<Box<[u8]>> {
621            #[inline]
622            fn borrow(&self) -> &$name<[u8]> {
623                self.as_borrowed()
624            }
625        }
626        // Note: Can't generalize this impl due to the `core::borrow::Borrow` blanket impl.
627        #[cfg(feature = "alloc")]
628        impl Borrow<$name<[u8]>> for $name<Vec<u8>> {
629            #[inline]
630            fn borrow(&self) -> &$name<[u8]> {
631                self.as_borrowed()
632            }
633        }
634        #[cfg(feature = "alloc")]
635        impl alloc::borrow::ToOwned for $name<[u8]> {
636            type Owned = $name<Box<[u8]>>;
637            #[doc = concat!("This impl allows [`", stringify!($name), "`] to be used inside of a [`Cow`](alloc::borrow::Cow).")]
638            ///
639            #[doc = concat!("Note that it is also possible to use `", stringify!($name), "<ZeroVec<u8>>` for a similar result.")]
640            ///
641            /// ✨ *Enabled with the `alloc` Cargo feature.*
642            ///
643            /// # Examples
644            ///
645            /// ```
646            /// use std::borrow::Cow;
647            #[doc = concat!("use zerotrie::", stringify!($name), ";")]
648            ///
649            #[doc = concat!("let trie: Cow<", stringify!($name), "<[u8]>> = Cow::Borrowed(", stringify!($name), "::from_bytes(b\"abc\\x85\"));")]
650            /// assert_eq!(trie.get(b"abc"), Some(5));
651            /// ```
652            fn to_owned(&self) -> Self::Owned {
653                let bytes: &[u8] = self.store.as_ref();
654                $name::from_store(
655                    Vec::from(bytes).into_boxed_slice(),
656                )
657            }
658        }
659        // TODO(#2778): Auto-derive these impls based on the repr(transparent).
660        //
661        // Safety (based on the safety checklist on the VarULE trait):
662        //  1. `$name` does not include any uninitialized or padding bytes as it is `repr(transparent)`
663        //     over a `VarULE` type, `Store`, as evidenced by the existence of `transparent_ref_from_store()`
664        //  2. `$name` is aligned to 1 byte for the same reason
665        //  3. The impl of `validate_bytes()` returns an error if any byte is not valid (passed down to `VarULE` impl of `Store`)
666        //  4. The impl of `validate_bytes()` returns an error if the slice cannot be used in its entirety (passed down to `VarULE` impl of `Store`)
667        //  5. The impl of `from_bytes_unchecked()` returns a reference to the same data.
668        //  6. `parse_bytes()` is left to its default impl
669        //  7. byte equality is semantic equality
670        #[cfg(feature = "zerovec")]
671        #[allow(unsafe_code)] // ULE impls are documented
672        unsafe impl<Store> zerovec::ule::VarULE for $name<Store>
673        where
674            Store: zerovec::ule::VarULE,
675        {
676            #[inline]
677            fn validate_bytes(bytes: &[u8]) -> Result<(), zerovec::ule::UleError> {
678                Store::validate_bytes(bytes)
679            }
680            #[inline]
681            unsafe fn from_bytes_unchecked(bytes: &[u8]) -> &Self {
682                // Safety: we can pass down the validity invariant to Store
683                Self::transparent_ref_from_store(Store::from_bytes_unchecked(bytes))
684            }
685        }
686        #[cfg(feature = "zerofrom")]
687        impl<'zf, Store1, Store2> zerofrom::ZeroFrom<'zf, $name<Store1>> for $name<Store2>
688        where
689            Store2: zerofrom::ZeroFrom<'zf, Store1>,
690        {
691            #[inline]
692            fn zero_from(other: &'zf $name<Store1>) -> Self {
693                $name::from_store(zerofrom::ZeroFrom::zero_from(&other.store))
694            }
695        }
696    };
697}
698
699#[cfg(feature = "alloc")]
700fn string_to_box_u8(input: String) -> Box<[u8]> {
701    input.into_boxed_str().into_boxed_bytes()
702}
703
704#[doc(hidden)] // subject to change
705#[cfg(feature = "alloc")]
706pub type ZeroTrieStringIterator<'a> =
707    core::iter::Map<reader::ZeroTrieIterator<'a>, fn((Vec<u8>, usize)) -> (String, usize)>;
708
709impl_zerotrie_subtype!(
710    ZeroTrieSimpleAscii,
711    String,
712    reader::get_iter_ascii_or_panic,
713    ZeroTrieStringIterator<'_>,
714    string_to_box_u8
715);
716impl_zerotrie_subtype!(
717    ZeroAsciiIgnoreCaseTrie,
718    String,
719    reader::get_iter_ascii_or_panic,
720    ZeroTrieStringIterator<'_>,
721    string_to_box_u8
722);
723impl_zerotrie_subtype!(
724    ZeroTriePerfectHash,
725    Vec<u8>,
726    reader::get_iter_phf,
727    reader::ZeroTrieIterator<'_>,
728    Vec::into_boxed_slice
729);
730impl_zerotrie_subtype!(
731    ZeroTrieExtendedCapacity,
732    Vec<u8>,
733    reader::get_iter_phf,
734    reader::ZeroTrieIterator<'_>,
735    Vec::into_boxed_slice
736);
737
738#[allow(unused_macro_rules)] // feature
739macro_rules! impl_dispatch {
740    ($self:ident, $inner_fn:ident()) => {
741        match $self.0 {
742            ZeroTrieFlavor::SimpleAscii(subtype) => subtype.$inner_fn(),
743            ZeroTrieFlavor::PerfectHash(subtype) => subtype.$inner_fn(),
744            ZeroTrieFlavor::ExtendedCapacity(subtype) => subtype.$inner_fn(),
745        }
746    };
747    ($self:ident, $inner_fn:ident().into_zerotrie()) => {
748        match $self.0 {
749            ZeroTrieFlavor::SimpleAscii(subtype) => subtype.$inner_fn().into_zerotrie(),
750            ZeroTrieFlavor::PerfectHash(subtype) => subtype.$inner_fn().into_zerotrie(),
751            ZeroTrieFlavor::ExtendedCapacity(subtype) => subtype.$inner_fn().into_zerotrie(),
752        }
753    };
754    (&$self:ident, $inner_fn:ident()) => {
755        match &$self.0 {
756            ZeroTrieFlavor::SimpleAscii(subtype) => subtype.$inner_fn(),
757            ZeroTrieFlavor::PerfectHash(subtype) => subtype.$inner_fn(),
758            ZeroTrieFlavor::ExtendedCapacity(subtype) => subtype.$inner_fn(),
759        }
760    };
761    (&$self:ident, $inner_fn:ident($arg:ident)) => {
762        match &$self.0 {
763            ZeroTrieFlavor::SimpleAscii(subtype) => subtype.$inner_fn($arg),
764            ZeroTrieFlavor::PerfectHash(subtype) => subtype.$inner_fn($arg),
765            ZeroTrieFlavor::ExtendedCapacity(subtype) => subtype.$inner_fn($arg),
766        }
767    };
768    (&$self:ident, $trait:ident::$inner_fn:ident()) => {
769        match &$self.0 {
770            ZeroTrieFlavor::SimpleAscii(subtype) => {
771                ZeroTrie(ZeroTrieFlavor::SimpleAscii($trait::$inner_fn(subtype)))
772            }
773            ZeroTrieFlavor::PerfectHash(subtype) => {
774                ZeroTrie(ZeroTrieFlavor::PerfectHash($trait::$inner_fn(subtype)))
775            }
776            ZeroTrieFlavor::ExtendedCapacity(subtype) => {
777                ZeroTrie(ZeroTrieFlavor::ExtendedCapacity($trait::$inner_fn(subtype)))
778            }
779        }
780    };
781}
782
783impl<Store> ZeroTrie<Store> {
784    /// Takes the byte store from this trie.
785    pub fn into_store(self) -> Store {
786        impl_dispatch!(self, into_store())
787    }
788    /// Converts this trie's store to a different store implementing the `From` trait.
789    ///
790    /// For example, use this to change `ZeroTrie<Vec<u8>>` to `ZeroTrie<Cow<[u8]>>`.
791    pub fn convert_store<NewStore>(self) -> ZeroTrie<NewStore>
792    where
793        NewStore: From<Store>,
794    {
795        impl_dispatch!(self, convert_store().into_zerotrie())
796    }
797}
798
799impl<Store> ZeroTrie<Store>
800where
801    Store: AsRef<[u8]>,
802{
803    /// Queries the trie for a string.
804    pub fn get<K>(&self, key: K) -> Option<usize>
805    where
806        K: AsRef<[u8]>,
807    {
808        impl_dispatch!(&self, get(key))
809    }
810    /// Returns `true` if the trie is empty.
811    pub fn is_empty(&self) -> bool {
812        impl_dispatch!(&self, is_empty())
813    }
814    /// Returns the size of the trie in number of bytes.
815    ///
816    /// To get the number of keys in the trie, use `.iter().count()`.
817    pub fn byte_len(&self) -> usize {
818        impl_dispatch!(&self, byte_len())
819    }
820}
821
822#[cfg(feature = "alloc")]
823impl<Store> ZeroTrie<Store>
824where
825    Store: AsRef<[u8]>,
826{
827    /// Exports the data from this [`ZeroTrie`] into a [`BTreeMap`].
828    ///
829    /// ✨ *Enabled with the `alloc` Cargo feature.*
830    pub fn to_btreemap(&self) -> BTreeMap<Box<[u8]>, usize> {
831        impl_dispatch!(&self, to_btreemap_bytes())
832    }
833}
834
835#[cfg(feature = "litemap")]
836impl<Store> ZeroTrie<Store>
837where
838    Store: AsRef<[u8]>,
839{
840    /// Exports the data from this [`ZeroTrie`] into a `LiteMap`.
841    #[cfg(feature = "serde")]
842    pub fn to_litemap(&self) -> LiteMap<Box<[u8]>, usize> {
843        impl_dispatch!(&self, to_litemap_bytes())
844    }
845
846    /// Exports the data from this [`ZeroTrie`] into a `LiteMap` for Serde.
847    #[cfg(feature = "serde")]
848    pub(crate) fn to_litemap_serde(&self) -> LiteMap<crate::serde::SerdeByteStrOwned, usize> {
849        impl_dispatch!(&self, to_litemap_serde())
850    }
851}
852#[cfg(feature = "alloc")]
853impl ZeroTrie<Vec<u8>> {
854    pub(crate) fn try_from_tuple_slice(
855        items: ByteSliceWithIndices,
856    ) -> Result<Self, ZeroTrieBuildError> {
857        let is_all_ascii = items.is_all_ascii();
858        if is_all_ascii && items.len() < 512 {
859            ZeroTrieSimpleAscii::try_from_tuple_slice(items).map(|x| x.into_zerotrie())
860        } else {
861            ZeroTriePerfectHash::try_from_tuple_slice(items).map(|x| x.into_zerotrie())
862        }
863    }
864}
865
866#[cfg(feature = "alloc")]
867impl<K> FromIterator<(K, usize)> for ZeroTrie<Vec<u8>>
868where
869    K: AsRef<[u8]>,
870{
871    fn from_iter<T: IntoIterator<Item = (K, usize)>>(iter: T) -> Self {
872        // We need two Vecs because the first one anchors the `K`s that the second one borrows.
873        let items = Vec::from_iter(iter);
874        let mut items: Vec<(&[u8], usize)> = items.iter().map(|(k, v)| (k.as_ref(), *v)).collect();
875        items.sort();
876        let byte_str_slice = ByteSliceWithIndices::from_byte_slice(&items);
877        #[expect(clippy::unwrap_used)] // FromIterator is panicky
878        Self::try_from_tuple_slice(byte_str_slice).unwrap()
879    }
880}
881
882#[cfg(feature = "databake")]
883impl<Store> databake::Bake for ZeroTrie<Store>
884where
885    Store: databake::Bake,
886{
887    fn bake(&self, env: &databake::CrateEnv) -> databake::TokenStream {
888        use databake::*;
889        let inner = impl_dispatch!(&self, bake(env));
890        quote! { #inner.into_zerotrie() }
891    }
892}
893
894#[cfg(feature = "databake")]
895impl<Store> databake::BakeSize for ZeroTrie<Store>
896where
897    Store: databake::BakeSize,
898{
899    fn borrows_size(&self) -> usize {
900        impl_dispatch!(&self, borrows_size())
901    }
902}
903
904#[cfg(feature = "zerofrom")]
905impl<'zf, Store1, Store2> zerofrom::ZeroFrom<'zf, ZeroTrie<Store1>> for ZeroTrie<Store2>
906where
907    Store2: zerofrom::ZeroFrom<'zf, Store1>,
908{
909    fn zero_from(other: &'zf ZeroTrie<Store1>) -> Self {
910        use zerofrom::ZeroFrom;
911        impl_dispatch!(&other, ZeroFrom::zero_from())
912    }
913}
zerotrie/zerotrie.rs

zerotrie/
zerotrie.rs
