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
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139

use std::marker::PhantomData;

use crate::Normalize;

use super::ReadWeights;

/// Casting iterator for `Weights`.
#[derive(Clone, Debug)]
pub struct CastingIter<'a, T>(ReadWeights<'a>, PhantomData<T>);

/// Type which describes how to cast any weight into u8.
#[derive(Clone, Debug)]
pub struct U8;

/// Type which describes how to cast any weight into u16.
#[derive(Clone, Debug)]
pub struct U16;

/// Type which describes how to cast any weight into f32.
#[derive(Clone, Debug)]
pub struct F32;

/// Trait for types which describe casting behaviour.
pub trait Cast {
    /// Output type.
    type Output;

    /// Cast from u8.
    fn cast_u8(x: [u8; 4]) -> Self::Output;

    /// Cast from u16.
    fn cast_u16(x: [u16; 4]) -> Self::Output;

    /// Cast from f32.
    fn cast_f32(x: [f32; 4]) -> Self::Output;
}

impl<'a, A> CastingIter<'a, A> {
    pub(crate) fn new(iter: ReadWeights<'a>) -> Self {
        CastingIter(iter, PhantomData)
    }

    /// Unwrap underlying `Weights` object.
    pub fn unwrap(self) -> ReadWeights<'a> {
        self.0
    }
}

impl<'a, A: Cast> ExactSizeIterator for CastingIter<'a, A> {}
impl<'a, A: Cast> Iterator for CastingIter<'a, A> {
    type Item = A::Output;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        match self.0 {
            ReadWeights::U8(ref mut i) => i.next().map(A::cast_u8),
            ReadWeights::U16(ref mut i) => i.next().map(A::cast_u16),
            ReadWeights::F32(ref mut i) => i.next().map(A::cast_f32),
        }
    }

    #[inline]
    fn nth(&mut self, x: usize) -> Option<Self::Item> {
        match self.0 {
            ReadWeights::U8(ref mut i) => i.nth(x).map(A::cast_u8),
            ReadWeights::U16(ref mut i) => i.nth(x).map(A::cast_u16),
            ReadWeights::F32(ref mut i) => i.nth(x).map(A::cast_f32),
        }
    }

    fn last(self) -> Option<Self::Item> {
        match self.0 {
            ReadWeights::U8(i) => i.last().map(A::cast_u8),
            ReadWeights::U16(i) => i.last().map(A::cast_u16),
            ReadWeights::F32(i) => i.last().map(A::cast_f32),
        }
    }

    fn count(self) -> usize {
        self.size_hint().0
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        match self.0 {
            ReadWeights::U8(ref i) => i.size_hint(),
            ReadWeights::U16(ref i) => i.size_hint(),
            ReadWeights::F32(ref i) => i.size_hint(),
        }
    }
}

impl Cast for U8 {
    type Output = [u8; 4];

    fn cast_u8(x: [u8; 4]) -> Self::Output {
        x.normalize()
    }

    fn cast_u16(x: [u16; 4]) -> Self::Output {
        x.normalize()
    }

    fn cast_f32(x: [f32; 4]) -> Self::Output {
        x.normalize()
    }
}

impl Cast for U16 {
    type Output = [u16; 4];

    fn cast_u8(x: [u8; 4]) -> Self::Output {
        x.normalize()
    }

    fn cast_u16(x: [u16; 4]) -> Self::Output {
        x.normalize()
    }

    fn cast_f32(x: [f32; 4]) -> Self::Output {
        x.normalize()
    }
}

impl Cast for F32 {
    type Output = [f32; 4];

    fn cast_u8(x: [u8; 4]) -> Self::Output {
        x.normalize()
    }

    fn cast_u16(x: [u16; 4]) -> Self::Output {
        x.normalize()
    }

    fn cast_f32(x: [f32; 4]) -> Self::Output {
        x.normalize()
    }
}