poseidon_over_3_fields.rs

use blstrs::Scalar as Fr;
use halo::pasta::{Fp, Fq as Fv};
// more imports...

lazy_static! {
    pub static ref BLS_P_MINUS_1: [u8; 32] = (-Fr::one()).to_repr();
    pub static ref PALLAS_P_MINUS_1: [u8; 32] = (-Fp::one()).to_repr();
    pub static ref VESTA_P_MINUS_1: [u8; 32] = (-Fv::one()).to_repr();

    pub static ref POSEIDON_CONSTANTS_PALLAS_2: PoseidonConstants::<Fp, U2> = PoseidonConstants::new();
    pub static ref POSEIDON_CONSTANTS_PALLAS_4: PoseidonConstants::<Fp, U4> = PoseidonConstants::new();
    pub static ref POSEIDON_CONSTANTS_PALLAS_8: PoseidonConstants::<Fp, U8> = PoseidonConstants::new();
    pub static ref POSEIDON_CONSTANTS_PALLAS_16: PoseidonConstants::<Fp, U16> = PoseidonConstants::new();
    pub static ref POSEIDON_MD_CONSTANTS_PALLAS: PoseidonConstants::<Fp, PoseidonMDArity> =
        PoseidonConstants::new();

    pub static ref POSEIDON_CONSTANTS_VESTA_2: PoseidonConstants::<Fv, U2> = PoseidonConstants::new();
    pub static ref POSEIDON_CONSTANTS_VESTA_4: PoseidonConstants::<Fv, U4> = PoseidonConstants::new();
    pub static ref POSEIDON_CONSTANTS_VESTA_8: PoseidonConstants::<Fv, U8> = PoseidonConstants::new();
    pub static ref POSEIDON_CONSTANTS_VESTA_16: PoseidonConstants::<Fv, U16> = PoseidonConstants::new();
    pub static ref POSEIDON_MD_CONSTANTS_VESTA: PoseidonConstants::<Fv, PoseidonMDArity> =
        PoseidonConstants::new();
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum FieldName {
    Bls,
    Pallas,
    Vesta,
}

#[inline]
fn field_name<F: PrimeField<Repr = [u8; 32]>>() -> FieldName {
    let p_minus_1 = (-F::one()).to_repr();

    if p_minus_1 == *BLS_P_MINUS_1 {
        FieldName::Bls
    } else if p_minus_1 == *PALLAS_P_MINUS_1 {
        FieldName::Pallas
    } else if p_minus_1 == *VESTA_P_MINUS_1 {
        FieldName::Vesta
    } else {
        unimplemented!("provided field is not BLS12-381, Pallas, or Vesta scalar field");
    }
}

// Convert a field element from type `F1` to `F2`.
fn f1_into_f2<F1, F2>(f: F1) -> F2
where
    F1: PrimeField<Repr = [u8; 32]>,
    F2: PrimeField<Repr = [u8; 32]>,
{
    debug_assert_eq!(field_name::<F1>(), field_name::<F2>());
    let mut repr_out = <F2 as PrimeField>::Repr::default();
    repr_out.as_mut().copy_from_slice(f.to_repr().as_ref());
    F2::from_repr_vartime(repr_out).unwrap()
}

fn poseidon<F: PrimeField<Repr = [u8; 32]>>(preimage: &[F]) -> F {
    match field_name::<F>() {
        FieldName::Bls => {
            let preimage: Vec<Fr> = preimage.iter().copied().map(f1_into_f2::<F, Fr>).collect();
            let digest: Fr = match preimage.len() {
                2 => Poseidon::new_with_preimage(&preimage, &*POSEIDON_CONSTANTS_2).hash(),
                4 => Poseidon::new_with_preimage(&preimage, &*POSEIDON_CONSTANTS_4).hash(),
                8 => Poseidon::new_with_preimage(&preimage, &*POSEIDON_CONSTANTS_8).hash(),
                16 => Poseidon::new_with_preimage(&preimage, &*POSEIDON_CONSTANTS_16).hash(),
                _ => unimplemented!("preimage length is not supported"),
            };
            f1_into_f2::<Fr, F>(digest)
        }
        FieldName::Pallas => {
            let preimage: Vec<Fp> = preimage.iter().copied().map(f1_into_f2::<F, Fp>).collect();
            let digest: Fp = match preimage.len() {
                2 => Poseidon::new_with_preimage(&preimage, &*POSEIDON_CONSTANTS_PALLAS_2).hash(),
                4 => Poseidon::new_with_preimage(&preimage, &*POSEIDON_CONSTANTS_PALLAS_4).hash(),
                8 => Poseidon::new_with_preimage(&preimage, &*POSEIDON_CONSTANTS_PALLAS_8).hash(),
                16 => Poseidon::new_with_preimage(&preimage, &*POSEIDON_CONSTANTS_PALLAS_16).hash(),
                _ => unimplemented!("preimage length is not supported"),
            };
            f1_into_f2::<Fp, F>(digest)
        }
        FieldName::Vesta => {
            let preimage: Vec<Fv> = preimage.iter().copied().map(f1_into_f2::<F, Fv>).collect();
            let digest: Fv = match preimage.len() {
                2 => Poseidon::new_with_preimage(&preimage, &*POSEIDON_CONSTANTS_VESTA_2).hash(),
                4 => Poseidon::new_with_preimage(&preimage, &*POSEIDON_CONSTANTS_VESTA_4).hash(),
                8 => Poseidon::new_with_preimage(&preimage, &*POSEIDON_CONSTANTS_VESTA_8).hash(),
                16 => Poseidon::new_with_preimage(&preimage, &*POSEIDON_CONSTANTS_VESTA_16).hash(),
                _ => unimplemented!("preimage length is not supported"),
            };
            f1_into_f2::<Fv, F>(digest)
        }
    }
}