From 9990e36336beb8b7b2835a609ab68e032c774892 Mon Sep 17 00:00:00 2001
From: Jonathan Wang <31040440+jonathanpwang@users.noreply.github.com>
Date: Thu, 7 Sep 2023 23:14:31 -0700
Subject: [PATCH] chore: use `PrimeField` for `OptimizedPoseidonSpec`

---
 halo2-base/src/poseidon/hasher/mds.rs  | 10 +++++-----
 halo2-base/src/poseidon/hasher/spec.rs | 20 +++++++++++++-------
 2 files changed, 18 insertions(+), 12 deletions(-)

diff --git a/halo2-base/src/poseidon/hasher/mds.rs b/halo2-base/src/poseidon/hasher/mds.rs
index 536fd7b3..159b031f 100644
--- a/halo2-base/src/poseidon/hasher/mds.rs
+++ b/halo2-base/src/poseidon/hasher/mds.rs
@@ -1,5 +1,5 @@
 #![allow(clippy::needless_range_loop)]
-use crate::utils::ScalarField;
+use crate::ff::PrimeField;
 
 /// The type used to hold the MDS matrix
 pub(crate) type Mds<F, const T: usize> = [[F; T]; T];
@@ -8,7 +8,7 @@ pub(crate) type Mds<F, const T: usize> = [[F; T]; T];
 /// also called `pre_sparse_mds` and sparse matrices that enables us to reduce
 /// number of multiplications in apply MDS step
 #[derive(Debug, Clone)]
-pub struct MDSMatrices<F: ScalarField, const T: usize, const RATE: usize> {
+pub struct MDSMatrices<F: PrimeField, const T: usize, const RATE: usize> {
     pub(crate) mds: MDSMatrix<F, T, RATE>,
     pub(crate) pre_sparse_mds: MDSMatrix<F, T, RATE>,
     pub(crate) sparse_matrices: Vec<SparseMDSMatrix<F, T, RATE>>,
@@ -17,16 +17,16 @@ pub struct MDSMatrices<F: ScalarField, const T: usize, const RATE: usize> {
 /// `SparseMDSMatrix` are in `[row], [hat | identity]` form and used in linear
 /// layer of partial rounds instead of the original MDS
 #[derive(Debug, Clone)]
-pub struct SparseMDSMatrix<F: ScalarField, const T: usize, const RATE: usize> {
+pub struct SparseMDSMatrix<F: PrimeField, const T: usize, const RATE: usize> {
     pub(crate) row: [F; T],
     pub(crate) col_hat: [F; RATE],
 }
 
 /// `MDSMatrix` is applied to `State` to achive linear layer of Poseidon
 #[derive(Clone, Debug)]
-pub struct MDSMatrix<F: ScalarField, const T: usize, const RATE: usize>(pub(crate) Mds<F, T>);
+pub struct MDSMatrix<F: PrimeField, const T: usize, const RATE: usize>(pub(crate) Mds<F, T>);
 
-impl<F: ScalarField, const T: usize, const RATE: usize> MDSMatrix<F, T, RATE> {
+impl<F: PrimeField, const T: usize, const RATE: usize> MDSMatrix<F, T, RATE> {
     pub(crate) fn mul_vector(&self, v: &[F; T]) -> [F; T] {
         let mut res = [F::ZERO; T];
         for i in 0..T {
diff --git a/halo2-base/src/poseidon/hasher/spec.rs b/halo2-base/src/poseidon/hasher/spec.rs
index c0e7142c..1568935b 100644
--- a/halo2-base/src/poseidon/hasher/spec.rs
+++ b/halo2-base/src/poseidon/hasher/spec.rs
@@ -1,4 +1,7 @@
-use crate::{poseidon::hasher::mds::*, utils::ScalarField};
+use crate::{
+    ff::{FromUniformBytes, PrimeField},
+    poseidon::hasher::mds::*,
+};
 
 use poseidon_rs::poseidon::primitives::Spec as PoseidonSpec; // trait
 use std::marker::PhantomData;
@@ -6,7 +9,7 @@ use std::marker::PhantomData;
 // struct so we can use PoseidonSpec trait to generate round constants and MDS matrix
 #[derive(Debug)]
 pub(crate) struct Poseidon128Pow5Gen<
-    F: ScalarField,
+    F: PrimeField,
     const T: usize,
     const RATE: usize,
     const R_F: usize,
@@ -17,7 +20,7 @@ pub(crate) struct Poseidon128Pow5Gen<
 }
 
 impl<
-        F: ScalarField,
+        F: PrimeField,
         const T: usize,
         const RATE: usize,
         const R_F: usize,
@@ -51,7 +54,7 @@ impl<
 /// `OptimizedPoseidonSpec` holds construction parameters as well as constants that are used in
 /// permutation step.
 #[derive(Debug, Clone)]
-pub struct OptimizedPoseidonSpec<F: ScalarField, const T: usize, const RATE: usize> {
+pub struct OptimizedPoseidonSpec<F: PrimeField, const T: usize, const RATE: usize> {
     pub(crate) r_f: usize,
     pub(crate) mds_matrices: MDSMatrices<F, T, RATE>,
     pub(crate) constants: OptimizedConstants<F, T>,
@@ -61,15 +64,18 @@ pub struct OptimizedPoseidonSpec<F: ScalarField, const T: usize, const RATE: usi
 /// full rounds has T sized constants there is a single constant for each
 /// partial round
 #[derive(Debug, Clone)]
-pub struct OptimizedConstants<F: ScalarField, const T: usize> {
+pub struct OptimizedConstants<F: PrimeField, const T: usize> {
     pub(crate) start: Vec<[F; T]>,
     pub(crate) partial: Vec<F>,
     pub(crate) end: Vec<[F; T]>,
 }
 
-impl<F: ScalarField, const T: usize, const RATE: usize> OptimizedPoseidonSpec<F, T, RATE> {
+impl<F: PrimeField, const T: usize, const RATE: usize> OptimizedPoseidonSpec<F, T, RATE> {
     /// Generate new spec with specific number of full and partial rounds. `SECURE_MDS` is usually 0, but may need to be specified because insecure matrices may sometimes be generated
-    pub fn new<const R_F: usize, const R_P: usize, const SECURE_MDS: usize>() -> Self {
+    pub fn new<const R_F: usize, const R_P: usize, const SECURE_MDS: usize>() -> Self
+    where
+        F: FromUniformBytes<64> + Ord,
+    {
         let (round_constants, mds, mds_inv) =
             Poseidon128Pow5Gen::<F, T, RATE, R_F, R_P, SECURE_MDS>::constants();
         let mds = MDSMatrix(mds);