mdarray 0.7.2

Multidimensional array for Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118

#[cfg(feature = "nightly")]
use alloc::alloc::{Allocator, Global};

use core::mem::ManuallyDrop;
use core::ptr;

#[cfg(not(feature = "nightly"))]
use crate::allocator::{Allocator, Global};
use crate::array::Array;
use crate::dim::Const;
use crate::index::Axis;
use crate::mapping::Mapping;
use crate::shape::{ConstShape, Shape};
use crate::slice::Slice;
use crate::tensor::Tensor;
use crate::view::ViewMut;

/// Array buffer trait, for moving elements out of an array.
pub trait Buffer {
    /// Array element type.
    type Item;

    /// Array shape type.
    type Shape: Shape;

    #[doc(hidden)]
    fn as_mut_slice(&mut self) -> &mut Slice<ManuallyDrop<Self::Item>, Self::Shape>;

    #[doc(hidden)]
    fn as_slice(&self) -> &Slice<ManuallyDrop<Self::Item>, Self::Shape>;
}

/// Buffer for moving elements out of an array range.
pub struct Drain<'a, T, S: Shape, A: Allocator = Global> {
    tensor: &'a mut Tensor<T, S, A>,
    view: ViewMut<'a, ManuallyDrop<T>, S>,
    new_size: usize,
    tail: usize,
}

impl<'a, T, S: Shape, A: Allocator> Drain<'a, T, S, A> {
    #[inline]
    pub(crate) fn new(tensor: &'a mut Tensor<T, S, A>, start: usize, end: usize) -> Self {
        assert!(start <= end && end <= tensor.dim(0), "invalid range");

        let new_size = tensor.dim(0) - (end - start);
        let tail = Axis::resize(Const::<0>, tensor.mapping(), new_size - start).len();

        // Shrink the array, to be safe in case Drain is leaked.
        unsafe {
            tensor.set_mapping(Mapping::resize_dim(tensor.mapping(), 0, start));
        }

        let ptr = unsafe { tensor.as_mut_ptr().add(tensor.len()) as *mut ManuallyDrop<T> };
        let mapping = Mapping::resize_dim(tensor.mapping(), 0, end - start);

        let view = unsafe { ViewMut::new_unchecked(ptr, mapping) };

        Self { tensor, view, new_size, tail }
    }
}

impl<T, S: Shape, A: Allocator> Buffer for Drain<'_, T, S, A> {
    type Item = T;
    type Shape = S;

    #[inline]
    fn as_mut_slice(&mut self) -> &mut Slice<ManuallyDrop<T>, S> {
        &mut self.view
    }

    #[inline]
    fn as_slice(&self) -> &Slice<ManuallyDrop<T>, S> {
        &self.view
    }
}

impl<T, S: Shape, A: Allocator> Drop for Drain<'_, T, S, A> {
    #[inline]
    fn drop(&mut self) {
        let mapping = Mapping::resize_dim(self.tensor.mapping(), 0, self.new_size);

        unsafe {
            ptr::copy(self.view.as_ptr().add(self.view.len()), self.view.as_mut_ptr(), self.tail);
            self.tensor.set_mapping(mapping);
        }
    }
}

impl<T, S: ConstShape> Buffer for Array<ManuallyDrop<T>, S> {
    type Item = T;
    type Shape = S;

    #[inline]
    fn as_mut_slice(&mut self) -> &mut Slice<ManuallyDrop<T>, S> {
        self
    }

    #[inline]
    fn as_slice(&self) -> &Slice<ManuallyDrop<T>, S> {
        self
    }
}

impl<T, S: Shape, A: Allocator> Buffer for Tensor<ManuallyDrop<T>, S, A> {
    type Item = T;
    type Shape = S;

    #[inline]
    fn as_mut_slice(&mut self) -> &mut Slice<ManuallyDrop<T>, S> {
        self
    }

    #[inline]
    fn as_slice(&self) -> &Slice<ManuallyDrop<T>, S> {
        self
    }
}