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