acme_tensor/actions/create/
arange.rs

1/*
2    Appellation: arange <mod>
3    Contrib: FL03 <jo3mccain@icloud.com>
4*/
5use super::utils::steps;
6use num::traits::{FromPrimitive, Num, ToPrimitive};
7
8pub struct Arange<T> {
9    scope: usize,
10    start: T,
11    stop: T,
12    step: T,
13}
14
15impl<T> Arange<T> {
16    pub fn new(start: T, stop: T, step: T) -> Self {
17        Self {
18            scope: 0,
19            start,
20            stop,
21            step,
22        }
23    }
24
25    pub fn start(&self) -> &T {
26        &self.start
27    }
28
29    pub fn stop(&self) -> &T {
30        &self.stop
31    }
32
33    pub fn step(&self) -> &T {
34        &self.step
35    }
36
37    pub fn steps(&self) -> usize
38    where
39        T: Copy + Num + ToPrimitive,
40    {
41        steps(self.start, self.stop, self.step)
42    }
43}
44
45impl<T> Iterator for Arange<T>
46where
47    T: Copy + FromPrimitive + Num + ToPrimitive,
48{
49    type Item = T;
50
51    fn next(&mut self) -> Option<Self::Item> {
52        if self.scope < self.steps() {
53            let value = self.start + self.step * T::from_usize(self.scope).unwrap();
54            self.scope += 1;
55            Some(value)
56        } else {
57            None
58        }
59    }
60}
61
62#[cfg(test)]
63mod tests {
64    use super::*;
65
66    #[test]
67    fn test_arange() {
68        let mut arange = Arange::new(0, 10, 2);
69        assert_eq!(arange.next(), Some(0));
70        assert_eq!(arange.next(), Some(2));
71        assert_eq!(arange.next(), Some(4));
72        assert_eq!(arange.next(), Some(6));
73        assert_eq!(arange.next(), Some(8));
74        assert_eq!(arange.next(), None);
75    }
76}
acme_tensor/actions/create/arange.rs

acme_tensor/actions/create/
arange.rs
