Variable stake and randomized leader selection

hotshot-types/src/traits/election.rs

@@ -7,8 +7,12 @@
 //! The election trait, used to decide which node is the leader and determine if a vote is valid.
 use std::{collections::BTreeSet, fmt::Debug, num::NonZeroU64};
 
+use sha2::{Digest, Sha256, Sha512};
+
+use crate::traits::signature_key::StakeTableEntryType;
 use async_trait::async_trait;
 use hotshot_utils::anytrace::Result;
+use primitive_types::{U256, U512};
 
 use super::node_implementation::NodeType;
 use crate::{traits::signature_key::SignatureKey, PeerConfig};
@@ -153,3 +157,75 @@ pub trait Membership<TYPES: NodeType>: Debug + Send + Sync {
         None
     }
 }
+
+/// Calculate `xor(drb.cycle(), public_key)`, returning the result as a vector of bytes
+fn cyclic_xor(drb: [u8; 32], public_key: Vec<u8>) -> Vec<u8> {
+    let drb: Vec<u8> = drb.to_vec();
+
+    let mut result: Vec<u8> = vec![];
+
+    for (drb_byte, public_key_byte) in public_key.iter().zip(drb.iter().cycle()) {
+        result.push(drb_byte ^ public_key_byte);
+    }
+
+    result
+}
+
+/// Generate the stake table CDF, as well as a hash of the resulting stake table
+pub fn generate_stake_cdf<Key: SignatureKey, Entry: StakeTableEntryType<Key>>(
+    mut stake_table: Vec<Entry>,
+    drb: [u8; 32],
+) -> (Vec<(Entry, U256)>, [u8; 32]) {
+    // sort by xor(public_key, drb_result)
+    stake_table.sort_by(|a, b| {
+        cyclic_xor(drb, a.public_key().to_bytes()).cmp(&cyclic_xor(drb, b.public_key().to_bytes()))
+    });
+
+    let mut hasher = Sha256::new();
+
+    let mut cumulative_stake = U256::from(0);
+    let mut cdf = vec![];
+
+    for entry in stake_table {
+        cumulative_stake += entry.stake();
+        hasher.update(entry.public_key().to_bytes());
+
+        cdf.push((entry, cumulative_stake));
+    }
+
+    (cdf, hasher.finalize().into())
+}
+
+/// select the leader for a view
+///
+/// # Panics
+/// Panics if `cdf` is empty. Results in undefined behaviour if `cdf` is not ordered.
+///
+/// Note that we try to downcast a U512 to a U256,
+/// but this should never panic because the U512 should be strictly smaller than U256::MAX by construction.
+pub fn select_randomized_leader<SignatureKey, Entry: StakeTableEntryType<SignatureKey> + Clone>(
+    cdf: Vec<(Entry, U256)>,
+    stake_table_hash: [u8; 32],
+    drb: [u8; 32],
+    view: u64,
+) -> Entry {
+    // We hash the concatenated drb, view and stake table hash.
+    let mut hasher = Sha512::new();
+    hasher.update(drb);
+    hasher.update(view.to_le_bytes());
+    hasher.update(stake_table_hash);
+    let raw_breakpoint: [u8; 64] = hasher.finalize().into();
+
+    // then calculate the remainder modulo the total stake as a U512
+    let remainder: U512 =
+        U512::from_little_endian(&raw_breakpoint) % U512::from(cdf.last().unwrap().1);
+
+    // and drop the top 32 bytes, downcasting to a U256
+    let breakpoint: U256 = U256::try_from(remainder).unwrap();
+
+    // now find the first index where the breakpoint is strictly smaller than the cdf
+    let index = cdf.partition_point(|(_, cumulative_stake)| breakpoint < *cumulative_stake);
+
+    // and return the corresponding entry
+    cdf[index].0.clone()
+}