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}