firewood_storage/path/
component.rs

1// Copyright (C) 2025, Ava Labs, Inc. All rights reserved.
2// See the file LICENSE.md for licensing terms.
3
4use smallvec::SmallVec;
5
6use super::{PartialPath, TriePath, TriePathFromUnpackedBytes};
7
8/// A path component in a hexary trie; which is only 4 bits (aka a nibble).
9#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
10pub struct PathComponent(pub crate::u4::U4);
11
12/// An iterator over path components.
13pub type ComponentIter<'a> = std::iter::Copied<std::slice::Iter<'a, PathComponent>>;
14
15/// Extension methods for slices of path components.
16pub trait PathComponentSliceExt {
17    /// Casts this slice of path components to a byte slice.
18    fn as_byte_slice(&self) -> &[u8];
19}
20
21impl PathComponent {
22    /// All possible path components.
23    ///
24    /// This makes it easy to iterate over all possible children of a branch
25    /// in a type-safe way. It is preferrable to do:
26    ///
27    /// ```ignore
28    /// for (idx, slot) in PathComponent::ALL.into_iter().zip(branch.children.each_ref()) {
29    ///     /// use idx and slot
30    /// }
31    /// ```
32    ///
33    /// instead of using a raw range like (`0..16`) or  [`Iterator::enumerate`],
34    /// which does not give a type-safe path component.
35    pub const ALL: [Self; Self::LEN] = [
36        Self(crate::u4::U4::new_masked(0x0)),
37        Self(crate::u4::U4::new_masked(0x1)),
38        Self(crate::u4::U4::new_masked(0x2)),
39        Self(crate::u4::U4::new_masked(0x3)),
40        Self(crate::u4::U4::new_masked(0x4)),
41        Self(crate::u4::U4::new_masked(0x5)),
42        Self(crate::u4::U4::new_masked(0x6)),
43        Self(crate::u4::U4::new_masked(0x7)),
44        Self(crate::u4::U4::new_masked(0x8)),
45        Self(crate::u4::U4::new_masked(0x9)),
46        Self(crate::u4::U4::new_masked(0xA)),
47        Self(crate::u4::U4::new_masked(0xB)),
48        Self(crate::u4::U4::new_masked(0xC)),
49        Self(crate::u4::U4::new_masked(0xD)),
50        Self(crate::u4::U4::new_masked(0xE)),
51        Self(crate::u4::U4::new_masked(0xF)),
52    ];
53
54    /// The number of possible path components.
55    pub const LEN: usize = 16;
56}
57
58impl PathComponent {
59    /// Returns the path component as a [`u8`].
60    #[must_use]
61    pub const fn as_u8(self) -> u8 {
62        self.0.as_u8()
63    }
64
65    /// Returns the path component as a [`usize`].
66    #[must_use]
67    pub const fn as_usize(self) -> usize {
68        self.as_u8() as usize
69    }
70
71    /// Tries to create a path component from the given [`u8`].
72    ///
73    /// For hexary tries, the input must be in the range 0x00 to 0x0F inclusive.
74    /// Any value outside this range will result in [`None`].
75    ///
76    /// For 256-ary tries, any value is valid.
77    #[must_use]
78    pub const fn try_new(value: u8) -> Option<Self> {
79        match crate::u4::U4::try_new(value) {
80            Some(u4) => Some(Self(u4)),
81            None => None,
82        }
83    }
84
85    /// Creates a pair of path components from a single byte where the
86    /// upper 4 bits are the first component and the lower 4 bits are the
87    /// second component.
88    #[must_use]
89    pub const fn new_pair(v: u8) -> (Self, Self) {
90        let (upper, lower) = crate::u4::U4::new_pair(v);
91        (Self(upper), Self(lower))
92    }
93
94    /// Joins this [`PathComponent`] with another to create a single [`u8`] where
95    /// this component is the upper 4 bits and the provided component is the
96    /// lower 4 bits.
97    #[must_use]
98    pub const fn join(self, other: Self) -> u8 {
99        self.0.join(other.0)
100    }
101
102    pub(crate) const fn wrapping_next(self) -> Self {
103        match crate::u4::U4::try_new(self.0.as_u8().wrapping_add(1)) {
104            Some(next) => Self(next),
105            None => Self(crate::u4::U4::MIN),
106        }
107    }
108}
109
110impl std::fmt::Debug for PathComponent {
111    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
112        std::fmt::Debug::fmt(&self.0, f)
113    }
114}
115
116impl std::fmt::Display for PathComponent {
117    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
118        write!(f, "{self:X}")
119    }
120}
121
122impl std::fmt::Binary for PathComponent {
123    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
124        std::fmt::Binary::fmt(&self.0, f)
125    }
126}
127
128impl std::fmt::LowerHex for PathComponent {
129    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
130        std::fmt::LowerHex::fmt(&self.0, f)
131    }
132}
133
134impl std::fmt::UpperHex for PathComponent {
135    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
136        std::fmt::UpperHex::fmt(&self.0, f)
137    }
138}
139
140impl TriePath for PathComponent {
141    type Components<'a>
142        = std::option::IntoIter<Self>
143    where
144        Self: 'a;
145
146    fn len(&self) -> usize {
147        1
148    }
149
150    fn components(&self) -> Self::Components<'_> {
151        Some(*self).into_iter()
152    }
153
154    fn as_component_slice(&self) -> PartialPath<'_> {
155        PartialPath::Borrowed(std::slice::from_ref(self))
156    }
157}
158
159impl TriePath for Option<PathComponent> {
160    type Components<'a>
161        = std::option::IntoIter<PathComponent>
162    where
163        Self: 'a;
164
165    fn len(&self) -> usize {
166        usize::from(self.is_some())
167    }
168
169    fn components(&self) -> Self::Components<'_> {
170        (*self).into_iter()
171    }
172
173    fn as_component_slice(&self) -> PartialPath<'_> {
174        PartialPath::Borrowed(self.as_slice())
175    }
176}
177
178impl TriePath for [PathComponent] {
179    type Components<'a>
180        = ComponentIter<'a>
181    where
182        Self: 'a;
183
184    fn len(&self) -> usize {
185        self.len()
186    }
187
188    fn components(&self) -> Self::Components<'_> {
189        self.iter().copied()
190    }
191
192    fn as_component_slice(&self) -> PartialPath<'_> {
193        PartialPath::Borrowed(self)
194    }
195}
196
197impl<const N: usize> TriePath for [PathComponent; N] {
198    type Components<'a>
199        = ComponentIter<'a>
200    where
201        Self: 'a;
202
203    fn len(&self) -> usize {
204        N
205    }
206
207    fn components(&self) -> Self::Components<'_> {
208        self.iter().copied()
209    }
210
211    fn as_component_slice(&self) -> PartialPath<'_> {
212        PartialPath::Borrowed(self)
213    }
214}
215
216impl TriePath for Vec<PathComponent> {
217    type Components<'a>
218        = ComponentIter<'a>
219    where
220        Self: 'a;
221
222    fn len(&self) -> usize {
223        self.len()
224    }
225
226    fn components(&self) -> Self::Components<'_> {
227        self.iter().copied()
228    }
229
230    fn as_component_slice(&self) -> PartialPath<'_> {
231        PartialPath::Borrowed(self.as_slice())
232    }
233}
234
235impl<A: smallvec::Array<Item = PathComponent>> TriePath for SmallVec<A> {
236    type Components<'a>
237        = ComponentIter<'a>
238    where
239        Self: 'a;
240
241    fn len(&self) -> usize {
242        self.len()
243    }
244
245    fn components(&self) -> Self::Components<'_> {
246        self.iter().copied()
247    }
248
249    fn as_component_slice(&self) -> PartialPath<'_> {
250        PartialPath::Borrowed(self.as_slice())
251    }
252}
253
254impl<'input> TriePathFromUnpackedBytes<'input> for &'input [PathComponent] {
255    type Error = crate::u4::TryFromIntError;
256
257    fn path_from_unpacked_bytes(bytes: &'input [u8]) -> Result<Self, Self::Error> {
258        if bytes.iter().all(|&b| b <= 0x0F) {
259            #[expect(unsafe_code)]
260            // SAFETY: we have verified that all bytes are in the valid range for
261            // `U4` (0x00 to 0x0F inclusive); therefore, it is now safe for us
262            // to reinterpret a &[u8] as a &[PathComponent].
263            Ok(unsafe { byte_slice_as_path_components_unchecked(bytes) })
264        } else {
265            Err(crate::u4::TryFromIntError)
266        }
267    }
268}
269
270impl TriePathFromUnpackedBytes<'_> for Vec<PathComponent> {
271    type Error = crate::u4::TryFromIntError;
272
273    fn path_from_unpacked_bytes(bytes: &[u8]) -> Result<Self, Self::Error> {
274        try_from_maybe_u4(bytes.iter().copied())
275    }
276}
277
278impl<A: smallvec::Array<Item = PathComponent>> TriePathFromUnpackedBytes<'_> for SmallVec<A> {
279    type Error = crate::u4::TryFromIntError;
280
281    fn path_from_unpacked_bytes(bytes: &[u8]) -> Result<Self, Self::Error> {
282        try_from_maybe_u4(bytes.iter().copied())
283    }
284}
285
286impl super::TriePathFromPackedBytes<'_> for Vec<PathComponent> {
287    fn path_from_packed_bytes(bytes: &'_ [u8]) -> Self {
288        let path = super::PackedPathRef::path_from_packed_bytes(bytes);
289        let mut this = Vec::new();
290        // reserve_exact is used because we trust that `TriePath::len` returns the exact
291        // length but `Vec::extend` won't trust `Iterator::size_hint` and may
292        // over/under-allocate.
293        this.reserve_exact(path.len());
294        this.extend(path.components());
295        this
296    }
297}
298
299impl<A: smallvec::Array<Item = PathComponent>> super::TriePathFromPackedBytes<'_> for SmallVec<A> {
300    fn path_from_packed_bytes(bytes: &'_ [u8]) -> Self {
301        let path = super::PackedPathRef::path_from_packed_bytes(bytes);
302        let mut this = SmallVec::<A>::new();
303        // reserve_exact is used because we trust that `TriePath::len` returns the exact
304        // length but `SmallVec::extend` won't trust `Iterator::size_hint` and may
305        // over/under-allocate.
306        this.reserve_exact(path.len());
307        this.extend(path.components());
308        this
309    }
310}
311
312impl<'input> TriePathFromUnpackedBytes<'input> for Box<[PathComponent]> {
313    type Error = <Vec<PathComponent> as TriePathFromUnpackedBytes<'input>>::Error;
314
315    fn path_from_unpacked_bytes(bytes: &'input [u8]) -> Result<Self, Self::Error> {
316        Vec::<PathComponent>::path_from_unpacked_bytes(bytes).map(Into::into)
317    }
318}
319
320impl<'input> TriePathFromUnpackedBytes<'input> for std::rc::Rc<[PathComponent]> {
321    type Error = <Vec<PathComponent> as TriePathFromUnpackedBytes<'input>>::Error;
322
323    fn path_from_unpacked_bytes(bytes: &'input [u8]) -> Result<Self, Self::Error> {
324        Vec::<PathComponent>::path_from_unpacked_bytes(bytes).map(Into::into)
325    }
326}
327
328impl<'input> TriePathFromUnpackedBytes<'input> for std::sync::Arc<[PathComponent]> {
329    type Error = <Vec<PathComponent> as TriePathFromUnpackedBytes<'input>>::Error;
330
331    fn path_from_unpacked_bytes(bytes: &'input [u8]) -> Result<Self, Self::Error> {
332        Vec::<PathComponent>::path_from_unpacked_bytes(bytes).map(Into::into)
333    }
334}
335
336impl super::TriePathFromPackedBytes<'_> for Box<[PathComponent]> {
337    fn path_from_packed_bytes(bytes: &[u8]) -> Self {
338        Vec::<PathComponent>::path_from_packed_bytes(bytes).into()
339    }
340}
341
342impl super::TriePathFromPackedBytes<'_> for std::rc::Rc<[PathComponent]> {
343    fn path_from_packed_bytes(bytes: &[u8]) -> Self {
344        Vec::<PathComponent>::path_from_packed_bytes(bytes).into()
345    }
346}
347
348impl super::TriePathFromPackedBytes<'_> for std::sync::Arc<[PathComponent]> {
349    fn path_from_packed_bytes(bytes: &[u8]) -> Self {
350        Vec::<PathComponent>::path_from_packed_bytes(bytes).into()
351    }
352}
353
354impl PathComponentSliceExt for [PathComponent] {
355    fn as_byte_slice(&self) -> &[u8] {
356        path_components_as_byte_slice(self)
357    }
358}
359
360impl<T: std::ops::Deref<Target = [PathComponent]>> PathComponentSliceExt for T {
361    fn as_byte_slice(&self) -> &[u8] {
362        path_components_as_byte_slice(self)
363    }
364}
365
366#[inline]
367const unsafe fn byte_slice_as_path_components_unchecked(bytes: &[u8]) -> &[PathComponent] {
368    #![expect(unsafe_code)]
369
370    // SAFETY: The caller must ensure that all bytes are valid for `PathComponent`,
371    // which is trivially true for 256-ary tries. For hexary tries, the caller must
372    // ensure that each byte is in the range 0x00 to 0x0F inclusive.
373    //
374    // We also rely on the fact that `PathComponent` is a single element type
375    // over `u8` (or `u4` which looks like a `u8` for this purpose).
376    //
377    // borrow rules ensure that the pointer for `bytes` is not null and
378    // `bytes.len()` is always valid. The returned reference will have the same
379    // lifetime as `bytes` so it cannot outlive the original slice.
380    unsafe {
381        &*(std::ptr::slice_from_raw_parts(bytes.as_ptr().cast::<PathComponent>(), bytes.len()))
382    }
383}
384
385#[inline]
386const fn path_components_as_byte_slice(components: &[PathComponent]) -> &[u8] {
387    #![expect(unsafe_code)]
388
389    // SAFETY: We rely on the fact that `PathComponent` is a single element type
390    // over `u8` (or `u4` which looks like a `u8` for this purpose).
391    //
392    // borrow rules ensure that the pointer for `components` is not null and
393    // `components.len()` is always valid. The returned reference will have the same
394    // lifetime as `components` so it cannot outlive the original slice.
395    unsafe {
396        &*(std::ptr::slice_from_raw_parts(components.as_ptr().cast::<u8>(), components.len()))
397    }
398}
399
400#[inline]
401fn try_from_maybe_u4<I: FromIterator<PathComponent>>(
402    bytes: impl IntoIterator<Item = u8>,
403) -> Result<I, crate::u4::TryFromIntError> {
404    bytes
405        .into_iter()
406        .map(PathComponent::try_new)
407        .collect::<Option<I>>()
408        .ok_or(crate::u4::TryFromIntError)
409}
410
411#[cfg(test)]
412mod tests {
413    use std::marker::PhantomData;
414
415    use test_case::test_case;
416
417    use super::*;
418
419    #[test_case(PhantomData::<&[PathComponent]>; "slice")]
420    #[test_case(PhantomData::<Box<[PathComponent]>>; "boxed slice")]
421    #[test_case(PhantomData::<Vec<PathComponent>>; "vec")]
422    #[test_case(PhantomData::<SmallVec<[PathComponent; 32]>>; "smallvec")]
423    fn test_path_from_unpacked_bytes_hexary<T>(_: PhantomData<T>)
424    where
425        T: TriePathFromUnpackedBytes<'static, Error: std::fmt::Debug>,
426    {
427        let input: &[u8; _] = &[0x00, 0x01, 0x0A, 0x0F];
428        let output = <T>::path_from_unpacked_bytes(input).expect("valid input");
429
430        assert_eq!(output.len(), input.len());
431        assert_eq!(
432            output
433                .components()
434                .map(PathComponent::as_u8)
435                .zip(input.iter().copied())
436                .take_while(|&(pc, b)| pc == b)
437                .count(),
438            input.len(),
439        );
440    }
441
442    #[test_case(PhantomData::<&[PathComponent]>; "slice")]
443    #[test_case(PhantomData::<Box<[PathComponent]>>; "boxed slice")]
444    #[test_case(PhantomData::<Vec<PathComponent>>; "vec")]
445    #[test_case(PhantomData::<SmallVec<[PathComponent; 32]>>; "smallvec")]
446    fn test_path_from_unpacked_bytes_hexary_invalid<T>(_: PhantomData<T>)
447    where
448        T: TriePathFromUnpackedBytes<'static> + std::fmt::Debug,
449    {
450        let input: &[u8; _] = &[0x00, 0x10, 0x0A, 0x0F];
451        let _ = <T>::path_from_unpacked_bytes(input).expect_err("invalid input");
452    }
453
454    #[test]
455    fn test_joined_path() {
456        let path = <&[PathComponent] as TriePathFromUnpackedBytes>::path_from_unpacked_bytes(&[
457            0x0A, 0x0B, 0x0C,
458        ])
459        .expect("valid input");
460
461        let with_suffix = path.append(PathComponent::try_new(0x0D).expect("valid"));
462        assert_eq!(with_suffix.len(), 4);
463        assert_eq!(
464            with_suffix
465                .components()
466                .map(PathComponent::as_u8)
467                .collect::<Vec<_>>(),
468            vec![0x0A, 0x0B, 0x0C, 0x0D],
469        );
470
471        let with_prefix = with_suffix.prepend(PathComponent::try_new(0x09).expect("valid"));
472        assert_eq!(with_prefix.len(), 5);
473        assert_eq!(
474            with_prefix
475                .components()
476                .map(PathComponent::as_u8)
477                .collect::<Vec<_>>(),
478            vec![0x09, 0x0A, 0x0B, 0x0C, 0x0D],
479        );
480    }
481}
firewood_storage/path/component.rs

firewood_storage/path/
component.rs
