1//! A small module giving you a simple container that allows easy and cheap
2//! replacement of parts of its content, with the ability to prevent changing
3//! the same parts multiple times.
4
5use std::rc::Rc;
6
7use crate::error::Error;
8
9/// Indicates the change state of a [`Span`].
10#[derive(Debug, Clone, PartialEq, Eq)]
11enum State {
12 /// The initial state. No change applied.
13 Initial,
14 /// Has been replaced.
15 Replaced(Rc<[u8]>),
16 /// Has been inserted.
17 Inserted(Rc<[u8]>),
18}
19
20impl State {
21 fn is_inserted(&self) -> bool {
22 matches!(*self, State::Inserted(..))
23 }
24}
25
26/// Span with a change [`State`].
27#[derive(Debug, Clone, PartialEq, Eq)]
28struct Span {
29 /// Start of this span in parent data
30 start: usize,
31 /// up to end excluding
32 end: usize,
33 /// Whether the span is inserted, replaced or still fresh.
34 data: State,
35}
36
37/// A container that allows easily replacing chunks of its data
38#[derive(Debug, Clone, Default)]
39pub struct Data {
40 /// Original data.
41 original: Vec<u8>,
42 /// [`Span`]s covering the full range of the original data.
43 parts: Vec<Span>,
44}
45
46impl Data {
47 /// Create a new data container from a slice of bytes
48 pub fn new(data: &[u8]) -> Self {
49 Data {
50 original: data.into(),
51 parts: vec![Span {
52 data: State::Initial,
53 start: 0,
54 end: data.len(),
55 }],
56 }
57 }
58
59 /// Render this data as a vector of bytes
60 pub fn to_vec(&self) -> Vec<u8> {
61 if self.original.is_empty() {
62 return Vec::new();
63 }
64
65 self.parts.iter().fold(Vec::new(), |mut acc, d| {
66 match d.data {
67 State::Initial => acc.extend_from_slice(&self.original[d.start..d.end]),
68 State::Replaced(ref d) | State::Inserted(ref d) => acc.extend_from_slice(d),
69 };
70 acc
71 })
72 }
73
74 /// Replace a chunk of data with the given slice, erroring when this part
75 /// was already changed previously.
76 pub fn replace_range(
77 &mut self,
78 range: std::ops::Range<usize>,
79 data: &[u8],
80 ) -> Result<(), Error> {
81 if range.start > range.end {
82 return Err(Error::InvalidRange(range));
83 }
84
85 if range.end > self.original.len() {
86 return Err(Error::DataLengthExceeded(range, self.original.len()));
87 }
88
89 let insert_only = range.start == range.end;
90
91 // Since we error out when replacing an already replaced chunk of data,
92 // we can take some shortcuts here. For example, there can be no
93 // overlapping replacements -- we _always_ split a chunk of 'initial'
94 // data into three[^empty] parts, and there can't ever be two 'initial'
95 // parts touching.
96 //
97 // [^empty]: Leading and trailing ones might be empty if we replace
98 // the whole chunk. As an optimization and without loss of generality we
99 // don't add empty parts.
100 let new_parts = {
101 let Some(index_of_part_to_split) = self.parts.iter().position(|p| {
102 !p.data.is_inserted() && p.start <= range.start && p.end >= range.end
103 }) else {
104 if tracing::enabled!(tracing::Level::DEBUG) {
105 let slices = self
106 .parts
107 .iter()
108 .map(|p| {
109 (
110 p.start,
111 p.end,
112 match p.data {
113 State::Initial => "initial",
114 State::Replaced(..) => "replaced",
115 State::Inserted(..) => "inserted",
116 },
117 )
118 })
119 .collect::<Vec<_>>();
120 tracing::debug!(
121 "no single slice covering {}..{}, current slices: {:?}",
122 range.start,
123 range.end,
124 slices,
125 );
126 }
127
128 return Err(Error::MaybeAlreadyReplaced(range));
129 };
130
131 let part_to_split = &self.parts[index_of_part_to_split];
132
133 // If this replacement matches exactly the part that we would
134 // otherwise split then we ignore this for now. This means that you
135 // can replace the exact same range with the exact same content
136 // multiple times and we'll process and allow it.
137 //
138 // This is currently done to alleviate issues like
139 // rust-lang/rust#51211 although this clause likely wants to be
140 // removed if that's fixed deeper in the compiler.
141 if part_to_split.start == range.start && part_to_split.end == range.end {
142 if let State::Replaced(ref replacement) = part_to_split.data {
143 if &**replacement == data {
144 return Ok(());
145 }
146 }
147 }
148
149 if part_to_split.data != State::Initial {
150 return Err(Error::AlreadyReplaced);
151 }
152
153 let mut new_parts = Vec::with_capacity(self.parts.len() + 2);
154
155 // Previous parts
156 if let Some(ps) = self.parts.get(..index_of_part_to_split) {
157 new_parts.extend_from_slice(ps);
158 }
159
160 // Keep initial data on left side of part
161 if range.start > part_to_split.start {
162 new_parts.push(Span {
163 start: part_to_split.start,
164 end: range.start,
165 data: State::Initial,
166 });
167 }
168
169 // New part
170 new_parts.push(Span {
171 start: range.start,
172 end: range.end,
173 data: if insert_only {
174 State::Inserted(data.into())
175 } else {
176 State::Replaced(data.into())
177 },
178 });
179
180 // Keep initial data on right side of part
181 if range.end < part_to_split.end {
182 new_parts.push(Span {
183 start: range.end,
184 end: part_to_split.end,
185 data: State::Initial,
186 });
187 }
188
189 // Following parts
190 if let Some(ps) = self.parts.get(index_of_part_to_split + 1..) {
191 new_parts.extend_from_slice(ps);
192 }
193
194 new_parts
195 };
196
197 self.parts = new_parts;
198
199 Ok(())
200 }
201}
202
203#[cfg(test)]
204mod tests {
205 use super::*;
206 use proptest::prelude::*;
207
208 fn str(i: &[u8]) -> &str {
209 ::std::str::from_utf8(i).unwrap()
210 }
211
212 #[test]
213 fn insert_at_beginning() {
214 let mut d = Data::new(b"foo bar baz");
215 d.replace_range(0..0, b"oh no ").unwrap();
216 assert_eq!("oh no foo bar baz", str(&d.to_vec()));
217 }
218
219 #[test]
220 fn insert_at_end() {
221 let mut d = Data::new(b"foo bar baz");
222 d.replace_range(11..11, b" oh no").unwrap();
223 assert_eq!("foo bar baz oh no", str(&d.to_vec()));
224 }
225
226 #[test]
227 fn replace_some_stuff() {
228 let mut d = Data::new(b"foo bar baz");
229 d.replace_range(4..7, b"lol").unwrap();
230 assert_eq!("foo lol baz", str(&d.to_vec()));
231 }
232
233 #[test]
234 fn replace_a_single_char() {
235 let mut d = Data::new(b"let y = true;");
236 d.replace_range(4..5, b"mut y").unwrap();
237 assert_eq!("let mut y = true;", str(&d.to_vec()));
238 }
239
240 #[test]
241 fn replace_multiple_lines() {
242 let mut d = Data::new(b"lorem\nipsum\ndolor");
243
244 d.replace_range(6..11, b"lol").unwrap();
245 assert_eq!("lorem\nlol\ndolor", str(&d.to_vec()));
246
247 d.replace_range(12..17, b"lol").unwrap();
248 assert_eq!("lorem\nlol\nlol", str(&d.to_vec()));
249 }
250
251 #[test]
252 fn replace_multiple_lines_with_insert_only() {
253 let mut d = Data::new(b"foo!");
254
255 d.replace_range(3..3, b"bar").unwrap();
256 assert_eq!("foobar!", str(&d.to_vec()));
257
258 d.replace_range(0..3, b"baz").unwrap();
259 assert_eq!("bazbar!", str(&d.to_vec()));
260
261 d.replace_range(3..4, b"?").unwrap();
262 assert_eq!("bazbar?", str(&d.to_vec()));
263 }
264
265 #[test]
266 fn replace_invalid_range() {
267 let mut d = Data::new(b"foo!");
268
269 assert!(d.replace_range(2..1, b"bar").is_err());
270 assert!(d.replace_range(0..3, b"bar").is_ok());
271 }
272
273 #[test]
274 fn empty_to_vec_roundtrip() {
275 let s = "";
276 assert_eq!(s.as_bytes(), Data::new(s.as_bytes()).to_vec().as_slice());
277 }
278
279 #[test]
280 fn replace_overlapping_stuff_errs() {
281 let mut d = Data::new(b"foo bar baz");
282
283 d.replace_range(4..7, b"lol").unwrap();
284 assert_eq!("foo lol baz", str(&d.to_vec()));
285
286 assert!(matches!(
287 d.replace_range(4..7, b"lol2").unwrap_err(),
288 Error::AlreadyReplaced,
289 ));
290 }
291
292 #[test]
293 fn broken_replacements() {
294 let mut d = Data::new(b"foo");
295 assert!(matches!(
296 d.replace_range(4..8, b"lol").unwrap_err(),
297 Error::DataLengthExceeded(std::ops::Range { start: 4, end: 8 }, 3),
298 ));
299 }
300
301 #[test]
302 fn replace_same_twice() {
303 let mut d = Data::new(b"foo");
304 d.replace_range(0..1, b"b").unwrap();
305 d.replace_range(0..1, b"b").unwrap();
306 assert_eq!("boo", str(&d.to_vec()));
307 }
308
309 proptest! {
310 #[test]
311 fn new_to_vec_roundtrip(ref s in "\\PC*") {
312 assert_eq!(s.as_bytes(), Data::new(s.as_bytes()).to_vec().as_slice());
313 }
314
315 #[test]
316 fn replace_random_chunks(
317 ref data in "\\PC*",
318 ref replacements in prop::collection::vec(
319 (any::<::std::ops::Range<usize>>(), any::<Vec<u8>>()),
320 1..100,
321 )
322 ) {
323 let mut d = Data::new(data.as_bytes());
324 for &(ref range, ref bytes) in replacements {
325 let _ = d.replace_range(range.clone(), bytes);
326 }
327 }
328 }
329}
330