1use super::{Digest, Felt, StarkField, DIGEST_SIZE, ZERO};
  2use crate::utils::{
  3    string::String, ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable,
  4};
  5use core::{cmp::Ordering, fmt::Display, ops::Deref};
  6
  7// DIGEST TRAIT IMPLEMENTATIONS
  8// ================================================================================================
  9
 10#[derive(Debug, Default, Copy, Clone, Eq, PartialEq)]
 11pub struct RpoDigest([Felt; DIGEST_SIZE]);
 12
 13impl RpoDigest {
 14    pub const fn new(value: [Felt; DIGEST_SIZE]) -> Self {
 15        Self(value)
 16    }
 17
 18    pub fn as_elements(&self) -> &[Felt] {
 19        self.as_ref()
 20    }
 21
 22    pub fn as_bytes(&self) -> [u8; 32] {
 23        <Self as Digest>::as_bytes(self)
 24    }
 25
 26    pub fn digests_as_elements<'a, I>(digests: I) -> impl Iterator<Item = &'a Felt>
 27    where
 28        I: Iterator<Item = &'a Self>,
 29    {
 30        digests.flat_map(|d| d.0.iter())
 31    }
 32}
 33
 34impl Digest for RpoDigest {
 35    fn as_bytes(&self) -> [u8; 32] {
 36        let mut result = [0; 32];
 37
 38        result[..8].copy_from_slice(&self.0[0].as_int().to_le_bytes());
 39        result[8..16].copy_from_slice(&self.0[1].as_int().to_le_bytes());
 40        result[16..24].copy_from_slice(&self.0[2].as_int().to_le_bytes());
 41        result[24..].copy_from_slice(&self.0[3].as_int().to_le_bytes());
 42
 43        result
 44    }
 45}
 46
 47impl Serializable for RpoDigest {
 48    fn write_into<W: ByteWriter>(&self, target: &mut W) {
 49        target.write_bytes(&self.as_bytes());
 50    }
 51}
 52
 53impl Deserializable for RpoDigest {
 54    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
 55        let mut inner: [Felt; DIGEST_SIZE] = [ZERO; DIGEST_SIZE];
 56        for inner in inner.iter_mut() {
 57            let e = source.read_u64()?;
 58            if e >= Felt::MODULUS {
 59                return Err(DeserializationError::InvalidValue(String::from(
 60                    "Value not in the appropriate range",
 61                )));
 62            }
 63            *inner = Felt::new(e);
 64        }
 65
 66        Ok(Self(inner))
 67    }
 68}
 69
 70impl From<[Felt; DIGEST_SIZE]> for RpoDigest {
 71    fn from(value: [Felt; DIGEST_SIZE]) -> Self {
 72        Self(value)
 73    }
 74}
 75
 76impl From<&RpoDigest> for [Felt; DIGEST_SIZE] {
 77    fn from(value: &RpoDigest) -> Self {
 78        value.0
 79    }
 80}
 81
 82impl From<RpoDigest> for [Felt; DIGEST_SIZE] {
 83    fn from(value: RpoDigest) -> Self {
 84        value.0
 85    }
 86}
 87
 88impl From<&RpoDigest> for [u64; DIGEST_SIZE] {
 89    fn from(value: &RpoDigest) -> Self {
 90        [
 91            value.0[0].as_int(),
 92            value.0[1].as_int(),
 93            value.0[2].as_int(),
 94            value.0[3].as_int(),
 95        ]
 96    }
 97}
 98
 99impl From<RpoDigest> for [u64; DIGEST_SIZE] {
100    fn from(value: RpoDigest) -> Self {
101        [
102            value.0[0].as_int(),
103            value.0[1].as_int(),
104            value.0[2].as_int(),
105            value.0[3].as_int(),
106        ]
107    }
108}
109
110impl From<&RpoDigest> for [u8; 32] {
111    fn from(value: &RpoDigest) -> Self {
112        value.as_bytes()
113    }
114}
115
116impl From<RpoDigest> for [u8; 32] {
117    fn from(value: RpoDigest) -> Self {
118        value.as_bytes()
119    }
120}
121
122impl Deref for RpoDigest {
123    type Target = [Felt; DIGEST_SIZE];
124
125    fn deref(&self) -> &Self::Target {
126        &self.0
127    }
128}
129
130impl Ord for RpoDigest {
131    fn cmp(&self, other: &Self) -> Ordering {
132        // compare the inner u64 of both elements.
133        //
134        // it will iterate the elements and will return the first computation different than
135        // `Equal`. Otherwise, the ordering is equal.
136        //
137        // the endianness is irrelevant here because since, this being a cryptographically secure
138        // hash computation, the digest shouldn't have any ordered property of its input.
139        //
140        // finally, we use `Felt::inner` instead of `Felt::as_int` so we avoid performing a
141        // montgomery reduction for every limb. that is safe because every inner element of the
142        // digest is guaranteed to be in its canonical form (that is, `x in [0,p)`).
143        self.0.iter().map(Felt::inner).zip(other.0.iter().map(Felt::inner)).fold(
144            Ordering::Equal,
145            |ord, (a, b)| match ord {
146                Ordering::Equal => a.cmp(&b),
147                _ => ord,
148            },
149        )
150    }
151}
152
153impl PartialOrd for RpoDigest {
154    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
155        Some(self.cmp(other))
156    }
157}
158
159impl Display for RpoDigest {
160    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
161        for byte in self.as_bytes() {
162            write!(f, "{byte:02x}")?;
163        }
164        Ok(())
165    }
166}
167
168// TESTS
169// ================================================================================================
170
171#[cfg(test)]
172mod tests {
173
174    use super::{Deserializable, Felt, RpoDigest, Serializable};
175    use crate::utils::SliceReader;
176    use rand_utils::rand_value;
177
178    #[test]
179    fn digest_serialization() {
180        let e1 = Felt::new(rand_value());
181        let e2 = Felt::new(rand_value());
182        let e3 = Felt::new(rand_value());
183        let e4 = Felt::new(rand_value());
184
185        let d1 = RpoDigest([e1, e2, e3, e4]);
186
187        let mut bytes = vec![];
188        d1.write_into(&mut bytes);
189        assert_eq!(32, bytes.len());
190
191        let mut reader = SliceReader::new(&bytes);
192        let d2 = RpoDigest::read_from(&mut reader).unwrap();
193
194        assert_eq!(d1, d2);
195    }
196}