willow_data_model/
path.rs

1#[cfg(feature = "dev")]
2use arbitrary::{size_hint::and_all, Arbitrary, Error as ArbitraryError, Unstructured};
3
4// This struct is tested in `fuzz/path.rs`, `fuzz/path2.rs`, `fuzz/path3.rs`, `fuzz/path3.rs` by comparing against a non-optimised reference implementation.
5// Further, the `successor` and `greater_but_not_prefixed` methods of that reference implementation are tested in `fuzz/path_successor.rs` and friends, and `fuzz/path_successor_of_prefix.rs` and friends.
6
7use core::borrow::Borrow;
8use core::convert::AsRef;
9use core::fmt::Debug;
10use core::hash::Hash;
11use core::iter;
12use core::mem::size_of;
13use core::ops::Deref;
14
15use bytes::{BufMut, Bytes, BytesMut};
16
17/// A [component](https://willowprotocol.org/specs/data-model/index.html#Component) of a Willow Path.
18///
19/// This type is a thin wrapper around `&'a [u8]` that enforces a const-generic [maximum component length](https://willowprotocol.org/specs/data-model/index.html#max_component_length). Use the `AsRef`, `DeRef`, or `Borrow` implementation to access the immutable byte slice.
20#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
21pub struct Component<'a, const MAX_COMPONENT_LENGTH: usize>(&'a [u8]);
22
23impl<'a, const MAX_COMPONENT_LENGTH: usize> Component<'a, MAX_COMPONENT_LENGTH> {
24    /// Create a `Component` from a byte slice. Return `None` if the slice is longer than `MaxComponentLength`.
25    pub fn new(slice: &'a [u8]) -> Option<Self> {
26        if slice.len() <= MAX_COMPONENT_LENGTH {
27            return Some(unsafe { Self::new_unchecked(slice) }); // Safe because we just checked the length.
28        } else {
29            None
30        }
31    }
32
33    /// Create a `Component` from a byte slice, without verifying its length.
34    ///
35    /// #### Safety
36    ///
37    /// Supplying a slice of length strictly greater than `MAX_COMPONENT_LENGTH` may trigger undefined behavior,
38    /// either immediately, or on any subsequent function invocation that operates on the resulting [`Component`].
39    pub unsafe fn new_unchecked(slice: &'a [u8]) -> Self {
40        Self(slice)
41    }
42
43    /// Create an empty component.
44    pub fn new_empty() -> Self {
45        Self(&[])
46    }
47
48    pub fn into_inner(self) -> &'a [u8] {
49        self.0
50    }
51}
52
53impl<'a, const MAX_COMPONENT_LENGTH: usize> Deref for Component<'a, MAX_COMPONENT_LENGTH> {
54    type Target = [u8];
55
56    fn deref(&self) -> &Self::Target {
57        self.0
58    }
59}
60
61impl<'a, const MAX_COMPONENT_LENGTH: usize> AsRef<[u8]> for Component<'a, MAX_COMPONENT_LENGTH> {
62    fn as_ref(&self) -> &[u8] {
63        self.0
64    }
65}
66
67impl<'a, const MAX_COMPONENT_LENGTH: usize> Borrow<[u8]> for Component<'a, MAX_COMPONENT_LENGTH> {
68    fn borrow(&self) -> &[u8] {
69        self.0
70    }
71}
72
73/// An owned [component](https://willowprotocol.org/specs/data-model/index.html#Component) of a Willow Path that uses reference counting for cheap cloning.
74///
75/// This type enforces a const-generic [maximum component length](https://willowprotocol.org/specs/data-model/index.html#max_component_length). Use the `AsRef`, `DeRef`, or `Borrow` implementation to access the immutable byte slice.
76#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
77pub struct OwnedComponent<const MAX_COMPONENT_LENGTH: usize>(Bytes);
78
79impl<const MAX_COMPONENT_LENGTH: usize> OwnedComponent<MAX_COMPONENT_LENGTH> {
80    /// Create an `OwnedComponent` by copying data from a byte slice. Return `None` if the slice is longer than `MaxComponentLength`.
81    ///
82    /// #### Complexity
83    ///
84    /// Runs in `O(n)`, where `n` is the length of the slice. Performs a single allocation of `O(n)` bytes.
85    pub fn new(data: &[u8]) -> Option<Self> {
86        if data.len() <= MAX_COMPONENT_LENGTH {
87            Some(unsafe { Self::new_unchecked(data) }) // Safe because we just checked the length.
88        } else {
89            None
90        }
91    }
92
93    /// Create an `OwnedComponent` by copying data from a byte slice, without verifying its length.
94    ///
95    /// #### Safety
96    ///
97    /// Supplying a slice of length strictly greater than `MAX_COMPONENT_LENGTH` may trigger undefined behavior,
98    /// either immediately, or on any subsequent function invocation that operates on the resulting [`OwnedComponent`].
99    ///
100    /// #### Complexity
101    ///
102    /// Runs in `O(n)`, where `n` is the length of the slice. Performs a single allocation of `O(n)` bytes.
103    pub unsafe fn new_unchecked(data: &[u8]) -> Self {
104        Self(Bytes::copy_from_slice(data))
105    }
106
107    /// Create an empty component.
108    ///
109    /// #### Complexity
110    ///
111    /// Runs in `O(1)`, performs no allocations.
112    pub fn new_empty() -> Self {
113        Self(Bytes::new())
114    }
115}
116
117impl<const MAX_COMPONENT_LENGTH: usize> Deref for OwnedComponent<MAX_COMPONENT_LENGTH> {
118    type Target = [u8];
119
120    fn deref(&self) -> &Self::Target {
121        self.0.deref()
122    }
123}
124
125impl<const MAX_COMPONENT_LENGTH: usize> AsRef<[u8]> for OwnedComponent<MAX_COMPONENT_LENGTH> {
126    fn as_ref(&self) -> &[u8] {
127        self.0.as_ref()
128    }
129}
130
131impl<const MAX_COMPONENT_LENGTH: usize> Borrow<[u8]> for OwnedComponent<MAX_COMPONENT_LENGTH> {
132    fn borrow(&self) -> &[u8] {
133        self.0.borrow()
134    }
135}
136
137#[derive(Debug)]
138/// An error arising from trying to construct a invalid [`Path`] from valid components.
139pub enum InvalidPathError {
140    /// The path's total length in bytes is too large.
141    PathTooLong,
142    /// The path has too many components.
143    TooManyComponents,
144}
145
146impl core::fmt::Display for InvalidPathError {
147    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
148        match self {
149            InvalidPathError::PathTooLong => {
150                write!(
151                    f,
152                    "Total length of a path in bytes exceeded the maximum path length"
153                )
154            }
155            InvalidPathError::TooManyComponents => {
156                write!(
157                    f,
158                    "Number of components of a path exceeded the maximum component count"
159                )
160            }
161        }
162    }
163}
164
165impl std::error::Error for InvalidPathError {}
166
167/// An immutable Willow [path](https://willowprotocol.org/specs/data-model/index.html#Path). Thread-safe, cheap to clone, cheap to take prefixes of, expensive to append to.
168///
169/// Enforces that each component has a length of at most `MCL` ([**m**ax\_**c**omponent\_**l**ength](https://willowprotocol.org/specs/data-model/index.html#max_component_length)), that each path has at most `MCC` ([**m**ax\_**c**omponent\_**c**count](https://willowprotocol.org/specs/data-model/index.html#max_component_count)) components, and that the total size in bytes of all components is at most `MPL` ([**m**ax\_**p**ath\_**l**ength](https://willowprotocol.org/specs/data-model/index.html#max_path_length)).
170#[derive(Clone)]
171pub struct Path<const MCL: usize, const MCC: usize, const MPL: usize> {
172    /// The data of the underlying path.
173    data: HeapEncoding<MCL>,
174    /// Number of components of the `data` to consider for this particular path. Must be less than or equal to the total number of components.
175    /// This field enables cheap prefix creation by cloning the heap data and adjusting the `component_count`.
176    component_count: usize,
177}
178
179impl<const MCL: usize, const MCC: usize, const MPL: usize> Path<MCL, MCC, MPL> {
180    /// Construct an empty path, i.e., a path of zero components.
181    ///
182    /// #### Complexity
183    ///
184    /// Runs in `O(1)`, performs no allocations.
185    pub fn new_empty() -> Self {
186        Path {
187            // 16 zero bytes, to work even on platforms on which `usize` has a size of 16.
188            data: HeapEncoding(Bytes::from_static(&[
189                0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
190            ])),
191            component_count: 0,
192        }
193    }
194
195    /// Construct a singleton path, i.e., a path of exactly one component.
196    ///
197    /// Copies the bytes of the component into an owned allocation on the heap.
198    ///
199    /// #### Complexity
200    ///
201    /// Runs in `O(n)`, where `n` is the length of the component. Performs a single allocation of `O(n)` bytes.
202    pub fn new_singleton(comp: Component<MCL>) -> Result<Self, InvalidPathError> {
203        if 1 > MCC {
204            Err(InvalidPathError::TooManyComponents)
205        } else if comp.len() > MPL {
206            return Err(InvalidPathError::PathTooLong);
207        } else {
208            let mut buf = BytesMut::with_capacity((2 * size_of::<usize>()) + comp.len());
209            buf.extend_from_slice(&(1usize.to_ne_bytes())[..]);
210            buf.extend_from_slice(&(comp.len().to_ne_bytes())[..]);
211            buf.extend_from_slice(comp.as_ref());
212
213            return Ok(Path {
214                data: HeapEncoding(buf.freeze()),
215                component_count: 1,
216            });
217        }
218    }
219
220    /// Construct a path of known total length from an [`ExactSizeIterator`][core::iter::ExactSizeIterator] of components.
221    ///
222    /// Copies the bytes of the components into an owned allocation on the heap.
223    ///
224    /// Panics if the claimed `total_length` does not match the sum of the lengths of all the components.
225    ///
226    /// #### Complexity
227    ///
228    /// Runs in `O(n)`, where `n` is the total length of the path in bytes. Performs a single allocation of `O(n)` bytes.
229    pub fn new_from_iter<'a, I>(total_length: usize, iter: &mut I) -> Result<Self, InvalidPathError>
230    where
231        I: ExactSizeIterator<Item = Component<'a, MCL>>,
232    {
233        if total_length > MPL {
234            return Err(InvalidPathError::PathTooLong);
235        }
236
237        let component_count = iter.len();
238
239        if component_count > MCC {
240            return Err(InvalidPathError::TooManyComponents);
241        }
242
243        let mut buf =
244            BytesMut::with_capacity(((component_count + 1) * size_of::<usize>()) + total_length);
245        buf.extend_from_slice(&(component_count.to_ne_bytes())[..]);
246
247        // Fill up the accumulated component lengths with dummy values.
248        buf.put_bytes(0, component_count * size_of::<usize>());
249
250        let mut accumulated_component_length = 0;
251        for (i, comp) in iter.enumerate() {
252            // Overwrite the dummy accumulated component length for this component with the actual value.
253            accumulated_component_length += comp.len();
254            let start = (1 + i) * size_of::<usize>();
255            let end = start + size_of::<usize>();
256            buf.as_mut()[start..end]
257                .copy_from_slice(&accumulated_component_length.to_ne_bytes()[..]);
258
259            // Append the component to the path.
260            buf.extend_from_slice(comp.as_ref());
261        }
262
263        if accumulated_component_length != total_length {
264            panic!("Tried to construct a path of total length {}, but got components whose accumulated length was {}.", total_length, accumulated_component_length);
265        }
266
267        Ok(Path {
268            data: HeapEncoding(buf.freeze()),
269            component_count,
270        })
271    }
272
273    /// Construct a path from a slice of components.
274    ///
275    /// Copies the bytes of the components into an owned allocation on the heap.
276    ///
277    /// #### Complexity
278    ///
279    /// Runs in `O(n)`, where `n` is the total length of the path in bytes. Performs a single allocation of `O(n)` bytes.
280    pub fn new_from_slice(components: &[Component<MCL>]) -> Result<Self, InvalidPathError> {
281        let mut total_length = 0;
282        for comp in components {
283            total_length += comp.len();
284        }
285
286        return Self::new_from_iter(total_length, &mut components.iter().copied());
287    }
288
289    /// Construct a new path by appending a component to this one.
290    ///
291    /// Creates a fully separate copy of the new data on the heap; this function is not more efficient than constructing the new path from scratch.
292    ///
293    /// #### Complexity
294    ///
295    /// Runs in `O(n)`, where `n` is the total length of the new path in bytes. Performs a single allocation of `O(n)` bytes.
296    pub fn append(&self, comp: Component<MCL>) -> Result<Self, InvalidPathError> {
297        let total_length = self.get_path_length() + comp.len();
298        return Self::new_from_iter(
299            total_length,
300            &mut ExactLengthChain::new(self.components(), iter::once(comp)),
301        );
302    }
303
304    /// Construct a new path by appending a slice of components to this one.
305    ///
306    /// Creates a fully separate copy of the new data on the heap; this function is not more efficient than constructing the new path from scratch.
307    ///
308    /// #### Complexity
309    ///
310    /// Runs in `O(n)`, where `n` is the total length of the new path in bytes. Performs a single allocation of `O(n)` bytes.
311    pub fn append_slice(&self, components: &[Component<MCL>]) -> Result<Self, InvalidPathError> {
312        let mut total_length = self.get_path_length();
313        for comp in components {
314            total_length += comp.len();
315        }
316
317        return Self::new_from_iter(
318            total_length,
319            &mut ExactLengthChain::new(self.components(), components.iter().copied()),
320        );
321    }
322
323    /// Get the number of components in this path.
324    ///
325    /// Guaranteed to be at most `MCC`.
326    ///
327    /// #### Complexity
328    ///
329    /// Runs in `O(1)`, performs no allocations.
330    pub fn get_component_count(&self) -> usize {
331        self.component_count
332    }
333
334    /// Return whether this path has zero components.
335    ///
336    /// #### Complexity
337    ///
338    /// Runs in `O(1)`, performs no allocations.
339    pub fn is_empty(&self) -> bool {
340        self.get_component_count() == 0
341    }
342
343    /// Get the sum of the lengths of all components in this path.
344    ///
345    /// Guaranteed to be at most `MCC`.
346    ///
347    /// #### Complexity
348    ///
349    /// Runs in `O(1)`, performs no allocations.
350    pub fn get_path_length(&self) -> usize {
351        if self.component_count == 0 {
352            0
353        } else {
354            return HeapEncoding::<MCL>::get_sum_of_lengths_for_component(
355                self.data.as_ref(),
356                self.component_count - 1,
357            )
358            .unwrap();
359        }
360    }
361
362    /// Get the `i`-th [`Component`] of this path.
363    ///
364    /// #### Complexity
365    ///
366    /// Runs in `O(1)`, performs no allocations.
367    pub fn get_component(&self, i: usize) -> Option<Component<MCL>> {
368        return HeapEncoding::<MCL>::get_component(self.data.as_ref(), i);
369    }
370
371    /// Get an owned handle to the `i`-th [`Component`] of this path.
372    ///
373    /// #### Complexity
374    ///
375    /// Runs in `O(1)`, performs no allocations.
376    pub fn get_owned_component(&self, i: usize) -> Option<OwnedComponent<MCL>> {
377        let start = HeapEncoding::<MCL>::start_offset_of_component(self.data.0.as_ref(), i)?;
378        let end = HeapEncoding::<MCL>::end_offset_of_component(self.data.0.as_ref(), i)?;
379        Some(OwnedComponent(self.data.0.slice(start..end)))
380    }
381
382    /// Create an iterator over the components of this path.
383    ///
384    /// Stepping the iterator takes `O(1)` time and performs no memory allocations.
385    ///
386    /// #### Complexity
387    ///
388    /// Runs in `O(1)`, performs no allocations.
389    pub fn components(
390        &self,
391    ) -> impl DoubleEndedIterator<Item = Component<MCL>> + ExactSizeIterator<Item = Component<MCL>>
392    {
393        self.suffix_components(0)
394    }
395
396    /// Create an iterator over the components of this path, starting at the `i`-th component. If `i` is greater than or equal to the number of components, the iterator yields zero items.
397    ///
398    /// Stepping the iterator takes `O(1)` time and performs no memory allocations.
399    ///
400    /// #### Complexity
401    ///
402    /// Runs in `O(1)`, performs no allocations.
403    pub fn suffix_components(
404        &self,
405        i: usize,
406    ) -> impl DoubleEndedIterator<Item = Component<MCL>> + ExactSizeIterator<Item = Component<MCL>>
407    {
408        (i..self.get_component_count()).map(|i| {
409            self.get_component(i).unwrap() // Only `None` if `i >= self.get_component_count()`
410        })
411    }
412
413    /// Create an iterator over owned handles to the components of this path.
414    ///
415    /// Stepping the iterator takes `O(1)` time and performs no memory allocations.
416    ///
417    /// #### Complexity
418    ///
419    /// Runs in `O(1)`, performs no allocations.
420    pub fn owned_components(
421        &self,
422    ) -> impl DoubleEndedIterator<Item = OwnedComponent<MCL>>
423           + ExactSizeIterator<Item = OwnedComponent<MCL>>
424           + '_ {
425        self.suffix_owned_components(0)
426    }
427
428    /// Create an iterator over owned handles to the components of this path, starting at the `i`-th component. If `i` is greater than or equal to the number of components, the iterator yields zero items.
429    ///
430    /// Stepping the iterator takes `O(1)` time and performs no memory allocations.
431    ///
432    /// #### Complexity
433    ///
434    /// Runs in `O(1)`, performs no allocations.
435    pub fn suffix_owned_components(
436        &self,
437        i: usize,
438    ) -> impl DoubleEndedIterator<Item = OwnedComponent<MCL>>
439           + ExactSizeIterator<Item = OwnedComponent<MCL>>
440           + '_ {
441        (i..self.get_component_count()).map(|i| {
442            self.get_owned_component(i).unwrap() // Only `None` if `i >= self.get_component_count()`
443        })
444    }
445
446    /// Create a new path that consists of the first `length` components. More efficient than creating a new [`Path`] from scratch.
447    ///
448    /// Returns `None` if `length` is greater than `self.get_component_count()`.
449    ///
450    /// #### Complexity
451    ///
452    /// Runs in `O(1)`, performs no allocations.
453    pub fn create_prefix(&self, length: usize) -> Option<Self> {
454        if length > self.get_component_count() {
455            None
456        } else {
457            Some(unsafe { self.create_prefix_unchecked(length) })
458        }
459    }
460
461    /// Create a new path that consists of the first `length` components. More efficient than creating a new [`Path`] from scratch.
462    ///
463    /// #### Safety
464    ///
465    /// Undefined behaviour if `length` is greater than `self.get_component_count()`. May manifest directly, or at any later
466    /// function invocation that operates on the resulting [`Path`].
467    ///
468    /// #### Complexity
469    ///
470    /// Runs in `O(1)`, performs no allocations.
471    pub unsafe fn create_prefix_unchecked(&self, length: usize) -> Self {
472        Self {
473            data: self.data.clone(),
474            component_count: length,
475        }
476    }
477
478    /// Create an iterator over all prefixes of this path (including th empty path and the path itself).
479    ///
480    /// Stepping the iterator takes `O(1)` time and performs no memory allocations.
481    ///
482    /// #### Complexity
483    ///
484    /// Runs in `O(1)`, performs no allocations.
485    pub fn all_prefixes(&self) -> impl DoubleEndedIterator<Item = Self> + '_ {
486        (0..=self.get_component_count()).map(|i| {
487            unsafe {
488                self.create_prefix_unchecked(i) // safe to call for i <= self.get_component_count()
489            }
490        })
491    }
492
493    /// Test whether this path is a prefix of the given path.
494    /// Paths are always a prefix of themselves, and the empty path is a prefix of every path.
495    ///
496    /// #### Complexity
497    ///
498    /// Runs in `O(n)`, where `n` is the total length of the shorter of the two paths. Performs no allocations.
499    pub fn is_prefix_of(&self, other: &Self) -> bool {
500        for (comp_a, comp_b) in self.components().zip(other.components()) {
501            if comp_a != comp_b {
502                return false;
503            }
504        }
505
506        self.get_component_count() <= other.get_component_count()
507    }
508
509    /// Test whether this path is prefixed by the given path.
510    /// Paths are always a prefix of themselves.
511    ///
512    /// #### Complexity
513    ///
514    /// Runs in `O(n)`, where `n` is the total length of the shorter of the two paths. Performs no allocations.
515    pub fn is_prefixed_by(&self, other: &Self) -> bool {
516        other.is_prefix_of(self)
517    }
518
519    /// Return the longest common prefix of this path and the given path.
520    ///
521    /// #### Complexity
522    ///
523    /// Runs in `O(n)`, where `n` is the total length of the shorter of the two paths. Performs a single allocation to create the return value.
524    pub fn longest_common_prefix(&self, other: &Self) -> Self {
525        let mut lcp_len = 0;
526
527        for (comp_a, comp_b) in self.components().zip(other.components()) {
528            if comp_a != comp_b {
529                break;
530            }
531
532            lcp_len += 1
533        }
534
535        self.create_prefix(lcp_len).unwrap() // zip ensures that lcp_len <= self.get_component_count()
536    }
537
538    /// Return the least path which is strictly greater than `self`, or return `None` if `self` is the greatest possible path.
539    ///
540    /// #### Complexity
541    ///
542    /// Runs in `O(n)`, where `n` is the total length of the shorter of the two paths. Performs a single allocation to create the return value.
543    pub fn successor(&self) -> Option<Self> {
544        // If it is possible to append an empty component, then doing so yields the successor.
545        if let Ok(path) = self.append(Component::new_empty()) {
546            return Some(path);
547        }
548
549        // Otherwise, we try incrementing the final component. If that fails,
550        // we try to increment the second-to-final component, and so on.
551        // All components that come after the incremented component are discarded.
552        // If *no* component can be incremented, `self` is the maximal path and we return `None`.
553
554        for (i, component) in self.components().enumerate().rev() {
555            // It would be nice to call a `try_increment_component` function, but in order to avoid
556            // memory allocations, we write some lower-level but more efficient code.
557
558            // If it is possible to append a zero byte to a component, then doing so yields its successor.
559            if component.len() < MCL
560                && HeapEncoding::<MCL>::get_sum_of_lengths_for_component(self.data.as_ref(), i)
561                    .unwrap() // i < self.component_count
562                    < MPL
563            {
564                // We now know how to construct the path successor of `self`:
565                // Take the first `i` components (this *excludes* the current `component`),
566                // then append `component` with an additinoal zero byte at the end.
567                let mut buf = clone_prefix_and_lengthen_final_component(self, i, 1);
568                buf.put_u8(0);
569
570                return Some(Self::from_buffer_and_component_count(buf.freeze(), i + 1));
571            }
572
573            // We **cannot** append a zero byte, so instead we check whether we can treat the component as a fixed-width integer and increment it. The only failure case is if that component consists of 255-bytes only.
574            let can_increment = !component.iter().all(|byte| *byte == 255);
575
576            // If we cannot increment, we go to the next iteration of the loop. But if we can, we can create a copy of the
577            // prefix on the first `i + 1` components, and mutate its backing memory in-place.
578            if can_increment {
579                let mut buf = clone_prefix_and_lengthen_final_component(self, i, 0);
580
581                let start_component_offset =
582                    HeapEncoding::<MCL>::start_offset_of_component(buf.as_ref(), i).unwrap(); // i < self.component_count
583                let end_component_offset =
584                    HeapEncoding::<MCL>::end_offset_of_component(buf.as_ref(), i).unwrap(); // i < self.component_count
585                fixed_width_increment(
586                    &mut buf.as_mut()[start_component_offset..end_component_offset],
587                );
588
589                return Some(Self::from_buffer_and_component_count(buf.freeze(), i + 1));
590            }
591        }
592
593        // Failed to increment any component, so `self` is the maximal path.
594        None
595    }
596
597    /// Return the least path that is strictly greater than `self` and which is not prefixed by `self`, or `None` if no such path exists.
598    ///
599    /// #### Complexity
600    ///
601    /// Runs in `O(n)`, where `n` is the total length of the shorter of the two paths. Performs a single allocation to create the return value.
602    pub fn greater_but_not_prefixed(&self) -> Option<Self> {
603        // We iterate through all components in reverse order. For each component, we check whether we can replace it by another cmponent that is strictly greater but not prefixed by the original component. If that is possible, we do replace it with the least such component and drop all later components. If that is impossible, we try again with the previous component. If this impossible for all components, then this function returns `None`.
604
605        for (i, component) in self.components().enumerate().rev() {
606            // If it is possible to append a zero byte to a component, then doing so yields its successor.
607            if component.len() < MCL
608                && HeapEncoding::<MCL>::get_sum_of_lengths_for_component(self.data.as_ref(), i)
609                    .unwrap() // i < self.component_count
610                    < MPL
611            {
612                let mut buf = clone_prefix_and_lengthen_final_component(self, i, 1);
613                buf.put_u8(0);
614
615                return Some(Self::from_buffer_and_component_count(buf.freeze(), i + 1));
616            }
617
618            // Next, we check whether the i-th component can be changed into the least component that is greater but not prefixed by the original. If so, do that and cut off all later components.
619            let mut next_component_length = None;
620            for (j, comp_byte) in component.iter().enumerate().rev() {
621                if *comp_byte < 255 {
622                    next_component_length = Some(j + 1);
623                    break;
624                }
625            }
626
627            if let Some(next_component_length) = next_component_length {
628                // Yay, we can replace the i-th comopnent and then we are done.
629
630                let mut buf = clone_prefix_and_lengthen_final_component(self, i, 0);
631                let length_of_prefix =
632                    HeapEncoding::<MCL>::get_sum_of_lengths_for_component(&buf, i).unwrap();
633
634                // Update the length of the final component.
635                buf_set_final_component_length(
636                    buf.as_mut(),
637                    i,
638                    length_of_prefix - (component.len() - next_component_length),
639                );
640
641                // Increment the byte at position `next_component_length` of the final component.
642                let offset = HeapEncoding::<MCL>::start_offset_of_component(buf.as_ref(), i)
643                    .unwrap()
644                    + next_component_length
645                    - 1;
646                let byte = buf.as_ref()[offset]; // guaranteed < 255...
647                buf.as_mut()[offset] = byte + 1; // ... hence no overflow here.
648
649                return Some(Self::from_buffer_and_component_count(buf.freeze(), i + 1));
650            }
651        }
652
653        None
654    }
655
656    pub(crate) fn from_buffer_and_component_count(buf: Bytes, component_count: usize) -> Self {
657        Path {
658            data: HeapEncoding(buf),
659            component_count,
660        }
661    }
662
663    pub(crate) fn raw_buf(&self) -> &[u8] {
664        self.data.0.as_ref()
665    }
666}
667
668impl<const MCL: usize, const MCC: usize, const MPL: usize> PartialEq for Path<MCL, MCC, MPL> {
669    fn eq(&self, other: &Self) -> bool {
670        if self.component_count != other.component_count {
671            false
672        } else {
673            return self.components().eq(other.components());
674        }
675    }
676}
677
678impl<const MCL: usize, const MCC: usize, const MPL: usize> Eq for Path<MCL, MCC, MPL> {}
679
680impl<const MCL: usize, const MCC: usize, const MPL: usize> Hash for Path<MCL, MCC, MPL> {
681    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
682        self.component_count.hash(state);
683
684        for comp in self.components() {
685            comp.hash(state);
686        }
687    }
688}
689
690impl<const MCL: usize, const MCC: usize, const MPL: usize> PartialOrd for Path<MCL, MCC, MPL> {
691    fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
692        Some(self.cmp(other))
693    }
694}
695
696/// Compare paths lexicogrphically, since that is the path ordering that the Willow spec always uses.
697impl<const MCL: usize, const MCC: usize, const MPL: usize> Ord for Path<MCL, MCC, MPL> {
698    fn cmp(&self, other: &Self) -> core::cmp::Ordering {
699        return self.components().cmp(other.components());
700    }
701}
702
703impl<const MCL: usize, const MCC: usize, const MPL: usize> Debug for Path<MCL, MCC, MPL> {
704    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
705        let data_vec: Vec<_> = self.components().collect();
706
707        f.debug_tuple("Path").field(&data_vec).finish()
708    }
709}
710
711/// Efficient, heap-allocated storage of a path encoding:
712///
713/// - First, a usize that gives the total number of path components.
714/// - Second, that many usizes, where the i-th one gives the sum of the lengths of the first i components.
715/// - Third, the concatenation of all components.
716///
717/// Note that these are not guaranteed to fulfil alignment requirements of usize, so we need to be careful in how we access these.
718/// Always use the methods on this struct for that reason.
719#[derive(Clone)]
720struct HeapEncoding<const MAX_COMPONENT_LENGTH: usize>(Bytes);
721
722// All offsets are in bytes, unless otherwise specified.
723// Arguments named `i` are the index of a component in the string, *not* a byte offset.
724impl<const MAX_COMPONENT_LENGTH: usize> HeapEncoding<MAX_COMPONENT_LENGTH> {
725    fn get_component_count(buf: &[u8]) -> usize {
726        Self::get_usize_at_offset(buf, 0).unwrap() // We always store at least 8byte for the component count.
727    }
728
729    // None if i is outside the slice.
730    fn get_sum_of_lengths_for_component(buf: &[u8], i: usize) -> Option<usize> {
731        let start_offset_in_bytes = Self::start_offset_of_sum_of_lengths_for_component(i);
732        Self::get_usize_at_offset(buf, start_offset_in_bytes)
733    }
734
735    fn get_component(buf: &[u8], i: usize) -> Option<Component<MAX_COMPONENT_LENGTH>> {
736        let start = Self::start_offset_of_component(buf, i)?;
737        let end = Self::end_offset_of_component(buf, i)?;
738        return Some(unsafe { Component::new_unchecked(&buf[start..end]) });
739    }
740
741    fn start_offset_of_sum_of_lengths_for_component(i: usize) -> usize {
742        // First usize is the number of components, then the i usizes storing the lengths; hence at i + 1.
743        size_of::<usize>() * (i + 1)
744    }
745
746    fn start_offset_of_component(buf: &[u8], i: usize) -> Option<usize> {
747        let metadata_length = (Self::get_component_count(buf) + 1) * size_of::<usize>();
748        if i == 0 {
749            Some(metadata_length)
750        } else {
751            Self::get_sum_of_lengths_for_component(buf, i - 1) // Length of everything up until the previous component.
752                .map(|length| length + metadata_length)
753        }
754    }
755
756    fn end_offset_of_component(buf: &[u8], i: usize) -> Option<usize> {
757        let metadata_length = (Self::get_component_count(buf) + 1) * size_of::<usize>();
758        Self::get_sum_of_lengths_for_component(buf, i).map(|length| length + metadata_length)
759    }
760
761    fn get_usize_at_offset(buf: &[u8], offset_in_bytes: usize) -> Option<usize> {
762        let end = offset_in_bytes + size_of::<usize>();
763
764        // We cannot interpret the memory in the slice as a usize directly, because the alignment might not match.
765        // So we first copy the bytes onto the heap, then construct a usize from it.
766        let mut usize_bytes = [0u8; size_of::<usize>()];
767
768        if buf.len() < end {
769            None
770        } else {
771            usize_bytes.copy_from_slice(&buf[offset_in_bytes..end]);
772            Some(usize::from_ne_bytes(usize_bytes))
773        }
774    }
775}
776
777impl<const MAX_COMPONENT_LENGTH: usize> AsRef<[u8]> for HeapEncoding<MAX_COMPONENT_LENGTH> {
778    fn as_ref(&self) -> &[u8] {
779        self.0.as_ref()
780    }
781}
782
783use syncify::syncify;
784use syncify::syncify_replace;
785
786#[syncify(encoding_sync)]
787mod encoding {
788    use super::*;
789
790    #[syncify_replace(use ufotofu::sync::{BulkConsumer, BulkProducer};)]
791    use ufotofu::local_nb::{BulkConsumer, BulkProducer};
792
793    use willow_encoding::DecodeError;
794    #[syncify_replace(use willow_encoding::sync::{Decodable, Encodable};)]
795    use willow_encoding::{Decodable, Encodable};
796
797    #[syncify_replace(use willow_encoding::sync::{decode_max_power, encode_max_power};)]
798    use willow_encoding::{decode_max_power, encode_max_power};
799
800    impl<const MCL: usize, const MCC: usize, const MPL: usize> Encodable for Path<MCL, MCC, MPL> {
801        async fn encode<C>(&self, consumer: &mut C) -> Result<(), C::Error>
802        where
803            C: BulkConsumer<Item = u8>,
804        {
805            encode_max_power(self.get_component_count(), MCC, consumer).await?;
806
807            for component in self.components() {
808                encode_max_power(component.len(), MCL, consumer).await?;
809
810                consumer
811                    .bulk_consume_full_slice(component.as_ref())
812                    .await
813                    .map_err(|f| f.reason)?;
814            }
815
816            Ok(())
817        }
818    }
819
820    impl<const MCL: usize, const MCC: usize, const MPL: usize> Decodable for Path<MCL, MCC, MPL> {
821        async fn decode<P>(producer: &mut P) -> Result<Self, DecodeError<P::Error>>
822        where
823            P: BulkProducer<Item = u8>,
824        {
825            let component_count: usize = decode_max_power(MCC, producer).await?.try_into()?;
826            if component_count > MCC {
827                return Err(DecodeError::InvalidInput);
828            }
829
830            let mut buf = ScratchSpacePathDecoding::<MCC, MPL>::new();
831
832            let mut accumulated_component_length: usize = 0; // Always holds the acc length of all components we copied so far.
833            for i in 0..component_count {
834                let component_len: usize = decode_max_power(MCL, producer).await?.try_into()?;
835                if component_len > MCL {
836                    return Err(DecodeError::InvalidInput);
837                }
838
839                accumulated_component_length += component_len;
840                if accumulated_component_length > MPL {
841                    return Err(DecodeError::InvalidInput);
842                }
843
844                buf.set_component_accumulated_length(accumulated_component_length, i);
845
846                // Decode the component itself into the scratch buffer.
847                producer
848                    .bulk_overwrite_full_slice(unsafe {
849                        // Safe because we called set_component_Accumulated_length for all j <= i
850                        buf.path_data_as_mut(i)
851                    })
852                    .await?;
853            }
854
855            Ok(unsafe { buf.to_path(component_count) })
856        }
857    }
858}
859
860#[cfg(feature = "dev")]
861impl<'a, const MCL: usize, const MCC: usize, const MPL: usize> Arbitrary<'a>
862    for Path<MCL, MCC, MPL>
863{
864    fn arbitrary(u: &mut Unstructured<'a>) -> Result<Self, ArbitraryError> {
865        let mut total_length_in_bytes: usize = Arbitrary::arbitrary(u)?;
866        total_length_in_bytes %= MPL + 1;
867
868        let data: Vec<u8> = Arbitrary::arbitrary(u)?;
869        total_length_in_bytes = core::cmp::min(total_length_in_bytes, data.len());
870
871        let mut num_components: usize = Arbitrary::arbitrary(u)?;
872        num_components %= MCC + 1;
873
874        if num_components == 0 {
875            total_length_in_bytes = 0;
876        }
877
878        let buf_capacity = size_of::<usize>() * (1 + num_components) + total_length_in_bytes;
879        let mut buf = BytesMut::with_capacity(buf_capacity);
880
881        // Write the component count of the path as the first usize.
882        buf.extend_from_slice(&num_components.to_ne_bytes());
883
884        let mut length_total_so_far = 0;
885        for i in 0..num_components {
886            // The final component must be long enough to result in the total_length_in_bytes.
887            if i + 1 == num_components {
888                let final_component_length = total_length_in_bytes - length_total_so_far;
889
890                if final_component_length > MCL {
891                    return Err(ArbitraryError::IncorrectFormat);
892                } else {
893                    buf.extend_from_slice(&total_length_in_bytes.to_ne_bytes());
894                }
895            } else {
896                // Any non-final component can take on a random length, ...
897                let mut component_length: usize = Arbitrary::arbitrary(u)?;
898                // ... except it must be at most the MCL, and...
899                component_length %= MCL + 1;
900                // ... the total length of all components must not exceed the total path length.
901                component_length = core::cmp::min(
902                    component_length,
903                    total_length_in_bytes - length_total_so_far,
904                );
905
906                length_total_so_far += component_length;
907                buf.extend_from_slice(&length_total_so_far.to_ne_bytes());
908            }
909        }
910
911        // Finally, add the random path data.
912        buf.extend_from_slice(&data);
913
914        Ok(Path {
915            data: HeapEncoding(buf.freeze()),
916            component_count: num_components,
917        })
918    }
919
920    #[inline]
921    fn size_hint(depth: usize) -> (usize, Option<usize>) {
922        (
923            and_all(&[
924                usize::size_hint(depth),
925                usize::size_hint(depth),
926                Vec::<u8>::size_hint(depth),
927            ])
928            .0,
929            None,
930        )
931    }
932}
933
934/// Like core::iter::Chain, but implements ExactSizeIter if both components implement it. Panics if the resulting length overflows.
935///
936/// Code liberally copy-pasted from the standard library.
937struct ExactLengthChain<A, B> {
938    a: Option<A>,
939    b: Option<B>,
940}
941
942impl<A, B> ExactLengthChain<A, B> {
943    fn new(a: A, b: B) -> ExactLengthChain<A, B> {
944        ExactLengthChain {
945            a: Some(a),
946            b: Some(b),
947        }
948    }
949}
950
951impl<A, B> Iterator for ExactLengthChain<A, B>
952where
953    A: Iterator,
954    B: Iterator<Item = A::Item>,
955{
956    type Item = A::Item;
957
958    #[inline]
959    fn next(&mut self) -> Option<A::Item> {
960        and_then_or_clear(&mut self.a, Iterator::next).or_else(|| self.b.as_mut()?.next())
961    }
962
963    fn size_hint(&self) -> (usize, Option<usize>) {
964        let (lower_a, higher_a) = self.a.as_ref().map_or((0, None), |a| a.size_hint());
965        let (lower_b, higher_b) = self.b.as_ref().map_or((0, None), |b| b.size_hint());
966
967        let higher = match (higher_a, higher_b) {
968            (Some(a), Some(b)) => Some(
969                a.checked_add(b)
970                    .expect("Some of lengths of two iterators must not overflow."),
971            ),
972            _ => None,
973        };
974
975        (lower_a + lower_b, higher)
976    }
977}
978
979impl<A, B> DoubleEndedIterator for ExactLengthChain<A, B>
980where
981    A: DoubleEndedIterator,
982    B: DoubleEndedIterator<Item = A::Item>,
983{
984    #[inline]
985    fn next_back(&mut self) -> Option<A::Item> {
986        and_then_or_clear(&mut self.b, |b| b.next_back()).or_else(|| self.a.as_mut()?.next_back())
987    }
988}
989
990impl<A, B> ExactSizeIterator for ExactLengthChain<A, B>
991where
992    A: ExactSizeIterator,
993    B: ExactSizeIterator<Item = A::Item>,
994{
995}
996
997#[inline]
998fn and_then_or_clear<T, U>(opt: &mut Option<T>, f: impl FnOnce(&mut T) -> Option<U>) -> Option<U> {
999    let x = f(opt.as_mut()?);
1000    if x.is_none() {
1001        *opt = None;
1002    }
1003    x
1004}
1005
1006// In a buffer that stores a path on the heap, set the sum of all component lengths for the i-th component, which must be the final component.
1007fn buf_set_final_component_length(buf: &mut [u8], i: usize, new_sum_of_lengths: usize) {
1008    let comp_len_start = (1 + i) * size_of::<usize>();
1009    let comp_len_end = comp_len_start + size_of::<usize>();
1010    buf[comp_len_start..comp_len_end].copy_from_slice(&new_sum_of_lengths.to_ne_bytes()[..]);
1011}
1012
1013// Overflows to all zeroes if all bytes are 255.
1014fn fixed_width_increment(buf: &mut [u8]) {
1015    for byte_ref in buf.iter_mut().rev() {
1016        if *byte_ref == 255 {
1017            *byte_ref = 0;
1018        } else {
1019            *byte_ref += 1;
1020            return;
1021        }
1022    }
1023}
1024
1025/// Create a new BufMut that stores the heap encoding of the first i components of `original`, but increasing the length of the final component by `extra_capacity`. No data to fill that extra capacity is written into the buffer.
1026fn clone_prefix_and_lengthen_final_component<
1027    const MCL: usize,
1028    const MCC: usize,
1029    const MPL: usize,
1030>(
1031    original: &Path<MCL, MCC, MPL>,
1032    i: usize,
1033    extra_capacity: usize,
1034) -> BytesMut {
1035    let original_slice = original.data.as_ref();
1036    let successor_path_length =
1037        HeapEncoding::<MCL>::get_sum_of_lengths_for_component(original_slice, i).unwrap()
1038            + extra_capacity;
1039    let buf_capacity = size_of::<usize>() * (i + 2) + successor_path_length;
1040    let mut buf = BytesMut::with_capacity(buf_capacity);
1041
1042    // Write the length of the successor path as the first usize.
1043    buf.extend_from_slice(&(i + 1).to_ne_bytes());
1044
1045    // Next, copy the total path lengths for the first i prefixes.
1046    buf.extend_from_slice(&original_slice[size_of::<usize>()..size_of::<usize>() * (i + 2)]);
1047
1048    // Now, write the length of the final component, which is one greater than before.
1049    buf_set_final_component_length(buf.as_mut(), i, successor_path_length);
1050
1051    // Finally, copy the raw bytes of the first i+1 components.
1052    buf.extend_from_slice(
1053        &original_slice[HeapEncoding::<MCL>::start_offset_of_component(original_slice, 0).unwrap()
1054            ..HeapEncoding::<MCL>::start_offset_of_component(original_slice, i + 1).unwrap()],
1055    );
1056
1057    buf
1058}
1059
1060/// A memory region to use for decoding paths. Reused between many decodings.
1061#[derive(Debug)]
1062pub(crate) struct ScratchSpacePathDecoding<const MCC: usize, const MPL: usize> {
1063    // The i-th usize holds the total lengths of the first i components.
1064    component_accumulated_lengths: [usize; MCC],
1065    path_data: [u8; MPL],
1066}
1067
1068impl<const MCC: usize, const MPL: usize> ScratchSpacePathDecoding<MCC, MPL> {
1069    pub fn new() -> Self {
1070        ScratchSpacePathDecoding {
1071            component_accumulated_lengths: [0; MCC],
1072            path_data: [0; MPL],
1073        }
1074    }
1075
1076    /// Panic if i >= MCC.
1077    pub fn set_component_accumulated_length(
1078        &mut self,
1079        component_accumulated_length: usize,
1080        i: usize,
1081    ) {
1082        self.component_accumulated_lengths[i] = component_accumulated_length;
1083    }
1084
1085    /// # Safety
1086    ///
1087    /// UB if length of slice is greater than `size_of::<usize>() * MCC`.
1088    pub unsafe fn set_many_component_accumulated_lengths_from_ne(&mut self, lengths: &[u8]) {
1089        let slice: &mut [u8] = core::slice::from_raw_parts_mut(
1090            self.component_accumulated_lengths[..lengths.len() / size_of::<usize>()].as_mut_ptr()
1091                as *mut u8,
1092            lengths.len(),
1093        );
1094
1095        slice.copy_from_slice(lengths);
1096    }
1097
1098    /// # Safety
1099    ///
1100    /// Memory must have been initialised with a prior call to set_component_accumulated_length for the same `i`.
1101    unsafe fn get_component_accumulated_length(&self, i: usize) -> usize {
1102        self.component_accumulated_lengths[i]
1103    }
1104
1105    /// Return a slice of the accumulated component lengths up to but excluding the `i`-th component, encoded as native-endian u8s.
1106    ///
1107    /// # Safety
1108    ///
1109    /// Memory must have been initialised with prior call to set_component_accumulated_length for all `j <= i`
1110    pub unsafe fn get_accumumulated_component_lengths(&self, i: usize) -> &[u8] {
1111        core::slice::from_raw_parts(
1112            self.component_accumulated_lengths[..i].as_ptr() as *const u8,
1113            i * size_of::<usize>(),
1114        )
1115    }
1116
1117    /// Return a mutable slice of the i-th path_data.
1118    ///
1119    /// # Safety
1120    ///
1121    /// Accumulated component lengths for `i` and `i - 1` must have been set (only for `i` if `i == 0`).
1122    pub unsafe fn path_data_as_mut(&mut self, i: usize) -> &mut [u8] {
1123        let start = if i == 0 {
1124            0
1125        } else {
1126            self.get_component_accumulated_length(i - 1)
1127        };
1128        let end = self.get_component_accumulated_length(i);
1129        &mut self.path_data[start..end]
1130    }
1131
1132    /// Return a mutable slice of the path_data up to but excluding the i-th component.
1133    ///
1134    /// # Safety
1135    ///
1136    /// Accumulated component lengths for `i - 1` must have been set (unless `i == 0`).
1137    pub unsafe fn path_data_until_as_mut(&mut self, i: usize) -> &mut [u8] {
1138        let end = self.get_component_accumulated_length(i - 1);
1139        &mut self.path_data[0..end]
1140    }
1141
1142    /// Get the path data of the first `i` components.
1143    ///
1144    /// # Safety
1145    ///
1146    /// Memory must have been initialised with a prior call to set_component_accumulated_length for `i - 1` (ignored if `i == 0`).
1147    /// Also, the path data must have been initialised via a reference obtained from `self.path_data_as_mut()`.
1148    unsafe fn get_path_data(&self, i: usize) -> &[u8] {
1149        let end = if i == 0 {
1150            0
1151        } else {
1152            self.get_component_accumulated_length(i - 1)
1153        };
1154        &self.path_data[..end]
1155    }
1156
1157    /// Copy the data from this struct into a new Path of `i` components.
1158    ///
1159    /// # Safety
1160    ///
1161    /// The first `i` accumulated component lengths must have been set, and all corresponding path data must be initialised. MCL, MCC, and MCP are trusted blindly and must be adhered to by the data in the scratch buffer.
1162    pub unsafe fn to_path<const MCL: usize>(&self, i: usize) -> Path<MCL, MCC, MPL> {
1163        if i == 0 {
1164            Path::new_empty()
1165        } else {
1166            let total_length = if i == 0 {
1167                0
1168            } else {
1169                self.get_component_accumulated_length(i - 1)
1170            };
1171            let mut buf = BytesMut::with_capacity((size_of::<usize>() * (i + 1)) + total_length);
1172
1173            buf.extend_from_slice(&i.to_ne_bytes());
1174            buf.extend_from_slice(self.get_accumumulated_component_lengths(i));
1175            buf.extend_from_slice(self.get_path_data(i));
1176
1177            Path::from_buffer_and_component_count(buf.freeze(), i)
1178        }
1179    }
1180}
willow_data_model/path.rs

willow_data_model/
path.rs
