use super::{Mode, ParamType};
use ark_ff::{fields::PrimeField as _, UniformRand as _};
use ark_serialize::CanonicalSerialize as _;
use mina_curves::pasta::Fp;
use mina_poseidon::{
constants::{self, SpongeConstants},
pasta,
poseidon::{ArithmeticSponge as Poseidon, ArithmeticSpongeParams, Sponge as _},
};
use num_bigint::BigUint;
use rand::{prelude::*, Rng};
use serde::Serialize;
#[derive(Debug, Serialize)]
pub struct TestVectors {
name: String,
test_vectors: Vec<TestVector>,
}
#[derive(Debug, Serialize)]
pub struct TestVector {
input: Vec<String>,
output: String,
}
fn poseidon<SC: SpongeConstants>(input: &[Fp], params: &'static ArithmeticSpongeParams<Fp>) -> Fp {
let mut s = Poseidon::<Fp, SC>::new(params);
s.absorb(input);
s.squeeze()
}
fn rand_fields(rng: &mut impl Rng, length: u8) -> Vec<Fp> {
let mut fields = vec![];
for _ in 0..length {
let fe = Fp::rand(rng);
fields.push(fe)
}
fields
}
pub fn generate(mode: Mode, param_type: ParamType) -> TestVectors {
let mut rng = &mut rand::rngs::StdRng::from_seed([0u8; 32]);
let mut test_vectors = vec![];
for length in 0..6 {
let input = rand_fields(&mut rng, length);
let output = match param_type {
ParamType::Legacy => poseidon::<constants::PlonkSpongeConstantsLegacy>(
&input,
pasta::fp_legacy::static_params(),
),
ParamType::Kimchi => poseidon::<constants::PlonkSpongeConstantsKimchi>(
&input,
pasta::fp_kimchi::static_params(),
),
};
let input = input
.into_iter()
.map(|elem| {
let mut input_bytes = vec![];
elem.into_repr()
.serialize(&mut input_bytes)
.expect("canonical serialiation should work");
match mode {
Mode::Hex => hex::encode(&input_bytes),
Mode::B10 => BigUint::from_bytes_le(&input_bytes).to_string(),
}
})
.collect();
let mut output_bytes = vec![];
output
.into_repr()
.serialize(&mut output_bytes)
.expect("canonical serialization should work");
test_vectors.push(TestVector {
input,
output: match mode {
Mode::Hex => hex::encode(&output_bytes),
Mode::B10 => BigUint::from_bytes_le(&output_bytes).to_string(),
},
})
}
let name = match param_type {
ParamType::Legacy => "legacy",
ParamType::Kimchi => "kimchi",
}
.into();
TestVectors { name, test_vectors }
}