From 2ad8495bc22169c08b9da189c4e2f27186bced1e Mon Sep 17 00:00:00 2001
From: Al-Kindi-0 <82364884+Al-Kindi-0@users.noreply.github.com>
Date: Thu, 20 Oct 2022 08:35:09 +0200
Subject: [PATCH 1/9] feat:RPO implementation

---
 crypto/src/hash/rpo/digest.rs   | 112 ++++
 crypto/src/hash/rpo/mds_freq.rs | 195 +++++++
 crypto/src/hash/rpo/mod.rs      | 912 ++++++++++++++++++++++++++++++++
 crypto/src/hash/rpo/tests.rs    | 161 ++++++
 4 files changed, 1380 insertions(+)
 create mode 100644 crypto/src/hash/rpo/digest.rs
 create mode 100644 crypto/src/hash/rpo/mds_freq.rs
 create mode 100644 crypto/src/hash/rpo/mod.rs
 create mode 100644 crypto/src/hash/rpo/tests.rs

diff --git a/crypto/src/hash/rpo/digest.rs b/crypto/src/hash/rpo/digest.rs
new file mode 100644
index 0000000..1e07bb3
--- /dev/null
+++ b/crypto/src/hash/rpo/digest.rs
@@ -0,0 +1,112 @@
+use super::DIGEST_SIZE;
+use core::slice;
+use winterfell::crypto::Digest;
+use winterfell::math::{fields::f64::BaseElement, StarkField};
+use winterfell::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable};
+
+// DIGEST TRAIT IMPLEMENTATIONS
+// ================================================================================================
+
+#[derive(Debug, Copy, Clone, Eq, PartialEq)]
+pub struct ElementDigest([BaseElement; DIGEST_SIZE]);
+
+impl ElementDigest {
+    pub fn new(value: [BaseElement; DIGEST_SIZE]) -> Self {
+        Self(value)
+    }
+
+    pub fn as_elements(&self) -> &[BaseElement] {
+        &self.0
+    }
+
+    pub fn digests_as_elements(digests: &[Self]) -> &[BaseElement] {
+        let p = digests.as_ptr();
+        let len = digests.len() * DIGEST_SIZE;
+        unsafe { slice::from_raw_parts(p as *const BaseElement, len) }
+    }
+}
+
+impl Digest for ElementDigest {
+    fn as_bytes(&self) -> [u8; 32] {
+        let mut result = [0; 32];
+
+        result[..8].copy_from_slice(&self.0[0].as_int().to_le_bytes());
+        result[8..16].copy_from_slice(&self.0[1].as_int().to_le_bytes());
+        result[16..24].copy_from_slice(&self.0[2].as_int().to_le_bytes());
+        result[24..].copy_from_slice(&self.0[3].as_int().to_le_bytes());
+
+        result
+    }
+}
+
+impl Default for ElementDigest {
+    fn default() -> Self {
+        ElementDigest([BaseElement::default(); DIGEST_SIZE])
+    }
+}
+
+impl Serializable for ElementDigest {
+    fn write_into<W: ByteWriter>(&self, target: &mut W) {
+        target.write_u8_slice(&self.as_bytes());
+    }
+}
+
+impl Deserializable for ElementDigest {
+    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
+        // TODO: check if the field elements are valid?
+        let e1 = BaseElement::new(source.read_u64()?);
+        let e2 = BaseElement::new(source.read_u64()?);
+        let e3 = BaseElement::new(source.read_u64()?);
+        let e4 = BaseElement::new(source.read_u64()?);
+
+        Ok(Self([e1, e2, e3, e4]))
+    }
+}
+
+impl From<[BaseElement; DIGEST_SIZE]> for ElementDigest {
+    fn from(value: [BaseElement; DIGEST_SIZE]) -> Self {
+        Self(value)
+    }
+}
+
+impl From<ElementDigest> for [BaseElement; DIGEST_SIZE] {
+    fn from(value: ElementDigest) -> Self {
+        value.0
+    }
+}
+
+impl From<ElementDigest> for [u8; 32] {
+    fn from(value: ElementDigest) -> Self {
+        value.as_bytes()
+    }
+}
+
+// TESTS
+// ================================================================================================
+
+#[cfg(test)]
+mod tests {
+
+    use super::ElementDigest;
+    use rand_utils::rand_value;
+    use winterfell::{math::fields::f64::BaseElement, Deserializable, Serializable, SliceReader};
+
+    #[test]
+    fn digest_serialization() {
+        let e1 = BaseElement::new(rand_value());
+        let e2 = BaseElement::new(rand_value());
+        let e3 = BaseElement::new(rand_value());
+        let e4 = BaseElement::new(rand_value());
+
+        let d1 = ElementDigest([e1, e2, e3, e4]);
+
+        let mut bytes = vec![];
+        d1.write_into(&mut bytes);
+        assert_eq!(32, bytes.len());
+
+        let mut reader = SliceReader::new(&bytes);
+        let d2 = ElementDigest::read_from(&mut reader).unwrap();
+
+        assert_eq!(d1, d2);
+    }
+}
diff --git a/crypto/src/hash/rpo/mds_freq.rs b/crypto/src/hash/rpo/mds_freq.rs
new file mode 100644
index 0000000..8778c69
--- /dev/null
+++ b/crypto/src/hash/rpo/mds_freq.rs
@@ -0,0 +1,195 @@
+// FFT-BASED MDS MULTIPLICATION HELPER FUNCTIONS
+// ================================================================================================
+
+/// This module contains helper functions as well as constants used to perform the vector-matrix
+/// multiplication step of the Rescue prime permutation. The special form of our MDS matrix
+/// i.e. being circular, allows us to reduce the vector-matrix multiplication to a Hadamard product
+/// of two vectors in "frequency domain". This follows from the simple fact that every circulant
+/// matrix has the columns of the discrete Fourier transform matrix as orthogonal eigenvectors.
+/// The implementation also avoids the use of 3-point FFTs, and 3-point iFFTs, and substitutes that
+/// with explicit expressions. It also avoids, due to the form of our matrix in the frequency domain,
+/// divisions by 2 and repeated modular reductions. This is because of our explicit choice of
+/// an MDS matrix that has small powers of 2 entries in frequency domain.
+/// The following implementation has benefited greatly from the discussions and insights of
+/// Hamish Ivey-Law and Jacqueline Nabaglo of Polygon Zero.
+
+// Rescue MDS matrix in frequency domain.
+// More precisely, this is the output of the three 4-point (real) FFTs of the first column of
+// the MDS matrix i.e. just before the multiplication with the appropriate twiddle factors
+// and application of the final four 3-point FFT in order to get the full 12-point FFT.
+// The entries have been scaled appropriately in order to avoid divisions by 2 in iFFT2 and iFFT4.
+// The code to generate the matrix in frequency domain is based on an adaptation of a code, to generate
+// MDS matrices efficiently in original domain, that was developed by the Polygon Zero team.
+const MDS_FREQ_BLOCK_ONE: [i64; 3] = [16, 8, 16];
+const MDS_FREQ_BLOCK_TWO: [(i64, i64); 3] = [(-1, 2), (-1, 1), (4, 8)];
+const MDS_FREQ_BLOCK_THREE: [i64; 3] = [-8, 1, 1];
+
+// We use split 3 x 4 FFT transform in order to transform our vectors into the frequency domain.
+#[inline(always)]
+pub(crate) fn mds_multiply_freq(state: [u64; 12]) -> [u64; 12] {
+    let [s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11] = state;
+
+    let (u0, u1, u2) = fft4_real([s0, s3, s6, s9]);
+    let (u4, u5, u6) = fft4_real([s1, s4, s7, s10]);
+    let (u8, u9, u10) = fft4_real([s2, s5, s8, s11]);
+
+    // This where the multiplication in frequency domain is done. More precisely, and with
+    // the appropriate permuations in between, the sequence of
+    // 3-point FFTs --> multiplication by twiddle factors --> Hadamard multiplication -->
+    // 3 point iFFTs --> multiplication by (inverse) twiddle factors
+    // is "squashed" into one step composed of the functions "block1", "block2" and "block3".
+    // The expressions in the aformentioned functions are the result of explicit computations
+    // combined with the Karatsuba trick for the multiplication of Complex numbers.
+
+    let [v0, v4, v8] = block1([u0, u4, u8], MDS_FREQ_BLOCK_ONE);
+    let [v1, v5, v9] = block2([u1, u5, u9], MDS_FREQ_BLOCK_TWO);
+    let [v2, v6, v10] = block3([u2, u6, u10], MDS_FREQ_BLOCK_THREE);
+    // The 4th block is not computed as it is similar to the 2nd one, up to complex conjugation,
+    // and is, due to the use of the real FFT and iFFT, redundant.
+
+    let [s0, s3, s6, s9] = ifft4_real((v0, v1, v2));
+    let [s1, s4, s7, s10] = ifft4_real((v4, v5, v6));
+    let [s2, s5, s8, s11] = ifft4_real((v8, v9, v10));
+
+    [s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11]
+}
+
+// We use the real FFT to avoid redundant computations. See https://www.mdpi.com/2076-3417/12/9/4700
+#[inline(always)]
+fn fft2_real(x: [u64; 2]) -> [i64; 2] {
+    [(x[0] as i64 + x[1] as i64), (x[0] as i64 - x[1] as i64)]
+}
+
+#[inline(always)]
+fn ifft2_real(y: [i64; 2]) -> [u64; 2] {
+    // We avoid divisions by 2 by appropriately scaling the MDS matrix constants.
+    [(y[0] + y[1]) as u64, (y[0] - y[1]) as u64]
+}
+
+#[inline(always)]
+fn fft4_real(x: [u64; 4]) -> (i64, (i64, i64), i64) {
+    let [z0, z2] = fft2_real([x[0], x[2]]);
+    let [z1, z3] = fft2_real([x[1], x[3]]);
+    let y0 = z0 + z1;
+    let y1 = (z2, -z3);
+    let y2 = z0 - z1;
+    (y0, y1, y2)
+}
+
+#[inline(always)]
+fn ifft4_real(y: (i64, (i64, i64), i64)) -> [u64; 4] {
+    // In calculating 'z0' and 'z1', division by 2 is avoided by appropriately scaling
+    // the MDS matrix constants.
+    let z0 = y.0 + y.2;
+    let z1 = y.0 - y.2;
+    let z2 = y.1 .0;
+    let z3 = -y.1 .1;
+
+    let [x0, x2] = ifft2_real([z0, z2]);
+    let [x1, x3] = ifft2_real([z1, z3]);
+
+    [x0, x1, x2, x3]
+}
+
+#[inline(always)]
+fn block1(x: [i64; 3], y: [i64; 3]) -> [i64; 3] {
+    let [x0, x1, x2] = x;
+    let [y0, y1, y2] = y;
+    let z0 = x0 * y0 + x1 * y2 + x2 * y1;
+    let z1 = x0 * y1 + x1 * y0 + x2 * y2;
+    let z2 = x0 * y2 + x1 * y1 + x2 * y0;
+
+    [z0, z1, z2]
+}
+
+#[inline(always)]
+fn block2(x: [(i64, i64); 3], y: [(i64, i64); 3]) -> [(i64, i64); 3] {
+    let [(x0r, x0i), (x1r, x1i), (x2r, x2i)] = x;
+    let [(y0r, y0i), (y1r, y1i), (y2r, y2i)] = y;
+    let x0s = x0r + x0i;
+    let x1s = x1r + x1i;
+    let x2s = x2r + x2i;
+    let y0s = y0r + y0i;
+    let y1s = y1r + y1i;
+    let y2s = y2r + y2i;
+
+    // Compute x0​y0 ​− ix1​y2​ − ix2​y1​ using Karatsuba for complex numbers multiplication
+    let m0 = (x0r * y0r, x0i * y0i);
+    let m1 = (x1r * y2r, x1i * y2i);
+    let m2 = (x2r * y1r, x2i * y1i);
+    let z0r = (m0.0 - m0.1) + (x1s * y2s - m1.0 - m1.1) + (x2s * y1s - m2.0 - m2.1);
+    let z0i = (x0s * y0s - m0.0 - m0.1) + (-m1.0 + m1.1) + (-m2.0 + m2.1);
+    let z0 = (z0r, z0i);
+
+    // Compute x0​y1​ + x1​y0​ − ix2​y2 using Karatsuba for complex numbers multiplication
+    let m0 = (x0r * y1r, x0i * y1i);
+    let m1 = (x1r * y0r, x1i * y0i);
+    let m2 = (x2r * y2r, x2i * y2i);
+    let z1r = (m0.0 - m0.1) + (m1.0 - m1.1) + (x2s * y2s - m2.0 - m2.1);
+    let z1i = (x0s * y1s - m0.0 - m0.1) + (x1s * y0s - m1.0 - m1.1) + (-m2.0 + m2.1);
+    let z1 = (z1r, z1i);
+
+    // Compute x0​y2​ + x1​y1 ​+ x2​y0​ using Karatsuba for complex numbers multiplication
+    let m0 = (x0r * y2r, x0i * y2i);
+    let m1 = (x1r * y1r, x1i * y1i);
+    let m2 = (x2r * y0r, x2i * y0i);
+    let z2r = (m0.0 - m0.1) + (m1.0 - m1.1) + (m2.0 - m2.1);
+    let z2i = (x0s * y2s - m0.0 - m0.1) + (x1s * y1s - m1.0 - m1.1) + (x2s * y0s - m2.0 - m2.1);
+    let z2 = (z2r, z2i);
+
+    [z0, z1, z2]
+}
+
+#[inline(always)]
+fn block3(x: [i64; 3], y: [i64; 3]) -> [i64; 3] {
+    let [x0, x1, x2] = x;
+    let [y0, y1, y2] = y;
+    let z0 = x0 * y0 - x1 * y2 - x2 * y1;
+    let z1 = x0 * y1 + x1 * y0 - x2 * y2;
+    let z2 = x0 * y2 + x1 * y1 + x2 * y0;
+
+    [z0, z1, z2]
+}
+
+// TESTS
+// ================================================================================================
+
+#[cfg(test)]
+mod tests {
+    use super::super::Rpo;
+    use crate::hash::rpo::MDS;
+    use proptest::prelude::*;
+    use winterfell::math::{fields::f64::BaseElement, FieldElement};
+
+    const STATE_WIDTH: usize = 12;
+
+    #[inline(always)]
+    fn apply_mds_naive(state: &mut [BaseElement; STATE_WIDTH]) {
+        let mut result = [BaseElement::ZERO; STATE_WIDTH];
+        result.iter_mut().zip(MDS).for_each(|(r, mds_row)| {
+            state.iter().zip(mds_row).for_each(|(&s, m)| {
+                *r += m * s;
+            });
+        });
+        *state = result;
+    }
+
+    proptest! {
+        #[test]
+        fn mds_freq_proptest(a in any::<[u64;STATE_WIDTH]>()) {
+
+            let mut v1 = [BaseElement::ZERO;STATE_WIDTH];
+            let mut v2;
+
+            for i in 0..STATE_WIDTH {
+                v1[i] = BaseElement::new(a[i]);
+            }
+            v2 = v1.clone();
+
+            apply_mds_naive(&mut v1);
+            Rpo::apply_mds(&mut v2);
+
+            prop_assert_eq!(v1, v2);
+        }
+    }
+}
diff --git a/crypto/src/hash/rpo/mod.rs b/crypto/src/hash/rpo/mod.rs
new file mode 100644
index 0000000..e872974
--- /dev/null
+++ b/crypto/src/hash/rpo/mod.rs
@@ -0,0 +1,912 @@
+use super::{exp_acc, ElementHasher, HashFn};
+use core::{convert::TryInto, ops::Range};
+use winterfell::math::{fields::f64::BaseElement, FieldElement, StarkField};
+
+mod digest;
+pub use digest::ElementDigest;
+
+#[cfg(test)]
+mod tests;
+
+mod mds_freq;
+use mds_freq::mds_multiply_freq;
+
+// CONSTANTS
+// ================================================================================================
+
+/// Sponge state is set to 12 field elements or 768 bytes; 8 elements are reserved for rate and
+/// the remaining 4 elements are reserved for capacity.
+const STATE_WIDTH: usize = 12;
+
+/// The rate portion of the state is located in elements 4 through 11.
+const RATE_RANGE: Range<usize> = 4..12;
+const RATE_WIDTH: usize = RATE_RANGE.end - RATE_RANGE.start;
+
+const INPUT1_RANGE: Range<usize> = 4..8;
+const INPUT2_RANGE: Range<usize> = 8..12;
+
+/// The capacity portion of the state is located in elements 0, 1, 2, and 3.
+const CAPACITY_RANGE: Range<usize> = 0..4;
+
+/// The output of the hash function is a digest which consists of 4 field elements or 32 bytes.
+///
+/// The digest is returned from state elements 4, 5, 6, and 7 (the first four elements of the
+/// rate portion).
+const DIGEST_RANGE: Range<usize> = 4..8;
+const DIGEST_SIZE: usize = DIGEST_RANGE.end - DIGEST_RANGE.start;
+
+/// The number of rounds is set to 7 to target 128-bit security level
+const NUM_ROUNDS: usize = 7;
+
+/// S-Box and Inverse S-Box powers;
+///
+/// The constants are defined for tests only because the exponentiations in the code are unrolled
+/// for efficiency reasons.
+#[cfg(test)]
+const ALPHA: u64 = 7;
+#[cfg(test)]
+const INV_ALPHA: u64 = 10540996611094048183;
+
+// HASHER IMPLEMENTATION
+// ================================================================================================
+
+/// Implementation of [Hasher] trait for Rescue Prime Optimized (RPO) hash function with 256-bit output.
+///
+/// The hash function is implemented according to the Rescue Prime Optimized
+/// [specifications](https://github.org/aszepieniec/rpo/)
+///
+/// The parameters used to instantiate the function are:
+/// * Field: 64-bit prime field with modulus 2^64 - 2^32 + 1.
+/// * State width: 12 field elements.
+/// * Capacity size: 4 field elements.
+/// * Number of founds: 7.
+/// * S-Box degree: 7.
+///
+/// The above parameters target 128-bit security level. The digest consists of four field elements
+/// and it can be serialized into 32 bytes (256 bits).
+///
+/// ## Hash output consistency
+/// Functions [hash_elements()](Rpo::hash_elements), [merge()](Rpo::merge), and
+/// [merge_with_int()](Rpo::merge_with_int) are internally consistent. That is, computing
+/// a hash for the same set of elements using these functions will always produce the same
+/// result. For example, merging two digests using [merge()](Rpo::merge) will produce the
+/// same result as hashing 8 elements which make up these digests using
+/// [hash_elements()](Rpo::hash_elements) function.
+///
+/// However, [hash()](Rpo::hash) function is not consistent with functions mentioned above.
+/// For example, if we take two field elements, serialize them to bytes and hash them using
+/// [hash()](Rpo::hash), the result will differ from the result obtained by hashing these
+/// elements directly using [hash_elements()](Rpo::hash_elements) function. The reason for
+/// this difference is that [hash()](Rpo::hash) function needs to be able to handle
+/// arbitrary binary strings, which may or may not encode valid field elements - and thus,
+/// deserialization procedure used by this function is different from the procedure used to
+/// deserialize valid field elements.
+///
+/// Thus, if the underlying data consists of valid field elements, it might make more sense
+/// to deserialize them into field elements and then hash them using
+/// [hash_elements()](Rpo::hash_elements) function rather then hashing the serialized bytes
+/// using [hash()](Rpo::hash) function.
+pub struct Rpo();
+
+impl HashFn for Rpo {
+    type Digest = ElementDigest;
+
+    fn hash(bytes: &[u8]) -> Self::Digest {
+        // compute the number of elements required to represent the string; we will be processing
+        // the string in 7-byte chunks, thus the number of elements will be equal to the number
+        // of such chunks (including a potential partial chunk at the end).
+        let num_elements = if bytes.len() % 7 == 0 {
+            bytes.len() / 7
+        } else {
+            bytes.len() / 7 + 1
+        };
+
+        // initialize state to all zeros, except for the first element of the capacity part, which
+        // is set to the number of elements to be hashed. this is done so that adding zero elements
+        // at the end of the list always results in a different hash.
+        let mut state = [BaseElement::ZERO; STATE_WIDTH];
+        state[CAPACITY_RANGE.start] = BaseElement::new(num_elements as u64);
+
+        // break the string into 7-byte chunks, convert each chunk into a field element, and
+        // absorb the element into the rate portion of the state. we use 7-byte chunks because
+        // every 7-byte chunk is guaranteed to map to some field element.
+        let mut i = 0;
+        let mut buf = [0_u8; 8];
+        for chunk in bytes.chunks(7) {
+            if i < num_elements - 1 {
+                buf[..7].copy_from_slice(chunk);
+            } else {
+                // if we are dealing with the last chunk, it may be smaller than 7 bytes long, so
+                // we need to handle it slightly differently. we also append a byte with value 1
+                // to the end of the string; this pads the string in such a way that adding
+                // trailing zeros results in different hash
+                let chunk_len = chunk.len();
+                buf = [0_u8; 8];
+                buf[..chunk_len].copy_from_slice(chunk);
+                buf[chunk_len] = 1;
+            }
+
+            // convert the bytes into a field element and absorb it into the rate portion of the
+            // state; if the rate is filled up, apply the Rescue permutation and start absorbing
+            // again from zero index.
+            state[RATE_RANGE.start + i] += BaseElement::new(u64::from_le_bytes(buf));
+            i += 1;
+            if i % RATE_WIDTH == 0 {
+                Self::apply_permutation(&mut state);
+                i = 0;
+            }
+        }
+
+        // if we absorbed some elements but didn't apply a permutation to them (would happen when
+        // the number of elements is not a multiple of RATE_WIDTH), apply the Rescue permutation.
+        // we don't need to apply any extra padding because we injected total number of elements
+        // in the input list into the capacity portion of the state during initialization.
+        if i > 0 {
+            Self::apply_permutation(&mut state);
+        }
+
+        // return the first 4 elements of the state as hash result
+        ElementDigest::new(state[DIGEST_RANGE].try_into().unwrap())
+    }
+
+    fn merge(values: &[Self::Digest; 2]) -> Self::Digest {
+        // initialize the state by copying the digest elements into the rate portion of the state
+        // (8 total elements), and set the capacity elements to 0.
+        let mut state = [BaseElement::ZERO; STATE_WIDTH];
+        state[RATE_RANGE].copy_from_slice(Self::Digest::digests_as_elements(values));
+
+        // apply the RPO permutation and return the first four elements of the state
+        Self::apply_permutation(&mut state);
+        ElementDigest::new(state[DIGEST_RANGE].try_into().unwrap())
+    }
+
+    fn merge_with_int(seed: Self::Digest, value: u64) -> Self::Digest {
+        // initialize the state as follows:
+        // - seed is copied into the first 4 elements of the rate portion of the state.
+        // - if the value fits into a single field element, copy it into the fifth rate element
+        //   and set the sixth rate element and first capacity element to 1.
+        // - if the value doesn't fit into a single field element, split it into two field
+        //   elements, copy them into rate elements 5 and 6, and set the seventh rate element
+        //   and first capacity element to 1.
+        let mut state = [BaseElement::ZERO; STATE_WIDTH];
+        state[INPUT1_RANGE].copy_from_slice(seed.as_elements());
+        state[INPUT2_RANGE.start] = BaseElement::new(value);
+        if value < BaseElement::MODULUS {
+            state[INPUT2_RANGE.start + 1] = BaseElement::ONE;
+        } else {
+            state[INPUT2_RANGE.start + 1] = BaseElement::new(value / BaseElement::MODULUS);
+            state[INPUT2_RANGE.start + 2] = BaseElement::ONE;            
+        }
+
+        // common padding for both cases
+        state[CAPACITY_RANGE.start] = BaseElement::ONE;
+        
+        // apply the RPO permutation and return the first four elements of the state
+        Self::apply_permutation(&mut state);
+        ElementDigest::new(state[DIGEST_RANGE].try_into().unwrap())
+    }
+}
+
+impl ElementHasher for Rpo {
+    type BaseField = BaseElement;
+
+    fn hash_elements<E: FieldElement<BaseField = Self::BaseField>>(elements: &[E]) -> Self::Digest {
+        // convert the elements into a list of base field elements
+        let elements = E::as_base_elements(elements);
+
+        // initialize state to all zeros, except for the first element of the capacity part, which
+        // is set to 1 if the number of elements is not a multiple of RATE_WIDTH.
+        let mut state = [BaseElement::ZERO; STATE_WIDTH];
+        if elements.len() % RATE_WIDTH != 0 {
+            state[CAPACITY_RANGE.start] = BaseElement::ONE;
+        }
+
+        // absorb elements into the state one by one until the rate portion of the state is filled
+        // up; then apply the Rescue permutation and start absorbing again; repeat until all
+        // elements have been absorbed
+        let mut i = 0;
+        for &element in elements.iter() {
+            state[RATE_RANGE.start + i] = element;
+            i += 1;
+            if i % RATE_WIDTH == 0 {
+                Self::apply_permutation(&mut state);
+                i = 0;
+            }
+        }
+
+        // if we absorbed some elements but didn't apply a permutation to them (would happen when
+        // the number of elements is not a multiple of RATE_WIDTH), apply the RPO permutation after
+        // padding by appending a 1 followed by as many 0 as necessary to make the input length a
+        // multiple of the RATE_WIDTH.
+        if i > 0 {
+            state[RATE_RANGE.start + i] = BaseElement::ONE;
+            i += 1;
+            while i != RATE_WIDTH {
+                state[RATE_RANGE.start + i] = BaseElement::ZERO;
+                i += 1;
+            }
+            Self::apply_permutation(&mut state);
+        }
+
+        // return the first 4 elements of the state as hash result
+        ElementDigest::new(state[DIGEST_RANGE].try_into().unwrap())
+    }
+}
+
+// HASH FUNCTION IMPLEMENTATION
+// ================================================================================================
+
+impl Rpo {
+    // RESCUE PERMUTATION
+    // --------------------------------------------------------------------------------------------
+
+    /// Applies RPO permutation to the provided state.
+    pub fn apply_permutation(state: &mut [BaseElement; STATE_WIDTH]) {
+        for i in 0..NUM_ROUNDS {
+            Self::apply_round(state, i);
+        }
+    }
+
+    /// RPO round function.
+    #[inline(always)]
+    pub fn apply_round(state: &mut [BaseElement; STATE_WIDTH], round: usize) {
+        // apply first half of RPO round
+        Self::apply_mds(state);
+        Self::add_constants(state, &ARK1[round]);
+        Self::apply_sbox(state);
+
+        // apply second half of RPO round
+        Self::apply_mds(state);
+        Self::add_constants(state, &ARK2[round]);
+        Self::apply_inv_sbox(state);
+    }
+
+    // HELPER FUNCTIONS
+    // --------------------------------------------------------------------------------------------
+
+    #[inline(always)]
+    fn apply_mds(state: &mut [BaseElement; STATE_WIDTH]) {
+        let mut result = [BaseElement::ZERO; STATE_WIDTH];
+
+        // Using the linearity of the operations we can split the state into a low||high decomposition
+        // and operate on each with no overflow and then combine/reduce the result to a field element.
+        // The no overflow is guaranteed by the fact that the MDS matrix is a small powers of two in
+        // frequency domain.
+        let mut state_l = [0u64; STATE_WIDTH];
+        let mut state_h = [0u64; STATE_WIDTH];
+
+        for r in 0..STATE_WIDTH {
+            let s = state[r].inner();
+            state_h[r] = s >> 32;
+            state_l[r] = (s as u32) as u64;
+        }
+
+        let state_h = mds_multiply_freq(state_h);
+        let state_l = mds_multiply_freq(state_l);
+
+        for r in 0..STATE_WIDTH {
+            let s = state_l[r] as u128 + ((state_h[r] as u128) << 32);
+            let s_hi = (s >> 64) as u64;
+            let s_lo = s as u64;
+            let z = (s_hi << 32) - s_hi;
+            let (res, over) = s_lo.overflowing_add(z);
+
+            result[r] =
+                BaseElement::from_mont(res.wrapping_add(0u32.wrapping_sub(over as u32) as u64));
+        }
+        *state = result;
+    }
+
+    #[inline(always)]
+    fn add_constants(state: &mut [BaseElement; STATE_WIDTH], ark: &[BaseElement; STATE_WIDTH]) {
+        state.iter_mut().zip(ark).for_each(|(s, &k)| *s += k);
+    }
+
+    #[inline(always)]
+    fn apply_sbox(state: &mut [BaseElement; STATE_WIDTH]) {
+        state[0] = state[0].exp7();
+        state[1] = state[1].exp7();
+        state[2] = state[2].exp7();
+        state[3] = state[3].exp7();
+        state[4] = state[4].exp7();
+        state[5] = state[5].exp7();
+        state[6] = state[6].exp7();
+        state[7] = state[7].exp7();
+        state[8] = state[8].exp7();
+        state[9] = state[9].exp7();
+        state[10] = state[10].exp7();
+        state[11] = state[11].exp7();
+    }
+
+    #[inline(always)]
+    fn apply_inv_sbox(state: &mut [BaseElement; STATE_WIDTH]) {
+        // compute base^10540996611094048183 using 72 multiplications per array element
+        // 10540996611094048183 = b1001001001001001001001001001000110110110110110110110110110110111
+
+        // compute base^10
+        let mut t1 = *state;
+        t1.iter_mut().for_each(|t| *t = t.square());
+
+        // compute base^100
+        let mut t2 = t1;
+        t2.iter_mut().for_each(|t| *t = t.square());
+
+        // compute base^100100
+        let t3 = exp_acc::<BaseElement, STATE_WIDTH, 3>(t2, t2);
+
+        // compute base^100100100100
+        let t4 = exp_acc::<BaseElement, STATE_WIDTH, 6>(t3, t3);
+
+        // compute base^100100100100100100100100
+        let t5 = exp_acc::<BaseElement, STATE_WIDTH, 12>(t4, t4);
+
+        // compute base^100100100100100100100100100100
+        let t6 = exp_acc::<BaseElement, STATE_WIDTH, 6>(t5, t3);
+
+        // compute base^1001001001001001001001001001000100100100100100100100100100100
+        let t7 = exp_acc::<BaseElement, STATE_WIDTH, 31>(t6, t6);
+
+        // compute base^1001001001001001001001001001000110110110110110110110110110110111
+        for (i, s) in state.iter_mut().enumerate() {
+            let a = (t7[i].square() * t6[i]).square().square();
+            let b = t1[i] * t2[i] * *s;
+            *s = a * b;
+        }
+    }
+}
+
+// MDS
+// ================================================================================================
+/// RPO MDS matrix
+const MDS: [[BaseElement; STATE_WIDTH]; STATE_WIDTH] = [
+    [
+        BaseElement::new(7),
+        BaseElement::new(23),
+        BaseElement::new(8),
+        BaseElement::new(26),
+        BaseElement::new(13),
+        BaseElement::new(10),
+        BaseElement::new(9),
+        BaseElement::new(7),
+        BaseElement::new(6),
+        BaseElement::new(22),
+        BaseElement::new(21),
+        BaseElement::new(8),
+    ],
+    [
+        BaseElement::new(8),
+        BaseElement::new(7),
+        BaseElement::new(23),
+        BaseElement::new(8),
+        BaseElement::new(26),
+        BaseElement::new(13),
+        BaseElement::new(10),
+        BaseElement::new(9),
+        BaseElement::new(7),
+        BaseElement::new(6),
+        BaseElement::new(22),
+        BaseElement::new(21),
+    ],
+    [
+        BaseElement::new(21),
+        BaseElement::new(8),
+        BaseElement::new(7),
+        BaseElement::new(23),
+        BaseElement::new(8),
+        BaseElement::new(26),
+        BaseElement::new(13),
+        BaseElement::new(10),
+        BaseElement::new(9),
+        BaseElement::new(7),
+        BaseElement::new(6),
+        BaseElement::new(22),
+    ],
+    [
+        BaseElement::new(22),
+        BaseElement::new(21),
+        BaseElement::new(8),
+        BaseElement::new(7),
+        BaseElement::new(23),
+        BaseElement::new(8),
+        BaseElement::new(26),
+        BaseElement::new(13),
+        BaseElement::new(10),
+        BaseElement::new(9),
+        BaseElement::new(7),
+        BaseElement::new(6),
+    ],
+    [
+        BaseElement::new(6),
+        BaseElement::new(22),
+        BaseElement::new(21),
+        BaseElement::new(8),
+        BaseElement::new(7),
+        BaseElement::new(23),
+        BaseElement::new(8),
+        BaseElement::new(26),
+        BaseElement::new(13),
+        BaseElement::new(10),
+        BaseElement::new(9),
+        BaseElement::new(7),
+    ],
+    [
+        BaseElement::new(7),
+        BaseElement::new(6),
+        BaseElement::new(22),
+        BaseElement::new(21),
+        BaseElement::new(8),
+        BaseElement::new(7),
+        BaseElement::new(23),
+        BaseElement::new(8),
+        BaseElement::new(26),
+        BaseElement::new(13),
+        BaseElement::new(10),
+        BaseElement::new(9),
+    ],
+    [
+        BaseElement::new(9),
+        BaseElement::new(7),
+        BaseElement::new(6),
+        BaseElement::new(22),
+        BaseElement::new(21),
+        BaseElement::new(8),
+        BaseElement::new(7),
+        BaseElement::new(23),
+        BaseElement::new(8),
+        BaseElement::new(26),
+        BaseElement::new(13),
+        BaseElement::new(10),
+    ],
+    [
+        BaseElement::new(10),
+        BaseElement::new(9),
+        BaseElement::new(7),
+        BaseElement::new(6),
+        BaseElement::new(22),
+        BaseElement::new(21),
+        BaseElement::new(8),
+        BaseElement::new(7),
+        BaseElement::new(23),
+        BaseElement::new(8),
+        BaseElement::new(26),
+        BaseElement::new(13),
+    ],
+    [
+        BaseElement::new(13),
+        BaseElement::new(10),
+        BaseElement::new(9),
+        BaseElement::new(7),
+        BaseElement::new(6),
+        BaseElement::new(22),
+        BaseElement::new(21),
+        BaseElement::new(8),
+        BaseElement::new(7),
+        BaseElement::new(23),
+        BaseElement::new(8),
+        BaseElement::new(26),
+    ],
+    [
+        BaseElement::new(26),
+        BaseElement::new(13),
+        BaseElement::new(10),
+        BaseElement::new(9),
+        BaseElement::new(7),
+        BaseElement::new(6),
+        BaseElement::new(22),
+        BaseElement::new(21),
+        BaseElement::new(8),
+        BaseElement::new(7),
+        BaseElement::new(23),
+        BaseElement::new(8),
+    ],
+    [
+        BaseElement::new(8),
+        BaseElement::new(26),
+        BaseElement::new(13),
+        BaseElement::new(10),
+        BaseElement::new(9),
+        BaseElement::new(7),
+        BaseElement::new(6),
+        BaseElement::new(22),
+        BaseElement::new(21),
+        BaseElement::new(8),
+        BaseElement::new(7),
+        BaseElement::new(23),
+    ],
+    [
+        BaseElement::new(23),
+        BaseElement::new(8),
+        BaseElement::new(26),
+        BaseElement::new(13),
+        BaseElement::new(10),
+        BaseElement::new(9),
+        BaseElement::new(7),
+        BaseElement::new(6),
+        BaseElement::new(22),
+        BaseElement::new(21),
+        BaseElement::new(8),
+        BaseElement::new(7),
+    ],
+];
+
+/// RPO Inverse MDS matrix
+const INV_MDS: [[BaseElement; STATE_WIDTH]; STATE_WIDTH] = [
+    [
+        BaseElement::new(14868391535953158196),
+        BaseElement::new(13278298489594233127),
+        BaseElement::new(389999932707070822),
+        BaseElement::new(9782021734907796003),
+        BaseElement::new(4829905704463175582),
+        BaseElement::new(7567822018949214430),
+        BaseElement::new(14205019324568680367),
+        BaseElement::new(15489674211196160593),
+        BaseElement::new(17636013826542227504),
+        BaseElement::new(16254215311946436093),
+        BaseElement::new(3641486184877122796),
+        BaseElement::new(11069068059762973582),
+    ],
+    [
+        BaseElement::new(11069068059762973582),
+        BaseElement::new(14868391535953158196),
+        BaseElement::new(13278298489594233127),
+        BaseElement::new(389999932707070822),
+        BaseElement::new(9782021734907796003),
+        BaseElement::new(4829905704463175582),
+        BaseElement::new(7567822018949214430),
+        BaseElement::new(14205019324568680367),
+        BaseElement::new(15489674211196160593),
+        BaseElement::new(17636013826542227504),
+        BaseElement::new(16254215311946436093),
+        BaseElement::new(3641486184877122796),
+    ],
+    [
+        BaseElement::new(3641486184877122796),
+        BaseElement::new(11069068059762973582),
+        BaseElement::new(14868391535953158196),
+        BaseElement::new(13278298489594233127),
+        BaseElement::new(389999932707070822),
+        BaseElement::new(9782021734907796003),
+        BaseElement::new(4829905704463175582),
+        BaseElement::new(7567822018949214430),
+        BaseElement::new(14205019324568680367),
+        BaseElement::new(15489674211196160593),
+        BaseElement::new(17636013826542227504),
+        BaseElement::new(16254215311946436093),
+    ],
+    [
+        BaseElement::new(16254215311946436093),
+        BaseElement::new(3641486184877122796),
+        BaseElement::new(11069068059762973582),
+        BaseElement::new(14868391535953158196),
+        BaseElement::new(13278298489594233127),
+        BaseElement::new(389999932707070822),
+        BaseElement::new(9782021734907796003),
+        BaseElement::new(4829905704463175582),
+        BaseElement::new(7567822018949214430),
+        BaseElement::new(14205019324568680367),
+        BaseElement::new(15489674211196160593),
+        BaseElement::new(17636013826542227504),
+    ],
+    [
+        BaseElement::new(17636013826542227504),
+        BaseElement::new(16254215311946436093),
+        BaseElement::new(3641486184877122796),
+        BaseElement::new(11069068059762973582),
+        BaseElement::new(14868391535953158196),
+        BaseElement::new(13278298489594233127),
+        BaseElement::new(389999932707070822),
+        BaseElement::new(9782021734907796003),
+        BaseElement::new(4829905704463175582),
+        BaseElement::new(7567822018949214430),
+        BaseElement::new(14205019324568680367),
+        BaseElement::new(15489674211196160593),
+    ],
+    [
+        BaseElement::new(15489674211196160593),
+        BaseElement::new(17636013826542227504),
+        BaseElement::new(16254215311946436093),
+        BaseElement::new(3641486184877122796),
+        BaseElement::new(11069068059762973582),
+        BaseElement::new(14868391535953158196),
+        BaseElement::new(13278298489594233127),
+        BaseElement::new(389999932707070822),
+        BaseElement::new(9782021734907796003),
+        BaseElement::new(4829905704463175582),
+        BaseElement::new(7567822018949214430),
+        BaseElement::new(14205019324568680367),
+    ],
+    [
+        BaseElement::new(14205019324568680367),
+        BaseElement::new(15489674211196160593),
+        BaseElement::new(17636013826542227504),
+        BaseElement::new(16254215311946436093),
+        BaseElement::new(3641486184877122796),
+        BaseElement::new(11069068059762973582),
+        BaseElement::new(14868391535953158196),
+        BaseElement::new(13278298489594233127),
+        BaseElement::new(389999932707070822),
+        BaseElement::new(9782021734907796003),
+        BaseElement::new(4829905704463175582),
+        BaseElement::new(7567822018949214430),
+    ],
+    [
+        BaseElement::new(7567822018949214430),
+        BaseElement::new(14205019324568680367),
+        BaseElement::new(15489674211196160593),
+        BaseElement::new(17636013826542227504),
+        BaseElement::new(16254215311946436093),
+        BaseElement::new(3641486184877122796),
+        BaseElement::new(11069068059762973582),
+        BaseElement::new(14868391535953158196),
+        BaseElement::new(13278298489594233127),
+        BaseElement::new(389999932707070822),
+        BaseElement::new(9782021734907796003),
+        BaseElement::new(4829905704463175582),
+    ],
+    [
+        BaseElement::new(4829905704463175582),
+        BaseElement::new(7567822018949214430),
+        BaseElement::new(14205019324568680367),
+        BaseElement::new(15489674211196160593),
+        BaseElement::new(17636013826542227504),
+        BaseElement::new(16254215311946436093),
+        BaseElement::new(3641486184877122796),
+        BaseElement::new(11069068059762973582),
+        BaseElement::new(14868391535953158196),
+        BaseElement::new(13278298489594233127),
+        BaseElement::new(389999932707070822),
+        BaseElement::new(9782021734907796003),
+    ],
+    [
+        BaseElement::new(9782021734907796003),
+        BaseElement::new(4829905704463175582),
+        BaseElement::new(7567822018949214430),
+        BaseElement::new(14205019324568680367),
+        BaseElement::new(15489674211196160593),
+        BaseElement::new(17636013826542227504),
+        BaseElement::new(16254215311946436093),
+        BaseElement::new(3641486184877122796),
+        BaseElement::new(11069068059762973582),
+        BaseElement::new(14868391535953158196),
+        BaseElement::new(13278298489594233127),
+        BaseElement::new(389999932707070822),
+    ],
+    [
+        BaseElement::new(389999932707070822),
+        BaseElement::new(9782021734907796003),
+        BaseElement::new(4829905704463175582),
+        BaseElement::new(7567822018949214430),
+        BaseElement::new(14205019324568680367),
+        BaseElement::new(15489674211196160593),
+        BaseElement::new(17636013826542227504),
+        BaseElement::new(16254215311946436093),
+        BaseElement::new(3641486184877122796),
+        BaseElement::new(11069068059762973582),
+        BaseElement::new(14868391535953158196),
+        BaseElement::new(13278298489594233127),
+    ],
+    [
+        BaseElement::new(13278298489594233127),
+        BaseElement::new(389999932707070822),
+        BaseElement::new(9782021734907796003),
+        BaseElement::new(4829905704463175582),
+        BaseElement::new(7567822018949214430),
+        BaseElement::new(14205019324568680367),
+        BaseElement::new(15489674211196160593),
+        BaseElement::new(17636013826542227504),
+        BaseElement::new(16254215311946436093),
+        BaseElement::new(3641486184877122796),
+        BaseElement::new(11069068059762973582),
+        BaseElement::new(14868391535953158196),
+    ],
+];
+
+// ROUND CONSTANTS
+// ================================================================================================
+
+/// Rescue round constants;
+/// computed as in section ?? from [specifications](https://github.org/aszepieniec/rpo/)
+///
+/// The constants are broken up into two arrays ARK1 and ARK2; ARK1 contains the constants for the
+/// first half of RPO round, and ARK2 contains constants for the second half of RPO round.
+const ARK1: [[BaseElement; STATE_WIDTH]; NUM_ROUNDS] = [
+    [
+        BaseElement::new(5789762306288267392),
+        BaseElement::new(6522564764413701783),
+        BaseElement::new(17809893479458208203),
+        BaseElement::new(107145243989736508),
+        BaseElement::new(6388978042437517382),
+        BaseElement::new(15844067734406016715),
+        BaseElement::new(9975000513555218239),
+        BaseElement::new(3344984123768313364),
+        BaseElement::new(9959189626657347191),
+        BaseElement::new(12960773468763563665),
+        BaseElement::new(9602914297752488475),
+        BaseElement::new(16657542370200465908),
+    ],
+    [
+        BaseElement::new(12987190162843096997),
+        BaseElement::new(653957632802705281),
+        BaseElement::new(4441654670647621225),
+        BaseElement::new(4038207883745915761),
+        BaseElement::new(5613464648874830118),
+        BaseElement::new(13222989726778338773),
+        BaseElement::new(3037761201230264149),
+        BaseElement::new(16683759727265180203),
+        BaseElement::new(8337364536491240715),
+        BaseElement::new(3227397518293416448),
+        BaseElement::new(8110510111539674682),
+        BaseElement::new(2872078294163232137),
+    ],
+    [
+        BaseElement::new(18072785500942327487),
+        BaseElement::new(6200974112677013481),
+        BaseElement::new(17682092219085884187),
+        BaseElement::new(10599526828986756440),
+        BaseElement::new(975003873302957338),
+        BaseElement::new(8264241093196931281),
+        BaseElement::new(10065763900435475170),
+        BaseElement::new(2181131744534710197),
+        BaseElement::new(6317303992309418647),
+        BaseElement::new(1401440938888741532),
+        BaseElement::new(8884468225181997494),
+        BaseElement::new(13066900325715521532),
+    ],
+    [
+        BaseElement::new(5674685213610121970),
+        BaseElement::new(5759084860419474071),
+        BaseElement::new(13943282657648897737),
+        BaseElement::new(1352748651966375394),
+        BaseElement::new(17110913224029905221),
+        BaseElement::new(1003883795902368422),
+        BaseElement::new(4141870621881018291),
+        BaseElement::new(8121410972417424656),
+        BaseElement::new(14300518605864919529),
+        BaseElement::new(13712227150607670181),
+        BaseElement::new(17021852944633065291),
+        BaseElement::new(6252096473787587650),
+    ],
+    [
+        BaseElement::new(4887609836208846458),
+        BaseElement::new(3027115137917284492),
+        BaseElement::new(9595098600469470675),
+        BaseElement::new(10528569829048484079),
+        BaseElement::new(7864689113198939815),
+        BaseElement::new(17533723827845969040),
+        BaseElement::new(5781638039037710951),
+        BaseElement::new(17024078752430719006),
+        BaseElement::new(109659393484013511),
+        BaseElement::new(7158933660534805869),
+        BaseElement::new(2955076958026921730),
+        BaseElement::new(7433723648458773977),
+    ],
+    [
+        BaseElement::new(16308865189192447297),
+        BaseElement::new(11977192855656444890),
+        BaseElement::new(12532242556065780287),
+        BaseElement::new(14594890931430968898),
+        BaseElement::new(7291784239689209784),
+        BaseElement::new(5514718540551361949),
+        BaseElement::new(10025733853830934803),
+        BaseElement::new(7293794580341021693),
+        BaseElement::new(6728552937464861756),
+        BaseElement::new(6332385040983343262),
+        BaseElement::new(13277683694236792804),
+        BaseElement::new(2600778905124452676),
+    ],
+    [
+        BaseElement::new(7123075680859040534),
+        BaseElement::new(1034205548717903090),
+        BaseElement::new(7717824418247931797),
+        BaseElement::new(3019070937878604058),
+        BaseElement::new(11403792746066867460),
+        BaseElement::new(10280580802233112374),
+        BaseElement::new(337153209462421218),
+        BaseElement::new(13333398568519923717),
+        BaseElement::new(3596153696935337464),
+        BaseElement::new(8104208463525993784),
+        BaseElement::new(14345062289456085693),
+        BaseElement::new(17036731477169661256),
+    ],
+];
+
+const ARK2: [[BaseElement; STATE_WIDTH]; NUM_ROUNDS] = [
+    [
+        BaseElement::new(6077062762357204287),
+        BaseElement::new(15277620170502011191),
+        BaseElement::new(5358738125714196705),
+        BaseElement::new(14233283787297595718),
+        BaseElement::new(13792579614346651365),
+        BaseElement::new(11614812331536767105),
+        BaseElement::new(14871063686742261166),
+        BaseElement::new(10148237148793043499),
+        BaseElement::new(4457428952329675767),
+        BaseElement::new(15590786458219172475),
+        BaseElement::new(10063319113072092615),
+        BaseElement::new(14200078843431360086),
+    ],
+    [
+        BaseElement::new(6202948458916099932),
+        BaseElement::new(17690140365333231091),
+        BaseElement::new(3595001575307484651),
+        BaseElement::new(373995945117666487),
+        BaseElement::new(1235734395091296013),
+        BaseElement::new(14172757457833931602),
+        BaseElement::new(707573103686350224),
+        BaseElement::new(15453217512188187135),
+        BaseElement::new(219777875004506018),
+        BaseElement::new(17876696346199469008),
+        BaseElement::new(17731621626449383378),
+        BaseElement::new(2897136237748376248),
+    ],
+    [
+        BaseElement::new(8023374565629191455),
+        BaseElement::new(15013690343205953430),
+        BaseElement::new(4485500052507912973),
+        BaseElement::new(12489737547229155153),
+        BaseElement::new(9500452585969030576),
+        BaseElement::new(2054001340201038870),
+        BaseElement::new(12420704059284934186),
+        BaseElement::new(355990932618543755),
+        BaseElement::new(9071225051243523860),
+        BaseElement::new(12766199826003448536),
+        BaseElement::new(9045979173463556963),
+        BaseElement::new(12934431667190679898),
+    ],
+    [
+        BaseElement::new(18389244934624494276),
+        BaseElement::new(16731736864863925227),
+        BaseElement::new(4440209734760478192),
+        BaseElement::new(17208448209698888938),
+        BaseElement::new(8739495587021565984),
+        BaseElement::new(17000774922218161967),
+        BaseElement::new(13533282547195532087),
+        BaseElement::new(525402848358706231),
+        BaseElement::new(16987541523062161972),
+        BaseElement::new(5466806524462797102),
+        BaseElement::new(14512769585918244983),
+        BaseElement::new(10973956031244051118),
+    ],
+    [
+        BaseElement::new(6982293561042362913),
+        BaseElement::new(14065426295947720331),
+        BaseElement::new(16451845770444974180),
+        BaseElement::new(7139138592091306727),
+        BaseElement::new(9012006439959783127),
+        BaseElement::new(14619614108529063361),
+        BaseElement::new(1394813199588124371),
+        BaseElement::new(4635111139507788575),
+        BaseElement::new(16217473952264203365),
+        BaseElement::new(10782018226466330683),
+        BaseElement::new(6844229992533662050),
+        BaseElement::new(7446486531695178711),
+    ],
+    [
+        BaseElement::new(3736792340494631448),
+        BaseElement::new(577852220195055341),
+        BaseElement::new(6689998335515779805),
+        BaseElement::new(13886063479078013492),
+        BaseElement::new(14358505101923202168),
+        BaseElement::new(7744142531772274164),
+        BaseElement::new(16135070735728404443),
+        BaseElement::new(12290902521256031137),
+        BaseElement::new(12059913662657709804),
+        BaseElement::new(16456018495793751911),
+        BaseElement::new(4571485474751953524),
+        BaseElement::new(17200392109565783176),
+    ],
+    [
+        BaseElement::new(17130398059294018733),
+        BaseElement::new(519782857322261988),
+        BaseElement::new(9625384390925085478),
+        BaseElement::new(1664893052631119222),
+        BaseElement::new(7629576092524553570),
+        BaseElement::new(3485239601103661425),
+        BaseElement::new(9755891797164033838),
+        BaseElement::new(15218148195153269027),
+        BaseElement::new(16460604813734957368),
+        BaseElement::new(9643968136937729763),
+        BaseElement::new(3611348709641382851),
+        BaseElement::new(18256379591337759196),
+    ],
+];
diff --git a/crypto/src/hash/rpo/tests.rs b/crypto/src/hash/rpo/tests.rs
new file mode 100644
index 0000000..8eff6af
--- /dev/null
+++ b/crypto/src/hash/rpo/tests.rs
@@ -0,0 +1,161 @@
+use super::{
+    BaseElement, ElementDigest, ElementHasher, FieldElement, HashFn, Rpo, StarkField, ALPHA,
+    INV_ALPHA, INV_MDS, MDS, STATE_WIDTH,
+};
+
+use core::convert::TryInto;
+
+use rand_utils::rand_value;
+
+#[test]
+fn mds_inv_test() {
+    let mut mul_result = [[BaseElement::new(0); STATE_WIDTH]; STATE_WIDTH];
+    for i in 0..STATE_WIDTH {
+        for j in 0..STATE_WIDTH {
+            let result = {
+                let mut result = BaseElement::new(0);
+                for k in 0..STATE_WIDTH {
+                    result += MDS[i][k] * INV_MDS[k][j]
+                }
+                result
+            };
+            mul_result[i][j] = result;
+            if i == j {
+                assert_eq!(result, BaseElement::new(1));
+            } else {
+                assert_eq!(result, BaseElement::new(0));
+            }
+        }
+    }
+}
+#[test]
+fn test_alphas() {
+    let e: BaseElement = BaseElement::new(rand_value());
+    let e_exp = e.exp(ALPHA.into());
+    assert_eq!(e, e_exp.exp(INV_ALPHA));
+}
+
+#[test]
+fn test_sbox() {
+    let state = [BaseElement::new(rand_value()); STATE_WIDTH];
+
+    let mut expected = state;
+    expected.iter_mut().for_each(|v| *v = v.exp(ALPHA));
+
+    let mut actual = state;
+    Rpo::apply_sbox(&mut actual);
+
+    assert_eq!(expected, actual);
+}
+
+#[test]
+fn test_inv_sbox() {
+    let state = [BaseElement::new(rand_value()); STATE_WIDTH];
+
+    let mut expected = state;
+    expected.iter_mut().for_each(|v| *v = v.exp(INV_ALPHA));
+
+    let mut actual = state;
+    Rpo::apply_inv_sbox(&mut actual);
+
+    assert_eq!(expected, actual);
+}
+
+#[test]
+fn hash_elements_vs_merge() {
+    let elements = [BaseElement::new(rand_value()); 8];
+
+    let digests: [ElementDigest; 2] = [
+        ElementDigest::new(elements[..4].try_into().unwrap()),
+        ElementDigest::new(elements[4..].try_into().unwrap()),
+    ];
+
+    let m_result = Rpo::merge(&digests);
+    let h_result = Rpo::hash_elements(&elements);
+    assert_eq!(m_result, h_result);
+}
+
+#[test]
+fn hash_elements_vs_merge_with_int() {
+    let tmp = [BaseElement::new(rand_value()); 4];
+    let seed = ElementDigest::new(tmp);
+
+    // ----- value fits into a field element ------------------------------------------------------
+    let val: BaseElement = BaseElement::new(rand_value());
+    let m_result = Rpo::merge_with_int(seed, val.as_int());
+
+    let mut elements = seed.as_elements().to_vec();
+    elements.push(val);
+    let h_result = Rpo::hash_elements(&elements);
+
+    assert_eq!(m_result, h_result);
+
+    // ----- value does not fit into a field element ----------------------------------------------
+    let val = BaseElement::MODULUS + 2;
+    let m_result = Rpo::merge_with_int(seed, val);
+
+    let mut elements = seed.as_elements().to_vec();
+    elements.push(BaseElement::new(val));
+    elements.push(BaseElement::new(1));
+    let h_result = Rpo::hash_elements(&elements);
+
+    assert_eq!(m_result, h_result);
+}
+
+#[test]
+fn hash_padding() {
+    // adding a zero bytes at the end of a byte string should result in a different hash
+    let r1 = Rpo::hash(&[1_u8, 2, 3]);
+    let r2 = Rpo::hash(&[1_u8, 2, 3, 0]);
+    assert_ne!(r1, r2);
+
+    // same as above but with bigger inputs
+    let r1 = Rpo::hash(&[1_u8, 2, 3, 4, 5, 6]);
+    let r2 = Rpo::hash(&[1_u8, 2, 3, 4, 5, 6, 0]);
+    assert_ne!(r1, r2);
+
+    // same as above but with input splitting over two elements
+    let r1 = Rpo::hash(&[1_u8, 2, 3, 4, 5, 6, 7]);
+    let r2 = Rpo::hash(&[1_u8, 2, 3, 4, 5, 6, 7, 0]);
+    assert_ne!(r1, r2);
+
+    // same as above but with multiple zeros
+    let r1 = Rpo::hash(&[1_u8, 2, 3, 4, 5, 6, 7, 0, 0]);
+    let r2 = Rpo::hash(&[1_u8, 2, 3, 4, 5, 6, 7, 0, 0, 0, 0]);
+    assert_ne!(r1, r2);
+}
+
+#[test]
+fn hash_elements_padding() {
+    let e1 = [BaseElement::new(rand_value()); 2];
+    let e2 = [e1[0], e1[1], BaseElement::ZERO];
+
+    let r1 = Rpo::hash_elements(&e1);
+    let r2 = Rpo::hash_elements(&e2);
+    assert_ne!(r1, r2);
+}
+
+#[test]
+fn hash_elements() {
+    let elements = [
+        BaseElement::new(0),
+        BaseElement::new(1),
+        BaseElement::new(2),
+        BaseElement::new(3),
+        BaseElement::new(4),
+        BaseElement::new(5),
+        BaseElement::new(6),
+        BaseElement::new(7)
+    ];
+
+    let digests: [ElementDigest; 2] = [
+        ElementDigest::new(elements[..4].try_into().unwrap()),
+        ElementDigest::new(elements[4..8].try_into().unwrap()),
+    ];
+
+    let m_result = Rpo::merge(&digests);
+    let h_result = Rpo::hash_elements(&elements);
+    let m_elements :Vec<u64> = m_result.as_elements().iter().map(|a| {a.as_int()}).collect();
+    let h_elements :Vec<u64> = h_result.as_elements().iter().map(|a| {a.as_int()}).collect();
+    assert_eq!(m_elements, h_elements);
+}
\ No newline at end of file

From bbfddcfbd3ce5fa973f0fb9becb5d73bbd0bc8d0 Mon Sep 17 00:00:00 2001
From: Al-Kindi-0 <82364884+Al-Kindi-0@users.noreply.github.com>
Date: Thu, 20 Oct 2022 14:26:09 +0200
Subject: [PATCH 2/9] fix misc. nits

---
 crypto/src/hash/mod.rs               |   17 +-
 crypto/src/hash/rpo/digest.rs        |   62 +-
 crypto/src/hash/rpo/mod.rs           | 1161 +++++++++++++-------------
 crypto/src/hash/rpo/tests.rs         |   70 +-
 crypto/src/lib.rs                    |   13 +
 crypto/src/merkle/merkle_path_set.rs |   18 +-
 crypto/src/merkle/merkle_tree.rs     |    7 +-
 crypto/src/merkle/mod.rs             |    7 +
 8 files changed, 683 insertions(+), 672 deletions(-)

diff --git a/crypto/src/hash/mod.rs b/crypto/src/hash/mod.rs
index 0b0eed9..6674e22 100644
--- a/crypto/src/hash/mod.rs
+++ b/crypto/src/hash/mod.rs
@@ -1,6 +1,7 @@
-use crate::StarkField;
-pub use winterfell::crypto::hashers::Rp64_256 as Hasher;
-pub use winterfell::crypto::{ElementHasher, Hasher as HashFn};
+use crate::{ElementHasher, HashFn};
+
+mod rpo;
+pub use rpo::Rpo as Hasher;
 
 // TYPE ALIASES
 // ================================================================================================
@@ -10,16 +11,6 @@ pub type Digest = <Hasher as HashFn>::Digest;
 // HELPER FUNCTIONS
 // ================================================================================================
 
-#[inline(always)]
-fn _exp_acc<B: StarkField, const N: usize, const M: usize>(base: [B; N], tail: [B; N]) -> [B; N] {
-    let mut result = base;
-    for _ in 0..M {
-        result.iter_mut().for_each(|r| *r = r.square());
-    }
-    result.iter_mut().zip(tail).for_each(|(r, t)| *r *= t);
-    result
-}
-
 #[inline(always)]
 pub fn merge(values: &[Digest; 2]) -> Digest {
     Hasher::merge(values)
diff --git a/crypto/src/hash/rpo/digest.rs b/crypto/src/hash/rpo/digest.rs
index 1e07bb3..7c21298 100644
--- a/crypto/src/hash/rpo/digest.rs
+++ b/crypto/src/hash/rpo/digest.rs
@@ -1,32 +1,31 @@
 use super::DIGEST_SIZE;
+use crate::{Digest, Felt, StarkField};
 use core::slice;
-use winterfell::crypto::Digest;
-use winterfell::math::{fields::f64::BaseElement, StarkField};
 use winterfell::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable};
 
 // DIGEST TRAIT IMPLEMENTATIONS
 // ================================================================================================
 
 #[derive(Debug, Copy, Clone, Eq, PartialEq)]
-pub struct ElementDigest([BaseElement; DIGEST_SIZE]);
+pub struct RpoDigest256([Felt; DIGEST_SIZE]);
 
-impl ElementDigest {
-    pub fn new(value: [BaseElement; DIGEST_SIZE]) -> Self {
+impl RpoDigest256 {
+    pub fn new(value: [Felt; DIGEST_SIZE]) -> Self {
         Self(value)
     }
 
-    pub fn as_elements(&self) -> &[BaseElement] {
+    pub fn as_elements(&self) -> &[Felt] {
         &self.0
     }
 
-    pub fn digests_as_elements(digests: &[Self]) -> &[BaseElement] {
+    pub fn digests_as_elements(digests: &[Self]) -> &[Felt] {
         let p = digests.as_ptr();
         let len = digests.len() * DIGEST_SIZE;
-        unsafe { slice::from_raw_parts(p as *const BaseElement, len) }
+        unsafe { slice::from_raw_parts(p as *const Felt, len) }
     }
 }
 
-impl Digest for ElementDigest {
+impl Digest for RpoDigest256 {
     fn as_bytes(&self) -> [u8; 32] {
         let mut result = [0; 32];
 
@@ -39,44 +38,44 @@ impl Digest for ElementDigest {
     }
 }
 
-impl Default for ElementDigest {
+impl Default for RpoDigest256 {
     fn default() -> Self {
-        ElementDigest([BaseElement::default(); DIGEST_SIZE])
+        RpoDigest256([Felt::default(); DIGEST_SIZE])
     }
 }
 
-impl Serializable for ElementDigest {
+impl Serializable for RpoDigest256 {
     fn write_into<W: ByteWriter>(&self, target: &mut W) {
         target.write_u8_slice(&self.as_bytes());
     }
 }
 
-impl Deserializable for ElementDigest {
+impl Deserializable for RpoDigest256 {
     fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
         // TODO: check if the field elements are valid?
-        let e1 = BaseElement::new(source.read_u64()?);
-        let e2 = BaseElement::new(source.read_u64()?);
-        let e3 = BaseElement::new(source.read_u64()?);
-        let e4 = BaseElement::new(source.read_u64()?);
+        let e1 = Felt::new(source.read_u64()?);
+        let e2 = Felt::new(source.read_u64()?);
+        let e3 = Felt::new(source.read_u64()?);
+        let e4 = Felt::new(source.read_u64()?);
 
         Ok(Self([e1, e2, e3, e4]))
     }
 }
 
-impl From<[BaseElement; DIGEST_SIZE]> for ElementDigest {
-    fn from(value: [BaseElement; DIGEST_SIZE]) -> Self {
+impl From<[Felt; DIGEST_SIZE]> for RpoDigest256 {
+    fn from(value: [Felt; DIGEST_SIZE]) -> Self {
         Self(value)
     }
 }
 
-impl From<ElementDigest> for [BaseElement; DIGEST_SIZE] {
-    fn from(value: ElementDigest) -> Self {
+impl From<RpoDigest256> for [Felt; DIGEST_SIZE] {
+    fn from(value: RpoDigest256) -> Self {
         value.0
     }
 }
 
-impl From<ElementDigest> for [u8; 32] {
-    fn from(value: ElementDigest) -> Self {
+impl From<RpoDigest256> for [u8; 32] {
+    fn from(value: RpoDigest256) -> Self {
         value.as_bytes()
     }
 }
@@ -87,25 +86,26 @@ impl From<ElementDigest> for [u8; 32] {
 #[cfg(test)]
 mod tests {
 
-    use super::ElementDigest;
+    use super::RpoDigest256;
+    use crate::Felt;
     use rand_utils::rand_value;
-    use winterfell::{math::fields::f64::BaseElement, Deserializable, Serializable, SliceReader};
+    use winterfell::{Deserializable, Serializable, SliceReader};
 
     #[test]
     fn digest_serialization() {
-        let e1 = BaseElement::new(rand_value());
-        let e2 = BaseElement::new(rand_value());
-        let e3 = BaseElement::new(rand_value());
-        let e4 = BaseElement::new(rand_value());
+        let e1 = Felt::new(rand_value());
+        let e2 = Felt::new(rand_value());
+        let e3 = Felt::new(rand_value());
+        let e4 = Felt::new(rand_value());
 
-        let d1 = ElementDigest([e1, e2, e3, e4]);
+        let d1 = RpoDigest256([e1, e2, e3, e4]);
 
         let mut bytes = vec![];
         d1.write_into(&mut bytes);
         assert_eq!(32, bytes.len());
 
         let mut reader = SliceReader::new(&bytes);
-        let d2 = ElementDigest::read_from(&mut reader).unwrap();
+        let d2 = RpoDigest256::read_from(&mut reader).unwrap();
 
         assert_eq!(d1, d2);
     }
diff --git a/crypto/src/hash/rpo/mod.rs b/crypto/src/hash/rpo/mod.rs
index e872974..8c4b66a 100644
--- a/crypto/src/hash/rpo/mod.rs
+++ b/crypto/src/hash/rpo/mod.rs
@@ -1,9 +1,9 @@
-use super::{exp_acc, ElementHasher, HashFn};
+use super::{ElementHasher, HashFn};
+use crate::{Felt, FieldElement, StarkField, ONE, ZERO};
 use core::{convert::TryInto, ops::Range};
-use winterfell::math::{fields::f64::BaseElement, FieldElement, StarkField};
 
 mod digest;
-pub use digest::ElementDigest;
+pub use digest::RpoDigest256;
 
 #[cfg(test)]
 mod tests;
@@ -14,7 +14,7 @@ use mds_freq::mds_multiply_freq;
 // CONSTANTS
 // ================================================================================================
 
-/// Sponge state is set to 12 field elements or 768 bytes; 8 elements are reserved for rate and
+/// Sponge state is set to 12 field elements or 96 bytes; 8 elements are reserved for rate and
 /// the remaining 4 elements are reserved for capacity.
 const STATE_WIDTH: usize = 12;
 
@@ -89,7 +89,7 @@ const INV_ALPHA: u64 = 10540996611094048183;
 pub struct Rpo();
 
 impl HashFn for Rpo {
-    type Digest = ElementDigest;
+    type Digest = RpoDigest256;
 
     fn hash(bytes: &[u8]) -> Self::Digest {
         // compute the number of elements required to represent the string; we will be processing
@@ -104,8 +104,8 @@ impl HashFn for Rpo {
         // initialize state to all zeros, except for the first element of the capacity part, which
         // is set to the number of elements to be hashed. this is done so that adding zero elements
         // at the end of the list always results in a different hash.
-        let mut state = [BaseElement::ZERO; STATE_WIDTH];
-        state[CAPACITY_RANGE.start] = BaseElement::new(num_elements as u64);
+        let mut state = [ZERO; STATE_WIDTH];
+        state[CAPACITY_RANGE.start] = Felt::new(num_elements as u64);
 
         // break the string into 7-byte chunks, convert each chunk into a field element, and
         // absorb the element into the rate portion of the state. we use 7-byte chunks because
@@ -129,7 +129,7 @@ impl HashFn for Rpo {
             // convert the bytes into a field element and absorb it into the rate portion of the
             // state; if the rate is filled up, apply the Rescue permutation and start absorbing
             // again from zero index.
-            state[RATE_RANGE.start + i] += BaseElement::new(u64::from_le_bytes(buf));
+            state[RATE_RANGE.start + i] += Felt::new(u64::from_le_bytes(buf));
             i += 1;
             if i % RATE_WIDTH == 0 {
                 Self::apply_permutation(&mut state);
@@ -146,49 +146,50 @@ impl HashFn for Rpo {
         }
 
         // return the first 4 elements of the state as hash result
-        ElementDigest::new(state[DIGEST_RANGE].try_into().unwrap())
+        RpoDigest256::new(state[DIGEST_RANGE].try_into().unwrap())
     }
 
     fn merge(values: &[Self::Digest; 2]) -> Self::Digest {
         // initialize the state by copying the digest elements into the rate portion of the state
         // (8 total elements), and set the capacity elements to 0.
-        let mut state = [BaseElement::ZERO; STATE_WIDTH];
+        let mut state = [ZERO; STATE_WIDTH];
         state[RATE_RANGE].copy_from_slice(Self::Digest::digests_as_elements(values));
 
         // apply the RPO permutation and return the first four elements of the state
         Self::apply_permutation(&mut state);
-        ElementDigest::new(state[DIGEST_RANGE].try_into().unwrap())
+        RpoDigest256::new(state[DIGEST_RANGE].try_into().unwrap())
     }
 
     fn merge_with_int(seed: Self::Digest, value: u64) -> Self::Digest {
         // initialize the state as follows:
         // - seed is copied into the first 4 elements of the rate portion of the state.
         // - if the value fits into a single field element, copy it into the fifth rate element
-        //   and set the sixth rate element and first capacity element to 1.
+        //   and set the sixth rate element to 1.
         // - if the value doesn't fit into a single field element, split it into two field
         //   elements, copy them into rate elements 5 and 6, and set the seventh rate element
-        //   and first capacity element to 1.
-        let mut state = [BaseElement::ZERO; STATE_WIDTH];
+        //   to 1.
+        // - set the first capacity element to 1
+        let mut state = [ZERO; STATE_WIDTH];
         state[INPUT1_RANGE].copy_from_slice(seed.as_elements());
-        state[INPUT2_RANGE.start] = BaseElement::new(value);
-        if value < BaseElement::MODULUS {
-            state[INPUT2_RANGE.start + 1] = BaseElement::ONE;
+        state[INPUT2_RANGE.start] = Felt::new(value);
+        if value < Felt::MODULUS {
+            state[INPUT2_RANGE.start + 1] = ONE;
         } else {
-            state[INPUT2_RANGE.start + 1] = BaseElement::new(value / BaseElement::MODULUS);
-            state[INPUT2_RANGE.start + 2] = BaseElement::ONE;            
+            state[INPUT2_RANGE.start + 1] = Felt::new(value / Felt::MODULUS);
+            state[INPUT2_RANGE.start + 2] = ONE;
         }
 
         // common padding for both cases
-        state[CAPACITY_RANGE.start] = BaseElement::ONE;
-        
+        state[CAPACITY_RANGE.start] = ONE;
+
         // apply the RPO permutation and return the first four elements of the state
         Self::apply_permutation(&mut state);
-        ElementDigest::new(state[DIGEST_RANGE].try_into().unwrap())
+        RpoDigest256::new(state[DIGEST_RANGE].try_into().unwrap())
     }
 }
 
 impl ElementHasher for Rpo {
-    type BaseField = BaseElement;
+    type BaseField = Felt;
 
     fn hash_elements<E: FieldElement<BaseField = Self::BaseField>>(elements: &[E]) -> Self::Digest {
         // convert the elements into a list of base field elements
@@ -196,9 +197,9 @@ impl ElementHasher for Rpo {
 
         // initialize state to all zeros, except for the first element of the capacity part, which
         // is set to 1 if the number of elements is not a multiple of RATE_WIDTH.
-        let mut state = [BaseElement::ZERO; STATE_WIDTH];
+        let mut state = [ZERO; STATE_WIDTH];
         if elements.len() % RATE_WIDTH != 0 {
-            state[CAPACITY_RANGE.start] = BaseElement::ONE;
+            state[CAPACITY_RANGE.start] = ONE;
         }
 
         // absorb elements into the state one by one until the rate portion of the state is filled
@@ -219,17 +220,17 @@ impl ElementHasher for Rpo {
         // padding by appending a 1 followed by as many 0 as necessary to make the input length a
         // multiple of the RATE_WIDTH.
         if i > 0 {
-            state[RATE_RANGE.start + i] = BaseElement::ONE;
+            state[RATE_RANGE.start + i] = ONE;
             i += 1;
             while i != RATE_WIDTH {
-                state[RATE_RANGE.start + i] = BaseElement::ZERO;
+                state[RATE_RANGE.start + i] = ZERO;
                 i += 1;
             }
             Self::apply_permutation(&mut state);
         }
 
         // return the first 4 elements of the state as hash result
-        ElementDigest::new(state[DIGEST_RANGE].try_into().unwrap())
+        RpoDigest256::new(state[DIGEST_RANGE].try_into().unwrap())
     }
 }
 
@@ -241,7 +242,7 @@ impl Rpo {
     // --------------------------------------------------------------------------------------------
 
     /// Applies RPO permutation to the provided state.
-    pub fn apply_permutation(state: &mut [BaseElement; STATE_WIDTH]) {
+    pub fn apply_permutation(state: &mut [Felt; STATE_WIDTH]) {
         for i in 0..NUM_ROUNDS {
             Self::apply_round(state, i);
         }
@@ -249,7 +250,7 @@ impl Rpo {
 
     /// RPO round function.
     #[inline(always)]
-    pub fn apply_round(state: &mut [BaseElement; STATE_WIDTH], round: usize) {
+    pub fn apply_round(state: &mut [Felt; STATE_WIDTH], round: usize) {
         // apply first half of RPO round
         Self::apply_mds(state);
         Self::add_constants(state, &ARK1[round]);
@@ -265,8 +266,8 @@ impl Rpo {
     // --------------------------------------------------------------------------------------------
 
     #[inline(always)]
-    fn apply_mds(state: &mut [BaseElement; STATE_WIDTH]) {
-        let mut result = [BaseElement::ZERO; STATE_WIDTH];
+    fn apply_mds(state: &mut [Felt; STATE_WIDTH]) {
+        let mut result = [ZERO; STATE_WIDTH];
 
         // Using the linearity of the operations we can split the state into a low||high decomposition
         // and operate on each with no overflow and then combine/reduce the result to a field element.
@@ -291,19 +292,18 @@ impl Rpo {
             let z = (s_hi << 32) - s_hi;
             let (res, over) = s_lo.overflowing_add(z);
 
-            result[r] =
-                BaseElement::from_mont(res.wrapping_add(0u32.wrapping_sub(over as u32) as u64));
+            result[r] = Felt::from_mont(res.wrapping_add(0u32.wrapping_sub(over as u32) as u64));
         }
         *state = result;
     }
 
     #[inline(always)]
-    fn add_constants(state: &mut [BaseElement; STATE_WIDTH], ark: &[BaseElement; STATE_WIDTH]) {
+    fn add_constants(state: &mut [Felt; STATE_WIDTH], ark: &[Felt; STATE_WIDTH]) {
         state.iter_mut().zip(ark).for_each(|(s, &k)| *s += k);
     }
 
     #[inline(always)]
-    fn apply_sbox(state: &mut [BaseElement; STATE_WIDTH]) {
+    fn apply_sbox(state: &mut [Felt; STATE_WIDTH]) {
         state[0] = state[0].exp7();
         state[1] = state[1].exp7();
         state[2] = state[2].exp7();
@@ -319,7 +319,7 @@ impl Rpo {
     }
 
     #[inline(always)]
-    fn apply_inv_sbox(state: &mut [BaseElement; STATE_WIDTH]) {
+    fn apply_inv_sbox(state: &mut [Felt; STATE_WIDTH]) {
         // compute base^10540996611094048183 using 72 multiplications per array element
         // 10540996611094048183 = b1001001001001001001001001001000110110110110110110110110110110111
 
@@ -332,19 +332,19 @@ impl Rpo {
         t2.iter_mut().for_each(|t| *t = t.square());
 
         // compute base^100100
-        let t3 = exp_acc::<BaseElement, STATE_WIDTH, 3>(t2, t2);
+        let t3 = Self::exp_acc::<Felt, STATE_WIDTH, 3>(t2, t2);
 
         // compute base^100100100100
-        let t4 = exp_acc::<BaseElement, STATE_WIDTH, 6>(t3, t3);
+        let t4 = Self::exp_acc::<Felt, STATE_WIDTH, 6>(t3, t3);
 
         // compute base^100100100100100100100100
-        let t5 = exp_acc::<BaseElement, STATE_WIDTH, 12>(t4, t4);
+        let t5 = Self::exp_acc::<Felt, STATE_WIDTH, 12>(t4, t4);
 
         // compute base^100100100100100100100100100100
-        let t6 = exp_acc::<BaseElement, STATE_WIDTH, 6>(t5, t3);
+        let t6 = Self::exp_acc::<Felt, STATE_WIDTH, 6>(t5, t3);
 
         // compute base^1001001001001001001001001001000100100100100100100100100100100
-        let t7 = exp_acc::<BaseElement, STATE_WIDTH, 31>(t6, t6);
+        let t7 = Self::exp_acc::<Felt, STATE_WIDTH, 31>(t6, t6);
 
         // compute base^1001001001001001001001001001000110110110110110110110110110110111
         for (i, s) in state.iter_mut().enumerate() {
@@ -353,351 +353,364 @@ impl Rpo {
             *s = a * b;
         }
     }
+
+    #[inline(always)]
+    fn exp_acc<B: StarkField, const N: usize, const M: usize>(
+        base: [B; N],
+        tail: [B; N],
+    ) -> [B; N] {
+        let mut result = base;
+        for _ in 0..M {
+            result.iter_mut().for_each(|r| *r = r.square());
+        }
+        result.iter_mut().zip(tail).for_each(|(r, t)| *r *= t);
+        result
+    }
 }
 
 // MDS
 // ================================================================================================
 /// RPO MDS matrix
-const MDS: [[BaseElement; STATE_WIDTH]; STATE_WIDTH] = [
-    [
-        BaseElement::new(7),
-        BaseElement::new(23),
-        BaseElement::new(8),
-        BaseElement::new(26),
-        BaseElement::new(13),
-        BaseElement::new(10),
-        BaseElement::new(9),
-        BaseElement::new(7),
-        BaseElement::new(6),
-        BaseElement::new(22),
-        BaseElement::new(21),
-        BaseElement::new(8),
-    ],
-    [
-        BaseElement::new(8),
-        BaseElement::new(7),
-        BaseElement::new(23),
-        BaseElement::new(8),
-        BaseElement::new(26),
-        BaseElement::new(13),
-        BaseElement::new(10),
-        BaseElement::new(9),
-        BaseElement::new(7),
-        BaseElement::new(6),
-        BaseElement::new(22),
-        BaseElement::new(21),
-    ],
-    [
-        BaseElement::new(21),
-        BaseElement::new(8),
-        BaseElement::new(7),
-        BaseElement::new(23),
-        BaseElement::new(8),
-        BaseElement::new(26),
-        BaseElement::new(13),
-        BaseElement::new(10),
-        BaseElement::new(9),
-        BaseElement::new(7),
-        BaseElement::new(6),
-        BaseElement::new(22),
-    ],
-    [
-        BaseElement::new(22),
-        BaseElement::new(21),
-        BaseElement::new(8),
-        BaseElement::new(7),
-        BaseElement::new(23),
-        BaseElement::new(8),
-        BaseElement::new(26),
-        BaseElement::new(13),
-        BaseElement::new(10),
-        BaseElement::new(9),
-        BaseElement::new(7),
-        BaseElement::new(6),
-    ],
-    [
-        BaseElement::new(6),
-        BaseElement::new(22),
-        BaseElement::new(21),
-        BaseElement::new(8),
-        BaseElement::new(7),
-        BaseElement::new(23),
-        BaseElement::new(8),
-        BaseElement::new(26),
-        BaseElement::new(13),
-        BaseElement::new(10),
-        BaseElement::new(9),
-        BaseElement::new(7),
-    ],
-    [
-        BaseElement::new(7),
-        BaseElement::new(6),
-        BaseElement::new(22),
-        BaseElement::new(21),
-        BaseElement::new(8),
-        BaseElement::new(7),
-        BaseElement::new(23),
-        BaseElement::new(8),
-        BaseElement::new(26),
-        BaseElement::new(13),
-        BaseElement::new(10),
-        BaseElement::new(9),
-    ],
-    [
-        BaseElement::new(9),
-        BaseElement::new(7),
-        BaseElement::new(6),
-        BaseElement::new(22),
-        BaseElement::new(21),
-        BaseElement::new(8),
-        BaseElement::new(7),
-        BaseElement::new(23),
-        BaseElement::new(8),
-        BaseElement::new(26),
-        BaseElement::new(13),
-        BaseElement::new(10),
-    ],
-    [
-        BaseElement::new(10),
-        BaseElement::new(9),
-        BaseElement::new(7),
-        BaseElement::new(6),
-        BaseElement::new(22),
-        BaseElement::new(21),
-        BaseElement::new(8),
-        BaseElement::new(7),
-        BaseElement::new(23),
-        BaseElement::new(8),
-        BaseElement::new(26),
-        BaseElement::new(13),
-    ],
-    [
-        BaseElement::new(13),
-        BaseElement::new(10),
-        BaseElement::new(9),
-        BaseElement::new(7),
-        BaseElement::new(6),
-        BaseElement::new(22),
-        BaseElement::new(21),
-        BaseElement::new(8),
-        BaseElement::new(7),
-        BaseElement::new(23),
-        BaseElement::new(8),
-        BaseElement::new(26),
-    ],
-    [
-        BaseElement::new(26),
-        BaseElement::new(13),
-        BaseElement::new(10),
-        BaseElement::new(9),
-        BaseElement::new(7),
-        BaseElement::new(6),
-        BaseElement::new(22),
-        BaseElement::new(21),
-        BaseElement::new(8),
-        BaseElement::new(7),
-        BaseElement::new(23),
-        BaseElement::new(8),
-    ],
-    [
-        BaseElement::new(8),
-        BaseElement::new(26),
-        BaseElement::new(13),
-        BaseElement::new(10),
-        BaseElement::new(9),
-        BaseElement::new(7),
-        BaseElement::new(6),
-        BaseElement::new(22),
-        BaseElement::new(21),
-        BaseElement::new(8),
-        BaseElement::new(7),
-        BaseElement::new(23),
-    ],
-    [
-        BaseElement::new(23),
-        BaseElement::new(8),
-        BaseElement::new(26),
-        BaseElement::new(13),
-        BaseElement::new(10),
-        BaseElement::new(9),
-        BaseElement::new(7),
-        BaseElement::new(6),
-        BaseElement::new(22),
-        BaseElement::new(21),
-        BaseElement::new(8),
-        BaseElement::new(7),
+const MDS: [[Felt; STATE_WIDTH]; STATE_WIDTH] = [
+    [
+        Felt::new(7),
+        Felt::new(23),
+        Felt::new(8),
+        Felt::new(26),
+        Felt::new(13),
+        Felt::new(10),
+        Felt::new(9),
+        Felt::new(7),
+        Felt::new(6),
+        Felt::new(22),
+        Felt::new(21),
+        Felt::new(8),
+    ],
+    [
+        Felt::new(8),
+        Felt::new(7),
+        Felt::new(23),
+        Felt::new(8),
+        Felt::new(26),
+        Felt::new(13),
+        Felt::new(10),
+        Felt::new(9),
+        Felt::new(7),
+        Felt::new(6),
+        Felt::new(22),
+        Felt::new(21),
+    ],
+    [
+        Felt::new(21),
+        Felt::new(8),
+        Felt::new(7),
+        Felt::new(23),
+        Felt::new(8),
+        Felt::new(26),
+        Felt::new(13),
+        Felt::new(10),
+        Felt::new(9),
+        Felt::new(7),
+        Felt::new(6),
+        Felt::new(22),
+    ],
+    [
+        Felt::new(22),
+        Felt::new(21),
+        Felt::new(8),
+        Felt::new(7),
+        Felt::new(23),
+        Felt::new(8),
+        Felt::new(26),
+        Felt::new(13),
+        Felt::new(10),
+        Felt::new(9),
+        Felt::new(7),
+        Felt::new(6),
+    ],
+    [
+        Felt::new(6),
+        Felt::new(22),
+        Felt::new(21),
+        Felt::new(8),
+        Felt::new(7),
+        Felt::new(23),
+        Felt::new(8),
+        Felt::new(26),
+        Felt::new(13),
+        Felt::new(10),
+        Felt::new(9),
+        Felt::new(7),
+    ],
+    [
+        Felt::new(7),
+        Felt::new(6),
+        Felt::new(22),
+        Felt::new(21),
+        Felt::new(8),
+        Felt::new(7),
+        Felt::new(23),
+        Felt::new(8),
+        Felt::new(26),
+        Felt::new(13),
+        Felt::new(10),
+        Felt::new(9),
+    ],
+    [
+        Felt::new(9),
+        Felt::new(7),
+        Felt::new(6),
+        Felt::new(22),
+        Felt::new(21),
+        Felt::new(8),
+        Felt::new(7),
+        Felt::new(23),
+        Felt::new(8),
+        Felt::new(26),
+        Felt::new(13),
+        Felt::new(10),
+    ],
+    [
+        Felt::new(10),
+        Felt::new(9),
+        Felt::new(7),
+        Felt::new(6),
+        Felt::new(22),
+        Felt::new(21),
+        Felt::new(8),
+        Felt::new(7),
+        Felt::new(23),
+        Felt::new(8),
+        Felt::new(26),
+        Felt::new(13),
+    ],
+    [
+        Felt::new(13),
+        Felt::new(10),
+        Felt::new(9),
+        Felt::new(7),
+        Felt::new(6),
+        Felt::new(22),
+        Felt::new(21),
+        Felt::new(8),
+        Felt::new(7),
+        Felt::new(23),
+        Felt::new(8),
+        Felt::new(26),
+    ],
+    [
+        Felt::new(26),
+        Felt::new(13),
+        Felt::new(10),
+        Felt::new(9),
+        Felt::new(7),
+        Felt::new(6),
+        Felt::new(22),
+        Felt::new(21),
+        Felt::new(8),
+        Felt::new(7),
+        Felt::new(23),
+        Felt::new(8),
+    ],
+    [
+        Felt::new(8),
+        Felt::new(26),
+        Felt::new(13),
+        Felt::new(10),
+        Felt::new(9),
+        Felt::new(7),
+        Felt::new(6),
+        Felt::new(22),
+        Felt::new(21),
+        Felt::new(8),
+        Felt::new(7),
+        Felt::new(23),
+    ],
+    [
+        Felt::new(23),
+        Felt::new(8),
+        Felt::new(26),
+        Felt::new(13),
+        Felt::new(10),
+        Felt::new(9),
+        Felt::new(7),
+        Felt::new(6),
+        Felt::new(22),
+        Felt::new(21),
+        Felt::new(8),
+        Felt::new(7),
     ],
 ];
 
 /// RPO Inverse MDS matrix
-const INV_MDS: [[BaseElement; STATE_WIDTH]; STATE_WIDTH] = [
-    [
-        BaseElement::new(14868391535953158196),
-        BaseElement::new(13278298489594233127),
-        BaseElement::new(389999932707070822),
-        BaseElement::new(9782021734907796003),
-        BaseElement::new(4829905704463175582),
-        BaseElement::new(7567822018949214430),
-        BaseElement::new(14205019324568680367),
-        BaseElement::new(15489674211196160593),
-        BaseElement::new(17636013826542227504),
-        BaseElement::new(16254215311946436093),
-        BaseElement::new(3641486184877122796),
-        BaseElement::new(11069068059762973582),
-    ],
-    [
-        BaseElement::new(11069068059762973582),
-        BaseElement::new(14868391535953158196),
-        BaseElement::new(13278298489594233127),
-        BaseElement::new(389999932707070822),
-        BaseElement::new(9782021734907796003),
-        BaseElement::new(4829905704463175582),
-        BaseElement::new(7567822018949214430),
-        BaseElement::new(14205019324568680367),
-        BaseElement::new(15489674211196160593),
-        BaseElement::new(17636013826542227504),
-        BaseElement::new(16254215311946436093),
-        BaseElement::new(3641486184877122796),
-    ],
-    [
-        BaseElement::new(3641486184877122796),
-        BaseElement::new(11069068059762973582),
-        BaseElement::new(14868391535953158196),
-        BaseElement::new(13278298489594233127),
-        BaseElement::new(389999932707070822),
-        BaseElement::new(9782021734907796003),
-        BaseElement::new(4829905704463175582),
-        BaseElement::new(7567822018949214430),
-        BaseElement::new(14205019324568680367),
-        BaseElement::new(15489674211196160593),
-        BaseElement::new(17636013826542227504),
-        BaseElement::new(16254215311946436093),
-    ],
-    [
-        BaseElement::new(16254215311946436093),
-        BaseElement::new(3641486184877122796),
-        BaseElement::new(11069068059762973582),
-        BaseElement::new(14868391535953158196),
-        BaseElement::new(13278298489594233127),
-        BaseElement::new(389999932707070822),
-        BaseElement::new(9782021734907796003),
-        BaseElement::new(4829905704463175582),
-        BaseElement::new(7567822018949214430),
-        BaseElement::new(14205019324568680367),
-        BaseElement::new(15489674211196160593),
-        BaseElement::new(17636013826542227504),
-    ],
-    [
-        BaseElement::new(17636013826542227504),
-        BaseElement::new(16254215311946436093),
-        BaseElement::new(3641486184877122796),
-        BaseElement::new(11069068059762973582),
-        BaseElement::new(14868391535953158196),
-        BaseElement::new(13278298489594233127),
-        BaseElement::new(389999932707070822),
-        BaseElement::new(9782021734907796003),
-        BaseElement::new(4829905704463175582),
-        BaseElement::new(7567822018949214430),
-        BaseElement::new(14205019324568680367),
-        BaseElement::new(15489674211196160593),
-    ],
-    [
-        BaseElement::new(15489674211196160593),
-        BaseElement::new(17636013826542227504),
-        BaseElement::new(16254215311946436093),
-        BaseElement::new(3641486184877122796),
-        BaseElement::new(11069068059762973582),
-        BaseElement::new(14868391535953158196),
-        BaseElement::new(13278298489594233127),
-        BaseElement::new(389999932707070822),
-        BaseElement::new(9782021734907796003),
-        BaseElement::new(4829905704463175582),
-        BaseElement::new(7567822018949214430),
-        BaseElement::new(14205019324568680367),
-    ],
-    [
-        BaseElement::new(14205019324568680367),
-        BaseElement::new(15489674211196160593),
-        BaseElement::new(17636013826542227504),
-        BaseElement::new(16254215311946436093),
-        BaseElement::new(3641486184877122796),
-        BaseElement::new(11069068059762973582),
-        BaseElement::new(14868391535953158196),
-        BaseElement::new(13278298489594233127),
-        BaseElement::new(389999932707070822),
-        BaseElement::new(9782021734907796003),
-        BaseElement::new(4829905704463175582),
-        BaseElement::new(7567822018949214430),
-    ],
-    [
-        BaseElement::new(7567822018949214430),
-        BaseElement::new(14205019324568680367),
-        BaseElement::new(15489674211196160593),
-        BaseElement::new(17636013826542227504),
-        BaseElement::new(16254215311946436093),
-        BaseElement::new(3641486184877122796),
-        BaseElement::new(11069068059762973582),
-        BaseElement::new(14868391535953158196),
-        BaseElement::new(13278298489594233127),
-        BaseElement::new(389999932707070822),
-        BaseElement::new(9782021734907796003),
-        BaseElement::new(4829905704463175582),
-    ],
-    [
-        BaseElement::new(4829905704463175582),
-        BaseElement::new(7567822018949214430),
-        BaseElement::new(14205019324568680367),
-        BaseElement::new(15489674211196160593),
-        BaseElement::new(17636013826542227504),
-        BaseElement::new(16254215311946436093),
-        BaseElement::new(3641486184877122796),
-        BaseElement::new(11069068059762973582),
-        BaseElement::new(14868391535953158196),
-        BaseElement::new(13278298489594233127),
-        BaseElement::new(389999932707070822),
-        BaseElement::new(9782021734907796003),
-    ],
-    [
-        BaseElement::new(9782021734907796003),
-        BaseElement::new(4829905704463175582),
-        BaseElement::new(7567822018949214430),
-        BaseElement::new(14205019324568680367),
-        BaseElement::new(15489674211196160593),
-        BaseElement::new(17636013826542227504),
-        BaseElement::new(16254215311946436093),
-        BaseElement::new(3641486184877122796),
-        BaseElement::new(11069068059762973582),
-        BaseElement::new(14868391535953158196),
-        BaseElement::new(13278298489594233127),
-        BaseElement::new(389999932707070822),
-    ],
-    [
-        BaseElement::new(389999932707070822),
-        BaseElement::new(9782021734907796003),
-        BaseElement::new(4829905704463175582),
-        BaseElement::new(7567822018949214430),
-        BaseElement::new(14205019324568680367),
-        BaseElement::new(15489674211196160593),
-        BaseElement::new(17636013826542227504),
-        BaseElement::new(16254215311946436093),
-        BaseElement::new(3641486184877122796),
-        BaseElement::new(11069068059762973582),
-        BaseElement::new(14868391535953158196),
-        BaseElement::new(13278298489594233127),
-    ],
-    [
-        BaseElement::new(13278298489594233127),
-        BaseElement::new(389999932707070822),
-        BaseElement::new(9782021734907796003),
-        BaseElement::new(4829905704463175582),
-        BaseElement::new(7567822018949214430),
-        BaseElement::new(14205019324568680367),
-        BaseElement::new(15489674211196160593),
-        BaseElement::new(17636013826542227504),
-        BaseElement::new(16254215311946436093),
-        BaseElement::new(3641486184877122796),
-        BaseElement::new(11069068059762973582),
-        BaseElement::new(14868391535953158196),
+const INV_MDS: [[Felt; STATE_WIDTH]; STATE_WIDTH] = [
+    [
+        Felt::new(14868391535953158196),
+        Felt::new(13278298489594233127),
+        Felt::new(389999932707070822),
+        Felt::new(9782021734907796003),
+        Felt::new(4829905704463175582),
+        Felt::new(7567822018949214430),
+        Felt::new(14205019324568680367),
+        Felt::new(15489674211196160593),
+        Felt::new(17636013826542227504),
+        Felt::new(16254215311946436093),
+        Felt::new(3641486184877122796),
+        Felt::new(11069068059762973582),
+    ],
+    [
+        Felt::new(11069068059762973582),
+        Felt::new(14868391535953158196),
+        Felt::new(13278298489594233127),
+        Felt::new(389999932707070822),
+        Felt::new(9782021734907796003),
+        Felt::new(4829905704463175582),
+        Felt::new(7567822018949214430),
+        Felt::new(14205019324568680367),
+        Felt::new(15489674211196160593),
+        Felt::new(17636013826542227504),
+        Felt::new(16254215311946436093),
+        Felt::new(3641486184877122796),
+    ],
+    [
+        Felt::new(3641486184877122796),
+        Felt::new(11069068059762973582),
+        Felt::new(14868391535953158196),
+        Felt::new(13278298489594233127),
+        Felt::new(389999932707070822),
+        Felt::new(9782021734907796003),
+        Felt::new(4829905704463175582),
+        Felt::new(7567822018949214430),
+        Felt::new(14205019324568680367),
+        Felt::new(15489674211196160593),
+        Felt::new(17636013826542227504),
+        Felt::new(16254215311946436093),
+    ],
+    [
+        Felt::new(16254215311946436093),
+        Felt::new(3641486184877122796),
+        Felt::new(11069068059762973582),
+        Felt::new(14868391535953158196),
+        Felt::new(13278298489594233127),
+        Felt::new(389999932707070822),
+        Felt::new(9782021734907796003),
+        Felt::new(4829905704463175582),
+        Felt::new(7567822018949214430),
+        Felt::new(14205019324568680367),
+        Felt::new(15489674211196160593),
+        Felt::new(17636013826542227504),
+    ],
+    [
+        Felt::new(17636013826542227504),
+        Felt::new(16254215311946436093),
+        Felt::new(3641486184877122796),
+        Felt::new(11069068059762973582),
+        Felt::new(14868391535953158196),
+        Felt::new(13278298489594233127),
+        Felt::new(389999932707070822),
+        Felt::new(9782021734907796003),
+        Felt::new(4829905704463175582),
+        Felt::new(7567822018949214430),
+        Felt::new(14205019324568680367),
+        Felt::new(15489674211196160593),
+    ],
+    [
+        Felt::new(15489674211196160593),
+        Felt::new(17636013826542227504),
+        Felt::new(16254215311946436093),
+        Felt::new(3641486184877122796),
+        Felt::new(11069068059762973582),
+        Felt::new(14868391535953158196),
+        Felt::new(13278298489594233127),
+        Felt::new(389999932707070822),
+        Felt::new(9782021734907796003),
+        Felt::new(4829905704463175582),
+        Felt::new(7567822018949214430),
+        Felt::new(14205019324568680367),
+    ],
+    [
+        Felt::new(14205019324568680367),
+        Felt::new(15489674211196160593),
+        Felt::new(17636013826542227504),
+        Felt::new(16254215311946436093),
+        Felt::new(3641486184877122796),
+        Felt::new(11069068059762973582),
+        Felt::new(14868391535953158196),
+        Felt::new(13278298489594233127),
+        Felt::new(389999932707070822),
+        Felt::new(9782021734907796003),
+        Felt::new(4829905704463175582),
+        Felt::new(7567822018949214430),
+    ],
+    [
+        Felt::new(7567822018949214430),
+        Felt::new(14205019324568680367),
+        Felt::new(15489674211196160593),
+        Felt::new(17636013826542227504),
+        Felt::new(16254215311946436093),
+        Felt::new(3641486184877122796),
+        Felt::new(11069068059762973582),
+        Felt::new(14868391535953158196),
+        Felt::new(13278298489594233127),
+        Felt::new(389999932707070822),
+        Felt::new(9782021734907796003),
+        Felt::new(4829905704463175582),
+    ],
+    [
+        Felt::new(4829905704463175582),
+        Felt::new(7567822018949214430),
+        Felt::new(14205019324568680367),
+        Felt::new(15489674211196160593),
+        Felt::new(17636013826542227504),
+        Felt::new(16254215311946436093),
+        Felt::new(3641486184877122796),
+        Felt::new(11069068059762973582),
+        Felt::new(14868391535953158196),
+        Felt::new(13278298489594233127),
+        Felt::new(389999932707070822),
+        Felt::new(9782021734907796003),
+    ],
+    [
+        Felt::new(9782021734907796003),
+        Felt::new(4829905704463175582),
+        Felt::new(7567822018949214430),
+        Felt::new(14205019324568680367),
+        Felt::new(15489674211196160593),
+        Felt::new(17636013826542227504),
+        Felt::new(16254215311946436093),
+        Felt::new(3641486184877122796),
+        Felt::new(11069068059762973582),
+        Felt::new(14868391535953158196),
+        Felt::new(13278298489594233127),
+        Felt::new(389999932707070822),
+    ],
+    [
+        Felt::new(389999932707070822),
+        Felt::new(9782021734907796003),
+        Felt::new(4829905704463175582),
+        Felt::new(7567822018949214430),
+        Felt::new(14205019324568680367),
+        Felt::new(15489674211196160593),
+        Felt::new(17636013826542227504),
+        Felt::new(16254215311946436093),
+        Felt::new(3641486184877122796),
+        Felt::new(11069068059762973582),
+        Felt::new(14868391535953158196),
+        Felt::new(13278298489594233127),
+    ],
+    [
+        Felt::new(13278298489594233127),
+        Felt::new(389999932707070822),
+        Felt::new(9782021734907796003),
+        Felt::new(4829905704463175582),
+        Felt::new(7567822018949214430),
+        Felt::new(14205019324568680367),
+        Felt::new(15489674211196160593),
+        Felt::new(17636013826542227504),
+        Felt::new(16254215311946436093),
+        Felt::new(3641486184877122796),
+        Felt::new(11069068059762973582),
+        Felt::new(14868391535953158196),
     ],
 ];
 
@@ -709,204 +722,204 @@ const INV_MDS: [[BaseElement; STATE_WIDTH]; STATE_WIDTH] = [
 ///
 /// The constants are broken up into two arrays ARK1 and ARK2; ARK1 contains the constants for the
 /// first half of RPO round, and ARK2 contains constants for the second half of RPO round.
-const ARK1: [[BaseElement; STATE_WIDTH]; NUM_ROUNDS] = [
-    [
-        BaseElement::new(5789762306288267392),
-        BaseElement::new(6522564764413701783),
-        BaseElement::new(17809893479458208203),
-        BaseElement::new(107145243989736508),
-        BaseElement::new(6388978042437517382),
-        BaseElement::new(15844067734406016715),
-        BaseElement::new(9975000513555218239),
-        BaseElement::new(3344984123768313364),
-        BaseElement::new(9959189626657347191),
-        BaseElement::new(12960773468763563665),
-        BaseElement::new(9602914297752488475),
-        BaseElement::new(16657542370200465908),
-    ],
-    [
-        BaseElement::new(12987190162843096997),
-        BaseElement::new(653957632802705281),
-        BaseElement::new(4441654670647621225),
-        BaseElement::new(4038207883745915761),
-        BaseElement::new(5613464648874830118),
-        BaseElement::new(13222989726778338773),
-        BaseElement::new(3037761201230264149),
-        BaseElement::new(16683759727265180203),
-        BaseElement::new(8337364536491240715),
-        BaseElement::new(3227397518293416448),
-        BaseElement::new(8110510111539674682),
-        BaseElement::new(2872078294163232137),
-    ],
-    [
-        BaseElement::new(18072785500942327487),
-        BaseElement::new(6200974112677013481),
-        BaseElement::new(17682092219085884187),
-        BaseElement::new(10599526828986756440),
-        BaseElement::new(975003873302957338),
-        BaseElement::new(8264241093196931281),
-        BaseElement::new(10065763900435475170),
-        BaseElement::new(2181131744534710197),
-        BaseElement::new(6317303992309418647),
-        BaseElement::new(1401440938888741532),
-        BaseElement::new(8884468225181997494),
-        BaseElement::new(13066900325715521532),
-    ],
-    [
-        BaseElement::new(5674685213610121970),
-        BaseElement::new(5759084860419474071),
-        BaseElement::new(13943282657648897737),
-        BaseElement::new(1352748651966375394),
-        BaseElement::new(17110913224029905221),
-        BaseElement::new(1003883795902368422),
-        BaseElement::new(4141870621881018291),
-        BaseElement::new(8121410972417424656),
-        BaseElement::new(14300518605864919529),
-        BaseElement::new(13712227150607670181),
-        BaseElement::new(17021852944633065291),
-        BaseElement::new(6252096473787587650),
-    ],
-    [
-        BaseElement::new(4887609836208846458),
-        BaseElement::new(3027115137917284492),
-        BaseElement::new(9595098600469470675),
-        BaseElement::new(10528569829048484079),
-        BaseElement::new(7864689113198939815),
-        BaseElement::new(17533723827845969040),
-        BaseElement::new(5781638039037710951),
-        BaseElement::new(17024078752430719006),
-        BaseElement::new(109659393484013511),
-        BaseElement::new(7158933660534805869),
-        BaseElement::new(2955076958026921730),
-        BaseElement::new(7433723648458773977),
-    ],
-    [
-        BaseElement::new(16308865189192447297),
-        BaseElement::new(11977192855656444890),
-        BaseElement::new(12532242556065780287),
-        BaseElement::new(14594890931430968898),
-        BaseElement::new(7291784239689209784),
-        BaseElement::new(5514718540551361949),
-        BaseElement::new(10025733853830934803),
-        BaseElement::new(7293794580341021693),
-        BaseElement::new(6728552937464861756),
-        BaseElement::new(6332385040983343262),
-        BaseElement::new(13277683694236792804),
-        BaseElement::new(2600778905124452676),
-    ],
-    [
-        BaseElement::new(7123075680859040534),
-        BaseElement::new(1034205548717903090),
-        BaseElement::new(7717824418247931797),
-        BaseElement::new(3019070937878604058),
-        BaseElement::new(11403792746066867460),
-        BaseElement::new(10280580802233112374),
-        BaseElement::new(337153209462421218),
-        BaseElement::new(13333398568519923717),
-        BaseElement::new(3596153696935337464),
-        BaseElement::new(8104208463525993784),
-        BaseElement::new(14345062289456085693),
-        BaseElement::new(17036731477169661256),
+const ARK1: [[Felt; STATE_WIDTH]; NUM_ROUNDS] = [
+    [
+        Felt::new(5789762306288267392),
+        Felt::new(6522564764413701783),
+        Felt::new(17809893479458208203),
+        Felt::new(107145243989736508),
+        Felt::new(6388978042437517382),
+        Felt::new(15844067734406016715),
+        Felt::new(9975000513555218239),
+        Felt::new(3344984123768313364),
+        Felt::new(9959189626657347191),
+        Felt::new(12960773468763563665),
+        Felt::new(9602914297752488475),
+        Felt::new(16657542370200465908),
+    ],
+    [
+        Felt::new(12987190162843096997),
+        Felt::new(653957632802705281),
+        Felt::new(4441654670647621225),
+        Felt::new(4038207883745915761),
+        Felt::new(5613464648874830118),
+        Felt::new(13222989726778338773),
+        Felt::new(3037761201230264149),
+        Felt::new(16683759727265180203),
+        Felt::new(8337364536491240715),
+        Felt::new(3227397518293416448),
+        Felt::new(8110510111539674682),
+        Felt::new(2872078294163232137),
+    ],
+    [
+        Felt::new(18072785500942327487),
+        Felt::new(6200974112677013481),
+        Felt::new(17682092219085884187),
+        Felt::new(10599526828986756440),
+        Felt::new(975003873302957338),
+        Felt::new(8264241093196931281),
+        Felt::new(10065763900435475170),
+        Felt::new(2181131744534710197),
+        Felt::new(6317303992309418647),
+        Felt::new(1401440938888741532),
+        Felt::new(8884468225181997494),
+        Felt::new(13066900325715521532),
+    ],
+    [
+        Felt::new(5674685213610121970),
+        Felt::new(5759084860419474071),
+        Felt::new(13943282657648897737),
+        Felt::new(1352748651966375394),
+        Felt::new(17110913224029905221),
+        Felt::new(1003883795902368422),
+        Felt::new(4141870621881018291),
+        Felt::new(8121410972417424656),
+        Felt::new(14300518605864919529),
+        Felt::new(13712227150607670181),
+        Felt::new(17021852944633065291),
+        Felt::new(6252096473787587650),
+    ],
+    [
+        Felt::new(4887609836208846458),
+        Felt::new(3027115137917284492),
+        Felt::new(9595098600469470675),
+        Felt::new(10528569829048484079),
+        Felt::new(7864689113198939815),
+        Felt::new(17533723827845969040),
+        Felt::new(5781638039037710951),
+        Felt::new(17024078752430719006),
+        Felt::new(109659393484013511),
+        Felt::new(7158933660534805869),
+        Felt::new(2955076958026921730),
+        Felt::new(7433723648458773977),
+    ],
+    [
+        Felt::new(16308865189192447297),
+        Felt::new(11977192855656444890),
+        Felt::new(12532242556065780287),
+        Felt::new(14594890931430968898),
+        Felt::new(7291784239689209784),
+        Felt::new(5514718540551361949),
+        Felt::new(10025733853830934803),
+        Felt::new(7293794580341021693),
+        Felt::new(6728552937464861756),
+        Felt::new(6332385040983343262),
+        Felt::new(13277683694236792804),
+        Felt::new(2600778905124452676),
+    ],
+    [
+        Felt::new(7123075680859040534),
+        Felt::new(1034205548717903090),
+        Felt::new(7717824418247931797),
+        Felt::new(3019070937878604058),
+        Felt::new(11403792746066867460),
+        Felt::new(10280580802233112374),
+        Felt::new(337153209462421218),
+        Felt::new(13333398568519923717),
+        Felt::new(3596153696935337464),
+        Felt::new(8104208463525993784),
+        Felt::new(14345062289456085693),
+        Felt::new(17036731477169661256),
     ],
 ];
 
-const ARK2: [[BaseElement; STATE_WIDTH]; NUM_ROUNDS] = [
-    [
-        BaseElement::new(6077062762357204287),
-        BaseElement::new(15277620170502011191),
-        BaseElement::new(5358738125714196705),
-        BaseElement::new(14233283787297595718),
-        BaseElement::new(13792579614346651365),
-        BaseElement::new(11614812331536767105),
-        BaseElement::new(14871063686742261166),
-        BaseElement::new(10148237148793043499),
-        BaseElement::new(4457428952329675767),
-        BaseElement::new(15590786458219172475),
-        BaseElement::new(10063319113072092615),
-        BaseElement::new(14200078843431360086),
-    ],
-    [
-        BaseElement::new(6202948458916099932),
-        BaseElement::new(17690140365333231091),
-        BaseElement::new(3595001575307484651),
-        BaseElement::new(373995945117666487),
-        BaseElement::new(1235734395091296013),
-        BaseElement::new(14172757457833931602),
-        BaseElement::new(707573103686350224),
-        BaseElement::new(15453217512188187135),
-        BaseElement::new(219777875004506018),
-        BaseElement::new(17876696346199469008),
-        BaseElement::new(17731621626449383378),
-        BaseElement::new(2897136237748376248),
-    ],
-    [
-        BaseElement::new(8023374565629191455),
-        BaseElement::new(15013690343205953430),
-        BaseElement::new(4485500052507912973),
-        BaseElement::new(12489737547229155153),
-        BaseElement::new(9500452585969030576),
-        BaseElement::new(2054001340201038870),
-        BaseElement::new(12420704059284934186),
-        BaseElement::new(355990932618543755),
-        BaseElement::new(9071225051243523860),
-        BaseElement::new(12766199826003448536),
-        BaseElement::new(9045979173463556963),
-        BaseElement::new(12934431667190679898),
-    ],
-    [
-        BaseElement::new(18389244934624494276),
-        BaseElement::new(16731736864863925227),
-        BaseElement::new(4440209734760478192),
-        BaseElement::new(17208448209698888938),
-        BaseElement::new(8739495587021565984),
-        BaseElement::new(17000774922218161967),
-        BaseElement::new(13533282547195532087),
-        BaseElement::new(525402848358706231),
-        BaseElement::new(16987541523062161972),
-        BaseElement::new(5466806524462797102),
-        BaseElement::new(14512769585918244983),
-        BaseElement::new(10973956031244051118),
-    ],
-    [
-        BaseElement::new(6982293561042362913),
-        BaseElement::new(14065426295947720331),
-        BaseElement::new(16451845770444974180),
-        BaseElement::new(7139138592091306727),
-        BaseElement::new(9012006439959783127),
-        BaseElement::new(14619614108529063361),
-        BaseElement::new(1394813199588124371),
-        BaseElement::new(4635111139507788575),
-        BaseElement::new(16217473952264203365),
-        BaseElement::new(10782018226466330683),
-        BaseElement::new(6844229992533662050),
-        BaseElement::new(7446486531695178711),
-    ],
-    [
-        BaseElement::new(3736792340494631448),
-        BaseElement::new(577852220195055341),
-        BaseElement::new(6689998335515779805),
-        BaseElement::new(13886063479078013492),
-        BaseElement::new(14358505101923202168),
-        BaseElement::new(7744142531772274164),
-        BaseElement::new(16135070735728404443),
-        BaseElement::new(12290902521256031137),
-        BaseElement::new(12059913662657709804),
-        BaseElement::new(16456018495793751911),
-        BaseElement::new(4571485474751953524),
-        BaseElement::new(17200392109565783176),
-    ],
-    [
-        BaseElement::new(17130398059294018733),
-        BaseElement::new(519782857322261988),
-        BaseElement::new(9625384390925085478),
-        BaseElement::new(1664893052631119222),
-        BaseElement::new(7629576092524553570),
-        BaseElement::new(3485239601103661425),
-        BaseElement::new(9755891797164033838),
-        BaseElement::new(15218148195153269027),
-        BaseElement::new(16460604813734957368),
-        BaseElement::new(9643968136937729763),
-        BaseElement::new(3611348709641382851),
-        BaseElement::new(18256379591337759196),
+const ARK2: [[Felt; STATE_WIDTH]; NUM_ROUNDS] = [
+    [
+        Felt::new(6077062762357204287),
+        Felt::new(15277620170502011191),
+        Felt::new(5358738125714196705),
+        Felt::new(14233283787297595718),
+        Felt::new(13792579614346651365),
+        Felt::new(11614812331536767105),
+        Felt::new(14871063686742261166),
+        Felt::new(10148237148793043499),
+        Felt::new(4457428952329675767),
+        Felt::new(15590786458219172475),
+        Felt::new(10063319113072092615),
+        Felt::new(14200078843431360086),
+    ],
+    [
+        Felt::new(6202948458916099932),
+        Felt::new(17690140365333231091),
+        Felt::new(3595001575307484651),
+        Felt::new(373995945117666487),
+        Felt::new(1235734395091296013),
+        Felt::new(14172757457833931602),
+        Felt::new(707573103686350224),
+        Felt::new(15453217512188187135),
+        Felt::new(219777875004506018),
+        Felt::new(17876696346199469008),
+        Felt::new(17731621626449383378),
+        Felt::new(2897136237748376248),
+    ],
+    [
+        Felt::new(8023374565629191455),
+        Felt::new(15013690343205953430),
+        Felt::new(4485500052507912973),
+        Felt::new(12489737547229155153),
+        Felt::new(9500452585969030576),
+        Felt::new(2054001340201038870),
+        Felt::new(12420704059284934186),
+        Felt::new(355990932618543755),
+        Felt::new(9071225051243523860),
+        Felt::new(12766199826003448536),
+        Felt::new(9045979173463556963),
+        Felt::new(12934431667190679898),
+    ],
+    [
+        Felt::new(18389244934624494276),
+        Felt::new(16731736864863925227),
+        Felt::new(4440209734760478192),
+        Felt::new(17208448209698888938),
+        Felt::new(8739495587021565984),
+        Felt::new(17000774922218161967),
+        Felt::new(13533282547195532087),
+        Felt::new(525402848358706231),
+        Felt::new(16987541523062161972),
+        Felt::new(5466806524462797102),
+        Felt::new(14512769585918244983),
+        Felt::new(10973956031244051118),
+    ],
+    [
+        Felt::new(6982293561042362913),
+        Felt::new(14065426295947720331),
+        Felt::new(16451845770444974180),
+        Felt::new(7139138592091306727),
+        Felt::new(9012006439959783127),
+        Felt::new(14619614108529063361),
+        Felt::new(1394813199588124371),
+        Felt::new(4635111139507788575),
+        Felt::new(16217473952264203365),
+        Felt::new(10782018226466330683),
+        Felt::new(6844229992533662050),
+        Felt::new(7446486531695178711),
+    ],
+    [
+        Felt::new(3736792340494631448),
+        Felt::new(577852220195055341),
+        Felt::new(6689998335515779805),
+        Felt::new(13886063479078013492),
+        Felt::new(14358505101923202168),
+        Felt::new(7744142531772274164),
+        Felt::new(16135070735728404443),
+        Felt::new(12290902521256031137),
+        Felt::new(12059913662657709804),
+        Felt::new(16456018495793751911),
+        Felt::new(4571485474751953524),
+        Felt::new(17200392109565783176),
+    ],
+    [
+        Felt::new(17130398059294018733),
+        Felt::new(519782857322261988),
+        Felt::new(9625384390925085478),
+        Felt::new(1664893052631119222),
+        Felt::new(7629576092524553570),
+        Felt::new(3485239601103661425),
+        Felt::new(9755891797164033838),
+        Felt::new(15218148195153269027),
+        Felt::new(16460604813734957368),
+        Felt::new(9643968136937729763),
+        Felt::new(3611348709641382851),
+        Felt::new(18256379591337759196),
     ],
 ];
diff --git a/crypto/src/hash/rpo/tests.rs b/crypto/src/hash/rpo/tests.rs
index 8eff6af..a6f7956 100644
--- a/crypto/src/hash/rpo/tests.rs
+++ b/crypto/src/hash/rpo/tests.rs
@@ -1,6 +1,6 @@
 use super::{
-    BaseElement, ElementDigest, ElementHasher, FieldElement, HashFn, Rpo, StarkField, ALPHA,
-    INV_ALPHA, INV_MDS, MDS, STATE_WIDTH,
+    ElementHasher, Felt, FieldElement, HashFn, Rpo, RpoDigest256, StarkField, ALPHA, INV_ALPHA,
+    INV_MDS, MDS, STATE_WIDTH, ZERO,
 };
 
 use core::convert::TryInto;
@@ -9,11 +9,11 @@ use rand_utils::rand_value;
 
 #[test]
 fn mds_inv_test() {
-    let mut mul_result = [[BaseElement::new(0); STATE_WIDTH]; STATE_WIDTH];
+    let mut mul_result = [[Felt::new(0); STATE_WIDTH]; STATE_WIDTH];
     for i in 0..STATE_WIDTH {
         for j in 0..STATE_WIDTH {
             let result = {
-                let mut result = BaseElement::new(0);
+                let mut result = Felt::new(0);
                 for k in 0..STATE_WIDTH {
                     result += MDS[i][k] * INV_MDS[k][j]
                 }
@@ -21,23 +21,23 @@ fn mds_inv_test() {
             };
             mul_result[i][j] = result;
             if i == j {
-                assert_eq!(result, BaseElement::new(1));
+                assert_eq!(result, Felt::new(1));
             } else {
-                assert_eq!(result, BaseElement::new(0));
+                assert_eq!(result, Felt::new(0));
             }
         }
     }
 }
 #[test]
 fn test_alphas() {
-    let e: BaseElement = BaseElement::new(rand_value());
+    let e: Felt = Felt::new(rand_value());
     let e_exp = e.exp(ALPHA.into());
     assert_eq!(e, e_exp.exp(INV_ALPHA));
 }
 
 #[test]
 fn test_sbox() {
-    let state = [BaseElement::new(rand_value()); STATE_WIDTH];
+    let state = [Felt::new(rand_value()); STATE_WIDTH];
 
     let mut expected = state;
     expected.iter_mut().for_each(|v| *v = v.exp(ALPHA));
@@ -50,7 +50,7 @@ fn test_sbox() {
 
 #[test]
 fn test_inv_sbox() {
-    let state = [BaseElement::new(rand_value()); STATE_WIDTH];
+    let state = [Felt::new(rand_value()); STATE_WIDTH];
 
     let mut expected = state;
     expected.iter_mut().for_each(|v| *v = v.exp(INV_ALPHA));
@@ -63,11 +63,11 @@ fn test_inv_sbox() {
 
 #[test]
 fn hash_elements_vs_merge() {
-    let elements = [BaseElement::new(rand_value()); 8];
+    let elements = [Felt::new(rand_value()); 8];
 
-    let digests: [ElementDigest; 2] = [
-        ElementDigest::new(elements[..4].try_into().unwrap()),
-        ElementDigest::new(elements[4..].try_into().unwrap()),
+    let digests: [RpoDigest256; 2] = [
+        RpoDigest256::new(elements[..4].try_into().unwrap()),
+        RpoDigest256::new(elements[4..].try_into().unwrap()),
     ];
 
     let m_result = Rpo::merge(&digests);
@@ -77,11 +77,11 @@ fn hash_elements_vs_merge() {
 
 #[test]
 fn hash_elements_vs_merge_with_int() {
-    let tmp = [BaseElement::new(rand_value()); 4];
-    let seed = ElementDigest::new(tmp);
+    let tmp = [Felt::new(rand_value()); 4];
+    let seed = RpoDigest256::new(tmp);
 
     // ----- value fits into a field element ------------------------------------------------------
-    let val: BaseElement = BaseElement::new(rand_value());
+    let val: Felt = Felt::new(rand_value());
     let m_result = Rpo::merge_with_int(seed, val.as_int());
 
     let mut elements = seed.as_elements().to_vec();
@@ -91,12 +91,12 @@ fn hash_elements_vs_merge_with_int() {
     assert_eq!(m_result, h_result);
 
     // ----- value does not fit into a field element ----------------------------------------------
-    let val = BaseElement::MODULUS + 2;
+    let val = Felt::MODULUS + 2;
     let m_result = Rpo::merge_with_int(seed, val);
 
     let mut elements = seed.as_elements().to_vec();
-    elements.push(BaseElement::new(val));
-    elements.push(BaseElement::new(1));
+    elements.push(Felt::new(val));
+    elements.push(Felt::new(1));
     let h_result = Rpo::hash_elements(&elements);
 
     assert_eq!(m_result, h_result);
@@ -127,8 +127,8 @@ fn hash_padding() {
 
 #[test]
 fn hash_elements_padding() {
-    let e1 = [BaseElement::new(rand_value()); 2];
-    let e2 = [e1[0], e1[1], BaseElement::ZERO];
+    let e1 = [Felt::new(rand_value()); 2];
+    let e2 = [e1[0], e1[1], ZERO];
 
     let r1 = Rpo::hash_elements(&e1);
     let r2 = Rpo::hash_elements(&e2);
@@ -138,24 +138,22 @@ fn hash_elements_padding() {
 #[test]
 fn hash_elements() {
     let elements = [
-        BaseElement::new(0),
-        BaseElement::new(1),
-        BaseElement::new(2),
-        BaseElement::new(3),
-        BaseElement::new(4),
-        BaseElement::new(5),
-        BaseElement::new(6),
-        BaseElement::new(7)
+        Felt::new(0),
+        Felt::new(1),
+        Felt::new(2),
+        Felt::new(3),
+        Felt::new(4),
+        Felt::new(5),
+        Felt::new(6),
+        Felt::new(7),
     ];
 
-    let digests: [ElementDigest; 2] = [
-        ElementDigest::new(elements[..4].try_into().unwrap()),
-        ElementDigest::new(elements[4..8].try_into().unwrap()),
+    let digests: [RpoDigest256; 2] = [
+        RpoDigest256::new(elements[..4].try_into().unwrap()),
+        RpoDigest256::new(elements[4..8].try_into().unwrap()),
     ];
 
     let m_result = Rpo::merge(&digests);
     let h_result = Rpo::hash_elements(&elements);
-    let m_elements :Vec<u64> = m_result.as_elements().iter().map(|a| {a.as_int()}).collect();
-    let h_elements :Vec<u64> = h_result.as_elements().iter().map(|a| {a.as_int()}).collect();
-    assert_eq!(m_elements, h_elements);
-}
\ No newline at end of file
+    assert_eq!(m_result, h_result);
+}
diff --git a/crypto/src/lib.rs b/crypto/src/lib.rs
index 62954c2..dfc1ab2 100644
--- a/crypto/src/lib.rs
+++ b/crypto/src/lib.rs
@@ -1,3 +1,7 @@
+pub use winter_utils::{
+    ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable,
+};
+pub use winterfell::crypto::{Digest, ElementHasher, Hasher as HashFn};
 pub use winterfell::math::{
     fields::{f64::BaseElement as Felt, QuadExtension},
     ExtensionOf, FieldElement, StarkField,
@@ -10,3 +14,12 @@ pub mod merkle;
 // ================================================================================================
 
 pub type Word = [Felt; 4];
+
+// CONSTANTS
+// ================================================================================================
+
+/// Field element representing ZERO in the base field of the VM.
+pub const ZERO: Felt = Felt::ZERO;
+
+/// Field element representing ONE in the base field of the VM.
+pub const ONE: Felt = Felt::ONE;
diff --git a/crypto/src/merkle/merkle_path_set.rs b/crypto/src/merkle/merkle_path_set.rs
index 8010676..0741fdc 100644
--- a/crypto/src/merkle/merkle_path_set.rs
+++ b/crypto/src/merkle/merkle_path_set.rs
@@ -1,5 +1,5 @@
-use super::{Felt, FieldElement, MerkleError, Word};
-use crate::hash::merge;
+use super::{MerkleError, Word};
+use crate::{hash::merge, ZERO};
 use winter_utils::collections::{BTreeMap, Vec};
 
 // MERKLE PATH SET
@@ -21,7 +21,7 @@ impl MerklePathSet {
 
     /// Returns an empty MerklePathSet.
     pub fn new(depth: u32) -> Result<Self, MerkleError> {
-        let root = [Felt::ZERO; 4];
+        let root = [ZERO; 4];
         let paths = BTreeMap::<u64, Vec<Word>>::new();
 
         Ok(Self {
@@ -68,7 +68,7 @@ impl MerklePathSet {
         }
 
         let root_of_current_path = compute_path_root(&extended_path, depth, index);
-        if self.root == [Felt::ZERO; 4] {
+        if self.root == [ZERO; 4] {
             self.root = root_of_current_path;
         } else if self.root != root_of_current_path {
             return Err(MerkleError::InvalidPath(extended_path));
@@ -254,7 +254,8 @@ fn compute_path_root(path: &[Word], depth: u32, index: u64) -> Word {
 
 #[cfg(test)]
 mod tests {
-    use super::{calculate_parent_hash, Felt, FieldElement, Word};
+    use super::calculate_parent_hash;
+    use crate::merkle::int_to_node;
 
     #[test]
     fn get_root() {
@@ -344,11 +345,4 @@ mod tests {
         assert_eq!(new_hash_45, new_path_6[1]);
         assert_eq!(new_hash_5, new_path_4[0]);
     }
-
-    // HELPER FUNCTIONS
-    // --------------------------------------------------------------------------------------------
-
-    const fn int_to_node(value: u64) -> Word {
-        [Felt::new(value), Felt::ZERO, Felt::ZERO, Felt::ZERO]
-    }
 }
diff --git a/crypto/src/merkle/merkle_tree.rs b/crypto/src/merkle/merkle_tree.rs
index d784d30..eec6d02 100644
--- a/crypto/src/merkle/merkle_tree.rs
+++ b/crypto/src/merkle/merkle_tree.rs
@@ -145,8 +145,7 @@ impl MerkleTree {
 
 #[cfg(test)]
 mod tests {
-    use super::{Felt, FieldElement, Word};
-    use crate::hash::{ElementHasher, HashFn, Hasher};
+    use crate::{hash::Hasher, merkle::int_to_node, ElementHasher, HashFn, Word};
 
     const LEAVES4: [Word; 4] = [
         int_to_node(1),
@@ -253,8 +252,4 @@ mod tests {
 
         (root.into(), node2.into(), node3.into())
     }
-
-    const fn int_to_node(value: u64) -> Word {
-        [Felt::new(value), Felt::ZERO, Felt::ZERO, Felt::ZERO]
-    }
 }
diff --git a/crypto/src/merkle/mod.rs b/crypto/src/merkle/mod.rs
index 290064f..fe9e815 100644
--- a/crypto/src/merkle/mod.rs
+++ b/crypto/src/merkle/mod.rs
@@ -16,3 +16,10 @@ pub enum MerkleError {
     InvalidPath(Vec<Word>),
     NodeNotInSet(u64),
 }
+
+// HELPER FUNCTIONS
+// ================================================================================================
+
+const fn int_to_node(value: u64) -> Word {
+    [Felt::new(value), Felt::ZERO, Felt::ZERO, Felt::ZERO]
+}

From a23c5192d77afc9d1529c8b0d95d4de3dcbed608 Mon Sep 17 00:00:00 2001
From: Al-Kindi-0 <82364884+Al-Kindi-0@users.noreply.github.com>
Date: Thu, 20 Oct 2022 14:44:36 +0200
Subject: [PATCH 3/9] fix misc. nits

---
 crypto/src/hash/rpo/tests.rs | 2 --
 1 file changed, 2 deletions(-)

diff --git a/crypto/src/hash/rpo/tests.rs b/crypto/src/hash/rpo/tests.rs
index a6f7956..a839f3a 100644
--- a/crypto/src/hash/rpo/tests.rs
+++ b/crypto/src/hash/rpo/tests.rs
@@ -2,9 +2,7 @@ use super::{
     ElementHasher, Felt, FieldElement, HashFn, Rpo, RpoDigest256, StarkField, ALPHA, INV_ALPHA,
     INV_MDS, MDS, STATE_WIDTH, ZERO,
 };
-
 use core::convert::TryInto;
-
 use rand_utils::rand_value;
 
 #[test]

From 1dd923a19963936a33a25aabe07e2ba491fa0b8b Mon Sep 17 00:00:00 2001
From: Al-Kindi-0 <82364884+Al-Kindi-0@users.noreply.github.com>
Date: Fri, 21 Oct 2022 09:24:54 +0200
Subject: [PATCH 4/9] add test vectors & fix nits

---
 crypto/src/hash/mod.rs          |   2 +-
 crypto/src/hash/rpo/mds_freq.rs |   4 +-
 crypto/src/hash/rpo/mod.rs      |  34 +++---
 crypto/src/hash/rpo/tests.rs    | 190 ++++++++++++++++++++++++++++----
 crypto/src/merkle/mod.rs        |   5 +-
 5 files changed, 192 insertions(+), 43 deletions(-)

diff --git a/crypto/src/hash/mod.rs b/crypto/src/hash/mod.rs
index 6674e22..2337fe1 100644
--- a/crypto/src/hash/mod.rs
+++ b/crypto/src/hash/mod.rs
@@ -1,7 +1,7 @@
 use crate::{ElementHasher, HashFn};
 
 mod rpo;
-pub use rpo::Rpo as Hasher;
+pub use rpo::Rpo256 as Hasher;
 
 // TYPE ALIASES
 // ================================================================================================
diff --git a/crypto/src/hash/rpo/mds_freq.rs b/crypto/src/hash/rpo/mds_freq.rs
index 8778c69..a510730 100644
--- a/crypto/src/hash/rpo/mds_freq.rs
+++ b/crypto/src/hash/rpo/mds_freq.rs
@@ -156,7 +156,7 @@ fn block3(x: [i64; 3], y: [i64; 3]) -> [i64; 3] {
 
 #[cfg(test)]
 mod tests {
-    use super::super::Rpo;
+    use super::super::Rpo256;
     use crate::hash::rpo::MDS;
     use proptest::prelude::*;
     use winterfell::math::{fields::f64::BaseElement, FieldElement};
@@ -187,7 +187,7 @@ mod tests {
             v2 = v1.clone();
 
             apply_mds_naive(&mut v1);
-            Rpo::apply_mds(&mut v2);
+            Rpo256::apply_mds(&mut v2);
 
             prop_assert_eq!(v1, v2);
         }
diff --git a/crypto/src/hash/rpo/mod.rs b/crypto/src/hash/rpo/mod.rs
index 8c4b66a..c711416 100644
--- a/crypto/src/hash/rpo/mod.rs
+++ b/crypto/src/hash/rpo/mod.rs
@@ -50,10 +50,10 @@ const INV_ALPHA: u64 = 10540996611094048183;
 // HASHER IMPLEMENTATION
 // ================================================================================================
 
-/// Implementation of [Hasher] trait for Rescue Prime Optimized (RPO) hash function with 256-bit output.
+/// Implementation of [Hasher] trait for Rescue Prime Optimized (Rpo256) hash function with 256-bit output.
 ///
 /// The hash function is implemented according to the Rescue Prime Optimized
-/// [specifications](https://github.org/aszepieniec/rpo/)
+/// [specifications](https://github.com/ASDiscreteMathematics/rpo)
 ///
 /// The parameters used to instantiate the function are:
 /// * Field: 64-bit prime field with modulus 2^64 - 2^32 + 1.
@@ -66,29 +66,29 @@ const INV_ALPHA: u64 = 10540996611094048183;
 /// and it can be serialized into 32 bytes (256 bits).
 ///
 /// ## Hash output consistency
-/// Functions [hash_elements()](Rpo::hash_elements), [merge()](Rpo::merge), and
-/// [merge_with_int()](Rpo::merge_with_int) are internally consistent. That is, computing
+/// Functions [hash_elements()](Rpo256::hash_elements), [merge()](Rpo256::merge), and
+/// [merge_with_int()](Rpo256::merge_with_int) are internally consistent. That is, computing
 /// a hash for the same set of elements using these functions will always produce the same
-/// result. For example, merging two digests using [merge()](Rpo::merge) will produce the
+/// result. For example, merging two digests using [merge()](Rpo256::merge) will produce the
 /// same result as hashing 8 elements which make up these digests using
-/// [hash_elements()](Rpo::hash_elements) function.
+/// [hash_elements()](Rpo256::hash_elements) function.
 ///
-/// However, [hash()](Rpo::hash) function is not consistent with functions mentioned above.
+/// However, [hash()](Rpo256::hash) function is not consistent with functions mentioned above.
 /// For example, if we take two field elements, serialize them to bytes and hash them using
-/// [hash()](Rpo::hash), the result will differ from the result obtained by hashing these
-/// elements directly using [hash_elements()](Rpo::hash_elements) function. The reason for
-/// this difference is that [hash()](Rpo::hash) function needs to be able to handle
+/// [hash()](Rpo256::hash), the result will differ from the result obtained by hashing these
+/// elements directly using [hash_elements()](Rpo256::hash_elements) function. The reason for
+/// this difference is that [hash()](Rpo256::hash) function needs to be able to handle
 /// arbitrary binary strings, which may or may not encode valid field elements - and thus,
 /// deserialization procedure used by this function is different from the procedure used to
 /// deserialize valid field elements.
 ///
 /// Thus, if the underlying data consists of valid field elements, it might make more sense
 /// to deserialize them into field elements and then hash them using
-/// [hash_elements()](Rpo::hash_elements) function rather then hashing the serialized bytes
-/// using [hash()](Rpo::hash) function.
-pub struct Rpo();
+/// [hash_elements()](Rpo256::hash_elements) function rather then hashing the serialized bytes
+/// using [hash()](Rpo256::hash) function.
+pub struct Rpo256();
 
-impl HashFn for Rpo {
+impl HashFn for Rpo256 {
     type Digest = RpoDigest256;
 
     fn hash(bytes: &[u8]) -> Self::Digest {
@@ -188,7 +188,7 @@ impl HashFn for Rpo {
     }
 }
 
-impl ElementHasher for Rpo {
+impl ElementHasher for Rpo256 {
     type BaseField = Felt;
 
     fn hash_elements<E: FieldElement<BaseField = Self::BaseField>>(elements: &[E]) -> Self::Digest {
@@ -237,7 +237,7 @@ impl ElementHasher for Rpo {
 // HASH FUNCTION IMPLEMENTATION
 // ================================================================================================
 
-impl Rpo {
+impl Rpo256 {
     // RESCUE PERMUTATION
     // --------------------------------------------------------------------------------------------
 
@@ -718,7 +718,7 @@ const INV_MDS: [[Felt; STATE_WIDTH]; STATE_WIDTH] = [
 // ================================================================================================
 
 /// Rescue round constants;
-/// computed as in section ?? from [specifications](https://github.org/aszepieniec/rpo/)
+/// computed as in [specifications](https://github.com/ASDiscreteMathematics/rpo)
 ///
 /// The constants are broken up into two arrays ARK1 and ARK2; ARK1 contains the constants for the
 /// first half of RPO round, and ARK2 contains constants for the second half of RPO round.
diff --git a/crypto/src/hash/rpo/tests.rs b/crypto/src/hash/rpo/tests.rs
index a839f3a..7f3bc37 100644
--- a/crypto/src/hash/rpo/tests.rs
+++ b/crypto/src/hash/rpo/tests.rs
@@ -1,5 +1,5 @@
 use super::{
-    ElementHasher, Felt, FieldElement, HashFn, Rpo, RpoDigest256, StarkField, ALPHA, INV_ALPHA,
+    ElementHasher, Felt, FieldElement, HashFn, Rpo256, RpoDigest256, StarkField, ALPHA, INV_ALPHA,
     INV_MDS, MDS, STATE_WIDTH, ZERO,
 };
 use core::convert::TryInto;
@@ -41,7 +41,7 @@ fn test_sbox() {
     expected.iter_mut().for_each(|v| *v = v.exp(ALPHA));
 
     let mut actual = state;
-    Rpo::apply_sbox(&mut actual);
+    Rpo256::apply_sbox(&mut actual);
 
     assert_eq!(expected, actual);
 }
@@ -54,7 +54,7 @@ fn test_inv_sbox() {
     expected.iter_mut().for_each(|v| *v = v.exp(INV_ALPHA));
 
     let mut actual = state;
-    Rpo::apply_inv_sbox(&mut actual);
+    Rpo256::apply_inv_sbox(&mut actual);
 
     assert_eq!(expected, actual);
 }
@@ -68,8 +68,8 @@ fn hash_elements_vs_merge() {
         RpoDigest256::new(elements[4..].try_into().unwrap()),
     ];
 
-    let m_result = Rpo::merge(&digests);
-    let h_result = Rpo::hash_elements(&elements);
+    let m_result = Rpo256::merge(&digests);
+    let h_result = Rpo256::hash_elements(&elements);
     assert_eq!(m_result, h_result);
 }
 
@@ -80,22 +80,22 @@ fn hash_elements_vs_merge_with_int() {
 
     // ----- value fits into a field element ------------------------------------------------------
     let val: Felt = Felt::new(rand_value());
-    let m_result = Rpo::merge_with_int(seed, val.as_int());
+    let m_result = Rpo256::merge_with_int(seed, val.as_int());
 
     let mut elements = seed.as_elements().to_vec();
     elements.push(val);
-    let h_result = Rpo::hash_elements(&elements);
+    let h_result = Rpo256::hash_elements(&elements);
 
     assert_eq!(m_result, h_result);
 
     // ----- value does not fit into a field element ----------------------------------------------
     let val = Felt::MODULUS + 2;
-    let m_result = Rpo::merge_with_int(seed, val);
+    let m_result = Rpo256::merge_with_int(seed, val);
 
     let mut elements = seed.as_elements().to_vec();
     elements.push(Felt::new(val));
     elements.push(Felt::new(1));
-    let h_result = Rpo::hash_elements(&elements);
+    let h_result = Rpo256::hash_elements(&elements);
 
     assert_eq!(m_result, h_result);
 }
@@ -103,23 +103,23 @@ fn hash_elements_vs_merge_with_int() {
 #[test]
 fn hash_padding() {
     // adding a zero bytes at the end of a byte string should result in a different hash
-    let r1 = Rpo::hash(&[1_u8, 2, 3]);
-    let r2 = Rpo::hash(&[1_u8, 2, 3, 0]);
+    let r1 = Rpo256::hash(&[1_u8, 2, 3]);
+    let r2 = Rpo256::hash(&[1_u8, 2, 3, 0]);
     assert_ne!(r1, r2);
 
     // same as above but with bigger inputs
-    let r1 = Rpo::hash(&[1_u8, 2, 3, 4, 5, 6]);
-    let r2 = Rpo::hash(&[1_u8, 2, 3, 4, 5, 6, 0]);
+    let r1 = Rpo256::hash(&[1_u8, 2, 3, 4, 5, 6]);
+    let r2 = Rpo256::hash(&[1_u8, 2, 3, 4, 5, 6, 0]);
     assert_ne!(r1, r2);
 
     // same as above but with input splitting over two elements
-    let r1 = Rpo::hash(&[1_u8, 2, 3, 4, 5, 6, 7]);
-    let r2 = Rpo::hash(&[1_u8, 2, 3, 4, 5, 6, 7, 0]);
+    let r1 = Rpo256::hash(&[1_u8, 2, 3, 4, 5, 6, 7]);
+    let r2 = Rpo256::hash(&[1_u8, 2, 3, 4, 5, 6, 7, 0]);
     assert_ne!(r1, r2);
 
     // same as above but with multiple zeros
-    let r1 = Rpo::hash(&[1_u8, 2, 3, 4, 5, 6, 7, 0, 0]);
-    let r2 = Rpo::hash(&[1_u8, 2, 3, 4, 5, 6, 7, 0, 0, 0, 0]);
+    let r1 = Rpo256::hash(&[1_u8, 2, 3, 4, 5, 6, 7, 0, 0]);
+    let r2 = Rpo256::hash(&[1_u8, 2, 3, 4, 5, 6, 7, 0, 0, 0, 0]);
     assert_ne!(r1, r2);
 }
 
@@ -128,8 +128,8 @@ fn hash_elements_padding() {
     let e1 = [Felt::new(rand_value()); 2];
     let e2 = [e1[0], e1[1], ZERO];
 
-    let r1 = Rpo::hash_elements(&e1);
-    let r2 = Rpo::hash_elements(&e2);
+    let r1 = Rpo256::hash_elements(&e1);
+    let r2 = Rpo256::hash_elements(&e2);
     assert_ne!(r1, r2);
 }
 
@@ -151,7 +151,155 @@ fn hash_elements() {
         RpoDigest256::new(elements[4..8].try_into().unwrap()),
     ];
 
-    let m_result = Rpo::merge(&digests);
-    let h_result = Rpo::hash_elements(&elements);
+    let m_result = Rpo256::merge(&digests);
+    let h_result = Rpo256::hash_elements(&elements);
     assert_eq!(m_result, h_result);
 }
+
+#[test]
+fn hash_test_vectors() {
+    let elements = [
+        Felt::new(0),
+        Felt::new(1),
+        Felt::new(2),
+        Felt::new(3),
+        Felt::new(4),
+        Felt::new(5),
+        Felt::new(6),
+        Felt::new(7),
+        Felt::new(8),
+        Felt::new(9),
+        Felt::new(10),
+        Felt::new(11),
+        Felt::new(12),
+        Felt::new(13),
+        Felt::new(14),
+        Felt::new(15),
+        Felt::new(16),
+        Felt::new(17),
+        Felt::new(18),
+    ];
+
+    for i in 0..elements.len() {
+        let expected = RpoDigest256::new(EXPECTED[i].try_into().unwrap());
+        let result = Rpo256::hash_elements(&elements[..(i + 1)]);
+        assert_eq!(result, expected);
+    }
+}
+
+const EXPECTED: [[Felt; 4]; 19] = [
+    [
+        Felt::new(1502364727743950833),
+        Felt::new(5880949717274681448),
+        Felt::new(162790463902224431),
+        Felt::new(6901340476773664264),
+    ],
+    [
+        Felt::new(7478710183745780580),
+        Felt::new(3308077307559720969),
+        Felt::new(3383561985796182409),
+        Felt::new(17205078494700259815),
+    ],
+    [
+        Felt::new(17439912364295172999),
+        Felt::new(17979156346142712171),
+        Felt::new(8280795511427637894),
+        Felt::new(9349844417834368814),
+    ],
+    [
+        Felt::new(5105868198472766874),
+        Felt::new(13090564195691924742),
+        Felt::new(1058904296915798891),
+        Felt::new(18379501748825152268),
+    ],
+    [
+        Felt::new(9133662113608941286),
+        Felt::new(12096627591905525991),
+        Felt::new(14963426595993304047),
+        Felt::new(13290205840019973377),
+    ],
+    [
+        Felt::new(3134262397541159485),
+        Felt::new(10106105871979362399),
+        Felt::new(138768814855329459),
+        Felt::new(15044809212457404677),
+    ],
+    [
+        Felt::new(162696376578462826),
+        Felt::new(4991300494838863586),
+        Felt::new(660346084748120605),
+        Felt::new(13179389528641752698),
+    ],
+    [
+        Felt::new(2242391899857912644),
+        Felt::new(12689382052053305418),
+        Felt::new(235236990017815546),
+        Felt::new(5046143039268215739),
+    ],
+    [
+        Felt::new(9585630502158073976),
+        Felt::new(1310051013427303477),
+        Felt::new(7491921222636097758),
+        Felt::new(9417501558995216762),
+    ],
+    [
+        Felt::new(1994394001720334744),
+        Felt::new(10866209900885216467),
+        Felt::new(13836092831163031683),
+        Felt::new(10814636682252756697),
+    ],
+    [
+        Felt::new(17486854790732826405),
+        Felt::new(17376549265955727562),
+        Felt::new(2371059831956435003),
+        Felt::new(17585704935858006533),
+    ],
+    [
+        Felt::new(11368277489137713825),
+        Felt::new(3906270146963049287),
+        Felt::new(10236262408213059745),
+        Felt::new(78552867005814007),
+    ],
+    [
+        Felt::new(17899847381280262181),
+        Felt::new(14717912805498651446),
+        Felt::new(10769146203951775298),
+        Felt::new(2774289833490417856),
+    ],
+    [
+        Felt::new(3794717687462954368),
+        Felt::new(4386865643074822822),
+        Felt::new(8854162840275334305),
+        Felt::new(7129983987107225269),
+    ],
+    [
+        Felt::new(7244773535611633983),
+        Felt::new(19359923075859320),
+        Felt::new(10898655967774994333),
+        Felt::new(9319339563065736480),
+    ],
+    [
+        Felt::new(4935426252518736883),
+        Felt::new(12584230452580950419),
+        Felt::new(8762518969632303998),
+        Felt::new(18159875708229758073),
+    ],
+    [
+        Felt::new(14871230873837295931),
+        Felt::new(11225255908868362971),
+        Felt::new(18100987641405432308),
+        Felt::new(1559244340089644233),
+    ],
+    [
+        Felt::new(8348203744950016968),
+        Felt::new(4041411241960726733),
+        Felt::new(17584743399305468057),
+        Felt::new(16836952610803537051),
+    ],
+    [
+        Felt::new(16139797453633030050),
+        Felt::new(1090233424040889412),
+        Felt::new(10770255347785669036),
+        Felt::new(16982398877290254028),
+    ],
+];
diff --git a/crypto/src/merkle/mod.rs b/crypto/src/merkle/mod.rs
index fe9e815..6c8819b 100644
--- a/crypto/src/merkle/mod.rs
+++ b/crypto/src/merkle/mod.rs
@@ -1,4 +1,4 @@
-use crate::{Felt, FieldElement, Word};
+use crate::{Felt, Word, ZERO};
 
 pub mod merkle_path_set;
 pub mod merkle_tree;
@@ -20,6 +20,7 @@ pub enum MerkleError {
 // HELPER FUNCTIONS
 // ================================================================================================
 
+#[cfg(test)]
 const fn int_to_node(value: u64) -> Word {
-    [Felt::new(value), Felt::ZERO, Felt::ZERO, Felt::ZERO]
+    [Felt::new(value), ZERO, ZERO, ZERO]
 }

From 1d5d11efaf8bd4fe6b747d1047297c96d13a94f7 Mon Sep 17 00:00:00 2001
From: Al-Kindi-0 <82364884+Al-Kindi-0@users.noreply.github.com>
Date: Mon, 24 Oct 2022 09:32:39 +0200
Subject: [PATCH 5/9] fix misc. nits

---
 crypto/src/hash/rpo/digest.rs   | 50 ++++++++++++++++++++++++++++-----
 crypto/src/hash/rpo/mds_freq.rs | 16 +++++------
 crypto/src/hash/rpo/mod.rs      |  5 +++-
 3 files changed, 55 insertions(+), 16 deletions(-)

diff --git a/crypto/src/hash/rpo/digest.rs b/crypto/src/hash/rpo/digest.rs
index 7c21298..a0c20a4 100644
--- a/crypto/src/hash/rpo/digest.rs
+++ b/crypto/src/hash/rpo/digest.rs
@@ -1,6 +1,7 @@
 use super::DIGEST_SIZE;
 use crate::{Digest, Felt, StarkField};
-use core::slice;
+use core::ops::Deref;
+
 use winterfell::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable};
 
 // DIGEST TRAIT IMPLEMENTATIONS
@@ -15,13 +16,14 @@ impl RpoDigest256 {
     }
 
     pub fn as_elements(&self) -> &[Felt] {
-        &self.0
+        self.as_ref()
     }
 
-    pub fn digests_as_elements(digests: &[Self]) -> &[Felt] {
-        let p = digests.as_ptr();
-        let len = digests.len() * DIGEST_SIZE;
-        unsafe { slice::from_raw_parts(p as *const Felt, len) }
+    pub fn digests_as_elements<'a, I>(digests: I) -> impl Iterator<Item = &'a Felt>
+    where
+        I: Iterator<Item = &'a Self>,
+    {
+        digests.map(|d| d.0.iter()).flatten()
     }
 }
 
@@ -52,7 +54,6 @@ impl Serializable for RpoDigest256 {
 
 impl Deserializable for RpoDigest256 {
     fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
-        // TODO: check if the field elements are valid?
         let e1 = Felt::new(source.read_u64()?);
         let e2 = Felt::new(source.read_u64()?);
         let e3 = Felt::new(source.read_u64()?);
@@ -80,6 +81,41 @@ impl From<RpoDigest256> for [u8; 32] {
     }
 }
 
+impl Deref for RpoDigest256 {
+    type Target = [Felt; DIGEST_SIZE];
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl RpoDigest256 {
+    fn iter(&self) -> RpoDigest256Iter<'_> {
+        RpoDigest256Iter {
+            values: &self.0,
+            index: 0,
+        }
+    }
+}
+
+pub struct RpoDigest256Iter<'a> {
+    values: &'a [Felt; DIGEST_SIZE],
+    index: usize,
+}
+
+impl<'a> Iterator for RpoDigest256Iter<'a> {
+    type Item = &'a Felt;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.index >= self.values.len() {
+            return None;
+        }
+
+        self.index += 1;
+        Some(&self.values[self.index - 1])
+    }
+}
+
 // TESTS
 // ================================================================================================
 
diff --git a/crypto/src/hash/rpo/mds_freq.rs b/crypto/src/hash/rpo/mds_freq.rs
index a510730..b3a6c90 100644
--- a/crypto/src/hash/rpo/mds_freq.rs
+++ b/crypto/src/hash/rpo/mds_freq.rs
@@ -26,7 +26,7 @@ const MDS_FREQ_BLOCK_THREE: [i64; 3] = [-8, 1, 1];
 
 // We use split 3 x 4 FFT transform in order to transform our vectors into the frequency domain.
 #[inline(always)]
-pub(crate) fn mds_multiply_freq(state: [u64; 12]) -> [u64; 12] {
+pub(crate) const fn mds_multiply_freq(state: [u64; 12]) -> [u64; 12] {
     let [s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11] = state;
 
     let (u0, u1, u2) = fft4_real([s0, s3, s6, s9]);
@@ -56,18 +56,18 @@ pub(crate) fn mds_multiply_freq(state: [u64; 12]) -> [u64; 12] {
 
 // We use the real FFT to avoid redundant computations. See https://www.mdpi.com/2076-3417/12/9/4700
 #[inline(always)]
-fn fft2_real(x: [u64; 2]) -> [i64; 2] {
+const fn fft2_real(x: [u64; 2]) -> [i64; 2] {
     [(x[0] as i64 + x[1] as i64), (x[0] as i64 - x[1] as i64)]
 }
 
 #[inline(always)]
-fn ifft2_real(y: [i64; 2]) -> [u64; 2] {
+const fn ifft2_real(y: [i64; 2]) -> [u64; 2] {
     // We avoid divisions by 2 by appropriately scaling the MDS matrix constants.
     [(y[0] + y[1]) as u64, (y[0] - y[1]) as u64]
 }
 
 #[inline(always)]
-fn fft4_real(x: [u64; 4]) -> (i64, (i64, i64), i64) {
+const fn fft4_real(x: [u64; 4]) -> (i64, (i64, i64), i64) {
     let [z0, z2] = fft2_real([x[0], x[2]]);
     let [z1, z3] = fft2_real([x[1], x[3]]);
     let y0 = z0 + z1;
@@ -77,7 +77,7 @@ fn fft4_real(x: [u64; 4]) -> (i64, (i64, i64), i64) {
 }
 
 #[inline(always)]
-fn ifft4_real(y: (i64, (i64, i64), i64)) -> [u64; 4] {
+const fn ifft4_real(y: (i64, (i64, i64), i64)) -> [u64; 4] {
     // In calculating 'z0' and 'z1', division by 2 is avoided by appropriately scaling
     // the MDS matrix constants.
     let z0 = y.0 + y.2;
@@ -92,7 +92,7 @@ fn ifft4_real(y: (i64, (i64, i64), i64)) -> [u64; 4] {
 }
 
 #[inline(always)]
-fn block1(x: [i64; 3], y: [i64; 3]) -> [i64; 3] {
+const fn block1(x: [i64; 3], y: [i64; 3]) -> [i64; 3] {
     let [x0, x1, x2] = x;
     let [y0, y1, y2] = y;
     let z0 = x0 * y0 + x1 * y2 + x2 * y1;
@@ -103,7 +103,7 @@ fn block1(x: [i64; 3], y: [i64; 3]) -> [i64; 3] {
 }
 
 #[inline(always)]
-fn block2(x: [(i64, i64); 3], y: [(i64, i64); 3]) -> [(i64, i64); 3] {
+const fn block2(x: [(i64, i64); 3], y: [(i64, i64); 3]) -> [(i64, i64); 3] {
     let [(x0r, x0i), (x1r, x1i), (x2r, x2i)] = x;
     let [(y0r, y0i), (y1r, y1i), (y2r, y2i)] = y;
     let x0s = x0r + x0i;
@@ -141,7 +141,7 @@ fn block2(x: [(i64, i64); 3], y: [(i64, i64); 3]) -> [(i64, i64); 3] {
 }
 
 #[inline(always)]
-fn block3(x: [i64; 3], y: [i64; 3]) -> [i64; 3] {
+const fn block3(x: [i64; 3], y: [i64; 3]) -> [i64; 3] {
     let [x0, x1, x2] = x;
     let [y0, y1, y2] = y;
     let z0 = x0 * y0 - x1 * y2 - x2 * y1;
diff --git a/crypto/src/hash/rpo/mod.rs b/crypto/src/hash/rpo/mod.rs
index c711416..b9370cb 100644
--- a/crypto/src/hash/rpo/mod.rs
+++ b/crypto/src/hash/rpo/mod.rs
@@ -153,7 +153,10 @@ impl HashFn for Rpo256 {
         // initialize the state by copying the digest elements into the rate portion of the state
         // (8 total elements), and set the capacity elements to 0.
         let mut state = [ZERO; STATE_WIDTH];
-        state[RATE_RANGE].copy_from_slice(Self::Digest::digests_as_elements(values));
+        let it = Self::Digest::digests_as_elements(values.into_iter());
+        for (i, v) in it.enumerate() {
+            state[RATE_RANGE.start + i] = *v;
+        }
 
         // apply the RPO permutation and return the first four elements of the state
         Self::apply_permutation(&mut state);

From 84f7a73666754e569f2f697cc6611b738bfc1a07 Mon Sep 17 00:00:00 2001
From: Al-Kindi-0 <82364884+Al-Kindi-0@users.noreply.github.com>
Date: Mon, 24 Oct 2022 11:53:24 +0200
Subject: [PATCH 6/9] feat: canonical range check deserialize Digest

---
 crypto/src/hash/rpo/digest.rs | 19 ++++++++++++-------
 1 file changed, 12 insertions(+), 7 deletions(-)

diff --git a/crypto/src/hash/rpo/digest.rs b/crypto/src/hash/rpo/digest.rs
index a0c20a4..da32980 100644
--- a/crypto/src/hash/rpo/digest.rs
+++ b/crypto/src/hash/rpo/digest.rs
@@ -1,7 +1,6 @@
 use super::DIGEST_SIZE;
-use crate::{Digest, Felt, StarkField};
+use crate::{Digest, Felt, StarkField, ZERO};
 use core::ops::Deref;
-
 use winterfell::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable};
 
 // DIGEST TRAIT IMPLEMENTATIONS
@@ -54,12 +53,18 @@ impl Serializable for RpoDigest256 {
 
 impl Deserializable for RpoDigest256 {
     fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
-        let e1 = Felt::new(source.read_u64()?);
-        let e2 = Felt::new(source.read_u64()?);
-        let e3 = Felt::new(source.read_u64()?);
-        let e4 = Felt::new(source.read_u64()?);
+        let mut inner: [Felt; DIGEST_SIZE] = [ZERO; DIGEST_SIZE];
+        for inner in inner.iter_mut() {
+            let e = source.read_u64()?;
+            if e >= Felt::MODULUS {
+                return Err(DeserializationError::InvalidValue(
+                    "Value not in the appropriate range".to_owned(),
+                ));
+            }
+            *inner = Felt::new(e);
+        }
 
-        Ok(Self([e1, e2, e3, e4]))
+        Ok(Self(inner))
     }
 }
 

From 7cb9f7bfdf3994fc1323943646d3fa429c9052b7 Mon Sep 17 00:00:00 2001
From: Al-Kindi-0 <82364884+Al-Kindi-0@users.noreply.github.com>
Date: Tue, 25 Oct 2022 14:00:02 +0200
Subject: [PATCH 7/9] chore: update to winter 0.4.1

---
 crypto/Cargo.toml                    |  5 +++--
 crypto/src/hash/rpo/digest.rs        |  7 +++----
 crypto/src/hash/rpo/mds_freq.rs      | 10 +++++-----
 crypto/src/hash/rpo/mod.rs           |  2 +-
 crypto/src/lib.rs                    | 14 ++++++++------
 crypto/src/merkle/merkle_path_set.rs |  3 +--
 crypto/src/merkle/merkle_tree.rs     |  4 +---
 7 files changed, 22 insertions(+), 23 deletions(-)

diff --git a/crypto/Cargo.toml b/crypto/Cargo.toml
index 87560ab..7d5b908 100644
--- a/crypto/Cargo.toml
+++ b/crypto/Cargo.toml
@@ -11,8 +11,9 @@ keywords = ["miden", "crypto", "hash", "merkle"]
 edition = "2021"
 
 [dependencies]
-winterfell = { git = "https://github.com/novifinancial/winterfell"}
-winter_utils = { version = "0.4", package = "winter-utils" }
+winter_utils = { version = "0.4.1", package = "winter-utils" }
+winter_math = { version = "0.4.1", package = "winter-math" }
+winter_crypto = { version = "0.4.1", package = "winter-crypto" }
 
 [dev-dependencies]
 proptest = "1.0.0"
diff --git a/crypto/src/hash/rpo/digest.rs b/crypto/src/hash/rpo/digest.rs
index da32980..f1ae3c2 100644
--- a/crypto/src/hash/rpo/digest.rs
+++ b/crypto/src/hash/rpo/digest.rs
@@ -1,7 +1,7 @@
 use super::DIGEST_SIZE;
+use crate::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable};
 use crate::{Digest, Felt, StarkField, ZERO};
 use core::ops::Deref;
-use winterfell::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable};
 
 // DIGEST TRAIT IMPLEMENTATIONS
 // ================================================================================================
@@ -22,7 +22,7 @@ impl RpoDigest256 {
     where
         I: Iterator<Item = &'a Self>,
     {
-        digests.map(|d| d.0.iter()).flatten()
+        digests.flat_map(|d| d.0.iter())
     }
 }
 
@@ -128,9 +128,8 @@ impl<'a> Iterator for RpoDigest256Iter<'a> {
 mod tests {
 
     use super::RpoDigest256;
-    use crate::Felt;
+    use crate::{Deserializable, Felt, Serializable, SliceReader};
     use rand_utils::rand_value;
-    use winterfell::{Deserializable, Serializable, SliceReader};
 
     #[test]
     fn digest_serialization() {
diff --git a/crypto/src/hash/rpo/mds_freq.rs b/crypto/src/hash/rpo/mds_freq.rs
index b3a6c90..8766bed 100644
--- a/crypto/src/hash/rpo/mds_freq.rs
+++ b/crypto/src/hash/rpo/mds_freq.rs
@@ -158,14 +158,14 @@ const fn block3(x: [i64; 3], y: [i64; 3]) -> [i64; 3] {
 mod tests {
     use super::super::Rpo256;
     use crate::hash::rpo::MDS;
+    use crate::{Felt, FieldElement};
     use proptest::prelude::*;
-    use winterfell::math::{fields::f64::BaseElement, FieldElement};
 
     const STATE_WIDTH: usize = 12;
 
     #[inline(always)]
-    fn apply_mds_naive(state: &mut [BaseElement; STATE_WIDTH]) {
-        let mut result = [BaseElement::ZERO; STATE_WIDTH];
+    fn apply_mds_naive(state: &mut [Felt; STATE_WIDTH]) {
+        let mut result = [Felt::ZERO; STATE_WIDTH];
         result.iter_mut().zip(MDS).for_each(|(r, mds_row)| {
             state.iter().zip(mds_row).for_each(|(&s, m)| {
                 *r += m * s;
@@ -178,11 +178,11 @@ mod tests {
         #[test]
         fn mds_freq_proptest(a in any::<[u64;STATE_WIDTH]>()) {
 
-            let mut v1 = [BaseElement::ZERO;STATE_WIDTH];
+            let mut v1 = [Felt::ZERO;STATE_WIDTH];
             let mut v2;
 
             for i in 0..STATE_WIDTH {
-                v1[i] = BaseElement::new(a[i]);
+                v1[i] = Felt::new(a[i]);
             }
             v2 = v1.clone();
 
diff --git a/crypto/src/hash/rpo/mod.rs b/crypto/src/hash/rpo/mod.rs
index b9370cb..cb1f290 100644
--- a/crypto/src/hash/rpo/mod.rs
+++ b/crypto/src/hash/rpo/mod.rs
@@ -153,7 +153,7 @@ impl HashFn for Rpo256 {
         // initialize the state by copying the digest elements into the rate portion of the state
         // (8 total elements), and set the capacity elements to 0.
         let mut state = [ZERO; STATE_WIDTH];
-        let it = Self::Digest::digests_as_elements(values.into_iter());
+        let it = Self::Digest::digests_as_elements(values.iter());
         for (i, v) in it.enumerate() {
             state[RATE_RANGE.start + i] = *v;
         }
diff --git a/crypto/src/lib.rs b/crypto/src/lib.rs
index dfc1ab2..3f7a48e 100644
--- a/crypto/src/lib.rs
+++ b/crypto/src/lib.rs
@@ -1,10 +1,12 @@
-pub use winter_utils::{
-    ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable,
-};
-pub use winterfell::crypto::{Digest, ElementHasher, Hasher as HashFn};
-pub use winterfell::math::{
+pub use winter_crypto::{Digest, ElementHasher, Hasher as HashFn};
+pub use winter_math::{
     fields::{f64::BaseElement as Felt, QuadExtension},
-    ExtensionOf, FieldElement, StarkField,
+    log2, ExtensionOf, FieldElement, StarkField,
+};
+pub use winter_utils::{
+    collections::{BTreeMap, Vec},
+    uninit_vector, ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable,
+    SliceReader,
 };
 
 pub mod hash;
diff --git a/crypto/src/merkle/merkle_path_set.rs b/crypto/src/merkle/merkle_path_set.rs
index 0741fdc..6ab7dde 100644
--- a/crypto/src/merkle/merkle_path_set.rs
+++ b/crypto/src/merkle/merkle_path_set.rs
@@ -1,6 +1,5 @@
 use super::{MerkleError, Word};
-use crate::{hash::merge, ZERO};
-use winter_utils::collections::{BTreeMap, Vec};
+use crate::{hash::merge, BTreeMap, Vec, ZERO};
 
 // MERKLE PATH SET
 // ================================================================================================
diff --git a/crypto/src/merkle/merkle_tree.rs b/crypto/src/merkle/merkle_tree.rs
index eec6d02..4a2953a 100644
--- a/crypto/src/merkle/merkle_tree.rs
+++ b/crypto/src/merkle/merkle_tree.rs
@@ -1,11 +1,9 @@
 use super::MerkleError;
 use crate::{
     hash::{merge, Digest},
-    Felt, FieldElement, Word,
+    log2, uninit_vector, Felt, FieldElement, Word,
 };
 use core::slice;
-use winter_utils::uninit_vector;
-use winterfell::math::log2;
 
 // MERKLE TREE
 // ================================================================================================

From 636c92a78bff6e004a849d3948159a16202851a5 Mon Sep 17 00:00:00 2001
From: Al-Kindi-0 <82364884+Al-Kindi-0@users.noreply.github.com>
Date: Wed, 26 Oct 2022 07:01:46 +0200
Subject: [PATCH 8/9] chore: change to BINARY_CHUNK_SIZE

---
 crypto/src/hash/rpo/mod.rs | 36 ++++++++++++++++++++----------------
 1 file changed, 20 insertions(+), 16 deletions(-)

diff --git a/crypto/src/hash/rpo/mod.rs b/crypto/src/hash/rpo/mod.rs
index cb1f290..720e7d2 100644
--- a/crypto/src/hash/rpo/mod.rs
+++ b/crypto/src/hash/rpo/mod.rs
@@ -38,6 +38,9 @@ const DIGEST_SIZE: usize = DIGEST_RANGE.end - DIGEST_RANGE.start;
 /// The number of rounds is set to 7 to target 128-bit security level
 const NUM_ROUNDS: usize = 7;
 
+/// The number of byte chunks defining a field element when hashing a sequence of bytes
+const BINARY_CHUNK_SIZE: usize = 7;
+
 /// S-Box and Inverse S-Box powers;
 ///
 /// The constants are defined for tests only because the exponentiations in the code are unrolled
@@ -93,12 +96,12 @@ impl HashFn for Rpo256 {
 
     fn hash(bytes: &[u8]) -> Self::Digest {
         // compute the number of elements required to represent the string; we will be processing
-        // the string in 7-byte chunks, thus the number of elements will be equal to the number
-        // of such chunks (including a potential partial chunk at the end).
-        let num_elements = if bytes.len() % 7 == 0 {
-            bytes.len() / 7
+        // the string in BINARY_CHUNK_SIZE-byte chunks, thus the number of elements will be equal
+        // to the number of such chunks (including a potential partial chunk at the end).
+        let num_elements = if bytes.len() % BINARY_CHUNK_SIZE == 0 {
+            bytes.len() / BINARY_CHUNK_SIZE
         } else {
-            bytes.len() / 7 + 1
+            bytes.len() / BINARY_CHUNK_SIZE + 1
         };
 
         // initialize state to all zeros, except for the first element of the capacity part, which
@@ -107,19 +110,20 @@ impl HashFn for Rpo256 {
         let mut state = [ZERO; STATE_WIDTH];
         state[CAPACITY_RANGE.start] = Felt::new(num_elements as u64);
 
-        // break the string into 7-byte chunks, convert each chunk into a field element, and
-        // absorb the element into the rate portion of the state. we use 7-byte chunks because
-        // every 7-byte chunk is guaranteed to map to some field element.
+        // break the string into BINARY_CHUNK_SIZE-byte chunks, convert each chunk into a field
+        // element, and absorb the element into the rate portion of the state. we use
+        // BINARY_CHUNK_SIZE-byte chunks because every BINARY_CHUNK_SIZE-byte chunk is guaranteed
+        // to map to some field element.
         let mut i = 0;
         let mut buf = [0_u8; 8];
-        for chunk in bytes.chunks(7) {
+        for chunk in bytes.chunks(BINARY_CHUNK_SIZE) {
             if i < num_elements - 1 {
-                buf[..7].copy_from_slice(chunk);
+                buf[..BINARY_CHUNK_SIZE].copy_from_slice(chunk);
             } else {
-                // if we are dealing with the last chunk, it may be smaller than 7 bytes long, so
-                // we need to handle it slightly differently. we also append a byte with value 1
-                // to the end of the string; this pads the string in such a way that adding
-                // trailing zeros results in different hash
+                // if we are dealing with the last chunk, it may be smaller than BINARY_CHUNK_SIZE
+                // bytes long, so we need to handle it slightly differently. We also append a byte
+                // with value 1 to the end of the string; this pads the string in such a way that
+                // adding trailing zeros results in different hash
                 let chunk_len = chunk.len();
                 buf = [0_u8; 8];
                 buf[..chunk_len].copy_from_slice(chunk);
@@ -129,7 +133,7 @@ impl HashFn for Rpo256 {
             // convert the bytes into a field element and absorb it into the rate portion of the
             // state; if the rate is filled up, apply the Rescue permutation and start absorbing
             // again from zero index.
-            state[RATE_RANGE.start + i] += Felt::new(u64::from_le_bytes(buf));
+            state[RATE_RANGE.start + i] = Felt::new(u64::from_le_bytes(buf));
             i += 1;
             if i % RATE_WIDTH == 0 {
                 Self::apply_permutation(&mut state);
@@ -138,7 +142,7 @@ impl HashFn for Rpo256 {
         }
 
         // if we absorbed some elements but didn't apply a permutation to them (would happen when
-        // the number of elements is not a multiple of RATE_WIDTH), apply the Rescue permutation.
+        // the number of elements is not a multiple of RATE_WIDTH), apply the RPO permutation.
         // we don't need to apply any extra padding because we injected total number of elements
         // in the input list into the capacity portion of the state during initialization.
         if i > 0 {

From 09014a34d1c16bca8a021400f34d09ca765c7d59 Mon Sep 17 00:00:00 2001
From: Al-Kindi-0 <82364884+Al-Kindi-0@users.noreply.github.com>
Date: Wed, 26 Oct 2022 07:48:17 +0200
Subject: [PATCH 9/9] fix: clippy warnings

---
 crypto/src/hash/mod.rs        |  1 +
 crypto/src/hash/rpo/digest.rs | 27 ---------------------------
 crypto/src/hash/rpo/mod.rs    |  4 ++--
 crypto/src/merkle/mod.rs      |  5 ++++-
 4 files changed, 7 insertions(+), 30 deletions(-)

diff --git a/crypto/src/hash/mod.rs b/crypto/src/hash/mod.rs
index 2337fe1..dec2761 100644
--- a/crypto/src/hash/mod.rs
+++ b/crypto/src/hash/mod.rs
@@ -2,6 +2,7 @@ use crate::{ElementHasher, HashFn};
 
 mod rpo;
 pub use rpo::Rpo256 as Hasher;
+pub use rpo::{INV_MDS, MDS};
 
 // TYPE ALIASES
 // ================================================================================================
diff --git a/crypto/src/hash/rpo/digest.rs b/crypto/src/hash/rpo/digest.rs
index f1ae3c2..3eef96f 100644
--- a/crypto/src/hash/rpo/digest.rs
+++ b/crypto/src/hash/rpo/digest.rs
@@ -94,33 +94,6 @@ impl Deref for RpoDigest256 {
     }
 }
 
-impl RpoDigest256 {
-    fn iter(&self) -> RpoDigest256Iter<'_> {
-        RpoDigest256Iter {
-            values: &self.0,
-            index: 0,
-        }
-    }
-}
-
-pub struct RpoDigest256Iter<'a> {
-    values: &'a [Felt; DIGEST_SIZE],
-    index: usize,
-}
-
-impl<'a> Iterator for RpoDigest256Iter<'a> {
-    type Item = &'a Felt;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        if self.index >= self.values.len() {
-            return None;
-        }
-
-        self.index += 1;
-        Some(&self.values[self.index - 1])
-    }
-}
-
 // TESTS
 // ================================================================================================
 
diff --git a/crypto/src/hash/rpo/mod.rs b/crypto/src/hash/rpo/mod.rs
index 720e7d2..5a88eb8 100644
--- a/crypto/src/hash/rpo/mod.rs
+++ b/crypto/src/hash/rpo/mod.rs
@@ -378,7 +378,7 @@ impl Rpo256 {
 // MDS
 // ================================================================================================
 /// RPO MDS matrix
-const MDS: [[Felt; STATE_WIDTH]; STATE_WIDTH] = [
+pub const MDS: [[Felt; STATE_WIDTH]; STATE_WIDTH] = [
     [
         Felt::new(7),
         Felt::new(23),
@@ -550,7 +550,7 @@ const MDS: [[Felt; STATE_WIDTH]; STATE_WIDTH] = [
 ];
 
 /// RPO Inverse MDS matrix
-const INV_MDS: [[Felt; STATE_WIDTH]; STATE_WIDTH] = [
+pub const INV_MDS: [[Felt; STATE_WIDTH]; STATE_WIDTH] = [
     [
         Felt::new(14868391535953158196),
         Felt::new(13278298489594233127),
diff --git a/crypto/src/merkle/mod.rs b/crypto/src/merkle/mod.rs
index 6c8819b..63e5488 100644
--- a/crypto/src/merkle/mod.rs
+++ b/crypto/src/merkle/mod.rs
@@ -1,4 +1,7 @@
-use crate::{Felt, Word, ZERO};
+use crate::Word;
+
+#[cfg(test)]
+use crate::{Felt, ZERO};
 
 pub mod merkle_path_set;
 pub mod merkle_tree;