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