use libp2prs_core::Multiaddr;
use smallvec::SmallVec;
use std::{collections::VecDeque, num::NonZeroUsize};
#[derive(Clone)]
pub struct Addresses {
registry: SmallVec<[Record; 8]>,
limit: NonZeroUsize,
reports: VecDeque<Multiaddr>,
}
impl std::fmt::Debug for Addresses {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Addresses")
.field("registry", &self.registry)
.field("limit", &self.limit)
.finish()
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
struct Record {
score: u32,
addr: Multiaddr,
}
impl Default for Addresses {
fn default() -> Self {
Addresses::new(NonZeroUsize::new(200).expect("200 > 0"))
}
}
impl Addresses {
pub fn new(limit: NonZeroUsize) -> Self {
Addresses {
registry: SmallVec::new(),
limit,
reports: VecDeque::with_capacity(limit.get()),
}
}
pub fn add(&mut self, a: Multiaddr) {
let oldest = if self.reports.len() == self.limit.get() {
self.reports.pop_front()
} else {
None
};
if let Some(oldest) = oldest {
if let Some(in_registry) = self.registry.iter_mut().find(|r| r.addr == oldest) {
in_registry.score = in_registry.score.saturating_sub(1);
isort(&mut self.registry);
}
}
while self.registry.last().map(|e| e.score == 0).unwrap_or(false) {
self.registry.pop();
}
self.reports.push_back(a.clone());
for r in &mut self.registry {
if r.addr == a {
r.score = r.score.saturating_add(1);
isort(&mut self.registry);
return;
}
}
let r = Record { score: 1, addr: a };
self.registry.push(r)
}
pub fn iter(&self) -> AddressIter<'_> {
AddressIter {
items: &self.registry,
offset: 0,
}
}
#[allow(dead_code)]
pub fn into_iter(self) -> AddressIntoIter {
AddressIntoIter { items: self.registry }
}
}
#[derive(Clone)]
pub struct AddressIter<'a> {
items: &'a [Record],
offset: usize,
}
impl<'a> Iterator for AddressIter<'a> {
type Item = &'a Multiaddr;
fn next(&mut self) -> Option<Self::Item> {
if self.offset == self.items.len() {
return None;
}
let item = &self.items[self.offset];
self.offset += 1;
Some(&item.addr)
}
fn size_hint(&self) -> (usize, Option<usize>) {
let n = self.items.len() - self.offset;
(n, Some(n))
}
}
impl<'a> ExactSizeIterator for AddressIter<'a> {}
#[derive(Clone)]
pub struct AddressIntoIter {
items: SmallVec<[Record; 8]>,
}
impl Iterator for AddressIntoIter {
type Item = Multiaddr;
fn next(&mut self) -> Option<Self::Item> {
if !self.items.is_empty() {
Some(self.items.remove(0).addr)
} else {
None
}
}
fn size_hint(&self) -> (usize, Option<usize>) {
let n = self.items.len();
(n, Some(n))
}
}
impl ExactSizeIterator for AddressIntoIter {}
fn isort(xs: &mut [Record]) {
for i in 1..xs.len() {
for j in (1..=i).rev() {
if xs[j].score <= xs[j - 1].score {
break;
}
xs.swap(j, j - 1)
}
}
}
#[cfg(test)]
mod tests {
use super::{isort, Addresses, Record};
use libp2prs_core::multiaddr::{protocol::Protocol, Multiaddr};
use quickcheck::{Arbitrary, Gen, QuickCheck};
use rand::Rng;
use std::num::NonZeroUsize;
#[test]
fn isort_sorts() {
fn property(xs: Vec<u32>) -> bool {
let mut xs = xs
.into_iter()
.map(|s| Record {
score: s,
addr: Multiaddr::empty(),
})
.collect::<Vec<_>>();
isort(&mut xs);
for i in 1..xs.len() {
assert!(xs[i - 1].score >= xs[i].score)
}
true
}
QuickCheck::new().quickcheck(property as fn(Vec<u32>) -> bool)
}
#[test]
fn old_reports_disappear() {
let mut addresses = Addresses::default();
let single: Multiaddr = "/tcp/2108".parse().unwrap();
addresses.add(single.clone());
assert!(addresses.iter().any(|a| *a == single));
let other: Multiaddr = "/tcp/120".parse().unwrap();
for _ in 0..2000 {
addresses.add(other.clone());
}
assert!(addresses.iter().find(|a| **a == single).is_none());
}
#[test]
fn record_score_equals_last_n_reports() {
#[derive(PartialEq, Eq, Clone, Hash, Debug)]
struct Ma(Multiaddr);
impl Arbitrary for Ma {
fn arbitrary<G: Gen>(g: &mut G) -> Self {
Ma(Protocol::Tcp(g.gen::<u16>() % 16).into())
}
}
fn property(xs: Vec<Ma>, n: u8) -> bool {
let n = std::cmp::max(n, 1);
let mut addresses = Addresses::new(NonZeroUsize::new(usize::from(n)).unwrap());
for Ma(a) in &xs {
addresses.add(a.clone())
}
for r in &addresses.registry {
let count = xs.iter().rev().take(usize::from(n)).filter(|Ma(x)| x == &r.addr).count();
if r.score as usize != count {
return false;
}
}
true
}
QuickCheck::new().quickcheck(property as fn(Vec<Ma>, u8) -> bool)
}
}