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 Some(index_of_final_component) => {
109 // Overwrite the dummy accumulated component lengths for the first `prefix_component_count` components with the actual values.
110 let accumulated_component_lengths_start =
111 Representation::start_offset_of_sum_of_lengths_for_component(0);
112 let accumulated_component_lengths_end =
113 Representation::start_offset_of_sum_of_lengths_for_component(
114 index_of_final_component + 1,
115 );
116 builder.bytes.as_mut()
117 [accumulated_component_lengths_start..accumulated_component_lengths_end]
118 .copy_from_slice(
119 &reference.data[accumulated_component_lengths_start
120 ..accumulated_component_lengths_end],
121 );
122
123 let components_start_in_prefix =
124 Representation::start_offset_of_component(&reference.data, 0);
125 let components_end_in_prefix = Representation::end_offset_of_component(
126 &reference.data,
127 index_of_final_component,
128 );
129 builder.bytes.extend_from_slice(
130 &reference.data[components_start_in_prefix..components_end_in_prefix],
131 );
132
133 // Record that we added the components.
134 builder.initialised_components = prefix_component_count;
135 }
136 }
137
138 // For reference, here is a naive implementation of the same functionality:
139
140 // for comp in reference
141 // .create_prefix(prefix_component_count)
142 // .expect("`prefix_component_count` must be less than the number of components in `reference`")
143 // .components()
144 // {
145 // builder.append_component(comp);
146 // }
147
148 Ok(builder)
149 }
150
151 /// Appends the data for the next [`Component`].
152 ///
153 /// #### Complexity
154 ///
155 /// Runs in `O(component_length)` time. Performs no allocations.
156 ///
157 /// #### Examples
158 ///
159 /// ```
160 /// use willow_data_model::prelude::*;
161 /// let mut builder = PathBuilder::<4, 4, 4>::new(4, 2)?;
162 /// builder.append_component(Component::new(b"hi")?);
163 /// builder.append_component(Component::new(b"ho")?);
164 /// assert_eq!(builder.build(), Path::from_slices(&[b"hi", b"ho"])?);
165 /// # Ok::<(), PathError>(())
166 /// ```
167 pub fn append_component(&mut self, component: &Component<MCL>) {
168 // Compute the accumulated length for the new component.
169 let total_length_so_far = match self.initialised_components.checked_sub(1) {
170 Some(i) => Representation::sum_of_lengths_for_component(&self.bytes, i),
171 None => 0,
172 };
173 let acc_length = component.as_ref().len() + total_length_so_far;
174
175 // Overwrite the dummy accumulated component length for this component with the actual value.
176 let start = Representation::start_offset_of_sum_of_lengths_for_component(
177 self.initialised_components,
178 );
179 let end = start + size_of::<usize>();
180 self.bytes.as_mut()[start..end].copy_from_slice(&acc_length.to_ne_bytes()[..]);
181
182 // Append the component to the path.
183 self.bytes.extend_from_slice(component.as_ref());
184
185 // Record that we added a component.
186 self.initialised_components += 1;
187 }
188
189 /// Appends the data for the next [`Component`], from a slice of bytes.
190 ///
191 /// #### Complexity
192 ///
193 /// Runs in `O(component_length)` time. Performs no allocations.
194 ///
195 /// #### Examples
196 ///
197 /// ```
198 /// use willow_data_model::prelude::*;
199 /// let mut builder = PathBuilder::<4, 4, 4>::new(4, 2)?;
200 /// builder.append_slice(b"hi")?;
201 /// builder.append_slice(b"ho")?;
202 /// assert_eq!(builder.build(), Path::from_slices(&[b"hi", b"ho"])?);
203 /// # Ok::<(), PathError>(())
204 /// ```
205 pub fn append_slice(&mut self, component: &[u8]) -> Result<(), InvalidComponentError> {
206 self.append_component(Component::new(component)?);
207 Ok(())
208 }
209
210 /// Appends data for a component of known length by reading data from a [`BulkProducer`] of bytes. Panics if `component_length > MCL`.
211 ///
212 /// #### Complexity
213 ///
214 /// Runs in `O(component_length)` time (assuming the producer takes `O(1)` time per byte). Performs no allocations.
215 ///
216 /// #### Examples
217 ///
218 /// ```
219 /// use willow_data_model::prelude::*;
220 /// use ufotofu::{codec_prelude::*, producer::clone_from_slice};
221 ///
222 /// # pollster::block_on(async {
223 /// let mut builder = PathBuilder::<4, 4, 4>::new(4, 2)?;
224 /// builder.append_component_from_bulk_producer(2, &mut clone_from_slice(b"hi")).await.unwrap();
225 /// builder.append_component_from_bulk_producer(2, &mut clone_from_slice(b"ho")).await.unwrap();
226 /// assert_eq!(builder.build(), Path::from_slices(&[b"hi", b"ho"])?);
227 /// # Ok::<(), PathError>(())
228 /// # });
229 /// ```
230 pub async fn append_component_from_bulk_producer<P>(
231 &mut self,
232 component_length: usize,
233 p: &mut P,
234 ) -> Result<(), DecodeError<P::Final, P::Error, Blame>>
235 where
236 P: BulkProducer<Item = u8> + ?Sized,
237 {
238 assert!(component_length <= MCL);
239
240 // Compute the accumulated length for the new component.
241 let total_length_so_far = match self.initialised_components.checked_sub(1) {
242 Some(i) => Representation::sum_of_lengths_for_component(&self.bytes, i),
243 None => 0,
244 };
245 let acc_length = component_length + total_length_so_far;
246
247 // Overwrite the dummy accumulated component length for this component with the actual value.
248 let start = Representation::start_offset_of_sum_of_lengths_for_component(
249 self.initialised_components,
250 );
251 let end = start + size_of::<usize>();
252 self.bytes.as_mut()[start..end].copy_from_slice(&acc_length.to_ne_bytes()[..]);
253
254 // Ufotofu prohibits empty slices, so it is easier to handle the case that would require them explicitly.
255 if component_length == 0 {
256 // Record that we added a component.
257 self.initialised_components += 1;
258 return Ok(());
259 }
260
261 // Now, read bytes until we have the full component.
262 let mut produced_so_far = 0;
263 while produced_so_far < component_length {
264 if let Right(fin) = p
265 .expose_items_sync(|items| {
266 let remaining_len =
267 core::cmp::min(items.len(), component_length - produced_so_far);
268 self.bytes.extend_from_slice(&items[..remaining_len]);
269 produced_so_far += remaining_len;
270 (remaining_len, ())
271 })
272 .await?
273 {
274 return Err(DecodeError::UnexpectedEndOfInput(fin));
275 }
276 }
277
278 // Record that we added a component.
279 self.initialised_components += 1;
280
281 Ok(())
282 }
283
284 /// Turns this builder into an immutable [`Path`].
285 ///
286 /// Panics if the number of [`Component`]s or the total length does not match what was claimed in [`PathBuilder::new`].
287 ///
288 /// #### Complexity
289 ///
290 /// Runs in `O(1)` time. Performs no allocations.
291 ///
292 /// #### Examples
293 ///
294 /// ```
295 /// use willow_data_model::prelude::*;
296 /// let mut builder = PathBuilder::<4, 4, 4>::new(4, 2)?;
297 /// builder.append_component(Component::new(b"hi")?);
298 /// builder.append_slice(b"ho")?;
299 /// assert_eq!(builder.build(), Path::from_slices(&[b"hi", b"ho"])?);
300 /// # Ok::<(), PathError>(())
301 /// ```
302 pub fn build(self) -> Path<MCL, MCC, MPL> {
303 // Check whether we appended the correct number of components.
304 assert_eq!(
305 self.initialised_components,
306 Representation::component_count(self.bytes.as_ref())
307 );
308
309 assert_eq!(
310 self.target_length,
311 self.total_length(),
312 "Expected a target length of {}, but got an actual total_length of {}\nRaw representation:{:?}",
313 self.target_length,
314 self.total_length(),
315 self.bytes
316 );
317
318 Path {
319 data: self.bytes.freeze(),
320 component_count: self.initialised_components,
321 }
322 }
323
324 /// Returns the total length of all components added to the builder so far.
325 ///
326 /// ```
327 /// use willow_data_model::prelude::*;
328 ///
329 /// let mut builder = PathBuilder::<4, 4, 4>::new(4, 2)?;
330 /// assert_eq!(builder.total_length(), 0);
331 ///
332 /// builder.append_component(Component::new(b"hi")?);
333 /// assert_eq!(builder.total_length(), 2);
334 ///
335 /// builder.append_slice(b"ho")?;
336 /// assert_eq!(builder.total_length(), 4);
337 /// # Ok::<(), PathError>(())
338 /// ```
339 pub fn total_length(&self) -> usize {
340 Representation::total_length(&self.bytes, self.initialised_components)
341 }
342
343 /// Returns the number of components added to the builder so far.
344 ///
345 /// ```
346 /// use willow_data_model::prelude::*;
347 ///
348 /// let mut builder = PathBuilder::<4, 4, 4>::new(4, 2)?;
349 /// assert_eq!(builder.component_count(), 0);
350 ///
351 /// builder.append_component(Component::new(b"hi")?);
352 /// assert_eq!(builder.component_count(), 1);
353 ///
354 /// builder.append_slice(b"ho")?;
355 /// assert_eq!(builder.component_count(), 2);
356 /// # Ok::<(), PathError>(())
357 /// ```
358 pub fn component_count(&self) -> usize {
359 self.initialised_components
360 }
361}