-
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 73ad301
Showing
4 changed files
with
163 additions
and
26 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
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,10 @@ | ||
[package] | ||
name = "cuprate-p2p-bucket" | ||
version = "0.1.0" | ||
edition = "2021" | ||
|
||
[dependencies] | ||
arrayvec = { workspace = true } | ||
|
||
[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,117 @@ | ||
//! 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. | ||
//! | ||
|
||
use arrayvec::ArrayVec; | ||
use std::{collections::BTreeMap, net::Ipv4Addr}; | ||
|
||
/// A discriminant can be compute 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. Enable a fair distribution of the | ||
/// items based on their discriminants when popped. | ||
pub struct Bucket<const N: usize, I: Bucketable> { | ||
/// The storage of the bucket | ||
storage: BTreeMap<I::Discriminant, ArrayVec<I, N>>, | ||
/// Internal round-robin state | ||
round_robin: usize, | ||
} | ||
|
||
impl<const N: usize, I: Bucketable> Bucket<N, I> { | ||
/// Create a new Bucket | ||
pub const fn new() -> Self { | ||
Self { | ||
storage: BTreeMap::new(), | ||
round_robin: 0, | ||
} | ||
} | ||
|
||
/// 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. | ||
/// | ||
pub fn push(&mut self, item: I) { | ||
let discriminant = item.discriminant(); | ||
|
||
if let Some(vec) = self.storage.get_mut(&discriminant) { | ||
// Check if the element already exists | ||
let already_exist = vec.iter().any(|v| &item == v); | ||
if !already_exist { | ||
// Push the new element | ||
vec.push(item); | ||
} | ||
} else { | ||
// Initialize the vector if not found | ||
let mut vec = ArrayVec::<I, N>::new(); | ||
vec.push(item); | ||
self.storage.insert(discriminant, vec); | ||
} | ||
} | ||
|
||
/// 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| { | ||
let find = vec | ||
.iter() | ||
.enumerate() | ||
.filter(|(_, v)| &item == v) | ||
.map(|(i, _)| i) | ||
.last(); | ||
find.map(|index| vec.swap_remove(index)) | ||
}) | ||
} | ||
|
||
/// Will remove a new element from any bucket following a round-robin | ||
/// pattern. | ||
/// | ||
/// Repeated use of this function will provide a fair distribution of | ||
/// item based on their discriminants | ||
pub fn pop(&mut self) -> Option<I> { | ||
// Get the total amount of discriminants to explore | ||
let len = self.storage.len(); | ||
|
||
// Loop over every bucket for an element. | ||
for _ in 0..len { | ||
let (_, vec) = self.storage.iter_mut().nth(self.round_robin / len).unwrap(); | ||
self.round_robin = self.round_robin.wrapping_add(1); | ||
if let Some(item) = vec.pop() { | ||
return Some(item); | ||
} | ||
} | ||
|
||
None | ||
} | ||
} | ||
|
||
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]] | ||
} | ||
} |