1use std::array;
  2
  3use plonky2::{
  4    field::{extension::Extendable, goldilocks_field::GoldilocksField as F},
  5    hash::hash_types::RichField,
  6    iop::{
  7        target::Target,
  8        witness::{PartialWitness, WitnessWrite},
  9    },
 10    plonk::circuit_builder::CircuitBuilder,
 11};
 12
 13pub(crate) const D: usize = 2;
 14
 15#[derive(Debug, Copy, Clone)]
 16pub struct FeistelStateTarget<const STATE_LEN: usize>([Target; STATE_LEN]);
 17
 18pub trait CircuitBuilderFeistel<F: RichField + Extendable<D>, const D: usize> {
 19    fn add_feistel_state_target<const STATE_LEN: usize>(&mut self)
 20    -> FeistelStateTarget<STATE_LEN>;
 21
 22    fn feistel_cipher<const STATE_HALF_LEN: usize, const KEY_LEN: usize>(
 23        &mut self,
 24        state: FeistelStateTarget<{ 2 * STATE_HALF_LEN }>,
 25        key_schedule: &[[Target; KEY_LEN]],
 26        f: &impl Fn(&mut Self, Vec<Target>) -> [Target; STATE_HALF_LEN],
 27    ) -> FeistelStateTarget<{ 2 * STATE_HALF_LEN }>;
 28}
 29
 30impl CircuitBuilderFeistel<F, D> for CircuitBuilder<F, D> {
 31    fn add_feistel_state_target<const STATE_LEN: usize>(
 32        &mut self,
 33    ) -> FeistelStateTarget<STATE_LEN> {
 34        FeistelStateTarget(self.add_virtual_target_arr())
 35    }
 36
 37    fn feistel_cipher<const STATE_HALF_LEN: usize, const KEY_LEN: usize>(
 38        &mut self,
 39        state: FeistelStateTarget<{ 2 * STATE_HALF_LEN }>,
 40        key_schedule: &[[Target; KEY_LEN]],
 41        f: &impl Fn(&mut Self, Vec<Target>) -> [Target; STATE_HALF_LEN],
 42    ) -> FeistelStateTarget<{ 2 * STATE_HALF_LEN }> {
 43        if key_schedule.is_empty() {
 44            state
 45        } else {
 46            let k = key_schedule[0];
 47            let l = &state.0[..STATE_HALF_LEN];
 48            let r = &state.0[STATE_HALF_LEN..];
 49
 50            let new_l = r;
 51            let round_offset = f(self, [r, &k].concat());
 52            let new_r: [_; STATE_HALF_LEN] = array::from_fn(|i| self.add(l[i], round_offset[i]));
 53            self.feistel_cipher(
 54                FeistelStateTarget(array::from_fn(|i| {
 55                    if i < STATE_HALF_LEN {
 56                        new_l[i]
 57                    } else {
 58                        new_r[i - STATE_HALF_LEN]
 59                    }
 60                })),
 61                &key_schedule[1..],
 62                f,
 63            )
 64        }
 65    }
 66}
 67
 68pub trait PartialWitnessFeistel {
 69    fn set_feistel_state_target<const STATE_SIZE: usize>(
 70        &mut self,
 71        target: &FeistelStateTarget<STATE_SIZE>,
 72        value: &[F; STATE_SIZE],
 73    ) -> anyhow::Result<()>;
 74}
 75
 76impl PartialWitnessFeistel for PartialWitness<F> {
 77    fn set_feistel_state_target<const STATE_SIZE: usize>(
 78        &mut self,
 79        target: &FeistelStateTarget<STATE_SIZE>,
 80        value: &[F; STATE_SIZE],
 81    ) -> anyhow::Result<()> {
 82        self.set_target_arr(&target.0, value)
 83    }
 84}
 85
 86#[cfg(test)]
 87mod tests {
 88    use plonky2::{
 89        field::goldilocks_field::GoldilocksField as F,
 90        hash::{
 91            hashing::hash_n_to_hash_no_pad,
 92            keccak::KeccakHash,
 93            poseidon::{PoseidonHash, PoseidonPermutation},
 94        },
 95        iop::witness::{PartialWitness, WitnessWrite},
 96        plonk::{
 97            circuit_builder::CircuitBuilder, circuit_data::CircuitConfig,
 98            config::PoseidonGoldilocksConfig,
 99        },
100    };
101
102    use super::D;
103    use crate::{
104        circuit::{CircuitBuilderFeistel, FeistelStateTarget, PartialWitnessFeistel},
105        feistel_cipher,
106        tests::random_fields,
107    };
108
109    #[test]
110    fn feistel_poseidon_check() -> anyhow::Result<()> {
111        const NR: usize = 32;
112        const STATE_HALF_LEN: usize = 4;
113        const KEY_LEN: usize = 4;
114
115        let config = CircuitConfig::standard_recursion_config();
116        let mut builder = CircuitBuilder::<F, D>::new(config);
117
118        let state_target: FeistelStateTarget<{ 2 * STATE_HALF_LEN }> =
119            builder.add_feistel_state_target();
120        let key_schedule_target: [[_; KEY_LEN]; NR] =
121            std::array::from_fn(|_| builder.add_virtual_target_arr());
122
123        let f = |builder: &mut CircuitBuilder<F, D>, targets: Vec<_>| {
124            builder
125                .hash_n_to_hash_no_pad::<PoseidonHash>(targets)
126                .elements
127        };
128
129        let out_state_target = builder.feistel_cipher(state_target, &key_schedule_target, &f);
130
131        let data = builder.build::<PoseidonGoldilocksConfig>();
132        (0..10).try_for_each(|_| {
133            let state: [_; { 2 * STATE_HALF_LEN }] = random_fields();
134            let key_schedule: [[_; KEY_LEN]; NR] = std::array::from_fn(|_| random_fields());
135            let eff = |fields: Vec<_>| {
136                hash_n_to_hash_no_pad::<F, PoseidonPermutation<_>>(&fields).elements
137            };
138            let out_state = feistel_cipher(state, &key_schedule, &eff);
139
140            let mut pw = PartialWitness::new();
141            pw.set_feistel_state_target(&state_target, &state)?;
142            std::iter::zip(key_schedule_target, key_schedule)
143                .try_for_each(|(t, v)| pw.set_target_arr(&t, &v))?;
144            pw.set_feistel_state_target(&out_state_target, &out_state)?;
145
146            let proof = data.prove(pw)?;
147            data.verify(proof)
148        })
149    }
150}