Mercurial > repos > alg_tools / file comparison

comparison: src/maputil.rs

src/maputil.rs

changeset 91
db870f2a2cde
parent 55
7b2ee3e84c5f
equal deleted inserted replaced
90:b3c35d16affe 91:db870f2a2cde
1 /*! 1 /*!
2 Utilities for mapping over various container types. 2 Utilities for mapping over various container types.
3 */ 3 */
4 4
5 use itertools::izip;
5 #[cfg(feature = "nightly")] 6 #[cfg(feature = "nightly")]
6 use std::mem::MaybeUninit; 7 use std::mem::MaybeUninit;
7 use itertools::izip;
8 8
9 /// Trait for a fixed-length container type. 9 /// Trait for a fixed-length container type.
10 /// 10 ///
11 /// Implemented by [`Loc`][crate::loc::Loc] vectors, [`Cube`][crate::sets::Cube]s, 11 /// Implemented by [`Loc`][crate::loc::Loc] vectors, [`Cube`][crate::sets::Cube]s,
12 /// and basic arrays. 12 /// and basic arrays.
13 pub trait FixedLength<const N : usize> { 13 pub trait FixedLength<const N: usize> {
14 /// Type of elements of the container. 14 /// Type of elements of the container.
15 type Elem; 15 type Elem;
16 /// Type of iterators over the elements of the container. 16 /// Type of iterators over the elements of the container.
17 type Iter : Iterator<Item = Self::Elem>; 17 type Iter: Iterator<Item = Self::Elem>;
18 18
19 /// Returns an iteartor over the elements of the container. 19 /// Returns an iteartor over the elements of the container.
20 fn fl_iter(self) -> Self::Iter; 20 fn fl_iter(self) -> Self::Iter;
21 } 21 }
22 22
23 /// Trait for a mutable fixed-length container type. 23 /// Trait for a mutable fixed-length container type.
24 pub trait FixedLengthMut<const N : usize> : FixedLength<N> { 24 pub trait FixedLengthMut<const N: usize>: FixedLength<N> {
25 /// Type of iterators over references to mutable elements of the container. 25 /// Type of iterators over references to mutable elements of the container.
26 type IterMut<'a> : Iterator<Item=&'a mut Self::Elem> where Self : 'a; 26 type IterMut<'a>: Iterator<Item = &'a mut Self::Elem>
27 where
28 Self: 'a;
27 29
28 /// Returns an iterator over mutable references to elements of the container. 30 /// Returns an iterator over mutable references to elements of the container.
29 fn fl_iter_mut(&mut self) -> Self::IterMut<'_>; 31 fn fl_iter_mut(&mut self) -> Self::IterMut<'_>;
30 } 32 }
31 33
32 impl<A, const N : usize> FixedLength<N> for [A; N] { 34 impl<A, const N: usize> FixedLength<N> for [A; N] {
33 type Elem = A; 35 type Elem = A;
34 type Iter = std::array::IntoIter<A, N>; 36 type Iter = std::array::IntoIter<A, N>;
35 #[inline] 37 #[inline]
36 fn fl_iter(self) -> Self::Iter { 38 fn fl_iter(self) -> Self::Iter {
37 self.into_iter() 39 self.into_iter()
38 } 40 }
39 } 41 }
40 42
41 impl<A, const N : usize> FixedLengthMut<N> for [A; N] { 43 impl<A, const N: usize> FixedLengthMut<N> for [A; N] {
42 type IterMut<'a> = std::slice::IterMut<'a, A> where A : 'a; 44 type IterMut<'a>
45 = std::slice::IterMut<'a, A>
46 where
47 A: 'a;
43 #[inline] 48 #[inline]
44 fn fl_iter_mut(&mut self) -> Self::IterMut<'_> { 49 fn fl_iter_mut(&mut self) -> Self::IterMut<'_> {
45 self.iter_mut() 50 self.iter_mut()
46 } 51 }
47 } 52 }
48 53
49 impl<'a, A, const N : usize> FixedLength<N> for &'a [A; N] { 54 impl<'a, A, const N: usize> FixedLength<N> for &'a [A; N] {
50 type Elem = &'a A; 55 type Elem = &'a A;
51 type Iter = std::slice::Iter<'a, A>; 56 type Iter = std::slice::Iter<'a, A>;
52 #[inline] 57 #[inline]
53 fn fl_iter(self) -> Self::Iter { 58 fn fl_iter(self) -> Self::Iter {
54 self.iter() 59 self.iter()
155 make_mapmany_mut!(map3_mut, map3_indexed_mut, a b c; A B C, CA CB CC); 160 make_mapmany_mut!(map3_mut, map3_indexed_mut, a b c; A B C, CA CB CC);
156 make_mapmany_mut!(map4_mut, map4_indexed_mut, a b c d; A B C D, CA CB CC CD); 161 make_mapmany_mut!(map4_mut, map4_indexed_mut, a b c d; A B C D, CA CB CC CD);
157 make_mapmany_mut!(map5_mut, map5_indexed_mut, a b c d e; A B C D E, CA CB CC CD CE); 162 make_mapmany_mut!(map5_mut, map5_indexed_mut, a b c d e; A B C D E, CA CB CC CD CE);
158 make_mapmany_mut!(map6_mut, map6_indexed_mut, a b c d e f; A B C D E F, CA CB CC CD CE CF); 163 make_mapmany_mut!(map6_mut, map6_indexed_mut, a b c d e f; A B C D E F, CA CB CC CD CE CF);
159 164
160
161 /// Initialise an array of length `N` by calling `f` multiple times. 165 /// Initialise an array of length `N` by calling `f` multiple times.
162 #[inline] 166 #[inline]
163 pub fn array_init<A, F : Fn() -> A, const N : usize>(f : F) -> [A; N] { 167 pub fn array_init<A, F: Fn() -> A, const N: usize>(f: F) -> [A; N] {
164 //[(); N].map(|_| f()) 168 //[(); N].map(|_| f())
165 core::array::from_fn(|_| f()) 169 core::array::from_fn(|_| f())
166 } 170 }
167 171
168 // /// Initialise an array of length `N` by calling `f` with the index of each element. 172 // /// Initialise an array of length `N` by calling `f` with the index of each element.
170 // pub fn array_gen<A, F : Fn(usize) -> A, const N : usize>(f : F) -> [A; N] { 174 // pub fn array_gen<A, F : Fn(usize) -> A, const N : usize>(f : F) -> [A; N] {
171 // //[(); N].indexmap(|i, _| f(i)) 175 // //[(); N].indexmap(|i, _| f(i))
172 // core::array::from_fn(f) 176 // core::array::from_fn(f)
173 // } 177 // }
174 178
175
176
177 /// Iterator returned by [`foldmap`][FoldMappable::foldmap] applied to an iterator. 179 /// Iterator returned by [`foldmap`][FoldMappable::foldmap] applied to an iterator.
178 180
179 pub struct FoldMap<I : Iterator<Item=A>, A, B, J : Copy, F : Fn(J, A) -> (J, B)> { 181 pub struct FoldMap<I: Iterator<Item = A>, A, B, J: Copy, F: Fn(J, A) -> (J, B)> {
180 iter : I, 182 iter: I,
181 f : F, 183 f: F,
182 j : J, 184 j: J,
183 } 185 }
184 186
185 impl<A, B, I : Iterator<Item=A>, J : Copy, F : Fn(J, A) -> (J, B)> Iterator for FoldMap<I, A, B, J, F> { 187 impl<A, B, I: Iterator<Item = A>, J: Copy, F: Fn(J, A) -> (J, B)> Iterator
188 for FoldMap<I, A, B, J, F>
189 {
186 type Item = B; 190 type Item = B;
187 #[inline] 191 #[inline]
188 fn next(&mut self) -> Option<B> { 192 fn next(&mut self) -> Option<B> {
189 self.iter.next().map(|a| { 193 self.iter.next().map(|a| {
190 let (jnew, b) = (self.f)(self.j, a); 194 let (jnew, b) = (self.f)(self.j, a);
193 }) 197 })
194 } 198 }
195 } 199 }
196 200
197 /// Iterator returned by [`indexmap`][IndexMappable::indexmap] applied to an iterator. 201 /// Iterator returned by [`indexmap`][IndexMappable::indexmap] applied to an iterator.
198 pub struct IndexMap<I : Iterator<Item=A>, A, B, F : Fn(usize, A) -> B> { 202 pub struct IndexMap<I: Iterator<Item = A>, A, B, F: Fn(usize, A) -> B> {
199 iter : I, 203 iter: I,
200 f : F, 204 f: F,
201 j : usize, 205 j: usize,
202 } 206 }
203 207
204 impl<A, B, I : Iterator<Item=A>, F : Fn(usize, A) -> B> Iterator for IndexMap<I, A, B, F> { 208 impl<A, B, I: Iterator<Item = A>, F: Fn(usize, A) -> B> Iterator for IndexMap<I, A, B, F> {
205 type Item = B; 209 type Item = B;
206 #[inline] 210 #[inline]
207 fn next(&mut self) -> Option<B> { 211 fn next(&mut self) -> Option<B> {
208 self.iter.next().map(|a| { 212 self.iter.next().map(|a| {
209 let b = (self.f)(self.j, a); 213 let b = (self.f)(self.j, a);
210 self.j = self.j+1; 214 self.j = self.j + 1;
211 b 215 b
212 }) 216 })
213 } 217 }
214 } 218 }
215 219
216 /// Trait for things that can be foldmapped. 220 /// Trait for things that can be foldmapped.
217 /// 221 ///
218 /// `A` is the type of elements of `Self`, and `J` the accumulator type for the folding. 222 /// `A` is the type of elements of `Self`, and `J` the accumulator type for the folding.
219 pub trait FoldMappable<A, J> : Sized { 223 pub trait FoldMappable<A, J>: Sized {
220 type Output<B, F> where F : Fn(J, A) -> (J, B); 224 type Output<B, F>
225 where
226 F: Fn(J, A) -> (J, B);
221 /// Fold and map over `self` with `f`. `j` is the initial accumulator for folding. 227 /// Fold and map over `self` with `f`. `j` is the initial accumulator for folding.
222 /// 228 ///
223 /// The output type depends on the implementation, but will generally have elements of 229 /// The output type depends on the implementation, but will generally have elements of
224 /// type `B`. 230 /// type `B`.
225 fn foldmap<B, F : Fn(J, A) -> (J, B)>(self, j : J, f : F) -> Self::Output<B, F>; 231 fn foldmap<B, F: Fn(J, A) -> (J, B)>(self, j: J, f: F) -> Self::Output<B, F>;
226 } 232 }
227 233
228 /// Trait for things that can be indexmapped. 234 /// Trait for things that can be indexmapped.
229 /// 235 ///
230 /// `A` is the type of elements of `Self`. 236 /// `A` is the type of elements of `Self`.
231 pub trait IndexMappable<A> : Sized { 237 pub trait IndexMappable<A>: Sized {
232 type Output<B, F> where F : Fn(usize, A) -> B; 238 type Output<B, F>
239 where
240 F: Fn(usize, A) -> B;
233 /// Map over element indices and elements of `self`. 241 /// Map over element indices and elements of `self`.
234 /// 242 ///
235 /// The output type depends on the implementation, but will generally have elements of 243 /// The output type depends on the implementation, but will generally have elements of
236 /// type `B`. 244 /// type `B`.
237 fn indexmap<B, F : Fn(usize, A) -> B>(self, f : F) -> Self::Output<B, F>; 245 fn indexmap<B, F: Fn(usize, A) -> B>(self, f: F) -> Self::Output<B, F>;
238 } 246 }
239 247
240 impl<'a, A, J : Copy> FoldMappable<&'a A, J> 248 impl<'a, A, J: Copy> FoldMappable<&'a A, J> for std::slice::Iter<'a, A> {
241 for std::slice::Iter<'a, A> { 249 type Output<B, F>
242 type Output<B, F> = FoldMap<Self, &'a A, B, J, F> where F : Fn(J, &'a A) -> (J, B); 250 = FoldMap<Self, &'a A, B, J, F>
243 #[inline] 251 where
244 fn foldmap<B, F : Fn(J, &'a A) -> (J, B)>(self, j : J, f : F) -> Self::Output<B, F> { 252 F: Fn(J, &'a A) -> (J, B);
245 FoldMap { iter : self, j, f } 253 #[inline]
246 } 254 fn foldmap<B, F: Fn(J, &'a A) -> (J, B)>(self, j: J, f: F) -> Self::Output<B, F> {
247 } 255 FoldMap { iter: self, j, f }
248 256 }
249 impl<'a, A> IndexMappable<&'a A> 257 }
250 for std::slice::Iter<'a, A> { 258
251 type Output<B, F> = IndexMap<Self, &'a A, B, F> where F : Fn(usize, &'a A) -> B; 259 impl<'a, A> IndexMappable<&'a A> for std::slice::Iter<'a, A> {
252 #[inline] 260 type Output<B, F>
253 fn indexmap<B, F : Fn(usize, &'a A) -> B>(self, f : F) -> Self::Output<B, F> { 261 = IndexMap<Self, &'a A, B, F>
254 IndexMap { iter : self, j : 0, f } 262 where
255 } 263 F: Fn(usize, &'a A) -> B;
256 } 264 #[inline]
257 265 fn indexmap<B, F: Fn(usize, &'a A) -> B>(self, f: F) -> Self::Output<B, F> {
258 266 IndexMap {
259 impl<A, J : Copy, const N : usize> FoldMappable<A, J> 267 iter: self,
260 for std::array::IntoIter<A, N> { 268 j: 0,
261 type Output<B, F> = FoldMap<Self, A, B, J, F> where F : Fn(J, A) -> (J, B); 269 f,
262 #[inline] 270 }
263 fn foldmap<B, F : Fn(J, A) -> (J, B)>(self, j : J, f : F) -> Self::Output<B, F> { 271 }
264 FoldMap { iter : self, j, f } 272 }
265 } 273
266 } 274 impl<A, J: Copy, const N: usize> FoldMappable<A, J> for std::array::IntoIter<A, N> {
267 275 type Output<B, F>
268 impl<'a, A, const N : usize> IndexMappable<A> 276 = FoldMap<Self, A, B, J, F>
269 for std::array::IntoIter<A, N> { 277 where
270 type Output<B, F> = IndexMap<Self, A, B, F> where F : Fn(usize, A) -> B; 278 F: Fn(J, A) -> (J, B);
271 #[inline] 279 #[inline]
272 fn indexmap<B, F : Fn(usize, A) -> B>(self, f : F) -> Self::Output<B, F> { 280 fn foldmap<B, F: Fn(J, A) -> (J, B)>(self, j: J, f: F) -> Self::Output<B, F> {
273 IndexMap { iter : self, j : 0, f } 281 FoldMap { iter: self, j, f }
274 } 282 }
275 } 283 }
276 284
277 impl<A, J : Copy, const N : usize> FoldMappable<A, J> for [A; N] { 285 impl<'a, A, const N: usize> IndexMappable<A> for std::array::IntoIter<A, N> {
278 type Output<B, F> = [B; N] where F : Fn(J, A) -> (J, B); 286 type Output<B, F>
279 #[inline] 287 = IndexMap<Self, A, B, F>
280 fn foldmap<B, F : Fn(J, A) -> (J, B)>(self, j : J, f : F) -> [B; N] { 288 where
289 F: Fn(usize, A) -> B;
290 #[inline]
291 fn indexmap<B, F: Fn(usize, A) -> B>(self, f: F) -> Self::Output<B, F> {
292 IndexMap {
293 iter: self,
294 j: 0,
295 f,
296 }
297 }
298 }
299
300 impl<A, J: Copy, const N: usize> FoldMappable<A, J> for [A; N] {
301 type Output<B, F>
302 = [B; N]
303 where
304 F: Fn(J, A) -> (J, B);
305 #[inline]
306 fn foldmap<B, F: Fn(J, A) -> (J, B)>(self, j: J, f: F) -> [B; N] {
281 // //let mut res : [MaybeUninit<B>; N] = unsafe { MaybeUninit::uninit().assume_init() }; 307 // //let mut res : [MaybeUninit<B>; N] = unsafe { MaybeUninit::uninit().assume_init() };
282 // let mut res = MaybeUninit::uninit_array::<N>(); 308 // let mut res = MaybeUninit::uninit_array::<N>();
283 // for (a, i) in self.into_iter().zip(0..N) { 309 // for (a, i) in self.into_iter().zip(0..N) {
284 // let (jnew, b) = f(j, a); 310 // let (jnew, b) = f(j, a);
285 // unsafe { *(res.get_unchecked_mut(i)) = MaybeUninit::new(b) }; 311 // unsafe { *(res.get_unchecked_mut(i)) = MaybeUninit::new(b) };
290 let it = self.into_iter().foldmap(j, f); 316 let it = self.into_iter().foldmap(j, f);
291 collect_into_array_unchecked(it) 317 collect_into_array_unchecked(it)
292 } 318 }
293 } 319 }
294 320
295 impl<A, const N : usize> IndexMappable<A> for [A; N] { 321 impl<A, const N: usize> IndexMappable<A> for [A; N] {
296 type Output<B, F> = [B; N] where F : Fn(usize, A) -> B; 322 type Output<B, F>
297 #[inline] 323 = [B; N]
298 fn indexmap<B, F : Fn(usize, A) -> B>(self, f : F) -> [B; N] { 324 where
325 F: Fn(usize, A) -> B;
326 #[inline]
327 fn indexmap<B, F: Fn(usize, A) -> B>(self, f: F) -> [B; N] {
299 // //let mut res : [MaybeUninit<B>; N] = unsafe { MaybeUninit::uninit().assume_init() }; 328 // //let mut res : [MaybeUninit<B>; N] = unsafe { MaybeUninit::uninit().assume_init() };
300 // let mut res = MaybeUninit::uninit_array::<N>(); 329 // let mut res = MaybeUninit::uninit_array::<N>();
301 // for (a, i) in self.into_iter().zip(0..N) { 330 // for (a, i) in self.into_iter().zip(0..N) {
302 // let b = f(i, a); 331 // let b = f(i, a);
303 // unsafe { *(res.get_unchecked_mut(i)) = MaybeUninit::new(b) }; 332 // unsafe { *(res.get_unchecked_mut(i)) = MaybeUninit::new(b) };
323 /// 352 ///
324 /// If `iter.next()` panicks, all items already yielded by the iterator are 353 /// If `iter.next()` panicks, all items already yielded by the iterator are
325 /// dropped. 354 /// dropped.
326 #[cfg(feature = "nightly")] 355 #[cfg(feature = "nightly")]
327 #[inline] 356 #[inline]
328 pub(crate) fn collect_into_array_unchecked< 357 pub(crate) fn collect_into_array_unchecked<T, I: Iterator<Item = T>, const N: usize>(
329 T, 358 mut iter: I,
330 I : Iterator<Item=T>, 359 ) -> [T; N] {
331 const N: usize
332 >(mut iter: I) -> [T; N]
333 {
334 if N == 0 { 360 if N == 0 {
335 // SAFETY: An empty array is always inhabited and has no validity invariants. 361 // SAFETY: An empty array is always inhabited and has no validity invariants.
336 return unsafe { core::mem::zeroed() }; 362 return unsafe { core::mem::zeroed() };
337 } 363 }
338 364
346 fn drop(&mut self) { 372 fn drop(&mut self) {
347 debug_assert!(self.initialized <= N); 373 debug_assert!(self.initialized <= N);
348 374
349 // SAFETY: this slice will contain only initialized objects. 375 // SAFETY: this slice will contain only initialized objects.
350 unsafe { 376 unsafe {
351 core::ptr::drop_in_place(MaybeUninit::slice_assume_init_mut( 377 for elem in self.array_mut.get_unchecked_mut(..self.initialized) {
352 &mut self.array_mut.get_unchecked_mut(..self.initialized), 378 elem.assume_init_drop();
353 )); 379 }
354 } 380 }
355 } 381 }
356 } 382 }
357 383
358 let mut array = MaybeUninit::uninit_array::<N>(); 384 let mut array = MaybeUninit::uninit_array::<N>();
359 let mut guard = Guard { array_mut: &mut array, initialized: 0 }; 385 let mut guard = Guard {
386 array_mut: &mut array,
387 initialized: 0,
388 };
360 389
361 while let Some(item) = iter.next() { 390 while let Some(item) = iter.next() {
362 // SAFETY: `guard.initialized` starts at 0, is increased by one in the 391 // SAFETY: `guard.initialized` starts at 0, is increased by one in the
363 // loop and the loop is aborted once it reaches N (which is 392 // loop and the loop is aborted once it reaches N (which is
364 // `array.len()`). 393 // `array.len()`).
365 unsafe { 394 unsafe {
366 guard.array_mut.get_unchecked_mut(guard.initialized).write(item); 395 guard
396 .array_mut
397 .get_unchecked_mut(guard.initialized)
398 .write(item);
367 } 399 }
368 guard.initialized += 1; 400 guard.initialized += 1;
369 401
370 // Check if the whole array was initialized. 402 // Check if the whole array was initialized.
371 if guard.initialized == N { 403 if guard.initialized == N {
381 unreachable!("Something went wrong with iterator length") 413 unreachable!("Something went wrong with iterator length")
382 } 414 }
383 415
384 #[cfg(not(feature = "nightly"))] 416 #[cfg(not(feature = "nightly"))]
385 #[inline] 417 #[inline]
386 pub(crate) fn collect_into_array_unchecked< 418 pub(crate) fn collect_into_array_unchecked<T, I: Iterator<Item = T>, const N: usize>(
387 T, 419 iter: I,
388 I : Iterator<Item=T>, 420 ) -> [T; N] {
389 const N: usize
390 >(iter: I) -> [T; N]
391 {
392 match iter.collect::<Vec<T>>().try_into() { 421 match iter.collect::<Vec<T>>().try_into() {
393 Ok(a) => a, 422 Ok(a) => a,
394 Err(_) => panic!("collect_into_array failure: should not happen"), 423 Err(_) => panic!("collect_into_array failure: should not happen"),
395 } 424 }
396 } 425 }
399 mod tests { 428 mod tests {
400 use super::*; 429 use super::*;
401 430
402 #[test] 431 #[test]
403 fn mapx_test() { 432 fn mapx_test() {
404 let a = [0,1,2]; 433 let a = [0, 1, 2];
405 let mut b = [2,1,0]; 434 let mut b = [2, 1, 0];
406 assert_eq!(map1(a, |x| x+1), [1,2,3]); 435 assert_eq!(map1(a, |x| x + 1), [1, 2, 3]);
407 assert_eq!(map2(a, b, |x, y| x+y), [2,2,2]); 436 assert_eq!(map2(a, b, |x, y| x + y), [2, 2, 2]);
408 assert_eq!(map1_indexed(a, |i, y| y-i), [0,0,0]); 437 assert_eq!(map1_indexed(a, |i, y| y - i), [0, 0, 0]);
409 map1_indexed_mut(&mut b, |i, y| *y=i); 438 map1_indexed_mut(&mut b, |i, y| *y = i);
410 assert_eq!(b, a); 439 assert_eq!(b, a);
411 } 440 }
412 } 441 }

Mercurial Repository: alg_tools

mercurial