Merge pull request #1932 from EspressoSystems/vec-stake-table

feat: Stake table refactor & vector based implementation

Cargo.toml

@@ -6,25 +6,24 @@ rust-version = "1.65.0"
 homepage = "https://github.com/EspressoSystems/HotShot"
 documentation = "https://hotshot.docs.espressosys.com"
 repository = "https://github.com/EspressoSystems/HotShot"
-
 ### Workspace
 
 # The hotshot-types crate needs to be a seperate crate, as to not create a circular dependency
 # when implementing traits externally
 [workspace]
 members = [
-        "crates/libp2p-networking",
-        "crates/testing",
-        "crates/types",
-        "crates/utils",
         "crates/constants",
-        # "testing-macros",
-        "crates/task",
-        "crates/task-impls",
+        "crates/hotshot",
         "crates/hotshot-qc",
         "crates/hotshot-signature-key",
         "crates/hotshot-stake-table",
-        "crates/hotshot",
+        "crates/libp2p-networking",
+        # "testing-macros",
+        "crates/task",
+        "crates/task-impls",
+        "crates/testing",
+        "crates/types",
+        "crates/utils",
 ]
 resolver = "2"
 
@@ -115,9 +114,7 @@ tokio = { version = "1.32.0", features = [
         "time",
         "tracing",
 ] }
-
-### Profiles
-###
+### Profiles###
 ### Note: these only apply to example executables or tests built from within this crate. They have
 ### no effect on crates that depend on this crate.
 

crates/hotshot-qc/src/bit_vector.rs

@@ -219,9 +219,12 @@ mod tests {
             let key_pair3 = KeyPair::generate(&mut rng);
 
             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.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();
 

crates/hotshot-stake-table/Cargo.toml

@@ -8,6 +8,7 @@ rust-version = { workspace = true }
 
 [dependencies]
 ark-bn254 = "0.4.0"
+ark-ed-on-bn254 = "0.4.0"
 ark-ff = "0.4.0"
 ark-serialize = { workspace = true }
 ark-std = { workspace = true }

crates/hotshot-stake-table/src/config.rs

@@ -0,0 +1,4 @@
+//! Configuration file for stake table
+
+/// Capacity of a stake table
+pub const STAKE_TABLE_CAPACITY: usize = 1000;

crates/hotshot-stake-table/src/lib.rs

@@ -2,4 +2,7 @@
 #![deny(warnings)]
 #![deny(missing_docs)]
 
+pub mod config;
 pub mod mt_based;
+pub mod utils;
+pub mod vec_based;

crates/hotshot-stake-table/src/mt_based.rs

@@ -36,11 +36,13 @@ impl<K: Key> StakeTableScheme for StakeTable<K> {
     type Commitment = MerkleCommitment;
     type LookupProof = MerkleProof<K>;
     type IntoIter = internal::IntoIter<K>;
+    type Aux = ();
 
     fn register(
         &mut self,
         new_key: Self::Key,
         amount: Self::Amount,
+        _: Self::Aux,
     ) -> Result<(), StakeTableError> {
         match self.mapping.get(&new_key) {
             Some(_) => Err(StakeTableError::ExistingKey),
@@ -86,7 +88,18 @@ impl<K: Key> StakeTableScheme for StakeTable<K> {
         self.mapping.contains_key(key)
     }
 
-    fn lookup(
+    fn lookup(&self, version: SnapshotVersion, key: &K) -> Result<Self::Amount, StakeTableError> {
+        let root = Self::get_root(self, version)?;
+        match self.mapping.get(key) {
+            Some(index) => {
+                let branches = to_merkle_path(*index, self.height);
+                root.simple_lookup(self.height, &branches)
+            }
+            None => Err(StakeTableError::KeyNotFound),
+        }
+    }
+
+    fn lookup_with_proof(
         &self,
         version: SnapshotVersion,
         key: &Self::Key,
@@ -104,19 +117,13 @@ impl<K: Key> StakeTableScheme for StakeTable<K> {
         Ok((amount, proof))
     }
 
-    fn simple_lookup(
+    fn lookup_with_aux_and_proof(
         &self,
         version: SnapshotVersion,
-        key: &K,
-    ) -> Result<Self::Amount, StakeTableError> {
-        let root = Self::get_root(self, version)?;
-        match self.mapping.get(key) {
-            Some(index) => {
-                let branches = to_merkle_path(*index, self.height);
-                root.simple_lookup(self.height, &branches)
-            }
-            None => Err(StakeTableError::KeyNotFound),
-        }
+        key: &Self::Key,
+    ) -> Result<(Self::Amount, Self::Aux, Self::LookupProof), StakeTableError> {
+        let (amount, proof) = self.lookup_with_proof(version, key)?;
+        Ok((amount, (), proof))
     }
 
     fn update(
@@ -224,14 +231,14 @@ mod tests {
 
     #[test]
     fn test_stake_table() -> Result<(), StakeTableError> {
-        let mut st = StakeTable::new(3);
+        let mut st = StakeTable::<Key>::new(3);
         let keys = (0..10).map(Key::from).collect::<Vec<_>>();
         assert_eq!(st.total_stake(SnapshotVersion::Head)?, U256::from(0));
 
         // Registering keys
         keys.iter()
             .take(4)
-            .for_each(|&key| st.register(key, U256::from(100)).unwrap());
+            .for_each(|key| st.register(*key, U256::from(100), ()).unwrap());
         assert_eq!(st.total_stake(SnapshotVersion::Head)?, U256::from(400));
         assert_eq!(st.total_stake(SnapshotVersion::EpochStart)?, U256::from(0));
         assert_eq!(
@@ -247,7 +254,7 @@ mod tests {
         keys.iter()
             .skip(4)
             .take(3)
-            .for_each(|&key| st.register(key, U256::from(100)).unwrap());
+            .for_each(|key| st.register(*key, U256::from(100), ()).unwrap());
         assert_eq!(st.total_stake(SnapshotVersion::Head)?, U256::from(600));
         assert_eq!(
             st.total_stake(SnapshotVersion::EpochStart)?,
@@ -260,7 +267,7 @@ mod tests {
         st.advance();
         keys.iter()
             .skip(7)
-            .for_each(|&key| st.register(key, U256::from(100)).unwrap());
+            .for_each(|key| st.register(*key, U256::from(100), ()).unwrap());
         assert_eq!(st.total_stake(SnapshotVersion::Head)?, U256::from(900));
         assert_eq!(
             st.total_stake(SnapshotVersion::EpochStart)?,
@@ -272,7 +279,7 @@ mod tests {
         );
 
         // No duplicate register
-        assert!(st.register(keys[0], U256::from(100)).is_err());
+        assert!(st.register(keys[0], U256::from(100), ()).is_err());
         // The 9-th key is still in head stake table
         assert!(st.lookup(SnapshotVersion::EpochStart, &keys[9]).is_err());
         assert!(st.lookup(SnapshotVersion::EpochStart, &keys[5]).is_ok());
@@ -306,7 +313,9 @@ mod tests {
         );
 
         // Testing membership proof
-        let proof = st.lookup(SnapshotVersion::EpochStart, &keys[5])?.1;
+        let proof = st
+            .lookup_with_proof(SnapshotVersion::EpochStart, &keys[5])?
+            .1;
         assert!(proof
             .verify(&st.commitment(SnapshotVersion::EpochStart)?)
             .is_ok());

crates/hotshot-stake-table/src/mt_based/config.rs

@@ -1,8 +1,8 @@
 //! Config file for stake table
+use crate::utils::ToFields;
 use ark_ff::PrimeField;
 use ark_std::vec;
-use ethereum_types::U256;
-use jf_primitives::crhf::FixedLengthRescueCRHF;
+use jf_primitives::{crhf::FixedLengthRescueCRHF, signatures::bls_over_bn254};
 
 /// Branch of merkle tree.
 /// Set to 3 because we are currently using RATE-3 rescue hash function
@@ -13,9 +13,19 @@ 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>;
 
-/// convert a U256 to a field element.
-pub(crate) fn u256_to_field<F: PrimeField>(v: &U256) -> F {
-    let mut bytes = vec![0u8; 32];
-    v.to_little_endian(&mut bytes);
-    F::from_le_bytes_mod_order(&bytes)
+impl ToFields<FieldType> for FieldType {
+    const SIZE: usize = 1;
+    fn to_fields(&self) -> Vec<FieldType> {
+        vec![*self]
+    }
+}
+
+impl ToFields<FieldType> for bls_over_bn254::VerKey {
+    const SIZE: usize = 2;
+    fn to_fields(&self) -> Vec<FieldType> {
+        let bytes = jf_utils::to_bytes!(&self.to_affine()).unwrap();
+        let x = <ark_bn254::Fq as PrimeField>::from_le_bytes_mod_order(&bytes[..32]);
+        let y = <ark_bn254::Fq as PrimeField>::from_le_bytes_mod_order(&bytes[32..]);
+        vec![x, y]
+    }
 }