mdarray/expr/
buffer.rs

1#[cfg(feature = "nightly")]
2use alloc::alloc::{Allocator, Global};
3
4use core::mem::ManuallyDrop;
5use core::ptr;
6
7#[cfg(not(feature = "nightly"))]
8use crate::allocator::{Allocator, Global};
9use crate::array::Array;
10use crate::dim::Const;
11use crate::index::Axis;
12use crate::mapping::Mapping;
13use crate::shape::{ConstShape, Shape};
14use crate::slice::Slice;
15use crate::tensor::Tensor;
16use crate::view::ViewMut;
17
18/// Array buffer trait, for moving elements out of an array.
19pub trait Buffer {
20    /// Array element type.
21    type Item;
22
23    /// Array shape type.
24    type Shape: Shape;
25
26    #[doc(hidden)]
27    fn as_mut_slice(&mut self) -> &mut Slice<ManuallyDrop<Self::Item>, Self::Shape>;
28
29    #[doc(hidden)]
30    fn as_slice(&self) -> &Slice<ManuallyDrop<Self::Item>, Self::Shape>;
31}
32
33/// Buffer for moving elements out of an array range.
34pub struct Drain<'a, T, S: Shape, A: Allocator = Global> {
35    tensor: &'a mut Tensor<T, S, A>,
36    view: ViewMut<'a, ManuallyDrop<T>, S>,
37    new_size: usize,
38    tail: usize,
39}
40
41impl<'a, T, S: Shape, A: Allocator> Drain<'a, T, S, A> {
42    #[inline]
43    pub(crate) fn new(tensor: &'a mut Tensor<T, S, A>, start: usize, end: usize) -> Self {
44        assert!(start <= end && end <= tensor.dim(0), "invalid range");
45
46        let new_size = tensor.dim(0) - (end - start);
47        let tail = Axis::resize(Const::<0>, tensor.mapping(), new_size - start).len();
48
49        // Shrink the array, to be safe in case Drain is leaked.
50        unsafe {
51            tensor.set_mapping(Mapping::resize_dim(tensor.mapping(), 0, start));
52        }
53
54        let ptr = unsafe { tensor.as_mut_ptr().add(tensor.len()) as *mut ManuallyDrop<T> };
55        let mapping = Mapping::resize_dim(tensor.mapping(), 0, end - start);
56
57        let view = unsafe { ViewMut::new_unchecked(ptr, mapping) };
58
59        Self { tensor, view, new_size, tail }
60    }
61}
62
63impl<T, S: Shape, A: Allocator> Buffer for Drain<'_, T, S, A> {
64    type Item = T;
65    type Shape = S;
66
67    #[inline]
68    fn as_mut_slice(&mut self) -> &mut Slice<ManuallyDrop<T>, S> {
69        &mut self.view
70    }
71
72    #[inline]
73    fn as_slice(&self) -> &Slice<ManuallyDrop<T>, S> {
74        &self.view
75    }
76}
77
78impl<T, S: Shape, A: Allocator> Drop for Drain<'_, T, S, A> {
79    #[inline]
80    fn drop(&mut self) {
81        let mapping = Mapping::resize_dim(self.tensor.mapping(), 0, self.new_size);
82
83        unsafe {
84            ptr::copy(self.view.as_ptr().add(self.view.len()), self.view.as_mut_ptr(), self.tail);
85            self.tensor.set_mapping(mapping);
86        }
87    }
88}
89
90impl<T, S: ConstShape> Buffer for Array<ManuallyDrop<T>, S> {
91    type Item = T;
92    type Shape = S;
93
94    #[inline]
95    fn as_mut_slice(&mut self) -> &mut Slice<ManuallyDrop<T>, S> {
96        self
97    }
98
99    #[inline]
100    fn as_slice(&self) -> &Slice<ManuallyDrop<T>, S> {
101        self
102    }
103}
104
105impl<T, S: Shape, A: Allocator> Buffer for Tensor<ManuallyDrop<T>, S, A> {
106    type Item = T;
107    type Shape = S;
108
109    #[inline]
110    fn as_mut_slice(&mut self) -> &mut Slice<ManuallyDrop<T>, S> {
111        self
112    }
113
114    #[inline]
115    fn as_slice(&self) -> &Slice<ManuallyDrop<T>, S> {
116        self
117    }
118}
mdarray/expr/buffer.rs

mdarray/expr/
buffer.rs
