willow_data_model/paths/builder.rs
1use core::mem::size_of;
2
3use bytes::{BufMut, BytesMut};
4
5use ufotofu::codec_prelude::*;
6
7use crate::prelude::*;
8
9use super::Representation;
10
11/// A helper struct for creating a [`Path`] with exactly one memory allocation. Requires total length and component count to be known in advance.
12///
13/// 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)).
14///
15/// ```
16/// use willow_data_model::prelude::*;
17/// let mut builder = PathBuilder::<4, 4, 4>::new(4, 2)?;
18/// builder.append_component(Component::new(b"hi")?);
19/// builder.append_slice(b"ho")?;
20/// assert_eq!(builder.build(), Path::from_slices(&[b"hi", b"ho"])?);
21/// # Ok::<(), PathError>(())
22/// ```
23#[derive(Debug)]
24pub struct PathBuilder<const MCL: usize, const MCC: usize, const MPL: usize> {
25 bytes: BytesMut, // Turns into the [`HeapEncoding`] when building.
26 initialised_components: usize, // How many of the components have been appended to the `bytes` already?
27 target_length: usize, // The total length that the path should have when building.
28}
29
30impl<const MCL: usize, const MCC: usize, const MPL: usize> PathBuilder<MCL, MCC, MPL> {
31 /// Creates a builder for a [`Path`] of known total length and component count.
32 /// The component data must be filled in before building.
33 ///
34 /// #### Complexity
35 ///
36 /// Runs in `O(total_length + component_count)`, performs a single allocation of `O(total_length + component_count)` bytes.
37 ///
38 /// #### Examples
39 ///
40 /// ```
41 /// use willow_data_model::prelude::*;
42 /// let mut builder = PathBuilder::<4, 4, 4>::new(4, 2)?;
43 /// builder.append_component(Component::new(b"hi")?);
44 /// builder.append_slice(b"ho")?;
45 /// assert_eq!(builder.build(), Path::from_slices(&[b"hi", b"ho"])?);
46 /// # Ok::<(), PathError>(())
47 /// ```
48 pub fn new(
49 total_length: usize,
50 component_count: usize,
51 ) -> Result<Self, PathFromComponentsError> {
52 if total_length > MPL {
53 return Err(PathFromComponentsError::PathTooLong);
54 }
55
56 if component_count > MCC {
57 return Err(PathFromComponentsError::TooManyComponents);
58 }
59
60 // Allocate all storage in a single go.
61 let mut buf = BytesMut::with_capacity(Representation::allocation_size(
62 total_length,
63 component_count,
64 ));
65
66 // Place the number of components at the start of the buffer.
67 buf.extend_from_slice(&(component_count.to_ne_bytes())[..]);
68
69 // 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`.
70 buf.put_bytes(0, component_count * size_of::<usize>());
71
72 Ok(Self {
73 bytes: buf,
74 initialised_components: 0,
75 target_length: total_length,
76 })
77 }
78
79 /// 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`.
80 ///
81 /// The missing component data must be filled in before building.
82 ///
83 /// #### Complexity
84 ///
85 /// Runs in `O(target_total_length + target_component_count)`, performs a single allocation of `O(total_length + component_count)` bytes.
86 ///
87 /// ```
88 /// use willow_data_model::prelude::*;
89 /// let p = Path::<4, 4, 4>::from_slices(&[b"hi", b"he"])?;
90 /// let mut builder = PathBuilder::new_from_prefix(4, 2, &p, 1)?;
91 /// builder.append_component(Component::new(b"ho")?);
92 /// assert_eq!(builder.build(), Path::from_slices(&[b"hi", b"ho"])?);
93 /// # Ok::<(), PathError>(())
94 /// ```
95 pub fn new_from_prefix(
96 target_total_length: usize,
97 target_component_count: usize,
98 reference: &Path<MCL, MCC, MPL>,
99 prefix_component_count: usize,
100 ) -> Result<Self, PathFromComponentsError> {
101 let mut builder = Self::new(target_total_length, target_component_count)?;
102
103 assert!(reference.component_count() >= prefix_component_count);
104
105 // 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.
106 match prefix_component_count.checked_sub(1) {
107 None => return Ok(builder),
108
109 Some(index_of_final_component) => {
110 // Overwrite the dummy accumulated component lengths for the first `prefix_component_count` components with the actual values.
111
112 if reference.first_component == 0 {
113 // If our `reference` path doesn't have an offset first component (from a slice or suffix operation), we can copy the cumulative path length data with one `memcpy`...
114 let accumulated_component_lengths_start =
115 Representation::start_offset_of_sum_of_lengths_for_component(0);
116 let accumulated_component_lengths_end =
117 Representation::start_offset_of_sum_of_lengths_for_component(
118 index_of_final_component + 1,
119 );
120 builder.bytes.as_mut()
121 [accumulated_component_lengths_start..accumulated_component_lengths_end]
122 .copy_from_slice(
123 &reference.data[accumulated_component_lengths_start
124 ..accumulated_component_lengths_end],
125 );
126 } else {
127 // ...otherwise we need to adjust each element of the cumulative path length when we copy it across.
128
129 let mut start_of_component_length =
130 Representation::start_offset_of_sum_of_lengths_for_component(0);
131 for i in 0..prefix_component_count {
132 let end_of_component_length =
133 start_of_component_length + size_of::<usize>();
134 builder.bytes.as_mut()[start_of_component_length..end_of_component_length]
135 .copy_from_slice(
136 &reference.total_length_of_prefix(i + 1).to_ne_bytes(),
137 );
138 start_of_component_length = end_of_component_length;
139 }
140 }
141
142 let start_of_prefix_components = Representation::start_offset_of_component(
143 &reference.data,
144 reference.first_component,
145 );
146
147 let end_of_prefix_components = Representation::end_offset_of_component(
148 &reference.data,
149 reference.first_component + index_of_final_component,
150 );
151
152 builder.bytes.extend_from_slice(
153 &reference.data[start_of_prefix_components..end_of_prefix_components],
154 );
155
156 builder.initialised_components = prefix_component_count;
157 }
158 }
159
160 // For reference, here is a naive implementation of the same functionality:
161
162 // for comp in reference
163 // .create_prefix(prefix_component_count)
164 // .expect("`prefix_component_count` must be less than the number of components in `reference`")
165 // .components()
166 // {
167 // builder.append_component(comp);
168 // }
169
170 Ok(builder)
171 }
172
173 /// Appends the data for the next [`Component`].
174 ///
175 /// #### Complexity
176 ///
177 /// Runs in `O(component_length)` time. Performs no allocations.
178 ///
179 /// #### Examples
180 ///
181 /// ```
182 /// use willow_data_model::prelude::*;
183 /// let mut builder = PathBuilder::<4, 4, 4>::new(4, 2)?;
184 /// builder.append_component(Component::new(b"hi")?);
185 /// builder.append_component(Component::new(b"ho")?);
186 /// assert_eq!(builder.build(), Path::from_slices(&[b"hi", b"ho"])?);
187 /// # Ok::<(), PathError>(())
188 /// ```
189 pub fn append_component(&mut self, component: &Component<MCL>) {
190 // Compute the accumulated length for the new component.
191 let total_length_so_far = match self.initialised_components.checked_sub(1) {
192 Some(i) => Representation::sum_of_lengths_for_component(&self.bytes, i),
193 None => 0,
194 };
195 let acc_length = component.as_ref().len() + total_length_so_far;
196
197 // Overwrite the dummy accumulated component length for this component with the actual value.
198 let start = Representation::start_offset_of_sum_of_lengths_for_component(
199 self.initialised_components,
200 );
201 let end = start + size_of::<usize>();
202 self.bytes.as_mut()[start..end].copy_from_slice(&acc_length.to_ne_bytes()[..]);
203
204 // Append the component to the path.
205 self.bytes.extend_from_slice(component.as_ref());
206
207 // Record that we added a component.
208 self.initialised_components += 1;
209 }
210
211 /// Appends the data for the next [`Component`], from a slice of bytes.
212 ///
213 /// #### Complexity
214 ///
215 /// Runs in `O(component_length)` time. Performs no allocations.
216 ///
217 /// #### Examples
218 ///
219 /// ```
220 /// use willow_data_model::prelude::*;
221 /// let mut builder = PathBuilder::<4, 4, 4>::new(4, 2)?;
222 /// builder.append_slice(b"hi")?;
223 /// builder.append_slice(b"ho")?;
224 /// assert_eq!(builder.build(), Path::from_slices(&[b"hi", b"ho"])?);
225 /// # Ok::<(), PathError>(())
226 /// ```
227 pub fn append_slice(&mut self, component: &[u8]) -> Result<(), InvalidComponentError> {
228 self.append_component(Component::new(component)?);
229 Ok(())
230 }
231
232 /// Appends data for a component of known length by reading data from a [`BulkProducer`] of bytes. Panics if `component_length > MCL`.
233 ///
234 /// #### Complexity
235 ///
236 /// Runs in `O(component_length)` time (assuming the producer takes `O(1)` time per byte). Performs no allocations.
237 ///
238 /// #### Examples
239 ///
240 /// ```
241 /// use willow_data_model::prelude::*;
242 /// use ufotofu::{codec_prelude::*, producer::clone_from_slice};
243 ///
244 /// # pollster::block_on(async {
245 /// let mut builder = PathBuilder::<4, 4, 4>::new(4, 2)?;
246 /// builder.append_component_from_bulk_producer(2, &mut clone_from_slice(b"hi")).await.unwrap();
247 /// builder.append_component_from_bulk_producer(2, &mut clone_from_slice(b"ho")).await.unwrap();
248 /// assert_eq!(builder.build(), Path::from_slices(&[b"hi", b"ho"])?);
249 /// # Ok::<(), PathError>(())
250 /// # });
251 /// ```
252 pub async fn append_component_from_bulk_producer<P>(
253 &mut self,
254 component_length: usize,
255 p: &mut P,
256 ) -> Result<(), DecodeError<P::Final, P::Error, Blame>>
257 where
258 P: BulkProducer<Item = u8> + ?Sized,
259 {
260 assert!(component_length <= MCL);
261
262 // Compute the accumulated length for the new component.
263 let total_length_so_far = match self.initialised_components.checked_sub(1) {
264 Some(i) => Representation::sum_of_lengths_for_component(&self.bytes, i),
265 None => 0,
266 };
267 let acc_length = component_length + total_length_so_far;
268
269 // Overwrite the dummy accumulated component length for this component with the actual value.
270 let start = Representation::start_offset_of_sum_of_lengths_for_component(
271 self.initialised_components,
272 );
273 let end = start + size_of::<usize>();
274 self.bytes.as_mut()[start..end].copy_from_slice(&acc_length.to_ne_bytes()[..]);
275
276 // Ufotofu prohibits empty slices, so it is easier to handle the case that would require them explicitly.
277 if component_length == 0 {
278 // Record that we added a component.
279 self.initialised_components += 1;
280 return Ok(());
281 }
282
283 // Now, read bytes until we have the full component.
284 let mut produced_so_far = 0;
285 while produced_so_far < component_length {
286 if let Right(fin) = p
287 .expose_items_sync(|items| {
288 let remaining_len =
289 core::cmp::min(items.len(), component_length - produced_so_far);
290 self.bytes.extend_from_slice(&items[..remaining_len]);
291 produced_so_far += remaining_len;
292 (remaining_len, ())
293 })
294 .await?
295 {
296 return Err(DecodeError::UnexpectedEndOfInput(fin));
297 }
298 }
299
300 // Record that we added a component.
301 self.initialised_components += 1;
302
303 Ok(())
304 }
305
306 /// Turns this builder into an immutable [`Path`].
307 ///
308 /// Panics if the number of [`Component`]s or the total length does not match what was claimed in [`PathBuilder::new`].
309 ///
310 /// #### Complexity
311 ///
312 /// Runs in `O(1)` time. Performs no allocations.
313 ///
314 /// #### Examples
315 ///
316 /// ```
317 /// use willow_data_model::prelude::*;
318 /// let mut builder = PathBuilder::<4, 4, 4>::new(4, 2)?;
319 /// builder.append_component(Component::new(b"hi")?);
320 /// builder.append_slice(b"ho")?;
321 /// assert_eq!(builder.build(), Path::from_slices(&[b"hi", b"ho"])?);
322 /// # Ok::<(), PathError>(())
323 /// ```
324 pub fn build(self) -> Path<MCL, MCC, MPL> {
325 // Check whether we appended the correct number of components.
326 assert_eq!(
327 self.initialised_components,
328 Representation::component_count(self.bytes.as_ref())
329 );
330
331 assert_eq!(
332 self.target_length,
333 self.total_length(),
334 "Expected a target length of {}, but got an actual total_length of {}\nRaw representation:{:?}",
335 self.target_length,
336 self.total_length(),
337 self.bytes
338 );
339
340 Path {
341 data: self.bytes.freeze(),
342 component_count: self.initialised_components,
343 first_component: 0,
344 }
345 }
346
347 /// Returns the total length of all components added to the builder so far.
348 ///
349 /// ```
350 /// use willow_data_model::prelude::*;
351 ///
352 /// let mut builder = PathBuilder::<4, 4, 4>::new(4, 2)?;
353 /// assert_eq!(builder.total_length(), 0);
354 ///
355 /// builder.append_component(Component::new(b"hi")?);
356 /// assert_eq!(builder.total_length(), 2);
357 ///
358 /// builder.append_slice(b"ho")?;
359 /// assert_eq!(builder.total_length(), 4);
360 /// # Ok::<(), PathError>(())
361 /// ```
362 pub fn total_length(&self) -> usize {
363 Representation::total_length(&self.bytes, self.initialised_components)
364 }
365
366 /// Returns the number of components added to the builder so far.
367 ///
368 /// ```
369 /// use willow_data_model::prelude::*;
370 ///
371 /// let mut builder = PathBuilder::<4, 4, 4>::new(4, 2)?;
372 /// assert_eq!(builder.component_count(), 0);
373 ///
374 /// builder.append_component(Component::new(b"hi")?);
375 /// assert_eq!(builder.component_count(), 1);
376 ///
377 /// builder.append_slice(b"ho")?;
378 /// assert_eq!(builder.component_count(), 2);
379 /// # Ok::<(), PathError>(())
380 /// ```
381 pub fn component_count(&self) -> usize {
382 self.initialised_components
383 }
384}