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