willow_data_model/paths/builder.rs
1use core::mem::size_of;
2
3use bytes::{BufMut, BytesMut};
4
5use crate::prelude::*;
6
7use super::Representation;
8
9/// A helper struct for creating a [`Path`] with exactly one memory allocation. Requires total length and component count to be known in advance.
10///
11/// 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)) [`Component`]s, and that the total size in bytes of all [`Component`]s is at most `MPL` ([**m**ax\_**p**ath\_**l**ength](https://willowprotocol.org/specs/data-model/index.html#max_path_length)).
12///
13/// ```
14/// use willow_data_model::prelude::*;
15/// let mut builder = PathBuilder::<4, 4, 4>::new(4, 2)?;
16/// builder.append_component(Component::new(b"hi")?);
17/// builder.append_slice(b"ho")?;
18/// assert_eq!(builder.build(), Path::from_slices(&[b"hi", b"ho"])?);
19/// # Ok::<(), PathError>(())
20/// ```
21pub struct PathBuilder<const MCL: usize, const MCC: usize, const MPL: usize> {
22 bytes: BytesMut, // Turns into the [`HeapEncoding`] when building.
23 initialised_components: usize, // How many of the components have been appended to the `bytes` already?
24 target_length: usize, // The total length that the path should have when building.
25}
26
27impl<const MCL: usize, const MCC: usize, const MPL: usize> PathBuilder<MCL, MCC, MPL> {
28 /// Creates a builder for a [`Path`] of known total length and component count.
29 /// The component data must be filled in before building.
30 ///
31 /// #### Complexity
32 ///
33 /// Runs in `O(total_length + component_count)`, performs a single allocation of `O(total_length + component_count)` bytes.
34 ///
35 /// #### Examples
36 ///
37 /// ```
38 /// use willow_data_model::prelude::*;
39 /// let mut builder = PathBuilder::<4, 4, 4>::new(4, 2)?;
40 /// builder.append_component(Component::new(b"hi")?);
41 /// builder.append_slice(b"ho")?;
42 /// assert_eq!(builder.build(), Path::from_slices(&[b"hi", b"ho"])?);
43 /// # Ok::<(), PathError>(())
44 /// ```
45 pub fn new(
46 total_length: usize,
47 component_count: usize,
48 ) -> Result<Self, PathFromComponentsError> {
49 if total_length > MPL {
50 return Err(PathFromComponentsError::PathTooLong);
51 }
52
53 if component_count > MCC {
54 return Err(PathFromComponentsError::TooManyComponents);
55 }
56
57 // Allocate all storage in a single go.
58 let mut buf = BytesMut::with_capacity(Representation::allocation_size(
59 total_length,
60 component_count,
61 ));
62
63 // Place the number of components at the start of the buffer.
64 buf.extend_from_slice(&(component_count.to_ne_bytes())[..]);
65
66 // Fill up the accumulated component lengths with dummy zeroes, so we can use buf.extend_from_slice to append the actual component data in `append_component`. We overwrite the zeroes with actual values in `append_component`.
67 buf.put_bytes(0, component_count * size_of::<usize>());
68
69 Ok(Self {
70 bytes: buf,
71 initialised_components: 0,
72 target_length: total_length,
73 })
74 }
75
76 /// Creates a builder for a [`Path`] of known total length and component count, efficiently prefilled with the first `prefix_component_count` [`Component`]s of a given `reference` [`Path`]. Panics if there are not enough [`Component`]s in the `reference`.
77 ///
78 /// The missing component data must be filled in before building.
79 ///
80 /// #### Complexity
81 ///
82 /// Runs in `O(target_total_length + target_component_count)`, performs a single allocation of `O(total_length + component_count)` bytes.
83 ///
84 /// ```
85 /// use willow_data_model::prelude::*;
86 /// let p = Path::<4, 4, 4>::from_slices(&[b"hi", b"he"])?;
87 /// let mut builder = PathBuilder::new_from_prefix(4, 2, &p, 1)?;
88 /// builder.append_component(Component::new(b"ho")?);
89 /// assert_eq!(builder.build(), Path::from_slices(&[b"hi", b"ho"])?);
90 /// # Ok::<(), PathError>(())
91 /// ```
92 pub fn new_from_prefix(
93 target_total_length: usize,
94 target_component_count: usize,
95 reference: &Path<MCL, MCC, MPL>,
96 prefix_component_count: usize,
97 ) -> Result<Self, PathFromComponentsError> {
98 let mut builder = Self::new(target_total_length, target_component_count)?;
99
100 assert!(reference.component_count() >= prefix_component_count);
101
102 // The following is logically equivalent to successively appending the first `prefix_component_count` components of `reference` to the builder, but more efficient in terms of `memcpy` calls.
103 match prefix_component_count.checked_sub(1) {
104 None => return Ok(builder),
105 Some(index_of_final_component) => {
106 // Overwrite the dummy accumulated component lengths for the first `prefix_component_count` components with the actual values.
107 let accumulated_component_lengths_start =
108 Representation::start_offset_of_sum_of_lengths_for_component(0);
109 let accumulated_component_lengths_end =
110 Representation::start_offset_of_sum_of_lengths_for_component(
111 index_of_final_component + 1,
112 );
113 builder.bytes.as_mut()
114 [accumulated_component_lengths_start..accumulated_component_lengths_end]
115 .copy_from_slice(
116 &reference.data[accumulated_component_lengths_start
117 ..accumulated_component_lengths_end],
118 );
119
120 let components_start_in_prefix =
121 Representation::start_offset_of_component(&reference.data, 0);
122 let components_end_in_prefix = Representation::end_offset_of_component(
123 &reference.data,
124 index_of_final_component,
125 );
126 builder.bytes.extend_from_slice(
127 &reference.data[components_start_in_prefix..components_end_in_prefix],
128 );
129
130 // Record that we added the components.
131 builder.initialised_components = prefix_component_count;
132 }
133 }
134
135 // For reference, here is a naive implementation of the same functionality:
136
137 // for comp in reference
138 // .create_prefix(prefix_component_count)
139 // .expect("`prefix_component_count` must be less than the number of components in `reference`")
140 // .components()
141 // {
142 // builder.append_component(comp);
143 // }
144
145 Ok(builder)
146 }
147
148 /// Appends the data for the next [`Component`].
149 ///
150 /// #### Complexity
151 ///
152 /// Runs in `O(component_length)` time. Performs no allocations.
153 ///
154 /// #### Examples
155 ///
156 /// ```
157 /// use willow_data_model::prelude::*;
158 /// let mut builder = PathBuilder::<4, 4, 4>::new(4, 2)?;
159 /// builder.append_component(Component::new(b"hi")?);
160 /// builder.append_component(Component::new(b"ho")?);
161 /// assert_eq!(builder.build(), Path::from_slices(&[b"hi", b"ho"])?);
162 /// # Ok::<(), PathError>(())
163 /// ```
164 pub fn append_component(&mut self, component: &Component<MCL>) {
165 // Compute the accumulated length for the new component.
166 let total_length_so_far = match self.initialised_components.checked_sub(1) {
167 Some(i) => Representation::sum_of_lengths_for_component(&self.bytes, i),
168 None => 0,
169 };
170 let acc_length = component.as_ref().len() + total_length_so_far;
171
172 // Overwrite the dummy accumulated component length for this component with the actual value.
173 let start = Representation::start_offset_of_sum_of_lengths_for_component(
174 self.initialised_components,
175 );
176 let end = start + size_of::<usize>();
177 self.bytes.as_mut()[start..end].copy_from_slice(&acc_length.to_ne_bytes()[..]);
178
179 // Append the component to the path.
180 self.bytes.extend_from_slice(component.as_ref());
181
182 // Record that we added a component.
183 self.initialised_components += 1;
184 }
185
186 /// Appends the data for the next [`Component`], from a slice of bytes.
187 ///
188 /// #### Complexity
189 ///
190 /// Runs in `O(component_length)` time. Performs no allocations.
191 ///
192 /// #### Examples
193 ///
194 /// ```
195 /// use willow_data_model::prelude::*;
196 /// let mut builder = PathBuilder::<4, 4, 4>::new(4, 2)?;
197 /// builder.append_slice(b"hi")?;
198 /// builder.append_slice(b"ho")?;
199 /// assert_eq!(builder.build(), Path::from_slices(&[b"hi", b"ho"])?);
200 /// # Ok::<(), PathError>(())
201 /// ```
202 pub fn append_slice(&mut self, component: &[u8]) -> Result<(), InvalidComponentError> {
203 self.append_component(Component::new(component)?);
204 Ok(())
205 }
206
207 /// Turns this builder into an immutable [`Path`].
208 ///
209 /// Panics if the number of [`Component`]s or the total length does not match what was claimed in [`PathBuilder::new`].
210 ///
211 /// #### Complexity
212 ///
213 /// Runs in `O(1)` time. Performs no allocations.
214 ///
215 /// #### Examples
216 ///
217 /// ```
218 /// use willow_data_model::prelude::*;
219 /// let mut builder = PathBuilder::<4, 4, 4>::new(4, 2)?;
220 /// builder.append_component(Component::new(b"hi")?);
221 /// builder.append_slice(b"ho")?;
222 /// assert_eq!(builder.build(), Path::from_slices(&[b"hi", b"ho"])?);
223 /// # Ok::<(), PathError>(())
224 /// ```
225 pub fn build(self) -> Path<MCL, MCC, MPL> {
226 // Check whether we appended the correct number of components.
227 assert_eq!(
228 self.initialised_components,
229 Representation::component_count(self.bytes.as_ref())
230 );
231
232 assert_eq!(
233 self.target_length,
234 Representation::total_length(&self.bytes, self.initialised_components),
235 "Expected a target length of {}, but got an actual total_length of {}\nRaw representation:{:?}",
236 self.target_length,
237 Representation::total_length(&self.bytes, self.initialised_components),
238 self.bytes
239 );
240
241 Path {
242 data: self.bytes.freeze(),
243 component_count: self.initialised_components,
244 }
245 }
246}