1use std::{borrow::Borrow, mem, ops::Deref, sync::Arc, vec};
5
6use serde::{Deserialize, Deserializer, Serialize, Serializer};
7
8#[derive(Debug, PartialOrd, PartialEq, Ord, Eq, Hash)]
9pub struct CowVec<T>
10where
11 T: Clone + PartialEq,
12{
13 inner: Arc<Vec<T>>,
14}
15
16impl<T> CowVec<T>
17where
18 T: Clone + PartialEq,
19{
20 pub fn with_capacity(capacity: usize) -> Self {
21 let aligned_capacity = (capacity + 7) & !7;
22 Self {
23 inner: Arc::new(Vec::with_capacity(aligned_capacity)),
24 }
25 }
26
27 pub fn with_aligned_capacity(capacity: usize) -> Self {
28 let simd_alignment = 32 / mem::size_of::<T>().max(1);
29 let aligned_capacity = capacity.div_ceil(simd_alignment) * simd_alignment;
30 Self {
31 inner: Arc::new(Vec::with_capacity(aligned_capacity)),
32 }
33 }
34
35 pub fn len(&self) -> usize {
36 self.inner.len()
37 }
38
39 pub fn is_empty(&self) -> bool {
40 self.len() == 0
41 }
42
43 pub fn capacity(&self) -> usize {
44 self.inner.capacity()
45 }
46}
47
48#[macro_export]
49macro_rules! cow_vec {
50 () => {
51 $crate::util::cowvec::CowVec::new(Vec::new())
52 };
53 ($($elem:expr),+ $(,)?) => {
54 $crate::util::cowvec::CowVec::new(vec![$($elem),+])
55 };
56}
57
58impl<T> Default for CowVec<T>
59where
60 T: Clone + PartialEq,
61{
62 fn default() -> Self {
63 Self {
64 inner: Arc::new(Vec::new()),
65 }
66 }
67}
68
69impl<T: Clone + PartialEq> PartialEq<[T]> for &CowVec<T> {
70 fn eq(&self, other: &[T]) -> bool {
71 self.inner.as_slice() == other
72 }
73}
74
75impl<T: Clone + PartialEq> PartialEq<[T]> for CowVec<T> {
76 fn eq(&self, other: &[T]) -> bool {
77 self.inner.as_slice() == other
78 }
79}
80
81impl<T: Clone + PartialEq> PartialEq<CowVec<T>> for [T] {
82 fn eq(&self, other: &CowVec<T>) -> bool {
83 self == other.inner.as_slice()
84 }
85}
86
87impl<T: Clone + PartialEq> Clone for CowVec<T> {
88 fn clone(&self) -> Self {
89 CowVec {
90 inner: Arc::clone(&self.inner),
91 }
92 }
93}
94
95impl<T: Clone + PartialEq> CowVec<T> {
96 pub fn new(vec: Vec<T>) -> Self {
97 CowVec {
98 inner: Arc::new(vec),
99 }
100 }
101
102 pub fn from_rc(rc: Arc<Vec<T>>) -> Self {
103 CowVec {
104 inner: rc,
105 }
106 }
107
108 pub fn try_into_vec(self) -> Result<Vec<T>, Self> {
109 match Arc::try_unwrap(self.inner) {
110 Ok(vec) => Ok(vec),
111 Err(arc) => Err(CowVec {
112 inner: arc,
113 }),
114 }
115 }
116
117 pub fn into_inner(self) -> Vec<T> {
118 match Arc::try_unwrap(self.inner) {
119 Ok(vec) => vec,
120 Err(arc) => (*arc).clone(),
121 }
122 }
123
124 pub fn as_slice(&self) -> &[T] {
125 &self.inner
126 }
127
128 pub fn is_owned(&self) -> bool {
129 Arc::strong_count(&self.inner) == 1
130 }
131
132 pub fn is_shared(&self) -> bool {
133 Arc::strong_count(&self.inner) > 1
134 }
135
136 pub fn get(&self, idx: usize) -> Option<&T> {
137 self.inner.get(idx)
138 }
139
140 pub fn make_mut(&mut self) -> &mut Vec<T> {
141 Arc::make_mut(&mut self.inner)
142 }
143
144 pub fn set(&mut self, idx: usize, value: T) {
145 self.make_mut()[idx] = value;
146 }
147
148 pub fn push(&mut self, value: T) {
149 self.make_mut().push(value);
150 }
151
152 pub fn clear(&mut self) {
153 self.make_mut().clear();
154 }
155
156 pub fn extend(&mut self, iter: impl IntoIterator<Item = T>) {
157 self.make_mut().extend(iter);
158 }
159
160 pub fn extend_from_slice(&mut self, slice: &[T]) {
161 self.make_mut().extend_from_slice(slice);
162 }
163
164 pub fn reorder(&mut self, indices: &[usize]) {
165 let vec = self.make_mut();
166 let len = vec.len();
167 assert_eq!(len, indices.len());
168
169 let mut visited = vec![false; len];
170 for start in 0..len {
171 if visited[start] || indices[start] == start {
172 continue;
173 }
174 let mut current = start;
175 while !visited[current] {
176 visited[current] = true;
177 let next = indices[current];
178 if next == start {
179 break;
180 }
181 vec.swap(current, next);
182 current = next;
183 }
184 }
185 }
186
187 pub fn aligned_chunks(&self, chunk_size: usize) -> impl Iterator<Item = &[T]> {
188 self.inner.chunks(chunk_size)
189 }
190
191 pub fn aligned_chunks_mut(&mut self, chunk_size: usize) -> impl Iterator<Item = &mut [T]> {
192 self.make_mut().chunks_mut(chunk_size)
193 }
194
195 pub fn is_simd_aligned(&self) -> bool {
196 let alignment = 32;
197 let ptr = self.inner.as_ptr() as usize;
198 ptr.is_multiple_of(alignment)
199 }
200
201 pub fn take(&self, n: usize) -> Self {
202 let len = n.min(self.len());
203 CowVec::new(self.inner[..len].to_vec())
204 }
205}
206
207impl<T: Clone + PartialEq> IntoIterator for CowVec<T> {
208 type Item = T;
209 type IntoIter = vec::IntoIter<T>;
210
211 fn into_iter(self) -> Self::IntoIter {
212 match Arc::try_unwrap(self.inner) {
213 Ok(vec) => vec.into_iter(),
214 Err(arc) => (*arc).clone().into_iter(),
215 }
216 }
217}
218
219impl<T: Clone + PartialEq> Deref for CowVec<T> {
220 type Target = [T];
221
222 fn deref(&self) -> &Self::Target {
223 self.as_slice()
224 }
225}
226
227impl<T: Clone + PartialEq> Borrow<[T]> for CowVec<T> {
228 fn borrow(&self) -> &[T] {
229 self.as_slice()
230 }
231}
232
233impl<T> Serialize for CowVec<T>
234where
235 T: Clone + PartialEq + Serialize,
236{
237 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
238 where
239 S: Serializer,
240 {
241 self.inner.serialize(serializer)
242 }
243}
244
245impl<'de, T> Deserialize<'de> for CowVec<T>
246where
247 T: Clone + PartialEq + Deserialize<'de>,
248{
249 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
250 where
251 D: Deserializer<'de>,
252 {
253 let vec = Vec::<T>::deserialize(deserializer)?;
254 Ok(CowVec {
255 inner: Arc::new(vec),
256 })
257 }
258}
259
260#[cfg(test)]
261pub mod tests {
262 use super::CowVec;
263
264 #[test]
265 fn test_new() {
266 let cow = CowVec::new(vec![1, 2, 3]);
267 assert_eq!(cow.get(0), Some(&1));
268 assert_eq!(cow.get(1), Some(&2));
269 assert_eq!(cow.get(2), Some(&3));
270 }
271
272 #[test]
273 fn test_is_owned() {
274 let mut owned = CowVec::new(Vec::with_capacity(16));
275 owned.extend([1, 2]);
276
277 assert!(owned.is_owned());
278
279 let shared = owned.clone();
280 assert!(!owned.is_owned());
281 assert!(!shared.is_owned());
282
283 drop(shared);
284
285 assert!(owned.is_owned());
286 }
287
288 #[test]
289 fn test_is_shared() {
290 let mut owned = CowVec::new(Vec::with_capacity(16));
291 owned.extend([1, 2]);
292
293 assert!(!owned.is_shared());
294
295 let shared = owned.clone();
296 assert!(owned.is_shared());
297 assert!(shared.is_shared());
298
299 drop(shared);
300
301 assert!(!owned.is_shared());
302 }
303
304 #[test]
305 fn test_extend() {
306 let mut owned = CowVec::new(Vec::with_capacity(16));
307 owned.extend([1, 2]);
308
309 let ptr_before_owned = ptr_of(&owned);
310 owned.extend([9, 9, 24]);
311 assert_eq!(ptr_before_owned, ptr_of(&owned)); assert_eq!(owned.len(), 5);
313
314 let mut shared = owned.clone();
315
316 let ptr_before_shared = ptr_of(&shared);
317 shared.extend([9, 9, 24]);
318 assert_ne!(ptr_before_shared, ptr_of(&shared)); assert_eq!(owned.len(), 5);
320 }
321
322 #[test]
323 fn test_push() {
324 let mut owned = CowVec::new(Vec::with_capacity(16));
325 owned.extend([1, 2]);
326
327 let ptr_before_owned = ptr_of(&owned);
328 owned.push(99);
329 assert_eq!(ptr_before_owned, ptr_of(&owned)); assert_eq!(owned.len(), 3);
331
332 let mut shared = owned.clone();
333
334 let ptr_before_shared = ptr_of(&shared);
335 shared.push(99);
336 assert_ne!(ptr_before_shared, ptr_of(&shared)); assert_eq!(owned.len(), 3);
338 }
339
340 #[test]
341 fn test_set() {
342 let mut owned = CowVec::new(Vec::with_capacity(16));
343 owned.extend([1, 2]);
344
345 let ptr_before_owned = ptr_of(&owned);
346 owned.set(1, 99);
347 assert_eq!(ptr_before_owned, ptr_of(&owned)); assert_eq!(*owned, [1, 99]);
349
350 let mut shared = owned.clone();
351
352 let ptr_before_shared = ptr_of(&shared);
353 shared.set(1, 99);
354 assert_ne!(ptr_before_shared, ptr_of(&shared)); assert_eq!(*owned, [1, 99]);
356 }
357
358 #[test]
359 fn test_reorder() {
360 let mut owned = CowVec::new(Vec::with_capacity(16));
361 owned.extend([1, 2]);
362
363 let ptr_before_owned = ptr_of(&owned);
364 owned.reorder(&[1usize, 0]);
365 assert_eq!(ptr_before_owned, ptr_of(&owned)); assert_eq!(*owned, [2, 1]);
367
368 let mut shared = owned.clone();
369
370 let ptr_before_shared = ptr_of(&shared);
371 shared.reorder(&[1usize, 0]);
372 assert_ne!(ptr_before_shared, ptr_of(&shared)); assert_eq!(*shared, [1, 2]);
374 }
375
376 #[test]
377 fn test_reorder_identity() {
378 let mut cow = CowVec::new(vec![10, 20, 30]);
379 cow.reorder(&[0, 1, 2]); assert_eq!(cow.as_slice(), &[10, 20, 30]);
381 }
382
383 #[test]
384 fn test_reorder_basic() {
385 let mut cow = CowVec::new(vec![10, 20, 30]);
386 cow.reorder(&[2, 0, 1]);
387 assert_eq!(cow.as_slice(), &[30, 10, 20]);
388 }
389
390 fn ptr_of(v: &CowVec<i32>) -> *const i32 {
391 v.as_slice().as_ptr()
392 }
393}