cuprate_p2p_bucket/lib.rs
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172
//! Bucket data structure
//!
//! A collection data structure that discriminates its unique items and place them into "buckets".
//!
//! The item must implement the [`Bucketable`] trait that defines how to create the discriminant
//! from the item type. The data structure will internally contain any item into "buckets" or vectors
//! of sized capacity `N` that regroup all the stored items with this specific discriminant.
//!
//! A practical example of this data structure is for storing `N` amount of IP discriminated by their subnets.
//! You can store in each "buckets" corresponding to a `/16` subnet up to `N` IPs of that subnet.
//!
//! # Example
//!
//! ```
//! use cuprate_p2p_bucket::Bucket;
//! use std::net::Ipv4Addr;
//!
//! // Create a new bucket that can store at most 2 IPs in a particular `/16` subnet.
//! let mut bucket = Bucket::<2,Ipv4Addr>::new();
//!
//! // Fulfill the `96.96.0.0/16` bucket.
//! bucket.push("96.96.0.1".parse().unwrap());
//! bucket.push("96.96.0.2".parse().unwrap());
//! assert_eq!(2, bucket.len());
//! assert_eq!(2, bucket.len_bucket(&[96_u8,96_u8]).unwrap());
//!
//! // Push a new IP from another subnet
//! bucket.push("127.0.0.1".parse().unwrap());
//! assert_eq!(3, bucket.len());
//! assert_eq!(2, bucket.len_bucket(&[96_u8,96_u8]).unwrap());
//! assert_eq!(1, bucket.len_bucket(&[127_u8,0_u8]).unwrap());
//!
//! // Attempting to push a new IP within `96.96.0.0/16` bucket will return the IP back
//! // as this subnet is already full.
//! let pushed = bucket.push("96.96.0.3".parse().unwrap());
//! assert!(pushed.is_some());
//! assert_eq!(2, bucket.len_bucket(&[96_u8,96_u8]).unwrap());
//!
//! ```
use arrayvec::{ArrayVec, CapacityError};
use rand::random;
use std::{collections::BTreeMap, net::Ipv4Addr};
/// A discriminant that can be computed from the type.
pub trait Bucketable: Sized + Eq + Clone {
/// The type of the discriminant being used in the Binary tree.
type Discriminant: Ord + AsRef<[u8]>;
/// Method that can compute the discriminant from the item.
fn discriminant(&self) -> Self::Discriminant;
}
/// A collection data structure discriminating its unique items
/// with a specified method. Limiting the amount of items stored
/// with that discriminant to the const `N`.
pub struct Bucket<const N: usize, I: Bucketable> {
/// The storage of the bucket
storage: BTreeMap<I::Discriminant, ArrayVec<I, N>>,
}
impl<const N: usize, I: Bucketable> Bucket<N, I> {
/// Create a new Bucket
pub const fn new() -> Self {
Self {
storage: BTreeMap::new(),
}
}
/// Push a new element into the Bucket
///
/// Will internally create a new vector for each new discriminant being
/// generated from an item.
///
/// This function WILL NOT push the element if it already exists.
///
/// Return `None` if the item has been pushed or ignored. `Some(I)` if
/// the vector is full.
///
/// # Example
///
/// ```
/// use cuprate_p2p_bucket::Bucket;
/// use std::net::Ipv4Addr;
///
/// let mut bucket = Bucket::<8,Ipv4Addr>::new();
///
/// // Push a first IP address.
/// bucket.push("127.0.0.1".parse().unwrap());
/// assert_eq!(1, bucket.len());
///
/// // Push the same IP address a second time.
/// bucket.push("127.0.0.1".parse().unwrap());
/// assert_eq!(1, bucket.len());
/// ```
pub fn push(&mut self, item: I) -> Option<I> {
let discriminant = item.discriminant();
if let Some(vec) = self.storage.get_mut(&discriminant) {
// Push the item if it doesn't exist.
if !vec.contains(&item) {
return vec.try_push(item).err().map(CapacityError::element);
}
} else {
// Initialize the vector if not found.
let mut vec = ArrayVec::<I, N>::new();
vec.push(item);
self.storage.insert(discriminant, vec);
}
None
}
/// Will attempt to remove an item from the bucket.
pub fn remove(&mut self, item: &I) -> Option<I> {
self.storage.get_mut(&item.discriminant()).and_then(|vec| {
vec.iter()
.enumerate()
.find_map(|(i, v)| (item == v).then_some(i))
.map(|index| vec.swap_remove(index))
})
}
/// Return the number of item stored within the storage
pub fn len(&self) -> usize {
self.storage.values().map(ArrayVec::len).sum()
}
/// Return the number of item stored with a specific discriminant.
///
/// This method returns None if the bucket with this discriminant
/// doesn't exist.
pub fn len_bucket(&self, discriminant: &I::Discriminant) -> Option<usize> {
self.storage.get(discriminant).map(ArrayVec::len)
}
/// Return `true` if the storage contains no items
pub fn is_empty(&self) -> bool {
self.len() == 0
}
/// Return a reference to an item chosen at random.
///
/// Repeated use of this function will provide a normal distribution of
/// items based on their discriminants.
pub fn get_random(&mut self) -> Option<&I> {
// Get the total amount of discriminants to explore.
let len = self.storage.len();
// Get a random bucket.
let (_, vec) = self.storage.iter().nth(random::<usize>() / len).unwrap();
// Return a reference chose at random.
vec.get(random::<usize>() / vec.len())
}
}
impl<const N: usize, I: Bucketable> Default for Bucket<N, I> {
fn default() -> Self {
Self::new()
}
}
impl Bucketable for Ipv4Addr {
/// We are discriminating by `/16` subnets.
type Discriminant = [u8; 2];
fn discriminant(&self) -> Self::Discriminant {
[self.octets()[0], self.octets()[1]]
}
}