Poseidon Parameter Generation: Separate LFSR from Sampling Distribution/Algorithm

[bhgomes]

To improve the modularity and usefulness of the Poseidon Parameter Setup code, we want to separate the LFSR implementation from the field element sampling algorithm and the generation of the round constants and MDS matrix.

Old Design

Right now when generating the Poseidon parameters we have something like the following interfaces:

impl LFSR {
    /// Generates a new `LFSR` random number generator with a seed derived from a Poseidon specification.
    fn new(field_size: u32, width: u32, full_rounds: u32, partial_rounds: u32) -> Self { ... }

    /// Samples a random field point using rejection sampling.
    fn sample_field_point(&mut self) -> F { ... }
}

/// Generates the Poseidon hasher from its specification.
fn generate_hasher(field_size: u32, width: u32, full_rounds: u32, partial_rounds: u32) -> Hasher {
    // Generates a new LFSR from the Poseidon specification.
    let lfsr = LFSR::new(field_size, width, full_rounds, partial_rounds);

    // Generates the round constants using LFSR as the random field point sampler.
    let round_constants = generate_round_constants(lfsr);

    // Generates the MDS matrix using a Cauchy matrix over `(0..t) x (t..2t)`.
    let mds = generate_mds(width);
    
    // Initializes the Hasher parameters checking that they match the correct sizes.
    Self::new(round_constants, mds)
}

The issue with this design is that we are very restricted in our ability to utilize other sampling methods or PRNGs for constructing these types. We want to be able to leverage the existing random sampling APIs available in Rust and manta-crypto.

New Design

For each type we have a structure that defines its own sampling methods.

Additive Round Constants

/// Additive Round Constants
pub struct AdditiveRoundConstants<F> 
where
    F: Field,
{ ... }

impl<D, F> Sample<D> for AdditiveRoundConstants<F> 
where
    D: Clone,
    F: Field + Sample<D>,
{
    #[inline]
    fn sample<R>(distribution: D, rng: &mut R) -> Self
    where
        R: CryptoRng + RngCore + ?Sized,
    {
        Self::new(rng.sample_iter(core::iter::repeat(distribution)).collect())
    }
}

Maximum Distance Separable Matrix

/// Maximum Distance Separable Matrix
pub struct MaximumDistanceSeparableMatrix<F> 
where
    F: Field,
{ ... }

impl<F, X, Y> Sample<CauchyMatrixDistribution<X, Y>> for MaximumDistanceSeparableMatrix<F> 
where
    F: Field,
    X: IntoIterator,
    Y: IntoIterator,
{
    #[inline]
    fn sample<R>(distribution: CauchyMatrixDistribution, rng: &mut R) -> Self
    where
        R: CryptoRng + RngCore + ?Sized,
    {
        Self::new(distribution.to_matrix(|x, y| F::sub(x, y).inverse()))
    }
}

/// Cauchy Matrix Generator
#[derive(Clone, Copy, Debug, Default, Eq, Hash, PartialEq)]
pub struct CauchyMatrixGenerator<X, Y> {
    /// X<sub>i</sub> Injective Sequence
    pub xs: X,
    /// Y<sub>j</sub> Injective Sequence
    pub ys: Y,
}

impl<X, Y> CauchyMatrixGenerator<X, Y> {
    /// Generates the Cauchy matrix arising from the pairwise inverse
    /// differences of `self.xs` and `self.ys` using `inv_diff` to compute the
    /// inverse difference.
    #[inline]
    pub fn as_matrix<'s, T, F>(&'s self, mut inv_diff: F) -> Option<Vec<Vec<T>>>
    where
        &'s X: IntoIterator,
        &'s Y: IntoIterator,
        F: for<'t> FnMut(
            &'t <&'s X as IntoIterator>::Item,
            &'t <&'s Y as IntoIterator>::Item,
        ) -> Option<T>,
    {
        let ys = self.ys.into_iter().collect::<Vec<_>>();
        self.xs
            .into_iter()
            .map(|x| ys.iter().map(|y| inv_diff(&x, y)).collect())
            .collect()
    }

    /// Generates the Cauchy matrix arising from the pairwise inverse
    /// differences of `self.xs` and `self.ys` using `inv_diff` to compute the
    /// inverse difference.
    #[inline]
    pub fn to_matrix<T, F>(self, mut inv_diff: F) -> Option<Vec<Vec<T>>>
    where
        X: IntoIterator,
        Y: IntoIterator,
        F: for<'t> FnMut(&'t X::Item, &'t Y::Item) -> Option<T>,
    {
        let ys = self.ys.into_iter().collect::<Vec<_>>();
        self.xs
            .into_iter()
            .map(|x| ys.iter().map(|y| inv_diff(&x, y)).collect())
            .collect()
    }
}

Poseidon Hasher

/// Poseidon Hasher Sampling Distribution
#[derive(Clone, Copy, Debug, Default, Eq, Hash, PartialEq)]
pub struct HasherDistribution<A, M> {
    /// Additive Round Constants Distribution
    pub additive_round_constants: A,
    /// Maximum Distance Separable Matrix Distribution
    pub maximum_distance_separable_matrix: M,
}

impl<A, M, S> Sample<HasherDistribution<A,M>> for Hasher<S> 
where
    S: Specification,
    AdditiveRoundConstants<S::ParameterField>: Sample<A>,
    MaximumDistanceSeparableMatrix<S::ParameterField>: Sample<M>,
{
    #[inline]
    fn sample<R>(distribution: HasherDistribution<A, M>, rng: &mut R) -> Self
    where
        R: CryptoRng + RngCore + ?Sized,
    {
        Self::new(
            rng.sample(distribution.additive_round_constants),
            rng.sample(distribution.maximum_distance_separable_matrix)
        )
    }
}

Linear Feedback Shift Register PRNG

/// Linear Feedback Shift Register Pseudo-Random Number Generator
pub struct LinearFeedbackShiftRegister { ... }

impl RngCore for LinearFeedbackShiftRegister { ... }

impl SeedableRng for LinearFeedbackShiftRegister { ... }

and in the Poseidon module:

/// Generates the canonical LFSR for Poseidon constructed for the given [`Specification`].
#[inline]
pub fn canonical_sampling_rng<S>() -> LinearFeedbackShiftRegister 
where
    S: Specification,
    S::ParamterField: BinarySize,
{
    todo!("generate the canonical LFSR based on the specification")
}

where BinarySize is a trait that tells us how many bits the field is made of (trait name not final).

Rejection Sampling Utilities

NB: The following must be defined in manta_crypto::rand.

/// Returns the first value from `sample` applied to `rng` that returns `Some`.
#[inline]
pub fn find_sample<R, T, F>(rng: &mut R, mut sample: F) -> T
where
    F: FnMut(&mut R) -> Option<T>,
{
    loop {
        if let Some(value) = sample(rng) {
            return value;
        }
    }
}

/// Rejection Sampling Distribution
/// 
/// This distribution uses [`find_sample`] to convert a [`TrySample`] sampling algorithm into
/// a [`Sample`] one by performing rejection sampling.
#[derive(Clone, Copy, Debug, Default, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct RejectionSampling<D = ()>(pub D);

impl<D, T> Sample<RejectionSampling<D>> for T 
where
    D: Clone,
    T: TrySample<D>,
{
    #[inline]
    fn sample<R>(distribution: RejectionSampling<D>, rng: &mut R) -> Self {
        find_sample(rng, |rng| Self::try_sample(distribution.0.clone(), rng).ok())
    }
}

To plug these into the Poseidon parameter generation trait we just need to add the TrySample<()> implementation to the arkworks::Fp field object then just call the standard parameter generation as:

let hasher = canonical_sampling_rng::<S>().sample(
    HasherDistribution {
        additive_round_constants: RejectionSampling,
        maximum_distance_separable_matrix: CauchyMatrixGenerator { xs: (0..S::WIDTH), ys: (S::WIDTH..2*S::WIDTH) }
    }
);

[tsunrise]

It might have some issues if we require the RNG for AdditiveRoundConstants to be CryptoRng, because LFSR itself is not necessarily a crypto RNG. Are we going to implement CryptoRng trait for LFSR anyway?

Or shall we modify Sample trait to release the CryptoRng constraint? @bhgomes

[bhgomes]

It might have some issues if we require the RNG for AdditiveRoundConstants to be CryptoRng, because LFSR itself is not necessarily a crypto RNG. Are we going to implement CryptoRng trait for LFSR anyway?

Or shall we modify Sample trait to release the CryptoRng constraint? @bhgomes

@tsunrise We can drop CryptoRng from the requirements on Sample. I figured this was going to happen eventually.