Introduce trait StakeTableScheme and migrate previous implementations

Cargo.toml

@@ -38,6 +38,7 @@ typenum = { version = "1.16.0" }
 [dev-dependencies]
 criterion = { version = "0.5.1", features = ["html_reports"] }
 jf-primitives = { git = "https://github.com/espressosystems/jellyfish", features = ["test-srs"] }
+rand_chacha = { version = "0.3.1", default-features = false }
 sha2 = { version = "0.10" }
 
 [[bench]]

src/qc/bit_vector.rs

@@ -1,7 +1,10 @@
 //! Implementation for BitVectorQC that uses BLS signature + Bit vector.
 //! See more details in HotShot paper.
 
-use crate::qc::QuorumCertificate;
+use crate::{
+    qc::QuorumCertificate,
+    stake_table::{SnapshotVersion, StakeTableScheme},
+};
 use ark_std::{
     fmt::Debug,
     format,
@@ -20,32 +23,30 @@ use serde::{Deserialize, Serialize};
 use typenum::U32;
 
 /// An implementation of QC using BLS signature and a bit-vector.
-#[derive(Serialize, Deserialize)]
-pub struct BitVectorQC<A: AggregateableSignatureSchemes + Serialize + for<'a> Deserialize<'a>>(
+pub struct BitVectorQC<A: AggregateableSignatureSchemes, ST: StakeTableScheme>(
     PhantomData<A>,
+    PhantomData<ST>,
 );
 
 #[derive(Serialize, Deserialize, PartialEq, Debug)]
-pub struct StakeTableEntry<V> {
-    pub stake_key: V,
-    pub stake_amount: U256,
-}
-
-#[derive(Serialize, Deserialize, PartialEq, Debug)]
-pub struct QCParams<V, P> {
-    pub stake_entries: Vec<StakeTableEntry<V>>,
+pub struct QCParams<A: AggregateableSignatureSchemes, ST: StakeTableScheme> {
+    pub stake_table: ST,
     pub threshold: U256,
-    pub agg_sig_pp: P,
+    pub agg_sig_pp: A::PublicParameter,
 }
 
-impl<A> QuorumCertificate<A> for BitVectorQC<A>
+impl<A, ST> QuorumCertificate<A> for BitVectorQC<A, ST>
 where
-    A: AggregateableSignatureSchemes + Serialize + for<'a> Deserialize<'a>,
+    A: AggregateableSignatureSchemes + Serialize + for<'a> Deserialize<'a> + PartialEq,
+    ST: StakeTableScheme<Key = A::VerificationKey, Amount = U256>
+        + Serialize
+        + for<'a> Deserialize<'a>
+        + PartialEq,
 {
-    type QCProverParams = QCParams<A::VerificationKey, A::PublicParameter>;
+    type QCProverParams = QCParams<A, ST>;
 
     // TODO: later with SNARKs we'll use a smaller verifier parameter
-    type QCVerifierParams = QCParams<A::VerificationKey, A::PublicParameter>;
+    type QCVerifierParams = QCParams<A, ST>;
 
     type QC = (A::Signature, BitVec);
     type MessageLength = U32;
@@ -65,35 +66,42 @@ where
         signers: &BitSlice,
         sigs: &[A::Signature],
     ) -> Result<Self::QC, PrimitivesError> {
-        if signers.len() != qc_pp.stake_entries.len() {
+        let st_len = qc_pp.stake_table.len(SnapshotVersion::LastEpochStart)?;
+        if signers.len() != st_len {
             return Err(ParameterError(format!(
                 "bit vector len {} != the number of stake entries {}",
                 signers.len(),
-                qc_pp.stake_entries.len(),
+                st_len,
             )));
         }
-        let total_weight: U256 =
-            qc_pp
-                .stake_entries
-                .iter()
-                .zip(signers.iter())
-                .fold(U256::zero(), |acc, (entry, b)| {
+        let total_weight: U256 = qc_pp
+            .stake_table
+            .iter(SnapshotVersion::LastEpochStart)?
+            .zip(signers.iter())
+            .fold(
+                U256::zero(),
+                |acc, (entry, b)| {
                     if *b {
-                        acc + entry.stake_amount
+                        acc + entry.1
                     } else {
                         acc
                     }
-                });
+                },
+            );
         if total_weight < qc_pp.threshold {
             return Err(ParameterError(format!(
                 "total_weight {} less than threshold {}",
                 total_weight, qc_pp.threshold,
             )));
         }
         let mut ver_keys = vec![];
-        for (entry, b) in qc_pp.stake_entries.iter().zip(signers.iter()) {
+        for (entry, b) in qc_pp
+            .stake_table
+            .iter(SnapshotVersion::LastEpochStart)?
+            .zip(signers.iter())
+        {
             if *b {
-                ver_keys.push(entry.stake_key.clone());
+                ver_keys.push(entry.0.clone());
             }
         }
         if ver_keys.len() != sigs.len() {
@@ -114,35 +122,42 @@ where
         qc: &Self::QC,
     ) -> Result<Self::QuorumSize, PrimitivesError> {
         let (sig, signers) = qc;
-        if signers.len() != qc_vp.stake_entries.len() {
+        let st_len = qc_vp.stake_table.len(SnapshotVersion::LastEpochStart)?;
+        if signers.len() != st_len {
             return Err(ParameterError(format!(
                 "signers bit vector len {} != the number of stake entries {}",
                 signers.len(),
-                qc_vp.stake_entries.len(),
+                st_len,
             )));
         }
-        let total_weight: U256 =
-            qc_vp
-                .stake_entries
-                .iter()
-                .zip(signers.iter())
-                .fold(U256::zero(), |acc, (entry, b)| {
+        let total_weight: U256 = qc_vp
+            .stake_table
+            .iter(SnapshotVersion::LastEpochStart)?
+            .zip(signers.iter())
+            .fold(
+                U256::zero(),
+                |acc, (entry, b)| {
                     if *b {
-                        acc + entry.stake_amount
+                        acc + entry.1
                     } else {
                         acc
                     }
-                });
+                },
+            );
         if total_weight < qc_vp.threshold {
             return Err(ParameterError(format!(
                 "total_weight {} less than threshold {}",
                 total_weight, qc_vp.threshold,
             )));
         }
         let mut ver_keys = vec![];
-        for (entry, b) in qc_vp.stake_entries.iter().zip(signers.iter()) {
+        for (entry, b) in qc_vp
+            .stake_table
+            .iter(SnapshotVersion::LastEpochStart)?
+            .zip(signers.iter())
+        {
             if *b {
-                ver_keys.push(entry.stake_key.clone());
+                ver_keys.push(entry.0.clone());
             }
         }
         A::multi_sig_verify(&qc_vp.agg_sig_pp, &ver_keys[..], message, sig)?;
@@ -156,22 +171,23 @@ where
         qc: &Self::QC,
     ) -> Result<Vec<<A>::VerificationKey>, PrimitivesError> {
         let (_sig, signers) = qc;
-        if signers.len() != qc_vp.stake_entries.len() {
+        let st_len = qc_vp.stake_table.len(SnapshotVersion::LastEpochStart)?;
+        if signers.len() != st_len {
             return Err(ParameterError(format!(
                 "signers bit vector len {} != the number of stake entries {}",
                 signers.len(),
-                qc_vp.stake_entries.len(),
+                st_len,
             )));
         }
 
         Self::check(qc_vp, message, qc)?;
 
         let signer_pks: Vec<_> = qc_vp
-            .stake_entries
-            .iter()
+            .stake_table
+            .iter(SnapshotVersion::LastEpochStart)?
             .zip(signers.iter())
             .filter(|(_, b)| **b)
-            .map(|(pk, _)| pk.stake_key.clone())
+            .map(|(pk, _)| pk.0)
             .collect();
         Ok(signer_pks)
     }
@@ -180,49 +196,49 @@ where
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::stake_table::{StakeTable, StakeTableScheme};
     use jf_primitives::signatures::bls_over_bn254::{BLSOverBN254CurveSignatureScheme, KeyPair};
     use jf_primitives::signatures::SignatureScheme;
 
     macro_rules! test_quorum_certificate {
         ($aggsig:tt) => {
+            type ST = StakeTable<<$aggsig as SignatureScheme>::VerificationKey>;
             let mut rng = jf_utils::test_rng();
+
             let agg_sig_pp = $aggsig::param_gen(Some(&mut rng)).unwrap();
             let key_pair1 = KeyPair::generate(&mut rng);
             let key_pair2 = KeyPair::generate(&mut rng);
             let key_pair3 = KeyPair::generate(&mut rng);
-            let entry1 = StakeTableEntry {
-                stake_key: key_pair1.ver_key(),
-                stake_amount: U256::from(3u8),
-            };
-            let entry2 = StakeTableEntry {
-                stake_key: key_pair2.ver_key(),
-                stake_amount: U256::from(5u8),
-            };
-            let entry3 = StakeTableEntry {
-                stake_key: key_pair3.ver_key(),
-                stake_amount: U256::from(7u8),
-            };
+
+            let mut st = ST::new(3);
+            st.register(key_pair1.ver_key(), U256::from(3u8)).unwrap();
+            st.register(key_pair2.ver_key(), U256::from(5u8)).unwrap();
+            st.register(key_pair3.ver_key(), U256::from(7u8)).unwrap();
+            st.advance();
+            st.advance();
+
             let qc_pp = QCParams {
-                stake_entries: vec![entry1, entry2, entry3],
+                stake_table: st,
                 threshold: U256::from(10u8),
                 agg_sig_pp,
             };
+
             let msg = [72u8; 32];
-            let sig1 = BitVectorQC::<$aggsig>::sign(
+            let sig1 = BitVectorQC::<$aggsig, ST>::sign(
                 &agg_sig_pp,
                 &msg.into(),
                 key_pair1.sign_key_ref(),
                 &mut rng,
             )
             .unwrap();
-            let sig2 = BitVectorQC::<$aggsig>::sign(
+            let sig2 = BitVectorQC::<$aggsig, ST>::sign(
                 &agg_sig_pp,
                 &msg.into(),
                 key_pair2.sign_key_ref(),
                 &mut rng,
             )
             .unwrap();
-            let sig3 = BitVectorQC::<$aggsig>::sign(
+            let sig3 = BitVectorQC::<$aggsig, ST>::sign(
                 &agg_sig_pp,
                 &msg.into(),
                 key_pair3.sign_key_ref(),
@@ -232,15 +248,15 @@ mod tests {
 
             // happy path
             let signers = bitvec![0, 1, 1];
-            let qc = BitVectorQC::<$aggsig>::assemble(
+            let qc = BitVectorQC::<$aggsig, ST>::assemble(
                 &qc_pp,
                 signers.as_bitslice(),
                 &[sig2.clone(), sig3.clone()],
             )
             .unwrap();
-            assert!(BitVectorQC::<$aggsig>::check(&qc_pp, &msg.into(), &qc).is_ok());
+            assert!(BitVectorQC::<$aggsig, ST>::check(&qc_pp, &msg.into(), &qc).is_ok());
             assert_eq!(
-                BitVectorQC::<$aggsig>::trace(&qc_pp, &msg.into(), &qc).unwrap(),
+                BitVectorQC::<$aggsig, ST>::trace(&qc_pp, &msg.into(), &qc).unwrap(),
                 vec![key_pair2.ver_key(), key_pair3.ver_key()],
             );
 
@@ -250,56 +266,78 @@ mod tests {
                 bincode::deserialize(&bincode::serialize(&qc).unwrap()).unwrap()
             );
 
+            // (alex) since deserialized stake table's leaf would contain normalized projective
+            // points with Z=1, which differs from the original projective representation.
+            // We compare individual fields for equivalence instead.
+            let de_qc_pp: QCParams<$aggsig, ST> =
+                bincode::deserialize(&bincode::serialize(&qc_pp).unwrap()).unwrap();
+            assert_eq!(
+                qc_pp.stake_table.commitment(SnapshotVersion::Head).unwrap(),
+                de_qc_pp
+                    .stake_table
+                    .commitment(SnapshotVersion::Head)
+                    .unwrap(),
+            );
             assert_eq!(
-                qc_pp,
-                bincode::deserialize(&bincode::serialize(&qc_pp).unwrap()).unwrap()
+                qc_pp
+                    .stake_table
+                    .commitment(SnapshotVersion::LastEpochStart)
+                    .unwrap(),
+                de_qc_pp
+                    .stake_table
+                    .commitment(SnapshotVersion::LastEpochStart)
+                    .unwrap(),
             );
+            assert_eq!(qc_pp.threshold, de_qc_pp.threshold);
+            assert_eq!(qc_pp.agg_sig_pp, de_qc_pp.agg_sig_pp);
 
             // bad paths
             // number of signatures unmatch
-            assert!(BitVectorQC::<$aggsig>::assemble(
+            assert!(BitVectorQC::<$aggsig, ST>::assemble(
                 &qc_pp,
                 signers.as_bitslice(),
                 &[sig2.clone()]
             )
             .is_err());
             // total weight under threshold
             let active_bad = bitvec![1, 1, 0];
-            assert!(BitVectorQC::<$aggsig>::assemble(
+            assert!(BitVectorQC::<$aggsig, ST>::assemble(
                 &qc_pp,
                 active_bad.as_bitslice(),
                 &[sig1.clone(), sig2.clone()]
             )
             .is_err());
             // wrong bool vector length
             let active_bad_2 = bitvec![0, 1, 1, 0];
-            assert!(BitVectorQC::<$aggsig>::assemble(
+            assert!(BitVectorQC::<$aggsig, ST>::assemble(
                 &qc_pp,
                 active_bad_2.as_bitslice(),
                 &[sig2, sig3],
             )
             .is_err());
 
-            assert!(BitVectorQC::<$aggsig>::check(
+            assert!(BitVectorQC::<$aggsig, ST>::check(
                 &qc_pp,
                 &msg.into(),
                 &(qc.0.clone(), active_bad)
             )
             .is_err());
-            assert!(BitVectorQC::<$aggsig>::check(
+            assert!(BitVectorQC::<$aggsig, ST>::check(
                 &qc_pp,
                 &msg.into(),
                 &(qc.0.clone(), active_bad_2)
             )
             .is_err());
             let bad_msg = [70u8; 32];
-            assert!(BitVectorQC::<$aggsig>::check(&qc_pp, &bad_msg.into(), &qc).is_err());
+            assert!(BitVectorQC::<$aggsig, ST>::check(&qc_pp, &bad_msg.into(), &qc).is_err());
 
             let bad_sig = &sig1;
-            assert!(
-                BitVectorQC::<$aggsig>::check(&qc_pp, &msg.into(), &(bad_sig.clone(), qc.1))
-                    .is_err()
-            );
+            assert!(BitVectorQC::<$aggsig, ST>::check(
+                &qc_pp,
+                &msg.into(),
+                &(bad_sig.clone(), qc.1)
+            )
+            .is_err());
         };
     }
     #[test]

src/stake_table/config.rs

@@ -9,7 +9,7 @@ use jf_primitives::crhf::FixedLengthRescueCRHF;
 pub(crate) const TREE_BRANCH: usize = 3;
 
 /// Internal type of Merkle node value(commitment)
-pub(crate) type FieldType = ark_bn254::Fr;
+pub(crate) type FieldType = ark_bn254::Fq;
 /// Hash algorithm used in Merkle tree, using a RATE-3 rescue
 pub(crate) type Digest = FixedLengthRescueCRHF<FieldType, TREE_BRANCH, 1>;