mixnet/core/
scattered.rs

1// Copyright 2022 Parity Technologies (UK) Ltd.
2//
3// Permission is hereby granted, free of charge, to any person obtaining a
4// copy of this software and associated documentation files (the "Software"),
5// to deal in the Software without restriction, including without limitation
6// the rights to use, copy, modify, merge, publish, distribute, sublicense,
7// and/or sell copies of the Software, and to permit persons to whom the
8// Software is furnished to do so, subject to the following conditions:
9//
10// The above copyright notice and this permission notice shall be included in
11// all copies or substantial portions of the Software.
12//
13// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
14// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
15// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
16// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
17// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
18// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
19// DEALINGS IN THE SOFTWARE.
20
21/// A concatenation of multiple slices. The slices are not copied until
22/// [`copy_to_slice`](Self::copy_to_slice) or [`to_vec`](Self::to_vec) is called.
23#[derive(Clone, Copy)]
24pub struct Scattered<'a, T> {
25	len: usize,
26	first_slice: &'a [T],
27	mid_slices: &'a [&'a [T]],
28	last_slice: &'a [T],
29}
30
31impl<'a, T> Scattered<'a, T> {
32	/// Returns the total number of elements.
33	pub fn len(&self) -> usize {
34		self.len
35	}
36
37	/// Returns `true` if there are no elements.
38	pub fn is_empty(&self) -> bool {
39		self.len == 0
40	}
41
42	/// Just like [`slice::split_at`].
43	pub fn split_at(&self, mid: usize) -> (Self, Self) {
44		let right_len = self.len.checked_sub(mid).expect("mid must be <= len");
45
46		// Split first_slice case
47		let Some(mut mid_in_remaining) = mid.checked_sub(self.first_slice.len()) else {
48			let (first_slice_left, first_slice_right) = self.first_slice.split_at(mid);
49			return (
50				Self { len: mid, first_slice: first_slice_left, mid_slices: &[], last_slice: &[] },
51				Self {
52					len: right_len,
53					first_slice: first_slice_right,
54					mid_slices: self.mid_slices,
55					last_slice: self.last_slice,
56				},
57			)
58		};
59
60		// Split mid_slices case
61		for (i, mid_slice) in self.mid_slices.iter().enumerate() {
62			mid_in_remaining = match mid_in_remaining.checked_sub(mid_slice.len()) {
63				Some(mid_in_remaining) => mid_in_remaining,
64				None => {
65					let (mid_slices_left, mid_slices_right) = self.mid_slices.split_at(i);
66					let mid_slices_right =
67						mid_slices_right.split_first().expect("i < self.mid_slices.len()").1;
68					let (mid_slice_left, mid_slice_right) = mid_slice.split_at(mid_in_remaining);
69					return (
70						Self {
71							len: mid,
72							first_slice: self.first_slice,
73							mid_slices: mid_slices_left,
74							last_slice: mid_slice_left,
75						},
76						Self {
77							len: right_len,
78							first_slice: mid_slice_right,
79							mid_slices: mid_slices_right,
80							last_slice: self.last_slice,
81						},
82					)
83				},
84			};
85		}
86
87		// Split last_slice case
88		let (last_slice_left, last_slice_right) = self.last_slice.split_at(mid_in_remaining);
89		(
90			Self {
91				len: mid,
92				first_slice: self.first_slice,
93				mid_slices: self.mid_slices,
94				last_slice: last_slice_left,
95			},
96			Self {
97				len: right_len,
98				first_slice: last_slice_right,
99				mid_slices: &[],
100				last_slice: &[],
101			},
102		)
103	}
104}
105
106impl<'a, T: Copy> Scattered<'a, T> {
107	/// Copy all elements into `dst`. `dst.len()` must equal `self.len()`.
108	pub fn copy_to_slice(&self, dst: &mut [T]) {
109		let (dst_first_slice, mut dst) = dst.split_at_mut(self.first_slice.len());
110		dst_first_slice.copy_from_slice(self.first_slice);
111		for mid_slice in self.mid_slices {
112			let (dst_mid_slice, remaining_dst) = dst.split_at_mut(mid_slice.len());
113			dst_mid_slice.copy_from_slice(mid_slice);
114			dst = remaining_dst;
115		}
116		dst.copy_from_slice(self.last_slice);
117	}
118}
119
120impl<'a, T: Clone> Scattered<'a, T> {
121	/// Copy all elements to a new [`Vec`].
122	pub fn to_vec(&self) -> Vec<T> {
123		let mut vec = Vec::with_capacity(self.len);
124		vec.extend_from_slice(self.first_slice);
125		for mid_slice in self.mid_slices {
126			vec.extend_from_slice(mid_slice);
127		}
128		vec.extend_from_slice(self.last_slice);
129		vec
130	}
131}
132
133impl<'a, T> From<&'a [T]> for Scattered<'a, T> {
134	fn from(slice: &'a [T]) -> Self {
135		Self { len: slice.len(), first_slice: slice, mid_slices: &[], last_slice: &[] }
136	}
137}
138
139impl<'a, T> From<&'a [&'a [T]]> for Scattered<'a, T> {
140	fn from(slices: &'a [&'a [T]]) -> Self {
141		Self {
142			len: slices.iter().map(|slice| slice.len()).sum(),
143			first_slice: &[],
144			mid_slices: slices,
145			last_slice: &[],
146		}
147	}
148}
149
150#[cfg(test)]
151mod tests {
152	use super::*;
153	use rand::RngCore;
154
155	fn to_vec_using_copy_to_slice(scattered: &Scattered<u8>) -> Vec<u8> {
156		let mut vec = vec![0; scattered.len()];
157		scattered.copy_to_slice(&mut vec);
158		vec
159	}
160
161	fn test_splits(slice_lens: &[usize], mids: &[usize]) {
162		let mut contig = vec![0; slice_lens.iter().sum()];
163		rand::thread_rng().fill_bytes(&mut contig);
164		let mut contig = contig.as_slice();
165
166		let slices: Vec<_> = {
167			let mut remaining = contig;
168			slice_lens
169				.iter()
170				.map(|slice_len| {
171					let (left, right) = remaining.split_at(*slice_len);
172					remaining = right;
173					left
174				})
175				.collect()
176		};
177		let mut scattered: Scattered<u8> = slices.as_slice().into();
178
179		for mid in mids {
180			let (contig_left, contig_right) = contig.split_at(*mid);
181			let (scattered_left, scattered_right) = scattered.split_at(*mid);
182			assert_eq!(contig_left, scattered_left.to_vec());
183			assert_eq!(contig_right, scattered_right.to_vec());
184			assert_eq!(contig_left, to_vec_using_copy_to_slice(&scattered_left));
185			assert_eq!(contig_right, to_vec_using_copy_to_slice(&scattered_right));
186			contig = contig_right;
187			scattered = scattered_right;
188		}
189	}
190
191	#[test]
192	fn single_slice() {
193		test_splits(&[20], &[0, 9, 5, 6]);
194	}
195
196	#[test]
197	fn multiple_slices() {
198		test_splits(&[5, 7, 10, 7, 5], &[3, 2, 3, 4, 6, 4, 3, 4, 4, 1]);
199		test_splits(&[5, 7, 10, 7, 5], &[6, 9, 16, 3]);
200		test_splits(&[5, 7, 10, 7, 5], &[33, 1]);
201		test_splits(&[5, 7, 10, 7, 5], &[34]);
202	}
203}