# 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*}

Notice that the first element in the previous state is raised to the power of 5.

# 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);
        // configure a chip here
        StandardPlonkConfig::configure(meta)
    }

    fn synthesize(
        &self,
        config: Self::Config,
        mut layouter: impl Layouter<Fr>,
        meta: &mut ConstraintSystem<Fr>
    ) -> Result<(), Error> {
        layouter.assign_region(
            || "",
            |mut region| {
                //
                // BEGIN of SAMPLE code
                //
                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()))?;
                //
                // END of SAMPLE code
                //
                // You may comment out the sample code and write your code here
                
                // meta.create_gate to query the advice and fixed columns
                // return the vector with operations
                // vec![<a> * MDS_MATIX + <k>]
                Ok(())
            },
        )
    }
    
    fn matrix_mul(&self, config: Self::Config, matrix: [[F;4];4], curr_state: [F; 4], meta: &mut ConstraintSystem<Fr>) -> [F; 4]{
        let final_matrix: [F; 4];
    // We will fill in the curr_state array with [a[0], a[1], ...]
        for (let i = 0; i < curr_state.len; i++){
            curr_state[i] = curr_state[i].map(|column| meta.query_advice(column, Rotation::(i)));
        }
        for (let i = 0; i < curr_state.len; i++) {
            for (let j = 0; j < curr_state.len; j++)
            // query curr_state adivce column and multiply by rows in matrix
            final_matrix = curr_state[i] * matrix[i][j];
            }
        }
        Ok(final_matrix)
        
        
    }
}

// Scalar field of bn256 / Scalar
// Chip needs to do multiply and add (perhaps having a lookup table for a^5) a is a field element of a^256
//