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]]
    }
}