acme_tensor/actions/iter/
mod.rs

1/*
2    Appellation: iter <mod>
3    Contrib: FL03 <jo3mccain@icloud.com>
4*/
5//! # Iter
6//!
7//!
8pub use self::{indexed::*, iterator::*, layout::*, utils::*};
9
10#[allow(dead_code, unused)]
11pub(crate) mod axis;
12pub(crate) mod indexed;
13pub(crate) mod iterator;
14pub(crate) mod layout;
15
16pub(crate) mod utils {
17    use core::ptr;
18
19    pub(crate) fn zip<I, J>(i: I, j: J) -> core::iter::Zip<I::IntoIter, J::IntoIter>
20    where
21        I: IntoIterator,
22        J: IntoIterator,
23    {
24        i.into_iter().zip(j)
25    }
26
27    pub fn to_vec_mapped<I, F, B>(iter: I, mut f: F) -> Vec<B>
28    where
29        I: ExactSizeIterator, // + TrustedIterator
30        F: FnMut(I::Item) -> B,
31    {
32        // Use an `unsafe` block to do this efficiently.
33        // We know that iter will produce exactly .size() elements,
34        // and the loop can vectorize if it's clean (without branch to grow the vector).
35        let (size, _) = iter.size_hint();
36        let mut result = Vec::with_capacity(size);
37        let mut out_ptr = result.as_mut_ptr();
38        let mut len = 0;
39        iter.fold((), |(), elt| unsafe {
40            ptr::write(out_ptr, f(elt));
41            len += 1;
42            result.set_len(len);
43            out_ptr = out_ptr.offset(1);
44        });
45        debug_assert_eq!(size, result.len());
46        result
47    }
48}
49
50#[cfg(test)]
51mod tests {
52    use super::*;
53
54    use crate::shape::Layout;
55
56    #[test]
57    fn test_to_vec_mapped() {
58        let v = [1, 2, 3, 4, 5];
59        let f = |x| x * 2;
60        let res = to_vec_mapped(v.iter(), f);
61        assert_eq!(res, vec![2, 4, 6, 8, 10]);
62    }
63
64    #[test]
65    fn test_position() {
66        let shape = (2, 2);
67        let layout = Layout::contiguous(shape);
68        let position = Position::new(1, vec![0, 1]);
69        assert_eq!(position.index(), 1);
70        assert_eq!(
71            position.next(&layout).unwrap(),
72            Position::new(2, vec![1, 0])
73        );
74    }
75}
acme_tensor/actions/iter/mod.rs

acme_tensor/actions/iter/
mod.rs
