firewood_storage/path/
packed.rs

1// Copyright (C) 2025, Ava Labs, Inc. All rights reserved.
2// See the file LICENSE.md for licensing terms.
3
4use crate::TriePathAsPackedBytes;
5
6use super::{PathComponent, SplitPath, TriePath, TriePathFromPackedBytes};
7
8/// A packed representation of a trie path where each byte encodes two path components.
9///
10/// This is borrowed over the underlying byte slice and does not allocate.
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
12pub struct PackedPathRef<'a> {
13    prefix: Option<PathComponent>,
14    middle: &'a [u8],
15    suffix: Option<PathComponent>,
16}
17
18/// An iterator over the components of a [`PackedPathRef`].
19#[derive(Clone, Debug, PartialEq, Eq, Hash)]
20#[must_use = "iterators are lazy and do nothing unless consumed"]
21pub struct PackedPathComponents<'a> {
22    path: PackedPathRef<'a>,
23}
24
25/// An iterator over the components of a path that yields them packed into bytes.
26///
27/// If there is a trailing component without a pair, it is padded with zero.
28#[derive(Clone, Debug, PartialEq, Eq, Hash)]
29#[must_use = "iterators are lazy and do nothing unless consumed"]
30pub struct PackedBytes<I> {
31    iter: I,
32}
33
34impl<I: Iterator<Item = PathComponent>> PackedBytes<I> {
35    /// Creates a new iterator over the packed bytes of the given component iterator.
36    pub const fn new(iter: I) -> Self {
37        Self { iter }
38    }
39}
40
41impl TriePath for PackedPathRef<'_> {
42    type Components<'a>
43        = PackedPathComponents<'a>
44    where
45        Self: 'a;
46
47    fn len(&self) -> usize {
48        self.middle
49            .len()
50            .wrapping_mul(2)
51            .wrapping_add(usize::from(self.prefix.is_some()))
52            .wrapping_add(usize::from(self.suffix.is_some()))
53    }
54
55    fn components(&self) -> Self::Components<'_> {
56        PackedPathComponents { path: *self }
57    }
58
59    fn as_component_slice(&self) -> super::PartialPath<'_> {
60        if self.is_empty() {
61            super::PartialPath::Borrowed(&[])
62        } else {
63            let mut buf = super::PathBuf::with_capacity(self.len());
64            if let Some(prefix) = self.prefix {
65                buf.push(prefix);
66            }
67            for &byte in self.middle {
68                let (upper, lower) = PathComponent::new_pair(byte);
69                buf.push(upper);
70                buf.push(lower);
71            }
72            if let Some(suffix) = self.suffix {
73                buf.push(suffix);
74            }
75            super::PartialPath::Owned(buf)
76        }
77    }
78}
79
80impl SplitPath for PackedPathRef<'_> {
81    fn split_at(self, mid: usize) -> (Self, Self) {
82        assert!(mid <= self.len(), "mid > self.len()");
83
84        if let Some(mid) = mid.checked_sub(usize::from(self.prefix.is_some())) {
85            // the mid consumed the prefix, if it existed
86            if mid.is_multiple_of(2) {
87                // middle is split cleanly
88                let (a_middle, b_middle) = self.middle.split_at(mid / 2);
89                let prefix = Self {
90                    prefix: self.prefix,
91                    middle: a_middle,
92                    suffix: None,
93                };
94                let suffix = Self {
95                    prefix: None,
96                    middle: b_middle,
97                    suffix: self.suffix,
98                };
99                (prefix, suffix)
100            } else {
101                // mid splits a middle component into the suffix of `prefix` and the
102                // prefix of `suffix`
103                let (a_middle, b_middle) = self.middle.split_at(mid / 2);
104                let Some((&middle_byte, b_middle)) = b_middle.split_first() else {
105                    // `mid` is oob of `b_middle`, which happens if self.suffix is Some,
106                    // and `mid` is self.len()
107                    return (self, Self::default());
108                };
109
110                let (upper, lower) = PathComponent::new_pair(middle_byte);
111                let prefix = Self {
112                    prefix: self.prefix,
113                    middle: a_middle,
114                    suffix: Some(upper),
115                };
116                let suffix = Self {
117                    prefix: Some(lower),
118                    middle: b_middle,
119                    suffix: self.suffix,
120                };
121                (prefix, suffix)
122            }
123        } else {
124            // `prefix` is some and `mid` is zero
125            (Self::default(), self)
126        }
127    }
128
129    fn split_first(self) -> Option<(PathComponent, Self)> {
130        let mut iter = self.into_iter();
131        let first = iter.next()?;
132        Some((first, iter.path))
133    }
134}
135
136impl<'input> TriePathFromPackedBytes<'input> for PackedPathRef<'input> {
137    fn path_from_packed_bytes(bytes: &'input [u8]) -> Self {
138        Self {
139            prefix: None,
140            middle: bytes,
141            suffix: None,
142        }
143    }
144}
145
146impl<'a> IntoIterator for PackedPathRef<'a> {
147    type Item = PathComponent;
148
149    type IntoIter = PackedPathComponents<'a>;
150
151    fn into_iter(self) -> Self::IntoIter {
152        PackedPathComponents { path: self }
153    }
154}
155
156impl Iterator for PackedPathComponents<'_> {
157    type Item = PathComponent;
158
159    fn next(&mut self) -> Option<Self::Item> {
160        if let Some(next) = self.path.prefix.take() {
161            return Some(next);
162        }
163
164        if let Some((&next, rest)) = self.path.middle.split_first() {
165            let (upper, lower) = PathComponent::new_pair(next);
166            self.path.prefix = Some(lower);
167            self.path.middle = rest;
168            return Some(upper);
169        }
170
171        self.path.suffix.take()
172    }
173
174    fn size_hint(&self) -> (usize, Option<usize>) {
175        let len = self.path.len();
176        (len, Some(len))
177    }
178}
179
180impl DoubleEndedIterator for PackedPathComponents<'_> {
181    fn next_back(&mut self) -> Option<Self::Item> {
182        if let Some(next) = self.path.suffix.take() {
183            return Some(next);
184        }
185
186        if let Some((&last, rest)) = self.path.middle.split_last() {
187            let (upper, lower) = PathComponent::new_pair(last);
188            self.path.suffix = Some(upper);
189            self.path.middle = rest;
190            return Some(lower);
191        }
192
193        self.path.prefix.take()
194    }
195}
196
197impl ExactSizeIterator for PackedPathComponents<'_> {}
198
199impl std::iter::FusedIterator for PackedPathComponents<'_> {}
200
201impl<I: Iterator<Item = PathComponent>> Iterator for PackedBytes<I> {
202    type Item = u8;
203
204    fn next(&mut self) -> Option<Self::Item> {
205        let hi = self.iter.next()?;
206        let lo = self.iter.next().unwrap_or(const { PathComponent::ALL[0] });
207        Some(hi.join(lo))
208    }
209
210    fn size_hint(&self) -> (usize, Option<usize>) {
211        let (lower, upper) = self.iter.size_hint();
212        let lower = lower.div_ceil(2);
213        let upper = upper.map(|u| u.div_ceil(2));
214        (lower, upper)
215    }
216}
217
218impl<I: DoubleEndedIterator<Item = PathComponent>> DoubleEndedIterator for PackedBytes<I> {
219    fn next_back(&mut self) -> Option<Self::Item> {
220        let lo = self.iter.next_back()?;
221        let hi = self
222            .iter
223            .next_back()
224            .unwrap_or(const { PathComponent::ALL[0] });
225        Some(hi.join(lo))
226    }
227}
228
229impl<I: ExactSizeIterator<Item = PathComponent>> ExactSizeIterator for PackedBytes<I> {}
230
231impl<I: std::iter::FusedIterator<Item = PathComponent>> std::iter::FusedIterator
232    for PackedBytes<I>
233{
234}
235
236impl<T: TriePath + ?Sized> TriePathAsPackedBytes for T {
237    type PackedBytesIter<'a>
238        = PackedBytes<T::Components<'a>>
239    where
240        Self: 'a;
241
242    fn as_packed_bytes(&self) -> Self::PackedBytesIter<'_> {
243        PackedBytes::new(self.components())
244    }
245}
246
247#[cfg(test)]
248mod tests {
249    #![expect(clippy::unwrap_used)]
250
251    use test_case::test_case;
252
253    use super::*;
254
255    /// Yields a [`PathComponent`] failing to compile if the given expression is
256    /// not a valid path component.
257    macro_rules! pc {
258        ($elem:expr) => {
259            const { PathComponent::ALL[$elem] }
260        };
261    }
262
263    /// Yields an array of [`PathComponent`]s failing to compile if any of the
264    /// given expressions are not valid path components.
265    ///
266    /// The expression yields an array, not a slice, and a reference must be taken
267    /// to convert it to a slice.
268    macro_rules! path {
269        ($($elem:expr),* $(,)?) => {
270            [ $( pc!($elem), )* ]
271        };
272    }
273
274    struct FromPackedBytesTest<'a> {
275        bytes: &'a [u8],
276        components: &'a [PathComponent],
277    }
278
279    #[test_case(
280        FromPackedBytesTest{
281            bytes: &[],
282            components: &path![],
283        };
284        "empty path"
285    )]
286    #[test_case(
287        FromPackedBytesTest{
288            bytes: b"a",
289            components: &path![6, 1],
290        };
291        "single byte path"
292    )]
293    #[test_case(
294        FromPackedBytesTest{
295            bytes: b"abc",
296            components: &path![6, 1, 6, 2, 6, 3],
297        };
298        "multi-byte path"
299    )]
300    fn test_from_packed_bytes(case: FromPackedBytesTest<'_>) {
301        let path = PackedPathRef::path_from_packed_bytes(case.bytes);
302        assert_eq!(path.len(), case.components.len());
303
304        assert!(path.components().eq(case.components.iter().copied()));
305        assert!(
306            path.components()
307                .rev()
308                .eq(case.components.iter().copied().rev())
309        );
310
311        assert!(path.as_packed_bytes().eq(case.bytes.iter().copied()));
312        assert!(
313            path.as_packed_bytes()
314                .rev()
315                .eq(case.bytes.iter().copied().rev())
316        );
317
318        assert!(TriePath::path_eq(&path, &case.components));
319        assert!(TriePath::path_cmp(&path, &case.components).is_eq());
320    }
321
322    struct SplitAtTest<'a> {
323        path: PackedPathRef<'a>,
324        mid: usize,
325        a: &'a [PathComponent],
326        b: &'a [PathComponent],
327    }
328
329    #[test_case(
330        SplitAtTest{
331            path: PackedPathRef::path_from_packed_bytes(b""),
332            mid: 0,
333            a: &path![],
334            b: &path![],
335        };
336        "empty path"
337    )]
338    #[test_case(
339        SplitAtTest{
340            path: PackedPathRef::path_from_packed_bytes(b"a"),
341            mid: 0,
342            a: &path![],
343            b: &path![6, 1],
344        };
345        "single byte path, split at 0"
346    )]
347    #[test_case(
348        SplitAtTest{
349            path: PackedPathRef::path_from_packed_bytes(b"a"),
350            mid: 1,
351            a: &path![6],
352            b: &path![1],
353        };
354        "single byte path, split at 1"
355    )]
356    #[test_case(
357        SplitAtTest{
358            path: PackedPathRef::path_from_packed_bytes(b"a"),
359            mid: 2,
360            a: &path![6, 1],
361            b: &path![],
362        };
363        "single byte path, split at len"
364    )]
365    #[test_case(
366        SplitAtTest{
367            path: PackedPathRef::path_from_packed_bytes(b"abc"),
368            mid: 2,
369            a: &path![6, 1],
370            b: &path![6, 2, 6, 3],
371        };
372        "multi byte path, split at 2"
373    )]
374    fn test_split_at(case: SplitAtTest<'_>) {
375        let (a, b) = case.path.split_at(case.mid);
376        assert_eq!(a.len(), case.mid);
377        assert_eq!(a.len().wrapping_add(b.len()), case.path.len());
378        assert!(a.path_eq(&case.a));
379        assert!(b.path_eq(&case.b));
380        assert!(a.append(b).path_eq(&case.path));
381        if let Some(mid) = case.mid.checked_sub(1) {
382            let (_, path) = dbg!(case.path.split_first()).unwrap();
383            let (_, b) = dbg!(path.split_at(mid));
384            assert!(
385                b.path_eq(&case.b),
386                "{} != {} ({}) (mid = {mid})",
387                b.display(),
388                case.b.display(),
389                path.display(),
390            );
391        }
392        if let Some(len) = case.path.len().checked_sub(1)
393            && len >= case.mid
394        {
395            let (path, _) = dbg!(case.path.split_at(len));
396            let (a, _) = dbg!(path.split_at(case.mid));
397            assert!(
398                a.path_eq(&case.a),
399                "{} != {} ({}) (len = {len})",
400                a.display(),
401                case.a.display(),
402                path.display(),
403            );
404        }
405    }
406
407    struct AsPackedBytesTest<'a, T> {
408        path: T,
409        expected: &'a [u8],
410    }
411
412    #[test_case(
413        AsPackedBytesTest {
414            path: PackedPathRef::path_from_packed_bytes(&[]),
415            expected: &[],
416        };
417        "empty packed path"
418    )]
419    #[test_case(
420        AsPackedBytesTest::<&[PathComponent]> {
421            path: &[],
422            expected: &[],
423        };
424        "empty unpacked path"
425    )]
426    #[test_case(
427        AsPackedBytesTest {
428            path: PackedPathRef::path_from_packed_bytes(b"abc"),
429            expected: b"abc",
430        };
431        "multi-byte packed path"
432    )]
433    #[test_case(
434        AsPackedBytesTest::<&[PathComponent]> {
435            path: &[PathComponent::ALL[6]],
436            expected: &[0x60],
437        };
438        "odd-lengthed unpacked path"
439    )]
440    #[test_case(
441        AsPackedBytesTest::<&[PathComponent]> {
442            path: &[PathComponent::ALL[6], PathComponent::ALL[1]],
443            expected: &[0x61],
444        };
445        "even-lengthed unpacked path"
446    )]
447    #[test_case(
448        AsPackedBytesTest::<&[PathComponent]> {
449            path: &[PathComponent::ALL[6], PathComponent::ALL[1], PathComponent::ALL[2]],
450            expected: &[0x61, 0x20],
451        };
452        "three-length unpacked path"
453    )]
454    fn test_as_packed_bytes<T: TriePath>(case: AsPackedBytesTest<'_, T>) {
455        let as_packed = case.path.as_packed_bytes().collect::<Vec<u8>>();
456        assert_eq!(*as_packed, *case.expected);
457    }
458}
firewood_storage/path/packed.rs

firewood_storage/path/
packed.rs
