//! Concrete instantiation of a hash function.

use std::convert::TryInto;

use crate::field::extension_field::Extendable;
use crate::field::field_types::RichField;
use crate::hash::hash_types::{HashOut, HashOutTarget};
use crate::iop::target::Target;
use crate::plonk::circuit_builder::CircuitBuilder;
use crate::plonk::config::AlgebraicHasher;

pub(crate) const SPONGE_RATE: usize = 8;
pub(crate) const SPONGE_CAPACITY: usize = 4;
pub const SPONGE_WIDTH: usize = SPONGE_RATE + SPONGE_CAPACITY;

/// Hash the vector if necessary to reduce its length to ~256 bits. If it already fits, this is a
/// no-op.
pub fn hash_or_noop<F: RichField, P: PlonkyPermutation<F>>(inputs: Vec<F>) -> HashOut<F> {
    if inputs.len() <= 4 {
        HashOut::from_partial(inputs)
    } else {
        hash_n_to_hash::<F, P>(inputs, false)
    }
}

impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
    pub fn hash_or_noop<H: AlgebraicHasher<F>>(&mut self, inputs: Vec<Target>) -> HashOutTarget {
        let zero = self.zero();
        if inputs.len() <= 4 {
            HashOutTarget::from_partial(inputs, zero)
        } else {
            self.hash_n_to_hash::<H>(inputs, false)
        }
    }

    pub fn hash_n_to_hash<H: AlgebraicHasher<F>>(
        &mut self,
        inputs: Vec<Target>,
        pad: bool,
    ) -> HashOutTarget {
        HashOutTarget::from_vec(self.hash_n_to_m::<H>(inputs, 4, pad))
    }

    pub fn hash_n_to_m<H: AlgebraicHasher<F>>(
        &mut self,
        mut inputs: Vec<Target>,
        num_outputs: usize,
        pad: bool,
    ) -> Vec<Target> {
        let zero = self.zero();
        let one = self.one();

        if pad {
            inputs.push(zero);
            while (inputs.len() + 1) % SPONGE_WIDTH != 0 {
                inputs.push(one);
            }
            inputs.push(zero);
        }

        let mut state = [zero; SPONGE_WIDTH];

        // Absorb all input chunks.
        for input_chunk in inputs.chunks(SPONGE_RATE) {
            // Overwrite the first r elements with the inputs. This differs from a standard sponge,
            // where we would xor or add in the inputs. This is a well-known variant, though,
            // sometimes called "overwrite mode".
            state[..input_chunk.len()].copy_from_slice(input_chunk);
            state = self.permute::<H>(state);
        }

        // Squeeze until we have the desired number of outputs.
        let mut outputs = Vec::new();
        loop {
            for i in 0..SPONGE_RATE {
                outputs.push(state[i]);
                if outputs.len() == num_outputs {
                    return outputs;
                }
            }
            state = self.permute::<H>(state);
        }
    }
}

/// A one-way compression function which takes two ~256 bit inputs and returns a ~256 bit output.
pub fn compress<F: RichField, P: PlonkyPermutation<F>>(x: HashOut<F>, y: HashOut<F>) -> HashOut<F> {
    let mut perm_inputs = [F::ZERO; SPONGE_WIDTH];
    perm_inputs[..4].copy_from_slice(&x.elements);
    perm_inputs[4..8].copy_from_slice(&y.elements);
    HashOut {
        elements: P::permute(perm_inputs)[..4].try_into().unwrap(),
    }
}

/// Permutation that can be used in the sponge construction for an algebraic hash.
pub trait PlonkyPermutation<F: RichField> {
    fn permute(input: [F; SPONGE_WIDTH]) -> [F; SPONGE_WIDTH];
}

pub struct PoseidonPermutation;
impl<F: RichField> PlonkyPermutation<F> for PoseidonPermutation {
    fn permute(input: [F; SPONGE_WIDTH]) -> [F; SPONGE_WIDTH] {
        F::poseidon(input)
    }
}
pub struct GMiMCPermutation;
impl<F: RichField> PlonkyPermutation<F> for GMiMCPermutation {
    fn permute(input: [F; SPONGE_WIDTH]) -> [F; SPONGE_WIDTH] {
        F::gmimc_permute(input)
    }
}

/// If `pad` is enabled, the message is padded using the pad10*1 rule. In general this is required
/// for the hash to be secure, but it can safely be disabled in certain cases, like if the input
/// length is fixed.
pub fn hash_n_to_m<F: RichField, P: PlonkyPermutation<F>>(
    mut inputs: Vec<F>,
    num_outputs: usize,
    pad: bool,
) -> Vec<F> {
    if pad {
        inputs.push(F::ZERO);
        while (inputs.len() + 1) % SPONGE_WIDTH != 0 {
            inputs.push(F::ONE);
        }
        inputs.push(F::ZERO);
    }

    let mut state = [F::ZERO; SPONGE_WIDTH];

    // Absorb all input chunks.
    for input_chunk in inputs.chunks(SPONGE_RATE) {
        for i in 0..input_chunk.len() {
            state[i] = input_chunk[i];
        }
        state = P::permute(state);
    }

    // Squeeze until we have the desired number of outputs.
    let mut outputs = Vec::new();
    loop {
        for &item in state.iter().take(SPONGE_RATE) {
            outputs.push(item);
            if outputs.len() == num_outputs {
                return outputs;
            }
        }
        state = P::permute(state);
    }
}

pub fn hash_n_to_hash<F: RichField, P: PlonkyPermutation<F>>(
    inputs: Vec<F>,
    pad: bool,
) -> HashOut<F> {
    HashOut::from_vec(hash_n_to_m::<F, P>(inputs, 4, pad))
}