-
Notifications
You must be signed in to change notification settings - Fork 19
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Implements P2P Bucket data structure
This commit implements a "Bucket" data structure that is a collection of data that discriminates its items into "buckets" (vector of size N) following a defined function. - Implements Bucket data structure and Bucketable trait - Implements Bucketable for Ipv4Addr - Added the crate to the workspace dependencies - Added arrayvec as a dependency
- Loading branch information
1 parent
b57ee2f
commit c104d65
Showing
5 changed files
with
197 additions
and
0 deletions.
There are no files selected for viewing
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,13 @@ | ||
| [package] | ||
| name = "cuprate-p2p-bucket" | ||
| version = "0.1.0" | ||
| edition = "2021" | ||
| license = "MIT" | ||
| authors = ["SyntheticBird"] | ||
|
|
||
| [dependencies] | ||
| arrayvec = { workspace = true } | ||
| rand = { workspace = true, features = ["std", "std_rng"]} | ||
|
|
||
| [lints] | ||
| workspace = true |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,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]] | ||
| } | ||
| } |