arnaucube
arnaucube/miden-cryptoPublic Self-hosted repository
miden-crypto/src/hash/rescue/mds/freq.rs
src/hash/rescue/mds/freq.rsRaw
  1// FFT-BASED MDS MULTIPLICATION HELPER FUNCTIONS
  2// ================================================================================================
  3
  4//! This module contains helper functions as well as constants used to perform the vector-matrix
  5//! multiplication step of the Rescue prime permutation. The special form of our MDS matrix
  6//! i.e. being circular, allows us to reduce the vector-matrix multiplication to a Hadamard product
  7//! of two vectors in "frequency domain". This follows from the simple fact that every circulant
  8//! matrix has the columns of the discrete Fourier transform matrix as orthogonal eigenvectors.
  9//! The implementation also avoids the use of 3-point FFTs, and 3-point iFFTs, and substitutes that
 10//! with explicit expressions. It also avoids, due to the form of our matrix in the frequency
 11//! domain, divisions by 2 and repeated modular reductions. This is because of our explicit choice
 12//! of an MDS matrix that has small powers of 2 entries in frequency domain.
 13//! The following implementation has benefited greatly from the discussions and insights of
 14//! Hamish Ivey-Law and Jacqueline Nabaglo of Polygon Zero and is base on Nabaglo's Plonky2
 15//! implementation.
 16
 17// Rescue MDS matrix in frequency domain.
 18//
 19// More precisely, this is the output of the three 4-point (real) FFTs of the first column of
 20// the MDS matrix i.e. just before the multiplication with the appropriate twiddle factors
 21// and application of the final four 3-point FFT in order to get the full 12-point FFT.
 22// The entries have been scaled appropriately in order to avoid divisions by 2 in iFFT2 and iFFT4.
 23// The code to generate the matrix in frequency domain is based on an adaptation of a code, to
 24// generate MDS matrices efficiently in original domain, that was developed by the Polygon Zero
 25// team.
 26const MDS_FREQ_BLOCK_ONE: [i64; 3] = [16, 8, 16];
 27const MDS_FREQ_BLOCK_TWO: [(i64, i64); 3] = [(-1, 2), (-1, 1), (4, 8)];
 28const MDS_FREQ_BLOCK_THREE: [i64; 3] = [-8, 1, 1];
 29
 30// We use split 3 x 4 FFT transform in order to transform our vectors into the frequency domain.
 31#[inline(always)]
 32pub const fn mds_multiply_freq(state: [u64; 12]) -> [u64; 12] {
 33    let [s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11] = state;
 34
 35    let (u0, u1, u2) = fft4_real([s0, s3, s6, s9]);
 36    let (u4, u5, u6) = fft4_real([s1, s4, s7, s10]);
 37    let (u8, u9, u10) = fft4_real([s2, s5, s8, s11]);
 38
 39    // This where the multiplication in frequency domain is done. More precisely, and with
 40    // the appropriate permutations in between, the sequence of
 41    // 3-point FFTs --> multiplication by twiddle factors --> Hadamard multiplication -->
 42    // 3 point iFFTs --> multiplication by (inverse) twiddle factors
 43    // is "squashed" into one step composed of the functions "block1", "block2" and "block3".
 44    // The expressions in the aforementioned functions are the result of explicit computations
 45    // combined with the Karatsuba trick for the multiplication of Complex numbers.
 46
 47    let [v0, v4, v8] = block1([u0, u4, u8], MDS_FREQ_BLOCK_ONE);
 48    let [v1, v5, v9] = block2([u1, u5, u9], MDS_FREQ_BLOCK_TWO);
 49    let [v2, v6, v10] = block3([u2, u6, u10], MDS_FREQ_BLOCK_THREE);
 50    // The 4th block is not computed as it is similar to the 2nd one, up to complex conjugation,
 51    // and is, due to the use of the real FFT and iFFT, redundant.
 52
 53    let [s0, s3, s6, s9] = ifft4_real((v0, v1, v2));
 54    let [s1, s4, s7, s10] = ifft4_real((v4, v5, v6));
 55    let [s2, s5, s8, s11] = ifft4_real((v8, v9, v10));
 56
 57    [s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11]
 58}
 59
 60// We use the real FFT to avoid redundant computations. See https://www.mdpi.com/2076-3417/12/9/4700
 61#[inline(always)]
 62const fn fft2_real(x: [u64; 2]) -> [i64; 2] {
 63    [(x[0] as i64 + x[1] as i64), (x[0] as i64 - x[1] as i64)]
 64}
 65
 66#[inline(always)]
 67const fn ifft2_real(y: [i64; 2]) -> [u64; 2] {
 68    // We avoid divisions by 2 by appropriately scaling the MDS matrix constants.
 69    [(y[0] + y[1]) as u64, (y[0] - y[1]) as u64]
 70}
 71
 72#[inline(always)]
 73const fn fft4_real(x: [u64; 4]) -> (i64, (i64, i64), i64) {
 74    let [z0, z2] = fft2_real([x[0], x[2]]);
 75    let [z1, z3] = fft2_real([x[1], x[3]]);
 76    let y0 = z0 + z1;
 77    let y1 = (z2, -z3);
 78    let y2 = z0 - z1;
 79    (y0, y1, y2)
 80}
 81
 82#[inline(always)]
 83const fn ifft4_real(y: (i64, (i64, i64), i64)) -> [u64; 4] {
 84    // In calculating 'z0' and 'z1', division by 2 is avoided by appropriately scaling
 85    // the MDS matrix constants.
 86    let z0 = y.0 + y.2;
 87    let z1 = y.0 - y.2;
 88    let z2 = y.1 .0;
 89    let z3 = -y.1 .1;
 90
 91    let [x0, x2] = ifft2_real([z0, z2]);
 92    let [x1, x3] = ifft2_real([z1, z3]);
 93
 94    [x0, x1, x2, x3]
 95}
 96
 97#[inline(always)]
 98const fn block1(x: [i64; 3], y: [i64; 3]) -> [i64; 3] {
 99    let [x0, x1, x2] = x;
100    let [y0, y1, y2] = y;
101    let z0 = x0 * y0 + x1 * y2 + x2 * y1;
102    let z1 = x0 * y1 + x1 * y0 + x2 * y2;
103    let z2 = x0 * y2 + x1 * y1 + x2 * y0;
104
105    [z0, z1, z2]
106}
107
108#[inline(always)]
109const fn block2(x: [(i64, i64); 3], y: [(i64, i64); 3]) -> [(i64, i64); 3] {
110    let [(x0r, x0i), (x1r, x1i), (x2r, x2i)] = x;
111    let [(y0r, y0i), (y1r, y1i), (y2r, y2i)] = y;
112    let x0s = x0r + x0i;
113    let x1s = x1r + x1i;
114    let x2s = x2r + x2i;
115    let y0s = y0r + y0i;
116    let y1s = y1r + y1i;
117    let y2s = y2r + y2i;
118
119    // Compute x0​y0 ​− ix1​y2​ − ix2​y1​ using Karatsuba for complex numbers multiplication
120    let m0 = (x0r * y0r, x0i * y0i);
121    let m1 = (x1r * y2r, x1i * y2i);
122    let m2 = (x2r * y1r, x2i * y1i);
123    let z0r = (m0.0 - m0.1) + (x1s * y2s - m1.0 - m1.1) + (x2s * y1s - m2.0 - m2.1);
124    let z0i = (x0s * y0s - m0.0 - m0.1) + (-m1.0 + m1.1) + (-m2.0 + m2.1);
125    let z0 = (z0r, z0i);
126
127    // Compute x0​y1​ + x1​y0​ − ix2​y2 using Karatsuba for complex numbers multiplication
128    let m0 = (x0r * y1r, x0i * y1i);
129    let m1 = (x1r * y0r, x1i * y0i);
130    let m2 = (x2r * y2r, x2i * y2i);
131    let z1r = (m0.0 - m0.1) + (m1.0 - m1.1) + (x2s * y2s - m2.0 - m2.1);
132    let z1i = (x0s * y1s - m0.0 - m0.1) + (x1s * y0s - m1.0 - m1.1) + (-m2.0 + m2.1);
133    let z1 = (z1r, z1i);
134
135    // Compute x0​y2​ + x1​y1 ​+ x2​y0​ using Karatsuba for complex numbers multiplication
136    let m0 = (x0r * y2r, x0i * y2i);
137    let m1 = (x1r * y1r, x1i * y1i);
138    let m2 = (x2r * y0r, x2i * y0i);
139    let z2r = (m0.0 - m0.1) + (m1.0 - m1.1) + (m2.0 - m2.1);
140    let z2i = (x0s * y2s - m0.0 - m0.1) + (x1s * y1s - m1.0 - m1.1) + (x2s * y0s - m2.0 - m2.1);
141    let z2 = (z2r, z2i);
142
143    [z0, z1, z2]
144}
145
146#[inline(always)]
147const fn block3(x: [i64; 3], y: [i64; 3]) -> [i64; 3] {
148    let [x0, x1, x2] = x;
149    let [y0, y1, y2] = y;
150    let z0 = x0 * y0 - x1 * y2 - x2 * y1;
151    let z1 = x0 * y1 + x1 * y0 - x2 * y2;
152    let z2 = x0 * y2 + x1 * y1 + x2 * y0;
153
154    [z0, z1, z2]
155}
156
157// TESTS
158// ================================================================================================
159
160#[cfg(test)]
161mod tests {
162    use proptest::prelude::*;
163
164    use super::super::{apply_mds, Felt, MDS, ZERO};
165
166    const STATE_WIDTH: usize = 12;
167
168    #[inline(always)]
169    fn apply_mds_naive(state: &mut [Felt; STATE_WIDTH]) {
170        let mut result = [ZERO; STATE_WIDTH];
171        result.iter_mut().zip(MDS).for_each(|(r, mds_row)| {
172            state.iter().zip(mds_row).for_each(|(&s, m)| {
173                *r += m * s;
174            });
175        });
176        *state = result;
177    }
178
179    proptest! {
180        #[test]
181        fn mds_freq_proptest(a in any::<[u64; STATE_WIDTH]>()) {
182
183            let mut v1 = [ZERO; STATE_WIDTH];
184            let mut v2;
185
186            for i in 0..STATE_WIDTH {
187                v1[i] = Felt::new(a[i]);
188            }
189            v2 = v1;
190
191            apply_mds_naive(&mut v1);
192            apply_mds(&mut v2);
193
194            prop_assert_eq!(v1, v2);
195        }
196    }
197}