use crate::Worker;
use snarkvm::{
ledger::narwhal::{Transmission, TransmissionID},
prelude::{Network, ToBytes},
};
use anyhow::{Result, bail};
use sha2::{Digest, Sha256};
fn double_sha256(data: &[u8]) -> [u8; 32] {
let digest = Sha256::digest(Sha256::digest(data));
let mut ret = [0u8; 32];
ret.copy_from_slice(&digest);
ret
}
pub fn sha256d_to_u128(data: &[u8]) -> u128 {
let hash_slice = double_sha256(data);
let mut hash = [0u8; 16];
hash[..].copy_from_slice(&hash_slice[..16]);
u128::from_le_bytes(hash)
}
pub fn assign_to_worker<N: Network>(transmission_id: impl Into<TransmissionID<N>>, num_workers: u8) -> Result<u8> {
if num_workers == 1 {
return Ok(0);
}
let hash = sha256d_to_u128(&transmission_id.into().to_bytes_le()?);
let worker_id = (hash % num_workers as u128) as u8;
Ok(worker_id)
}
pub fn assign_to_workers<N: Network>(
workers: &[Worker<N>],
transmissions: impl Iterator<Item = (TransmissionID<N>, Transmission<N>)>,
op: impl Fn(&Worker<N>, TransmissionID<N>, Transmission<N>),
) -> Result<()> {
let num_workers = u8::try_from(workers.len()).expect("Too many workers");
for (transmission_id, transmission) in transmissions.into_iter() {
let Ok(worker_id) = assign_to_worker(transmission_id, num_workers) else {
bail!("Unable to assign transmission ID '{transmission_id}' to a worker")
};
match workers.get(worker_id as usize) {
Some(worker) => op(worker, transmission_id, transmission),
None => bail!("Unable to find worker {worker_id}"),
};
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use snarkvm::prelude::puzzle::SolutionID;
type CurrentNetwork = snarkvm::prelude::MainnetV0;
#[test]
fn test_assign_to_worker() {
let data = "Hello Aleo".as_bytes();
let sha = double_sha256(data);
assert_eq!(sha, [
113, 157, 210, 34, 60, 51, 220, 8, 63, 213, 79, 8, 117, 190, 134, 206, 127, 197, 21, 180, 116, 49, 218,
150, 49, 116, 116, 38, 244, 135, 215, 14
]);
let hash = sha256d_to_u128(data);
assert_eq!(hash, 274520597840828436951879875061540363633u128);
let transmission_id: TransmissionID<CurrentNetwork> =
TransmissionID::Solution(SolutionID::from(123456789), 12345);
let worker_id = assign_to_worker(transmission_id, 5).unwrap();
assert_eq!(worker_id, 4);
}
}