# coding task

__Implement a circuit for poseidon hash using a vanilla plonk gate.__

The plonk is configured as follows. It uses a custom gate
$$q_a·a + q_b·b + q_c·c + q_{ab}·a·b + constant + instance = 0$$
and use the following struct to store the witnesses and selectors.

#[derive(Clone, Copy)]
pub struct StandardPlonkConfig {
    a: Column<Advice>,
    b: Column<Advice>,
    c: Column<Advice>,
    q_a: Column<Fixed>,
    q_b: Column<Fixed>,
    q_c: Column<Fixed>,
    q_ab: Column<Fixed>,
    constant: Column<Fixed>,
    #[allow(dead_code)]
    instance: Column<Instance>,
}

Then we enforce the custom gates over the witnesses and selectors via the following configuration.

impl<F:Prim> StandardPlonkConfig<F> {
    fn configure(meta: &mut ConstraintSystem<F>) -> Self {
        let [a, b, c] = [(); 3].map(|_| meta.advice_column());
        let [q_a, q_b, q_c, q_ab, constant] = [(); 5].map(|_| meta.fixed_column());
        let instance = meta.instance_column();

        [a, b, c].map(|column| meta.enable_equality(column));

        meta.create_gate(
            "q_a·a + q_b·b + q_c·c + q_ab·a·b + constant + instance = 0",
            |meta| {
                let [a, b, c] = [a, b, c].map(|column| meta.query_advice(column, Rotation::cur()));
                let [q_a, q_b, q_c, q_ab, constant] = [q_a, q_b, q_c, q_ab, constant]
                    .map(|column| meta.query_fixed(column, Rotation::cur()));
                let instance = meta.query_instance(instance, Rotation::cur());
                Some(
                    q_a * a.clone()
                        + q_b * b.clone()
                        + q_c * c
                        + q_ab * a * b
                        + constant
                        + instance,
                )
            },
        );

        StandardPlonkConfig { a, b, c, q_a, q_b, q_c, q_ab, constant, instance }
    }
}

The task is to design a circuit for the following subroutine of a 4-to-1 poseidon hash function.

struct PoseidonState<F: PrimeField> {
    state: [F; 4],
    index: usize
}

You may assume that the state is already initialized. It will go through a `MAX_ITER` number of iterations. `index` is used to track which iteration it is current at. Each iterations is parameterized with an iteration key `k_i: [F; 4]`. The hash function also takes a global parameter M which is a 4-by-4 matrix, i.e. `M: [[F; 4];4]`.

Suppose the current state is $a_{i,1}, a_{i,2}, a_{i,3}, a_{i,4}$, then the next state is defined as

\begin{align*}
\left<a_{i+1,1}, a_{i+1,2}, a_{i+1,3}, a_{i+1,4}\right>&:=\left<a_{i,1}^5, a_{i,2}, a_{i,3}, a_{i,4}\right>\times\begin{pmatrix} M_{1,1}, M_{1,2}, M_{1,3}, M_{1,4}\\M_{2,1}, M_{2,2}, M_{2,3}, M_{2,4}\\M_{3,1}, M_{3,2}, M_{3,3}, M_{3,4}\\M_{4,1}, M_{4,2}, M_{4,3}, M_{4,4} \end{pmatrix} \\ &+ \left<k_{i,1}, k_{i,2}, k_{i,3}, k_{i,4}\right>
\end{align*}

# Skelton of the code

/// Assume we do a maximum of 5 iterations
const MAX_ITER: usize = 5;
/// Assume the round key is given as an array of size MAX_ITER;
/// each element of the array is the round key for the corresponding iteration
/// You DO NOT need to complete this part.
const ROUND_KEY: [[F;4];MAX_ITER] = ... 
/// Assume the matrix is given as a 4-by-4 2-dimentional array;
/// You DO NOT need to complete this part.
const MDS_MATRIX: [[F;4];4] = ...

/// The state that holds the exsitng 
#[derive(Clone, Default)]
struct PoseidonState([F; 4]);

impl Circuit<Fr> for PoseidonState {
    type Config = StandardPlonkConfig;
    type FloorPlanner = SimpleFloorPlanner;

    fn without_witnesses(&self) -> Self {
        Self::default()
    }

    fn configure(meta: &mut ConstraintSystem<Fr>) -> Self::Config {
        meta.set_minimum_degree(4);
        StandardPlonkConfig::configure(meta)
    }

    fn synthesize(
        &self,
        config: Self::Config,
        mut layouter: impl Layouter<Fr>,
    ) -> Result<(), Error> {
        layouter.assign_region(
            || "",
            |mut region| {
                //
                // BEGIN of SAMPLE code
                //
                // In this sample code we constraint that  
                // 
                //     -1 * State[0] + 5 * State[1] + State[2] = 0
                //
                region.assign_advice(|| "", config.a, 0, || Value::known(self.0[0]))?;
                region.assign_fixed(|| "", config.q_a, 0, || Value::known(-F::one()))?;
                region.assign_advice(|| "", config.b, 0, || Value::known(self.0[1]))?;
                region.assign_fixed(|| "", config.q_b, 0, || Value::known(F::from(5)))?;
                region.assign_advice(|| "", config.c, 0, || Value::known(self.0[2]))?;                
                region.assign_fixed(|| "", config.q_c, 0, || Value::known(F::one()))?;
                
                
// q_a·a + q_b·b + q_c·c + q_{ab}·a·b + constant + instance = 0
/* multiply x and y
                region.assign_advice(|| "", config.a, 0, || Value::known(x))?;
                region.assign_advice(|| "", config.b, 0, || Value::known(y))?;  
                region.assign_fixed(|| "", config.q_ab, 0, || Value::known(F::one()))?;
                              
*/

                // start of code
                region.assign_advice(|| "", config.a, 1, || Value::known(ai,1))?;
                region.assign_advice(|| "", config.b, 1, || Value::known(ai,1))?;
                region.assign_fixed(|| "", config.q_ab, 1, || Value::known(F::one()))?;
                
                region.assign_advice(|| "", config.c, 1, || Value::known(ai2,1))?;
                region.assign_fixed(|| "", config.q_ab, 1, || Value::known(-F::one()))?;
                
                // ai4
                region.assign_advice(|| "", config.a, 2, || Value::known(ai2,1))?;
                region.assign_advice(|| "", config.b, 2, || Value::known(ai2,1))?;
                region.assign_fixed(|| "", config.q_ab, 2, || Value::known(F::one()))?;
                
                region.assign_advice(|| "", config.c, 2, || Value::known(ai4,1))?;
                region.assign_fixed(|| "", config.q_ab, 2, || Value::known(-F::one()))?;
                
                // ai5
                region.assign_advice(|| "", config.a, 3, || Value::known(ai4,1))?;
                region.assign_advice(|| "", config.b, 3, || Value::known(ai,1))?;
                region.assign_fixed(|| "", config.q_ab, 3, || Value::known(F::one()))?;
                
                region.assign_advice(|| "", config.c, 3, || Value::known(ai5,1))?;
                region.assign_fixed(|| "", config.q_ab, 3, || Value::known(-F::one()))?;
                
                
                // get a single multiplication with M coefficient
                region.assign_advice(|| "", config.a, 4, || Value::known(ai5,1))?;
                region.assign_fixed(|| "", config.q_a, 4, || Value::known(M1,1))?;
                
                region.assign_advice(|| "", config.b, 4, || Value::known(ai,2))?;
                region.assign_fixed(|| "", config.q_b, 4, || Value::known(M2,1))?;
                
                region.assign_advice(|| "", config.c, 4, || Value::known(a,1))?;
                region.assign_fixed(|| "", config.q_c, 4, || Value::known(-F::one()))?;
                
             
// q_a·a + q_b·b + q_c·c + q_{ab}·a·b + constant + instance = 0


                // q_dot * (q_m1 * a1 * a1 * a1 * a1 * a1 + q_m2 * a2 + q_m3 * a3 + q_m4 * a4 + q_k)
                
                

                
                //
                // END of SAMPLE code
                //
                // You may comment out the sample code and write your code here
                

                Ok(())
            },
        )
    }
}