M Cargo.toml → Cargo.toml

@@ -1,8 +1,11 @@
 [workspace]
 members = [
-    "aes-gcm"
+    "aes-gcm",
+    "feistel"
 ]
 
 resolver = "2"
 
 [workspace.dependencies]
+plonky2 = { git = "https://github.com/0xPARC/plonky2.git", rev = "109d517d09c210ae4c2cee381d3e3fbc04aa3812" }
+anyhow = "1.0.56"

A feistel/Cargo.toml

@@ -0,0 +1,9 @@
+[package]
+name = "plonky2-feistel"
+version = "0.1.0"
+edition = "2024"
+
+[dependencies]
+plonky2 = { workspace = true }
+anyhow = { workspace = true }
+rand = "0.8.5"

A feistel/src/circuit.rs

@@ -0,0 +1,150 @@
+use std::array;
+
+use plonky2::{
+    field::{extension::Extendable, goldilocks_field::GoldilocksField as F},
+    hash::hash_types::RichField,
+    iop::{
+        target::Target,
+        witness::{PartialWitness, WitnessWrite},
+    },
+    plonk::circuit_builder::CircuitBuilder,
+};
+
+pub(crate) const D: usize = 2;
+
+#[derive(Debug, Copy, Clone)]
+pub struct FeistelStateTarget<const STATE_LEN: usize>([Target; STATE_LEN]);
+
+pub trait CircuitBuilderFeistel<F: RichField + Extendable<D>, const D: usize> {
+    fn add_feistel_state_target<const STATE_LEN: usize>(&mut self)
+    -> FeistelStateTarget<STATE_LEN>;
+
+    fn feistel_cipher<const STATE_HALF_LEN: usize, const KEY_LEN: usize>(
+        &mut self,
+        state: FeistelStateTarget<{ 2 * STATE_HALF_LEN }>,
+        key_schedule: &[[Target; KEY_LEN]],
+        f: &impl Fn(&mut Self, Vec<Target>) -> [Target; STATE_HALF_LEN],
+    ) -> FeistelStateTarget<{ 2 * STATE_HALF_LEN }>;
+}
+
+impl CircuitBuilderFeistel<F, D> for CircuitBuilder<F, D> {
+    fn add_feistel_state_target<const STATE_LEN: usize>(
+        &mut self,
+    ) -> FeistelStateTarget<STATE_LEN> {
+        FeistelStateTarget(self.add_virtual_target_arr())
+    }
+
+    fn feistel_cipher<const STATE_HALF_LEN: usize, const KEY_LEN: usize>(
+        &mut self,
+        state: FeistelStateTarget<{ 2 * STATE_HALF_LEN }>,
+        key_schedule: &[[Target; KEY_LEN]],
+        f: &impl Fn(&mut Self, Vec<Target>) -> [Target; STATE_HALF_LEN],
+    ) -> FeistelStateTarget<{ 2 * STATE_HALF_LEN }> {
+        if key_schedule.is_empty() {
+            state
+        } else {
+            let k = key_schedule[0];
+            let l = &state.0[..STATE_HALF_LEN];
+            let r = &state.0[STATE_HALF_LEN..];
+
+            let new_l = r;
+            let round_offset = f(self, [r, &k].concat());
+            let new_r: [_; STATE_HALF_LEN] = array::from_fn(|i| self.add(l[i], round_offset[i]));
+            self.feistel_cipher(
+                FeistelStateTarget(array::from_fn(|i| {
+                    if i < STATE_HALF_LEN {
+                        new_l[i]
+                    } else {
+                        new_r[i - STATE_HALF_LEN]
+                    }
+                })),
+                &key_schedule[1..],
+                f,
+            )
+        }
+    }
+}
+
+pub trait PartialWitnessFeistel {
+    fn set_feistel_state_target<const STATE_SIZE: usize>(
+        &mut self,
+        target: &FeistelStateTarget<STATE_SIZE>,
+        value: &[F; STATE_SIZE],
+    ) -> anyhow::Result<()>;
+}
+
+impl PartialWitnessFeistel for PartialWitness<F> {
+    fn set_feistel_state_target<const STATE_SIZE: usize>(
+        &mut self,
+        target: &FeistelStateTarget<STATE_SIZE>,
+        value: &[F; STATE_SIZE],
+    ) -> anyhow::Result<()> {
+        self.set_target_arr(&target.0, value)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use plonky2::{
+        field::goldilocks_field::GoldilocksField as F,
+        hash::{
+            hashing::hash_n_to_hash_no_pad,
+            keccak::KeccakHash,
+            poseidon::{PoseidonHash, PoseidonPermutation},
+        },
+        iop::witness::{PartialWitness, WitnessWrite},
+        plonk::{
+            circuit_builder::CircuitBuilder, circuit_data::CircuitConfig,
+            config::PoseidonGoldilocksConfig,
+        },
+    };
+
+    use super::D;
+    use crate::{
+        circuit::{CircuitBuilderFeistel, FeistelStateTarget, PartialWitnessFeistel},
+        feistel_cipher,
+        tests::random_fields,
+    };
+
+    #[test]
+    fn feistel_poseidon_check() -> anyhow::Result<()> {
+        const NR: usize = 32;
+        const STATE_HALF_LEN: usize = 4;
+        const KEY_LEN: usize = 4;
+
+        let config = CircuitConfig::standard_recursion_config();
+        let mut builder = CircuitBuilder::<F, D>::new(config);
+
+        let state_target: FeistelStateTarget<{ 2 * STATE_HALF_LEN }> =
+            builder.add_feistel_state_target();
+        let key_schedule_target: [[_; KEY_LEN]; NR] =
+            std::array::from_fn(|_| builder.add_virtual_target_arr());
+
+        let f = |builder: &mut CircuitBuilder<F, D>, targets: Vec<_>| {
+            builder
+                .hash_n_to_hash_no_pad::<PoseidonHash>(targets)
+                .elements
+        };
+
+        let out_state_target = builder.feistel_cipher(state_target, &key_schedule_target, &f);
+
+        let data = builder.build::<PoseidonGoldilocksConfig>();
+        (0..10).try_for_each(|_| {
+            let state: [_; { 2 * STATE_HALF_LEN }] = random_fields();
+            let key_schedule: [[_; KEY_LEN]; NR] = std::array::from_fn(|_| random_fields());
+            let eff = |fields: Vec<_>| {
+                hash_n_to_hash_no_pad::<F, PoseidonPermutation<_>>(&fields).elements
+            };
+            let out_state = feistel_cipher(state, &key_schedule, &eff);
+
+            let mut pw = PartialWitness::new();
+            pw.set_feistel_state_target(&state_target, &state)?;
+            std::iter::zip(key_schedule_target, key_schedule)
+                .try_for_each(|(t, v)| pw.set_target_arr(&t, &v))?;
+            pw.set_feistel_state_target(&out_state_target, &out_state)?;
+
+            let proof = data.prove(pw)?;
+            data.verify(proof)
+        })
+    }
+}

A feistel/src/lib.rs

@@ -0,0 +1,114 @@
+#![allow(incomplete_features)]
+#![feature(generic_const_exprs)]
+
+pub mod circuit;
+
+use std::array;
+
+use plonky2::field::types::Field;
+
+pub fn feistel_cipher<const STATE_HALF_LEN: usize, const KEY_LEN: usize, F: Field>(
+    state: [F; 2 * STATE_HALF_LEN],
+    key_schedule: &[[F; KEY_LEN]],
+    f: &impl Fn(Vec<F>) -> [F; STATE_HALF_LEN],
+) -> [F; 2 * STATE_HALF_LEN] {
+    if key_schedule.is_empty() {
+        state
+    } else {
+        let k = key_schedule[0];
+        let l = &state[..STATE_HALF_LEN];
+        let r = &state[STATE_HALF_LEN..];
+
+        let new_l = r;
+        let round_offset = f([r, &k].concat());
+        let new_r: [_; STATE_HALF_LEN] = array::from_fn(|i| l[i] + round_offset[i]);
+
+        feistel_cipher(
+            array::from_fn(|i| {
+                if i < STATE_HALF_LEN {
+                    new_l[i]
+                } else {
+                    new_r[i - STATE_HALF_LEN]
+                }
+            }),
+            &key_schedule[1..],
+            f,
+        )
+    }
+}
+
+pub fn feistel_inv_cipher<const STATE_HALF_LEN: usize, const KEY_LEN: usize, F: Field>(
+    state: [F; 2 * STATE_HALF_LEN],
+    reverse_key_schedule: &[[F; KEY_LEN]],
+    f: &impl Fn(Vec<F>) -> [F; STATE_HALF_LEN],
+) -> [F; 2 * STATE_HALF_LEN] {
+    if reverse_key_schedule.is_empty() {
+        state
+    } else {
+        let k = reverse_key_schedule[0];
+        let l = &state[..STATE_HALF_LEN];
+        let r = &state[STATE_HALF_LEN..];
+
+        let new_r = l;
+        let round_offset = f([l, &k].concat());
+        let new_l: [_; STATE_HALF_LEN] = array::from_fn(|i| r[i] - round_offset[i]);
+
+        feistel_inv_cipher(
+            array::from_fn(|i| {
+                if i < STATE_HALF_LEN {
+                    new_l[i]
+                } else {
+                    new_r[i - STATE_HALF_LEN]
+                }
+            }),
+            &reverse_key_schedule[1..],
+            f,
+        )
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use plonky2::{
+        field::{
+            goldilocks_field::GoldilocksField as F,
+            types::{Field, Field64},
+        },
+        hash::{
+            hashing::hash_n_to_hash_no_pad, keccak::KeccakPermutation,
+            poseidon::PoseidonPermutation,
+        },
+    };
+    use rand::{Rng, rngs::OsRng};
+
+    use crate::{feistel_cipher, feistel_inv_cipher};
+
+    pub(crate) fn random_fields<const N: usize>() -> [F; N] {
+        std::array::from_fn(|_| F::from_canonical_u64(OsRng.gen_range(0..F::ORDER)))
+    }
+
+    #[test]
+    fn feistel_keccak_poseidon_roundtrip() {
+        const NR: usize = 32;
+        const STATE_HALF_LEN: usize = 4;
+        const KEY_LEN: usize = 4;
+
+        let f_keccak =
+            |fields: Vec<_>| hash_n_to_hash_no_pad::<F, KeccakPermutation<_>>(&fields).elements;
+        let f_poseidon =
+            |fields: Vec<_>| hash_n_to_hash_no_pad::<F, PoseidonPermutation<_>>(&fields).elements;
+
+        [f_keccak, f_poseidon].into_iter().for_each(|f| {
+            (0..10).for_each(|_| {
+                let state: [_; { 2 * STATE_HALF_LEN }] = random_fields();
+                let key_schedule: [[_; KEY_LEN]; NR] = std::array::from_fn(|_| random_fields());
+
+                let out_state = feistel_cipher(state, &key_schedule, &f);
+                let reverse_key_schedule = key_schedule.into_iter().rev().collect::<Vec<_>>();
+                let maybe_original_state = feistel_inv_cipher(out_state, &reverse_key_schedule, &f);
+
+                assert_eq!(state, maybe_original_state);
+            })
+        })
+    }
+}