manual_share/shared_vec.rs
1//! # Manually shared vector
2//!
3//! `SharedVec` is the `Vec`-based counterpart to `SharedBox`.
4//! It owns the original allocation and lets you create multiple immutable `SharedVecRef`
5//! values that can be sent to other threads while keeping a single shared owner.
6//!
7//! The API is similar to `SharedBox`:
8//! - use `SharedVec::from_vec` to create a shared vector from a `Vec`
9//! - use `SharedVec::borrow` to create a `SharedVecRef`
10//! - use `SharedVec::try_return` to give a borrowed reference back
11//! - use `SharedVec::try_into_vec` to recover the original `Vec` once no references remain
12//!
13//! ```
14//! use std::thread;
15//! use manual_share::SharedVec;
16//!
17//! let values = vec![1, 2, 3];
18//! let mut shared = SharedVec::from_vec(values);
19//!
20//! let shared_ref = shared.borrow();
21//! let handle = thread::spawn(move || {
22//! assert_eq!(shared_ref.as_slice(), &[1, 2, 3]);
23//! shared_ref
24//! });
25//!
26//! let shared_ref = shared.borrow();
27//! let handle2 = thread::spawn(move || {
28//! assert_eq!(shared_ref.as_slice(), &[1, 2, 3]);
29//! shared_ref
30//! });
31//!
32//! shared.try_return(handle.join().unwrap()).unwrap();
33//! shared.try_return(handle2.join().unwrap()).unwrap();
34//!
35//! let values = shared.try_into_vec().unwrap();
36//! assert_eq!(values, vec![1, 2, 3]);
37//! ```
38//!
39
40/// A structure owning the original `Vec` which can be used to create multiple immutable
41/// `SharedVecRef` values to send to other threads.
42/// It uses a counter to record how many references have been created and not returned yet.
43///
44/// Dropping `SharedVec` without returning all `SharedVecRef` values leaks the underlying allocation.
45/// When the `panic-on-drop` feature is enabled, it will panic:
46/// ```should_panic
47/// let r = {
48/// let mut values = manual_share::SharedVec::from_vec(vec![0]);
49/// values.borrow()
50/// };
51/// println!("{:?}", r.as_slice());
52/// ```
53///
54/// Once all `SharedVecRef` values have been returned, `SharedVec` can be converted back into
55/// a `Vec` and its allocation will be released when dropped:
56/// ```
57/// let mut values = manual_share::SharedVec::from_vec(vec![0]);
58/// let reference = values.borrow();
59/// values.try_return(reference).unwrap();
60/// let values = values.try_into_vec().unwrap();
61/// assert_eq!(values, vec![0]);
62/// ```
63#[derive(Debug)]
64pub struct SharedVec<T> {
65 borrow_count: usize,
66 ptr: *mut T,
67 len: usize,
68 cap: usize,
69}
70impl<T> SharedVec<T> {
71 /// Create a `SharedVec` by consuming a `Vec`.
72 pub fn from_vec(vec: Vec<T>) -> Self {
73 let (ptr, len, cap) = vec.into_raw_parts();
74 Self {
75 borrow_count: 0,
76 ptr,
77 len,
78 cap,
79 }
80 }
81 /// Create a `SharedVecRef` and increase the borrow count.
82 ///
83 /// ```
84 /// let mut values = manual_share::SharedVec::from_vec(vec![1, 2, 3]);
85 /// let reference = values.borrow();
86 /// assert_eq!(reference.as_slice(), &[1, 2, 3]);
87 /// values.try_return(reference).unwrap();
88 /// ```
89 ///
90 /// # panics
91 /// Panics when borrow count overflows `usize`.
92 pub fn borrow(&mut self) -> SharedVecRef<T> {
93 self.borrow_count = self.borrow_count.checked_add(1).unwrap();
94 SharedVecRef {
95 ptr: self.ptr,
96 len: self.len,
97 }
98 }
99 /// Try to return back a `SharedVecRef`.
100 /// Returns `Err` if the `SharedVecRef` does not originate from the same `SharedVec`.
101 ///
102 /// ```
103 /// let mut first = manual_share::SharedVec::from_vec(vec![8]);
104 /// let first_ref = first.borrow();
105 /// first.try_return(first_ref).unwrap();
106 ///
107 /// let mut second = manual_share::SharedVec::from_vec(vec![9]);
108 /// let second_ref = second.borrow();
109 /// let err = first.try_return(second_ref).unwrap_err();
110 ///
111 /// assert_eq!(err.as_slice(), &[9]);
112 /// second.try_return(err).unwrap();
113 /// ```
114 pub fn try_return(&mut self, reference: SharedVecRef<T>) -> Result<(), SharedVecRef<T>> {
115 if !core::ptr::eq(self.ptr, reference.ptr) {
116 return Err(reference);
117 }
118
119 if size_of::<T>() == 0 {
120 if self.len != reference.len {
121 return Err(reference);
122 }
123
124 if let Some(new_count) = self.borrow_count.checked_sub(1) {
125 self.borrow_count = new_count;
126 let _ = core::mem::ManuallyDrop::new(reference);
127 Ok(())
128 } else {
129 Err(reference)
130 }
131 } else {
132 self.borrow_count -= 1;
133 let _ = core::mem::ManuallyDrop::new(reference);
134 Ok(())
135 }
136 }
137 /// Try to convert `Self` into a `Vec` once all borrowed references are returned.
138 ///
139 /// ```
140 /// let mut values = manual_share::SharedVec::from_vec(vec![0]);
141 /// let reference = values.borrow();
142 ///
143 /// // Try to convert to Vec without returning all SharedVecRef returns Err.
144 /// let mut values = values.try_into_vec().unwrap_err();
145 ///
146 /// values.try_return(reference).unwrap();
147 /// let values = values.try_into_vec().unwrap();
148 ///
149 /// assert_eq!(values, vec![0]);
150 /// ```
151 pub fn try_into_vec(self) -> Result<Vec<T>, Self> {
152 if self.borrow_count > 0 {
153 Err(self)
154 } else {
155 let r = core::mem::ManuallyDrop::new(self);
156 Ok(unsafe { Vec::from_raw_parts(r.ptr, r.len, r.cap) })
157 }
158 }
159 /// Directly get a slice to the values inside the `SharedVec`.
160 /// This use rust built-in lifetime check to ensure the slice is valid as long as the `SharedVec` is alive,
161 /// and has no runtime overhead.
162 pub fn get(&self) -> &[T] {
163 // SAFETY:
164 // The pointer is valid as long as the SharedVec is alive.
165 // All other references can only get immutable reference.
166 unsafe { core::slice::from_raw_parts(self.ptr, self.len) }
167 }
168}
169
170unsafe impl<T: Send> Send for SharedVec<T> {}
171unsafe impl<T: Sync> Sync for SharedVec<T> {}
172
173impl<T> Drop for SharedVec<T> {
174 fn drop(&mut self) {
175 #[cfg(feature = "panic-on-drop")]
176 {
177 // Let user deal with other panics.
178 #[cfg(feature = "do-not-panic-when-panicking")]
179 if std::thread::panicking() {
180 return;
181 }
182
183 if self.borrow_count > 0 {
184 panic!("Dropping a SharedVec without giving back all SharedVecRef")
185 }
186 }
187 // Only drops when there are no outstanding SharedBoxRef values to prevent use-after-free.
188 if self.borrow_count == 0 {
189 unsafe {
190 drop(Vec::from_raw_parts(self.ptr, self.len, self.cap));
191 }
192 }
193 }
194}
195
196/// A reference to `SharedVec` that can be sent to other threads.
197///
198/// Dropping a `SharedVecRef` leaks the heap allocation it points to.
199/// When the `panic-on-drop` feature is enabled, dropping it will panic:
200/// ```should_panic
201/// let mut values = manual_share::SharedVec::from_vec(vec![1]);
202/// values.borrow();
203///
204/// // forget SharedVecMut to make sure the panic is not caused by dropping it first.
205/// std::mem::forget(values);
206///
207/// // panic here due to dropping ShareVecRef
208/// ```
209///
210/// Use `SharedVec::try_return` to consume it without causing panic.
211/// ```
212/// let mut values = manual_share::SharedVec::from_vec(vec![1]);
213/// let reference = values.borrow();
214/// values.try_return(reference).unwrap();
215/// ```
216#[derive(Debug)]
217pub struct SharedVecRef<T> {
218 ptr: *const T,
219 len: usize,
220}
221
222impl<T> SharedVecRef<T> {
223 /// View the referenced data as a slice.
224 ///
225 /// ```
226 /// let mut values = manual_share::SharedVec::from_vec(vec![1, 2, 3]);
227 /// let reference = values.borrow();
228 /// assert_eq!(reference.as_slice(), &[1, 2, 3]);
229 /// values.try_return(reference).unwrap();
230 /// ```
231 pub fn as_slice(&self) -> &[T] {
232 unsafe { core::slice::from_raw_parts(self.ptr, self.len) }
233 }
234}
235
236impl<T> Drop for SharedVecRef<T> {
237 fn drop(&mut self) {
238 #[cfg(feature = "panic-on-drop")]
239 {
240 // Let user deal with other panics.
241 #[cfg(feature = "do-not-panic-when-panicking")]
242 if std::thread::panicking() {
243 return;
244 }
245
246 panic!("Dropping a SharedVecRef without returning it to the SharedVec")
247 }
248 }
249}
250
251unsafe impl<T: Sync + Send> Send for SharedVecRef<T> {}
252unsafe impl<T: Sync> Sync for SharedVecRef<T> {}
253
254/// A container of a `Vec` allocation that can be split into multiple `SharedVecPart` values.
255///
256/// This type is useful when a single `Vec` needs to be partitioned into multiple independently
257/// owned segments that still refer to the same underlying allocation.
258/// ```
259/// let mut values = manual_share::SharedVecMut::from_vec(vec![1, 2, 3, 4]);
260/// let mut part = values.split_off(2).unwrap();
261///
262/// let join_handle = std::thread::spawn(|| {
263/// part.as_slice_mut().iter_mut().for_each(|v| *v += 1);
264/// part
265/// });
266///
267/// let part = join_handle.join().unwrap();
268///
269/// assert_eq!(part.as_slice(), &[4, 5]);
270/// assert!(values.try_unsplit_off(part).is_ok());
271/// ```
272///
273/// Dropping the `SharedVecMut` without returning all `SharedVecPart` values leaks the underlying allocation.
274/// When the `panic-on-drop` feature is enabled, it will panic:
275/// ```should_panic
276/// let r = {
277/// let mut values = manual_share::SharedVecMut::from_vec(vec![0]);
278/// values.split_off(1).unwrap()
279///
280/// // panic here because the part was not returned to the original SharedVecMut.
281/// };
282/// println!("{:?}", r.as_slice());
283/// ```
284#[derive(Debug)]
285pub struct SharedVecMut<T> {
286 borrow_count: usize,
287
288 ptr: *mut T,
289 len: usize,
290 cap: usize,
291
292 remain_start: usize,
293 remain_len: usize,
294}
295
296impl<T> SharedVecMut<T> {
297 /// Create a `SharedVecMut` by consuming a `Vec`.
298 pub fn from_vec(vec: Vec<T>) -> Self {
299 let (ptr, len, cap) = vec.into_raw_parts();
300 Self {
301 borrow_count: 0,
302 ptr,
303 len,
304 cap,
305 remain_start: 0,
306 remain_len: len,
307 }
308 }
309 /// Split off the suffix of the vector starting at `at`.
310 /// This method is similar to `bytes::BytesMut::split_off`.
311 ///
312 /// Returns None when:
313 /// 1. `at` is greater than the length of the vector.
314 /// 2. `borrow_count` overflows `usize`.
315 ///
316 /// If successful, the returned part will contain [at, len) and self will contain [0, at).
317 ///
318 /// Here is an example of splitting a `SharedVecMut` into 3 parts:
319 /// ```
320 /// let mut values = manual_share::SharedVecMut::from_vec(vec![1, 2, 3]);
321 ///
322 /// let part1 = values.split_off(2).unwrap();
323 /// let part2 = values.split_off(1).unwrap();
324 /// let part3 = values.split_off(0).unwrap();
325 ///
326 /// assert_eq!(part1.as_slice(), &[3]);
327 /// assert_eq!(part2.as_slice(), &[2]);
328 /// assert_eq!(part3.as_slice(), &[1]);
329 ///
330 /// values.try_unsplit_off(part3).unwrap();
331 /// values.try_unsplit_off(part2).unwrap();
332 /// values.try_unsplit_off(part1).unwrap();
333 /// ```
334 ///
335 pub fn split_off(&mut self, at: usize) -> Option<SharedVecPart<T>> {
336 if at > self.remain_len {
337 return None;
338 }
339 self.borrow_count = self.borrow_count.checked_add(1)?;
340
341 let last_len = self.remain_len;
342 self.remain_len = at;
343
344 Some(SharedVecPart {
345 ptr: self.ptr,
346 start: self.remain_start + at,
347 len: last_len - at,
348 })
349 }
350 /// Split off the prefix of the vector ending at `at`.
351 /// This method is similar to `bytes::BytesMut::split_to`.
352 ///
353 /// Returns None when:
354 /// 1. `at` is greater than the length of the vector.
355 /// 2. `borrow_count` overflows `usize`.
356 ///
357 /// If successful, the returned part will contain [0, at) and self will contain [at, len).
358 ///
359 /// Here is an example of splitting a `SharedVecMut` into 3 parts:
360 /// ```
361 /// let mut values = manual_share::SharedVecMut::from_vec(vec![1, 2, 3]);
362 ///
363 /// let part1 = values.split_to(1).unwrap();
364 /// let part2 = values.split_to(1).unwrap();
365 /// let part3 = values.split_to(1).unwrap();
366 ///
367 /// assert_eq!(part1.as_slice(), &[1]);
368 /// assert_eq!(part2.as_slice(), &[2]);
369 /// assert_eq!(part3.as_slice(), &[3]);
370 ///
371 /// values.try_unsplit_to(part3).unwrap();
372 /// values.try_unsplit_to(part2).unwrap();
373 /// values.try_unsplit_to(part1).unwrap();
374 /// ```
375 pub fn split_to(&mut self, at: usize) -> Option<SharedVecPart<T>> {
376 if at > self.remain_len {
377 return None;
378 }
379 self.borrow_count = self.borrow_count.checked_add(1)?;
380
381 let last_start = self.remain_start;
382 self.remain_start += at;
383 self.remain_len -= at;
384
385 Some(SharedVecPart {
386 ptr: self.ptr,
387 start: last_start,
388 len: at,
389 })
390 }
391 /// Try to unsplit a part that was previously split off with `split_off`.
392 pub fn try_unsplit_off(&mut self, part: SharedVecPart<T>) -> Result<(), SharedVecPart<T>> {
393 if !core::ptr::eq(self.ptr, part.ptr) {
394 return Err(part);
395 }
396 if part.start != self.remain_start + self.remain_len {
397 return Err(part);
398 }
399
400 self.remain_len += part.len;
401
402 self.consume_part(part)
403 }
404 /// Try to unsplit a part that was previously split off with `split_to`.
405 pub fn try_unsplit_to(&mut self, part: SharedVecPart<T>) -> Result<(), SharedVecPart<T>> {
406 if !core::ptr::eq(self.ptr, part.ptr) {
407 return Err(part);
408 }
409 if self.remain_start != part.start + part.len {
410 return Err(part);
411 }
412
413 self.remain_start = part.start;
414 self.remain_len += part.len;
415
416 self.consume_part(part)
417 }
418 fn consume_part(&mut self, part: SharedVecPart<T>) -> Result<(), SharedVecPart<T>> {
419 if size_of::<T>() == 0 {
420 // ZST types can have multiple allocations to the same address, so we need to check for overflow.
421 if let Some(new_count) = self.borrow_count.checked_sub(1) {
422 self.borrow_count = new_count;
423 let _ = core::mem::ManuallyDrop::new(part);
424 Ok(())
425 } else {
426 Err(part)
427 }
428 } else {
429 self.borrow_count -= 1;
430 let _ = core::mem::ManuallyDrop::new(part);
431 Ok(())
432 }
433 }
434 fn can_convert_back(&self) -> bool {
435 self.borrow_count == 0
436 && self.remain_start == 0
437 && if size_of::<T>() == 0 {
438 self.remain_len == self.len
439 } else {
440 true
441 }
442 }
443 /// Try to convert the mutable view back into a `Vec` when no parts remain outstanding.
444 pub fn try_into_vec(self) -> Result<Vec<T>, Self> {
445 if self.can_convert_back() {
446 let r = core::mem::ManuallyDrop::new(self);
447 let vec = unsafe { Vec::from_raw_parts(r.ptr, r.remain_len, r.cap) };
448
449 Ok(vec)
450 } else {
451 Err(self)
452 }
453 }
454 /// Directly get a slice of the remaining part of the `SharedVecMut`.
455 /// ```
456 /// let mut values = manual_share::SharedVecMut::from_vec(vec![1, 2, 3]);
457 ///
458 /// let part1 = values.split_to(1).unwrap();
459 /// let part2 = values.split_off(1).unwrap();
460 ///
461 /// assert_eq!(values.as_slice(), &[2]);
462 /// assert_eq!(part1.as_slice(), &[1]);
463 /// assert_eq!(part2.as_slice(), &[3]);
464 ///
465 /// values.try_unsplit_off(part2).unwrap();
466 /// values.try_unsplit_to(part1).unwrap();
467 /// assert_eq!(values.as_slice(), &[1, 2, 3]);
468 /// ```
469 ///
470 /// Further splitting is no longer possible as long as the returned slice is held alive:
471 /// ```compile_fail
472 /// let mut values = manual_share::SharedVecMut::from_vec(vec![1, 2, 3]);
473 /// let slice = values.as_slice();
474 /// let part = values.split_off(1).unwrap();
475 ///
476 /// println!("{:?}", slice);
477 /// ```
478 pub fn as_slice(&self) -> &[T] {
479 // SAFETY:
480 // The pointer is valid as long as the SharedVecMut is alive.
481 // SharedVecPart cannot point to the same or overlapping region as self.
482 // Also, splitting methods can't be called when the returned slice is alive.
483 unsafe { core::slice::from_raw_parts(self.ptr.add(self.remain_start), self.remain_len) }
484 }
485 /// Directly get a mutable slice of the remaining part of the `SharedVecMut`.
486 pub fn as_slice_mut(&mut self) -> &mut [T] {
487 // SAFETY:
488 // The pointer is valid as long as the SharedVecMut is alive.
489 // SharedVecPart cannot point to the same or overlapping region as self.
490 // Also, splitting methods can't be called when the returned slice is alive.
491 unsafe { core::slice::from_raw_parts_mut(self.ptr.add(self.remain_start), self.remain_len) }
492 }
493}
494
495unsafe impl<T: Send> Send for SharedVecMut<T> {}
496unsafe impl<T: Sync> Sync for SharedVecMut<T> {}
497
498impl<T> Drop for SharedVecMut<T> {
499 fn drop(&mut self) {
500 #[cfg(feature = "panic-on-drop")]
501 {
502 // Let user deal with other panics.
503 #[cfg(feature = "do-not-panic-when-panicking")]
504 if std::thread::panicking() {
505 return;
506 }
507
508 if self.borrow_count > 0 {
509 panic!("Dropping a SharedVecMut without giving back all SharedVecRef")
510 }
511 }
512
513 if self.can_convert_back() {
514 unsafe {
515 drop(Vec::from_raw_parts(self.ptr, self.len, self.cap));
516 }
517 }
518 }
519}
520
521/// A slice-like view into a segment of a `SharedVecMut` allocation.
522///
523/// It can be read as a slice or mutated in place while the underlying allocation is still owned
524/// by the original `SharedVecMut`.
525///
526/// Dropping a `SharedVecPart` leaks the underlying allocation.
527/// When the **`panic-on-drop`** feature is enabled, dropping it will panic:
528/// ```should_panic
529/// let mut values = manual_share::SharedVecMut::from_vec(vec![1, 2, 3, 4]);
530/// let mut part = values.split_off(2).unwrap();
531///
532/// // forget SharedVecMut to make sure the panic is not caused by dropping it first.
533/// std::mem::forget(values);
534///
535/// // panic here due to dropping ShareVecPart
536/// ```
537#[derive(Debug)]
538pub struct SharedVecPart<T> {
539 ptr: *mut T,
540 start: usize,
541 len: usize,
542}
543
544impl<T> SharedVecPart<T> {
545 /// View the part as an immutable slice.
546 pub fn as_slice(&self) -> &[T] {
547 unsafe { core::slice::from_raw_parts(self.ptr.add(self.start), self.len) }
548 }
549 /// View the part as a mutable slice.
550 pub fn as_slice_mut(&mut self) -> &mut [T] {
551 unsafe { core::slice::from_raw_parts_mut(self.ptr.add(self.start), self.len) }
552 }
553}
554
555impl<T> Drop for SharedVecPart<T> {
556 fn drop(&mut self) {
557 #[cfg(feature = "panic-on-drop")]
558 {
559 // Let user deal with other panics.
560 #[cfg(feature = "do-not-panic-when-panicking")]
561 if std::thread::panicking() {
562 return;
563 }
564
565 panic!("Dropping a SharedVecPart without returning it to the SharedVecMut")
566 }
567 }
568}
569
570unsafe impl<T: Send> Send for SharedVecPart<T> {}
571unsafe impl<T: Sync> Sync for SharedVecPart<T> {}
572
573#[cfg(test)]
574mod test {
575 use super::*;
576
577 #[test]
578 fn zst() {
579 let mut b1: SharedVec<()> = SharedVec::from_vec(Vec::new());
580 let mut b2: SharedVec<()> = SharedVec::from_vec(Vec::new());
581
582 let r11 = b1.borrow();
583 let r12 = b1.borrow();
584
585 let r2 = b2.borrow();
586
587 b1.try_return(r2).unwrap();
588
589 b1.try_return(r11).unwrap();
590 let r12 = b1.try_return(r12).unwrap_err();
591
592 b2.try_return(r12).unwrap();
593 }
594
595 #[test]
596 fn mut_zst() {
597 let mut b1: SharedVecMut<()> = SharedVecMut::from_vec(Vec::new());
598 let mut b2: SharedVecMut<()> = SharedVecMut::from_vec(Vec::new());
599
600 let r11 = b1.split_off(0).unwrap();
601 let r12 = b1.split_off(0).unwrap();
602
603 let r2 = b2.split_off(0).unwrap();
604
605 b1.try_unsplit_off(r2).unwrap();
606
607 b1.try_unsplit_off(r11).unwrap();
608 let r12 = b1.try_unsplit_off(r12).unwrap_err();
609
610 b2.try_unsplit_off(r12).unwrap();
611 }
612}