plonky2/src/gadgets/permutation.rs

use std::convert::TryInto;

use crate::field::{extension_field::Extendable, field_types::Field};
use crate::gates::switch::SwitchGate;
use crate::iop::generator::{GeneratedValues, SimpleGenerator};
use crate::iop::target::{BoolTarget, Target};
use crate::iop::witness::PartialWitness;
use crate::plonk::circuit_builder::CircuitBuilder;

impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
    /// Assert that two lists of expressions evaluate to permutations of one another.
    pub fn assert_permutation<const CHUNK_SIZE: usize>(
        &mut self,
        a: Vec<[Target; CHUNK_SIZE]>,
        b: Vec<[Target; CHUNK_SIZE]>,
    ) {
        assert_eq!(
            a.len(),
            b.len(),
            "Permutation must have same number of inputs and outputs"
        );
        assert_eq!(a[0].len(), b[0].len(), "Chunk sizes must be the same");

        match a.len() {
            // Two empty lists are permutations of one another, trivially.
            0 => (),
            // Two singleton lists are permutations of one another as long as their items are equal.
            1 => {
                for e in 0..CHUNK_SIZE {
                    self.assert_equal(a[0][e], b[0][e])
                }
            }
            2 => self.assert_permutation_2x2(a[0], a[1], b[0], b[1]),
            // For larger lists, we recursively use two smaller permutation networks.
            //_ => self.assert_permutation_recursive(a, b)
            _ => self.assert_permutation_recursive(a, b),
        }
    }

    /// Assert that [a, b] is a permutation of [c, d].
    fn assert_permutation_2x2<const CHUNK_SIZE: usize>(
        &mut self,
        a: [Target; CHUNK_SIZE],
        b: [Target; CHUNK_SIZE],
        c: [Target; CHUNK_SIZE],
        d: [Target; CHUNK_SIZE],
    ) {
        let (_, gate_c, gate_d) = self.create_switch(a, b);
        for e in 0..CHUNK_SIZE {
            self.route(c[e], gate_c[e]);
            self.route(d[e], gate_d[e]);
        }
    }

    fn create_switch<const CHUNK_SIZE: usize>(
        &mut self,
        a: [Target; CHUNK_SIZE],
        b: [Target; CHUNK_SIZE],
    ) -> (usize, [Target; CHUNK_SIZE], [Target; CHUNK_SIZE]) {
        if self.current_switch_gates.len() < CHUNK_SIZE {
            self.current_switch_gates
                .extend(vec![None; CHUNK_SIZE - self.current_switch_gates.len()]);
        }

        let (gate_index, mut next_copy) = match self.current_switch_gates[CHUNK_SIZE - 1] {
            None => {
                let gate = SwitchGate::<F, D, CHUNK_SIZE>::new_from_config(self.config.clone());
                let gate_index = self.add_gate(gate.clone(), vec![]);
                (gate_index, 0)
            }
            Some((idx, next_copy)) => (idx, next_copy),
        };

        let num_copies =
            SwitchGate::<F, D, CHUNK_SIZE>::max_num_copies(self.config.num_routed_wires);

        let mut c = Vec::new();
        let mut d = Vec::new();
        for e in 0..CHUNK_SIZE {
            self.route(
                a[e],
                Target::wire(
                    gate_index,
                    SwitchGate::<F, D, CHUNK_SIZE>::wire_first_input(0, e),
                ),
            );
            self.route(
                b[e],
                Target::wire(
                    gate_index,
                    SwitchGate::<F, D, CHUNK_SIZE>::wire_second_input(0, e),
                ),
            );
            c.push(Target::wire(
                gate_index,
                SwitchGate::<F, D, CHUNK_SIZE>::wire_first_output(0, e),
            ));
            d.push(Target::wire(
                gate_index,
                SwitchGate::<F, D, CHUNK_SIZE>::wire_second_output(0, e),
            ));
        }

        let c_arr: [Target; CHUNK_SIZE] = c.try_into().unwrap();
        let d_arr: [Target; CHUNK_SIZE] = d.try_into().unwrap();

        next_copy += 1;
        if next_copy == num_copies {
            let new_gate = SwitchGate::<F, D, CHUNK_SIZE>::new_from_config(self.config.clone());
            let new_gate_index = self.add_gate(new_gate.clone(), vec![]);
            self.current_switch_gates[CHUNK_SIZE - 1] = Some((new_gate_index, 0));
        } else {
            self.current_switch_gates[CHUNK_SIZE - 1] = Some((gate_index, next_copy));
        }

        (gate_index, c_arr, d_arr)
    }

    fn assert_permutation_recursive<const CHUNK_SIZE: usize>(
        &mut self,
        a: Vec<[Target; CHUNK_SIZE]>,
        b: Vec<[Target; CHUNK_SIZE]>,
    ) {
        let n = a.len();
        let even = n % 2 == 0;

        let mut child_1_a = Vec::new();
        let mut child_1_b = Vec::new();
        let mut child_2_a = Vec::new();
        let mut child_2_b = Vec::new();

        // See Figure 8 in the AS-Waksman paper.
        let a_num_switches = n / 2;
        let b_num_switches = if even {
            a_num_switches - 1
        } else {
            a_num_switches
        };

        for i in 0..a_num_switches {
            let (_, out_1, out_2) = self.create_switch(a[i * 2], a[i * 2 + 1]);
            child_1_a.push(out_1);
            child_2_a.push(out_2);
        }
        for i in 0..b_num_switches {
            let (_, out_1, out_2) = self.create_switch(b[i * 2], b[i * 2 + 1]);
            child_1_b.push(out_1);
            child_2_b.push(out_2);
        }

        // See Figure 8 in the AS-Waksman paper.
        if even {
            child_1_b.push(b[n - 2].clone());
            child_2_b.push(b[n - 1].clone());
        } else {
            child_2_a.push(a[n - 1].clone());
            child_2_b.push(b[n - 1].clone());
        }

        self.assert_permutation(child_1_a, child_1_b);
        self.assert_permutation(child_2_a, child_2_b);
    }
}

fn route_one_layer<F: Field, const CHUNK_SIZE: usize>(
    a_values: Vec<F>,
    b_values: Vec<F>,
    a_wires: Vec<[Target; CHUNK_SIZE]>,
    b_wires: Vec<[Target; CHUNK_SIZE]>,
) {
    todo!()
}

struct PermutationGenerator<F: Field, const CHUNK_SIZE: usize> {
    a_values: Vec<F>,
    b_values: Vec<F>,
    a_wires: Vec<[Target; CHUNK_SIZE]>,
    b_wires: Vec<[Target; CHUNK_SIZE]>,
}

impl<F: Field, const CHUNK_SIZE: usize> SimpleGenerator<F> for PermutationGenerator<F, CHUNK_SIZE> {
    fn dependencies(&self) -> Vec<Target> {
        self.a_wires
            .iter()
            .map(|arr| arr.to_vec())
            .flatten()
            .collect()
        //.chain(self.b_wires.iter()).collect()
    }

    fn run_once(&self, witness: &PartialWitness<F>, out_buffer: &mut GeneratedValues<F>) {
        todo!()
    }
}

#[cfg(test)]
mod tests {
    use anyhow::Result;

    use super::*;
    use crate::field::crandall_field::CrandallField;
    use crate::field::extension_field::quartic::QuarticCrandallField;
    use crate::field::field_types::Field;
    use crate::iop::witness::PartialWitness;
    use crate::plonk::circuit_data::CircuitConfig;
    use crate::plonk::verifier::verify;

    fn test_permutation(size: usize) -> Result<()> {
        type F = CrandallField;
        type FF = QuarticCrandallField;
        let len = 1 << len_log;
        let config = CircuitConfig::large_config();
        let pw = PartialWitness::new(config.num_wires);
        let mut builder = CircuitBuilder::<F, 4>::new(config);
        let vec = FF::rand_vec(len);
        let v: Vec<_> = vec.iter().map(|x| builder.constant_extension(*x)).collect();

        for i in 0..len {
            let it = builder.constant(F::from_canonical_usize(i));
            let elem = builder.constant_extension(vec[i]);
            builder.random_access(it, elem, v.clone());
        }

        let data = builder.build();
        let proof = data.prove(pw)?;

        verify(proof, &data.verifier_only, &data.common)
    }
}