1#![allow(incomplete_features)]
  2#![feature(generic_const_exprs)]
  3
  4pub mod circuit;
  5
  6use std::array;
  7
  8use plonky2::field::types::Field;
  9
 10pub fn feistel_cipher<const STATE_HALF_LEN: usize, const KEY_LEN: usize, F: Field>(
 11    state: [F; 2 * STATE_HALF_LEN],
 12    key_schedule: &[[F; KEY_LEN]],
 13    f: &impl Fn(Vec<F>) -> [F; STATE_HALF_LEN],
 14) -> [F; 2 * STATE_HALF_LEN] {
 15    if key_schedule.is_empty() {
 16        state
 17    } else {
 18        let k = key_schedule[0];
 19        let l = &state[..STATE_HALF_LEN];
 20        let r = &state[STATE_HALF_LEN..];
 21
 22        let new_l = r;
 23        let round_offset = f([r, &k].concat());
 24        let new_r: [_; STATE_HALF_LEN] = array::from_fn(|i| l[i] + round_offset[i]);
 25
 26        feistel_cipher(
 27            array::from_fn(|i| {
 28                if i < STATE_HALF_LEN {
 29                    new_l[i]
 30                } else {
 31                    new_r[i - STATE_HALF_LEN]
 32                }
 33            }),
 34            &key_schedule[1..],
 35            f,
 36        )
 37    }
 38}
 39
 40pub fn feistel_inv_cipher<const STATE_HALF_LEN: usize, const KEY_LEN: usize, F: Field>(
 41    state: [F; 2 * STATE_HALF_LEN],
 42    reverse_key_schedule: &[[F; KEY_LEN]],
 43    f: &impl Fn(Vec<F>) -> [F; STATE_HALF_LEN],
 44) -> [F; 2 * STATE_HALF_LEN] {
 45    if reverse_key_schedule.is_empty() {
 46        state
 47    } else {
 48        let k = reverse_key_schedule[0];
 49        let l = &state[..STATE_HALF_LEN];
 50        let r = &state[STATE_HALF_LEN..];
 51
 52        let new_r = l;
 53        let round_offset = f([l, &k].concat());
 54        let new_l: [_; STATE_HALF_LEN] = array::from_fn(|i| r[i] - round_offset[i]);
 55
 56        feistel_inv_cipher(
 57            array::from_fn(|i| {
 58                if i < STATE_HALF_LEN {
 59                    new_l[i]
 60                } else {
 61                    new_r[i - STATE_HALF_LEN]
 62                }
 63            }),
 64            &reverse_key_schedule[1..],
 65            f,
 66        )
 67    }
 68}
 69
 70#[cfg(test)]
 71mod tests {
 72    use plonky2::{
 73        field::{
 74            goldilocks_field::GoldilocksField as F,
 75            types::{Field, Field64},
 76        },
 77        hash::{
 78            hashing::hash_n_to_hash_no_pad, keccak::KeccakPermutation,
 79            poseidon::PoseidonPermutation,
 80        },
 81    };
 82    use rand::{Rng, rngs::OsRng};
 83
 84    use crate::{feistel_cipher, feistel_inv_cipher};
 85
 86    pub(crate) fn random_fields<const N: usize>() -> [F; N] {
 87        std::array::from_fn(|_| F::from_canonical_u64(OsRng.gen_range(0..F::ORDER)))
 88    }
 89
 90    #[test]
 91    fn feistel_keccak_poseidon_roundtrip() {
 92        const NR: usize = 32;
 93        const STATE_HALF_LEN: usize = 4;
 94        const KEY_LEN: usize = 4;
 95
 96        let f_keccak =
 97            |fields: Vec<_>| hash_n_to_hash_no_pad::<F, KeccakPermutation<_>>(&fields).elements;
 98        let f_poseidon =
 99            |fields: Vec<_>| hash_n_to_hash_no_pad::<F, PoseidonPermutation<_>>(&fields).elements;
100
101        [f_keccak, f_poseidon].into_iter().for_each(|f| {
102            (0..10).for_each(|_| {
103                let state: [_; { 2 * STATE_HALF_LEN }] = random_fields();
104                let key_schedule: [[_; KEY_LEN]; NR] = std::array::from_fn(|_| random_fields());
105
106                let out_state = feistel_cipher(state, &key_schedule, &f);
107                let reverse_key_schedule = key_schedule.into_iter().rev().collect::<Vec<_>>();
108                let maybe_original_state = feistel_inv_cipher(out_state, &reverse_key_schedule, &f);
109
110                assert_eq!(state, maybe_original_state);
111            })
112        })
113    }
114}