use {
crate::bank::{Bank, PartitionIndex, PartitionsPerCycle},
solana_sdk::{epoch_schedule::EpochSchedule, pubkey::Pubkey},
std::collections::HashSet,
};
lazy_static! {
pub static ref EMPTY_HASHSET: HashSet<Pubkey> = HashSet::default();
}
#[derive(Debug, Default)]
pub struct RentPayingAccountsByPartition {
pub accounts: Vec<HashSet<Pubkey>>,
partition_count: PartitionsPerCycle,
}
impl RentPayingAccountsByPartition {
pub fn new(epoch_schedule: &EpochSchedule) -> Self {
let partition_count = epoch_schedule.slots_per_epoch;
Self {
partition_count,
accounts: (0..=partition_count)
.map(|_| HashSet::<Pubkey>::default())
.collect(),
}
}
pub fn add_account(&mut self, pubkey: &Pubkey) {
let partition_end_index = Bank::partition_from_pubkey(pubkey, self.partition_count);
let list = &mut self.accounts[partition_end_index as usize];
list.insert(*pubkey);
}
pub fn get_pubkeys_in_partition_index(
&self,
partition_end_index: PartitionIndex,
) -> &HashSet<Pubkey> {
self.accounts
.get(partition_end_index as usize)
.unwrap_or(&EMPTY_HASHSET)
}
pub fn is_initialized(&self) -> bool {
self.partition_count != 0
}
}
#[cfg(test)]
pub(crate) mod tests {
use super::*;
#[test]
fn test_add() {
let mut test = RentPayingAccountsByPartition::new(&EpochSchedule::custom(32, 0, false));
let pk = Pubkey::new(&[1; 32]);
test.add_account(&pk);
test.add_account(&pk);
assert_eq!(test.get_pubkeys_in_partition_index(0).len(), 1);
assert!(test.get_pubkeys_in_partition_index(1).is_empty());
assert!(test.is_initialized());
let test = RentPayingAccountsByPartition::default();
assert!(!test.is_initialized());
}
}