use snarkvm::{
console::prelude::*,
ledger::{
block::Transaction,
narwhal::{Data, Transmission, TransmissionID},
puzzle::{Solution, SolutionID},
},
};
use indexmap::{IndexMap, IndexSet};
use std::collections::{HashMap, VecDeque, hash_map::Entry::Vacant};
#[derive(Clone, Debug)]
pub struct Ready<N: Network> {
transmission_ids: HashMap<TransmissionID<N>, i64>,
transmissions: VecDeque<(TransmissionID<N>, Transmission<N>)>,
offset: i64,
}
impl<N: Network> Default for Ready<N> {
fn default() -> Self {
Self::new()
}
}
impl<N: Network> Ready<N> {
pub fn new() -> Self {
Self { transmission_ids: Default::default(), transmissions: Default::default(), offset: Default::default() }
}
pub fn is_empty(&self) -> bool {
self.transmissions.is_empty()
}
pub fn num_transmissions(&self) -> usize {
self.transmissions.len()
}
pub fn num_ratifications(&self) -> usize {
self.transmission_ids.keys().filter(|id| matches!(id, TransmissionID::Ratification)).count()
}
pub fn num_solutions(&self) -> usize {
self.transmission_ids.keys().filter(|id| matches!(id, TransmissionID::Solution(..))).count()
}
pub fn num_transactions(&self) -> usize {
self.transmission_ids.keys().filter(|id| matches!(id, TransmissionID::Transaction(..))).count()
}
pub fn transmission_ids(&self) -> IndexSet<TransmissionID<N>> {
self.transmission_ids.keys().copied().collect()
}
pub fn transmissions(&self) -> IndexMap<TransmissionID<N>, Transmission<N>> {
self.transmissions.iter().cloned().collect()
}
pub fn solutions(&self) -> Vec<(SolutionID<N>, Data<Solution<N>>)> {
self.transmissions
.iter()
.filter_map(|(id, transmission)| match (id, transmission) {
(TransmissionID::Solution(id, _), Transmission::Solution(solution)) => Some((*id, solution.clone())),
_ => None,
})
.collect()
}
pub fn transactions(&self) -> Vec<(N::TransactionID, Data<Transaction<N>>)> {
self.transmissions
.iter()
.filter_map(|(id, transmission)| match (id, transmission) {
(TransmissionID::Transaction(id, _), Transmission::Transaction(tx)) => Some((*id, tx.clone())),
_ => None,
})
.collect()
}
}
impl<N: Network> Ready<N> {
pub fn contains(&self, transmission_id: impl Into<TransmissionID<N>>) -> bool {
self.transmission_ids.contains_key(&transmission_id.into())
}
pub fn get(&self, transmission_id: impl Into<TransmissionID<N>>) -> Option<Transmission<N>> {
self.transmission_ids
.get(&transmission_id.into())
.and_then(|&index| self.transmissions.get((index - self.offset) as usize))
.map(|(_, transmission)| transmission.clone())
}
pub fn insert(&mut self, transmission_id: impl Into<TransmissionID<N>>, transmission: Transmission<N>) -> bool {
let physical_index = self.transmissions.len();
let transmission_id = transmission_id.into();
if let Vacant(entry) = self.transmission_ids.entry(transmission_id) {
entry.insert(physical_index as i64 + self.offset);
self.transmissions.push_back((transmission_id, transmission));
true
} else {
false
}
}
pub fn insert_front(
&mut self,
transmission_id: impl Into<TransmissionID<N>>,
transmission: Transmission<N>,
) -> bool {
let transmission_id = transmission_id.into();
if let Vacant(entry) = self.transmission_ids.entry(transmission_id) {
self.offset -= 1;
let index = self.offset;
entry.insert(index);
self.transmissions.push_front((transmission_id, transmission));
true
} else {
false
}
}
pub fn remove_front(&mut self) -> Option<(TransmissionID<N>, Transmission<N>)> {
if let Some((transmission_id, transmission)) = self.transmissions.pop_front() {
self.transmission_ids.remove(&transmission_id);
if self.transmission_ids.is_empty() {
debug_assert!(self.transmissions.is_empty());
self.offset = 0;
} else {
self.offset += 1;
}
Some((transmission_id, transmission))
} else {
None
}
}
pub fn clear_solutions(&mut self) {
self.transmissions.retain(|(_, transmission)| !matches!(transmission, Transmission::Solution(_)));
self.transmission_ids.clear();
self.offset = 0;
for (i, (id, _)) in self.transmissions.iter().enumerate() {
self.transmission_ids.insert(*id, i as i64);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use snarkvm::{ledger::narwhal::Data, prelude::Field};
use ::bytes::Bytes;
type CurrentNetwork = snarkvm::prelude::MainnetV0;
#[test]
fn test_ready() {
let rng = &mut TestRng::default();
let data =
|rng: &mut TestRng| Data::Buffer(Bytes::from((0..512).map(|_| rng.r#gen::<u8>()).collect::<Vec<_>>()));
let mut ready = Ready::<CurrentNetwork>::new();
let solution_id_1 = TransmissionID::Solution(
rng.r#gen::<u64>().into(),
rng.r#gen::<<CurrentNetwork as Network>::TransmissionChecksum>(),
);
let solution_id_2 = TransmissionID::Solution(
rng.r#gen::<u64>().into(),
rng.r#gen::<<CurrentNetwork as Network>::TransmissionChecksum>(),
);
let solution_id_3 = TransmissionID::Solution(
rng.r#gen::<u64>().into(),
rng.r#gen::<<CurrentNetwork as Network>::TransmissionChecksum>(),
);
let solution_1 = Transmission::Solution(data(rng));
let solution_2 = Transmission::Solution(data(rng));
let solution_3 = Transmission::Solution(data(rng));
assert!(ready.insert(solution_id_1, solution_1.clone()));
assert!(ready.insert(solution_id_2, solution_2.clone()));
assert!(ready.insert(solution_id_3, solution_3.clone()));
assert_eq!(ready.num_transmissions(), 3);
let transmission_ids = vec![solution_id_1, solution_id_2, solution_id_3].into_iter().collect::<IndexSet<_>>();
assert_eq!(ready.transmission_ids(), transmission_ids);
transmission_ids.iter().for_each(|id| assert!(ready.contains(*id)));
let solution_id_unknown = TransmissionID::Solution(
rng.r#gen::<u64>().into(),
rng.r#gen::<<CurrentNetwork as Network>::TransmissionChecksum>(),
);
assert!(!ready.contains(solution_id_unknown));
assert_eq!(ready.get(solution_id_1), Some(solution_1.clone()));
assert_eq!(ready.get(solution_id_2), Some(solution_2.clone()));
assert_eq!(ready.get(solution_id_3), Some(solution_3.clone()));
assert_eq!(ready.get(solution_id_unknown), None);
let mut transmissions = Vec::with_capacity(3);
for _ in 0..3 {
transmissions.push(ready.remove_front().unwrap())
}
assert!(ready.is_empty());
assert_eq!(ready.transmission_ids(), IndexSet::new());
assert_eq!(transmissions, vec![
(solution_id_1, solution_1),
(solution_id_2, solution_2),
(solution_id_3, solution_3)
]);
}
#[test]
fn test_ready_duplicate() {
use rand::RngCore;
let rng = &mut TestRng::default();
let mut vec = vec![0u8; 512];
rng.fill_bytes(&mut vec);
let data = Data::Buffer(Bytes::from(vec));
let mut ready = Ready::<CurrentNetwork>::new();
let solution_id = TransmissionID::Solution(
rng.r#gen::<u64>().into(),
rng.r#gen::<<CurrentNetwork as Network>::TransmissionChecksum>(),
);
let solution = Transmission::Solution(data);
assert!(ready.insert(solution_id, solution.clone()));
assert!(!ready.insert(solution_id, solution));
assert_eq!(ready.num_transmissions(), 1);
}
#[test]
fn test_insert_front() {
let rng = &mut TestRng::default();
let data =
|rng: &mut TestRng| Data::Buffer(Bytes::from((0..512).map(|_| rng.r#gen::<u8>()).collect::<Vec<_>>()));
let mut ready = Ready::<CurrentNetwork>::new();
let solution_id_1 = TransmissionID::Solution(
rng.r#gen::<u64>().into(),
rng.r#gen::<<CurrentNetwork as Network>::TransmissionChecksum>(),
);
let solution_id_2 = TransmissionID::Solution(
rng.r#gen::<u64>().into(),
rng.r#gen::<<CurrentNetwork as Network>::TransmissionChecksum>(),
);
let solution_1 = Transmission::Solution(data(rng));
let solution_2 = Transmission::Solution(data(rng));
assert!(ready.insert_front(solution_id_1, solution_1.clone()));
assert_eq!(ready.offset, -1);
assert!(ready.insert_front(solution_id_2, solution_2.clone()));
assert_eq!(ready.offset, -2);
assert_eq!(ready.get(solution_id_1), Some(solution_1.clone()));
assert_eq!(ready.get(solution_id_2), Some(solution_2.clone()));
let removed_solution = ready.remove_front().unwrap();
assert_eq!(removed_solution, (solution_id_2, solution_2));
assert_eq!(ready.offset, -1);
let removed_solution = ready.remove_front().unwrap();
assert_eq!(removed_solution, (solution_id_1, solution_1));
assert_eq!(ready.offset, 0);
}
#[test]
fn test_clear_solutions() {
let rng = &mut TestRng::default();
let solution_data =
|rng: &mut TestRng| Data::Buffer(Bytes::from((0..512).map(|_| rng.r#gen::<u8>()).collect::<Vec<_>>()));
let transaction_data =
|rng: &mut TestRng| Data::Buffer(Bytes::from((0..512).map(|_| rng.r#gen::<u8>()).collect::<Vec<_>>()));
let mut ready = Ready::<CurrentNetwork>::new();
let solution_id_1 = TransmissionID::Solution(
rng.r#gen::<u64>().into(),
rng.r#gen::<<CurrentNetwork as Network>::TransmissionChecksum>(),
);
let solution_id_2 = TransmissionID::Solution(
rng.r#gen::<u64>().into(),
rng.r#gen::<<CurrentNetwork as Network>::TransmissionChecksum>(),
);
let transaction_id = TransmissionID::Transaction(
<CurrentNetwork as Network>::TransactionID::from(Field::rand(rng)),
<CurrentNetwork as Network>::TransmissionChecksum::from(rng.r#gen::<u128>()),
);
let solution_1 = Transmission::Solution(solution_data(rng));
let solution_2 = Transmission::Solution(solution_data(rng));
let transaction = Transmission::Transaction(transaction_data(rng));
assert!(ready.insert_front(solution_id_1, solution_1.clone()));
assert_eq!(ready.offset, -1);
assert!(ready.insert(transaction_id, transaction.clone()));
assert_eq!(ready.offset, -1);
assert!(ready.insert(solution_id_2, solution_2.clone()));
assert_eq!(ready.offset, -1);
ready.clear_solutions();
assert_eq!(ready.num_transmissions(), 1);
assert_eq!(ready.offset, 0);
assert_eq!(ready.get(transaction_id), Some(transaction));
}
}