ndarray/partial.rs
1// Copyright 2020 bluss and ndarray developers.
2//
3// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
4// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
5// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
6// option. This file may not be copied, modified, or distributed
7// except according to those terms.
8
9use std::ptr;
10
11/// Partial is a partially written contiguous slice of data;
12/// it is the owner of the elements, but not the allocation,
13/// and will drop the elements on drop.
14#[must_use]
15pub(crate) struct Partial<T> {
16 /// Data pointer
17 ptr: *mut T,
18 /// Current length
19 pub(crate) len: usize,
20}
21
22impl<T> Partial<T> {
23 /// Create an empty partial for this data pointer
24 ///
25 /// ## Safety
26 ///
27 /// Unless ownership is released, the Partial acts as an owner of the slice of data (not the
28 /// allocation); and will free the elements on drop; the pointer must be dereferenceable and
29 /// the `len` elements following it valid.
30 ///
31 /// The Partial has an accessible length field which must only be modified in trusted code.
32 pub(crate) unsafe fn new(ptr: *mut T) -> Self {
33 Self {
34 ptr,
35 len: 0,
36 }
37 }
38
39 #[cfg(feature = "rayon")]
40 pub(crate) fn stub() -> Self {
41 Self { len: 0, ptr: ptr::null_mut() }
42 }
43
44 #[cfg(feature = "rayon")]
45 pub(crate) fn is_stub(&self) -> bool {
46 self.ptr.is_null()
47 }
48
49 /// Release Partial's ownership of the written elements, and return the current length
50 pub(crate) fn release_ownership(mut self) -> usize {
51 let ret = self.len;
52 self.len = 0;
53 ret
54 }
55
56 #[cfg(feature = "rayon")]
57 /// Merge if they are in order (left to right) and contiguous.
58 /// Skips merge if T does not need drop.
59 pub(crate) fn try_merge(mut left: Self, right: Self) -> Self {
60 if !std::mem::needs_drop::<T>() {
61 return left;
62 }
63 // Merge the partial collect results; the final result will be a slice that
64 // covers the whole output.
65 if left.is_stub() {
66 right
67 } else if left.ptr.wrapping_add(left.len) == right.ptr {
68 left.len += right.release_ownership();
69 left
70 } else {
71 // failure to merge; this is a bug in collect, so we will never reach this
72 debug_assert!(false, "Partial: failure to merge left and right parts");
73 left
74 }
75 }
76}
77
78unsafe impl<T> Send for Partial<T> where T: Send { }
79
80impl<T> Drop for Partial<T> {
81 fn drop(&mut self) {
82 if !self.ptr.is_null() {
83 unsafe {
84 ptr::drop_in_place(alloc::slice::from_raw_parts_mut(self.ptr, self.len));
85 }
86 }
87 }
88}