use std::collections::HashMap;
use std::time::{Duration, Instant};
use saorsa_core::identity::PeerId;
use crate::ant_protocol::XorName;
pub const MAX_PROVERS_PER_KEY: usize = 16;
pub const PROVER_ENTRY_TTL: Duration = Duration::from_secs(40 * 60);
#[derive(Debug, Clone, Copy)]
pub struct ProverEntry {
pub peer_id: PeerId,
pub proved_at: Instant,
pub commitment_hash: [u8; 32],
}
#[derive(Debug, Default, Clone)]
pub struct RecentProvers {
entries: HashMap<XorName, Vec<ProverEntry>>,
}
impl RecentProvers {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn record_proof(
&mut self,
key: XorName,
peer_id: PeerId,
commitment_hash: [u8; 32],
proved_at: Instant,
) {
let bucket = self.entries.entry(key).or_default();
for e in bucket.iter_mut() {
if e.peer_id == peer_id && e.commitment_hash == commitment_hash {
e.proved_at = proved_at;
bucket.sort_by_key(|e| e.proved_at);
return;
}
}
if bucket.len() >= MAX_PROVERS_PER_KEY {
bucket.remove(0);
}
bucket.push(ProverEntry {
peer_id,
proved_at,
commitment_hash,
});
bucket.sort_by_key(|e| e.proved_at);
}
#[must_use]
pub fn is_credited_holder(
&self,
key: &XorName,
peer_id: &PeerId,
current_commitment_hash: &[u8; 32],
) -> bool {
let now = Instant::now();
self.entries.get(key).is_some_and(|bucket| {
bucket.iter().any(|e| {
&e.peer_id == peer_id
&& &e.commitment_hash == current_commitment_hash
&& now.saturating_duration_since(e.proved_at) < PROVER_ENTRY_TTL
})
})
}
pub fn sweep_expired(&mut self, now: Instant) -> usize {
let mut dropped = 0;
for bucket in self.entries.values_mut() {
let before = bucket.len();
bucket.retain(|e| now.saturating_duration_since(e.proved_at) < PROVER_ENTRY_TTL);
dropped += before - bucket.len();
}
self.entries.retain(|_, b| !b.is_empty());
dropped
}
pub fn forget_peer(&mut self, peer_id: &PeerId) {
for bucket in self.entries.values_mut() {
bucket.retain(|e| &e.peer_id != peer_id);
}
self.entries.retain(|_, b| !b.is_empty());
}
pub fn forget_commitment(&mut self, stale_hash: &[u8; 32]) {
for bucket in self.entries.values_mut() {
bucket.retain(|e| &e.commitment_hash != stale_hash);
}
self.entries.retain(|_, b| !b.is_empty());
}
#[must_use]
pub fn provers_for(&self, key: &XorName) -> usize {
self.entries.get(key).map_or(0, Vec::len)
}
#[must_use]
pub fn total_entries(&self) -> usize {
self.entries.values().map(Vec::len).sum()
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
use super::*;
use std::time::Duration;
fn peer(byte: u8) -> PeerId {
let mut bytes = [0u8; 32];
bytes[0] = byte;
PeerId::from_bytes(bytes)
}
fn key(byte: u8) -> XorName {
let mut k = [0u8; 32];
k[0] = byte;
k
}
fn hash(byte: u8) -> [u8; 32] {
[byte; 32]
}
#[test]
fn empty_cache_credits_no_one() {
let cache = RecentProvers::new();
assert!(!cache.is_credited_holder(&key(1), &peer(1), &hash(1)));
assert_eq!(cache.total_entries(), 0);
}
#[test]
fn recorded_proof_credits_under_same_hash() {
let mut cache = RecentProvers::new();
cache.record_proof(key(1), peer(7), hash(0xAB), Instant::now());
assert!(cache.is_credited_holder(&key(1), &peer(7), &hash(0xAB)));
}
#[test]
fn rotated_hash_loses_credit() {
let mut cache = RecentProvers::new();
cache.record_proof(key(1), peer(7), hash(0xAB), Instant::now());
assert!(!cache.is_credited_holder(&key(1), &peer(7), &hash(0xCD)));
}
#[test]
fn other_peer_under_same_hash_not_credited() {
let mut cache = RecentProvers::new();
cache.record_proof(key(1), peer(7), hash(0xAB), Instant::now());
assert!(!cache.is_credited_holder(&key(1), &peer(8), &hash(0xAB)));
}
#[test]
fn per_key_cap_evicts_oldest() {
let mut cache = RecentProvers::new();
let now = Instant::now();
let max_u8 = u8::try_from(MAX_PROVERS_PER_KEY).unwrap_or(u8::MAX);
for i in 0..=max_u8 {
let t = now + Duration::from_millis(u64::from(i));
cache.record_proof(key(1), peer(i), hash(0xAB), t);
}
assert_eq!(cache.provers_for(&key(1)), MAX_PROVERS_PER_KEY);
assert!(!cache.is_credited_holder(&key(1), &peer(0), &hash(0xAB)));
assert!(cache.is_credited_holder(&key(1), &peer(max_u8), &hash(0xAB)));
}
#[test]
fn refresh_in_place_does_not_grow_bucket() {
let mut cache = RecentProvers::new();
let now = Instant::now();
cache.record_proof(key(1), peer(1), hash(0xAB), now);
cache.record_proof(key(1), peer(1), hash(0xAB), now + Duration::from_secs(1));
cache.record_proof(key(1), peer(1), hash(0xAB), now + Duration::from_secs(2));
assert_eq!(cache.provers_for(&key(1)), 1);
}
#[test]
fn forget_peer_drops_all_entries() {
let mut cache = RecentProvers::new();
let now = Instant::now();
cache.record_proof(key(1), peer(1), hash(0xAB), now);
cache.record_proof(key(2), peer(1), hash(0xAB), now);
cache.record_proof(key(1), peer(2), hash(0xAB), now);
assert_eq!(cache.total_entries(), 3);
cache.forget_peer(&peer(1));
assert_eq!(cache.total_entries(), 1);
assert!(!cache.is_credited_holder(&key(1), &peer(1), &hash(0xAB)));
assert!(cache.is_credited_holder(&key(1), &peer(2), &hash(0xAB)));
}
#[test]
fn forget_commitment_drops_only_matching_entries() {
let mut cache = RecentProvers::new();
let now = Instant::now();
cache.record_proof(key(1), peer(1), hash(0xAB), now);
cache.record_proof(key(1), peer(1), hash(0xCD), now);
cache.record_proof(key(2), peer(2), hash(0xAB), now);
assert_eq!(cache.total_entries(), 3);
cache.forget_commitment(&hash(0xAB));
assert_eq!(cache.total_entries(), 1);
assert!(cache.is_credited_holder(&key(1), &peer(1), &hash(0xCD)));
assert!(!cache.is_credited_holder(&key(1), &peer(1), &hash(0xAB)));
assert!(!cache.is_credited_holder(&key(2), &peer(2), &hash(0xAB)));
}
#[test]
fn lazy_rotation_via_unknown_commitment_hash_drops_credit() {
let mut cache = RecentProvers::new();
cache.record_proof(key(1), peer(7), hash(0xAB), Instant::now());
assert!(cache.is_credited_holder(&key(1), &peer(7), &hash(0xAB)));
cache.forget_commitment(&hash(0xAB));
assert!(!cache.is_credited_holder(&key(1), &peer(7), &hash(0xAB)));
assert!(!cache.is_credited_holder(&key(1), &peer(7), &hash(0xCD)));
}
}