1use nd::iter::{Iter, IterMut};
6use nd::prelude::*;
7use nd::{Data, DataMut, OwnedRepr, RawData, RawDataClone};
8
9pub struct Mask<S = OwnedRepr<bool>, D = Ix2>(ArrayBase<S, D>)
10where
11 D: Dimension,
12 S: RawData<Elem = bool>;
13
14impl<S, D> Mask<S, D>
15where
16 D: Dimension,
17 S: RawData<Elem = bool>,
18{
19 pub fn from_arr(data: ArrayBase<S, D>) -> Self {
20 Self(data)
21 }
22
23 pub fn apply<A, T, F>(&mut self, data: &ArrayBase<T, D>, fill: A) -> ArrayBase<T, D>
24 where
25 A: Clone,
26 S: Data,
27 T: DataMut<Elem = A> + RawDataClone,
28 {
29 let mut res = data.clone();
30 res.zip_mut_with(self.as_mut(), |x, &m| {
31 if m {
32 *x = fill.clone();
33 }
34 });
35 res
36 }
37
38 pub fn mask_inplace<'a, A, T, F>(
39 &mut self,
40 data: &'a mut ArrayBase<T, D>,
41 fill: A,
42 ) -> &'a mut ArrayBase<T, D>
43 where
44 A: Clone,
45 S: Data,
46 T: DataMut<Elem = A>,
47 {
48 data.zip_mut_with(&mut self.0, |x, &m| {
49 if m {
50 *x = fill.clone();
51 }
52 });
53 data
54 }
55
56 pub fn as_slice(&self) -> &[bool]
57 where
58 S: Data,
59 {
60 self.get().as_slice().unwrap()
61 }
62
63 pub fn as_mut_slice(&mut self) -> &mut [bool]
64 where
65 S: DataMut,
66 {
67 self.get_mut().as_slice_mut().unwrap()
68 }
69
70 pub fn dim(&self) -> D::Pattern {
71 self.get().dim()
72 }
73
74 pub fn iter(&self) -> Iter<'_, bool, D>
75 where
76 S: Data,
77 {
78 self.get().iter()
79 }
80
81 pub fn iter_mut(&mut self) -> IterMut<'_, bool, D>
82 where
83 S: DataMut,
84 {
85 self.get_mut().iter_mut()
86 }
87
88 pub fn get(&self) -> &ArrayBase<S, D> {
89 &self.0
90 }
91
92 pub fn get_mut(&mut self) -> &mut ArrayBase<S, D> {
93 &mut self.0
94 }
95
96 pub fn into_inner(self) -> ArrayBase<S, D> {
97 self.0
98 }
99
100 pub fn ndim(&self) -> usize {
101 self.get().ndim()
102 }
103
104 pub fn raw_dim(&self) -> D {
105 self.get().raw_dim()
106 }
107
108 pub fn set(&mut self, data: ArrayBase<S, D>) {
109 self.0 = data;
110 }
111
112 pub fn shape(&self) -> D {
113 self.get().raw_dim()
114 }
115}
116
117mod impls {
121 use super::Mask;
122 use core::borrow::{Borrow, BorrowMut};
123 use core::ops::{Deref, DerefMut, Index, IndexMut};
124 use nd::{ArrayBase, Data, DataMut, Dimension, NdIndex, RawData};
125
126 impl<S, D> AsRef<ArrayBase<S, D>> for Mask<S, D>
127 where
128 D: Dimension,
129 S: RawData<Elem = bool>,
130 {
131 fn as_ref(&self) -> &ArrayBase<S, D> {
132 &self.0
133 }
134 }
135
136 impl<S, D> AsMut<ArrayBase<S, D>> for Mask<S, D>
137 where
138 D: Dimension,
139 S: RawData<Elem = bool>,
140 {
141 fn as_mut(&mut self) -> &mut ArrayBase<S, D> {
142 &mut self.0
143 }
144 }
145
146 impl<S, D> Borrow<ArrayBase<S, D>> for Mask<S, D>
147 where
148 D: Dimension,
149 S: RawData<Elem = bool>,
150 {
151 fn borrow(&self) -> &ArrayBase<S, D> {
152 &self.0
153 }
154 }
155
156 impl<S, D> BorrowMut<ArrayBase<S, D>> for Mask<S, D>
157 where
158 D: Dimension,
159 S: RawData<Elem = bool>,
160 {
161 fn borrow_mut(&mut self) -> &mut ArrayBase<S, D> {
162 &mut self.0
163 }
164 }
165
166 impl<S, D> Deref for Mask<S, D>
167 where
168 D: Dimension,
169 S: RawData<Elem = bool>,
170 {
171 type Target = ArrayBase<S, D>;
172
173 fn deref(&self) -> &Self::Target {
174 &self.0
175 }
176 }
177
178 impl<S, D> DerefMut for Mask<S, D>
179 where
180 D: Dimension,
181 S: RawData<Elem = bool>,
182 {
183 fn deref_mut(&mut self) -> &mut Self::Target {
184 &mut self.0
185 }
186 }
187
188 impl<S, D, I> Index<I> for Mask<S, D>
189 where
190 D: Dimension,
191 I: NdIndex<D>,
192 S: Data<Elem = bool>,
193 {
194 type Output = <ArrayBase<S, D> as Index<I>>::Output;
195
196 fn index(&self, index: I) -> &Self::Output {
197 &self.0[index]
198 }
199 }
200
201 impl<S, D, I> IndexMut<I> for Mask<S, D>
202 where
203 D: Dimension,
204 I: NdIndex<D>,
205 S: DataMut<Elem = bool>,
206 {
207 fn index_mut(&mut self, index: I) -> &mut Self::Output {
208 &mut self.0[index]
209 }
210 }
211}
212
213mod impl_from {
214 use super::Mask;
215 use nd::{ArrayBase, Dimension, RawData};
216
217 impl<S, D> From<ArrayBase<S, D>> for Mask<S, D>
218 where
219 D: Dimension,
220 S: RawData<Elem = bool>,
221 {
222 fn from(mask: ArrayBase<S, D>) -> Self {
223 Mask(mask)
224 }
225 }
226
227 impl<S, D> From<Mask<S, D>> for ArrayBase<S, D>
228 where
229 D: Dimension,
230 S: RawData<Elem = bool>,
231 {
232 fn from(mask: Mask<S, D>) -> Self {
233 mask.0
234 }
235 }
236}