use crate::dht::node::Id;
use core::ops::RangeInclusive;
#[cfg(all(feature = "alloc", not(feature = "std")))]
use alloc::{vec, vec::Vec};
#[cfg(feature = "std")]
use std::{vec, vec::Vec};
pub trait Node {
fn id(&self) -> Id;
}
impl<T> Node for &T
where
T: Node,
{
fn id(&self) -> Id {
(*self).id()
}
}
#[derive(Debug, Clone)]
pub struct Bucket<T, Instant> {
range: RangeInclusive<Id>,
nodes: Vec<T>,
refresh_deadline: Instant,
}
impl<T, Instant> Bucket<T, Instant>
where
T: Node,
Instant: crate::time::Instant,
{
#[must_use]
#[inline]
pub fn new(range: RangeInclusive<Id>, refresh_deadline: Instant) -> Self {
Bucket {
range,
nodes: Vec::new(),
refresh_deadline,
}
}
#[must_use]
#[inline]
pub fn range(&self) -> &RangeInclusive<Id> {
&self.range
}
#[must_use]
#[inline]
pub fn rand_id<R>(&self, rng: &mut R) -> Id
where
R: rand::RngCore,
{
internal::rand_in_inclusive_range(&self.range, rng)
}
#[must_use]
#[inline]
pub fn refresh_deadline(&self) -> &Instant {
&self.refresh_deadline
}
#[inline]
pub fn set_refresh_deadline(&mut self, refresh_deadline: Instant) {
self.refresh_deadline = refresh_deadline;
}
#[must_use]
#[inline]
pub fn len(&self) -> usize {
self.nodes.len()
}
#[must_use]
#[inline]
pub fn is_empty(&self) -> bool {
self.nodes.is_empty()
}
#[inline]
pub fn iter(&self) -> impl Iterator<Item = &'_ T> {
self.nodes.iter()
}
#[inline]
pub fn iter_mut(&mut self) -> impl Iterator<Item = &'_ mut T> {
self.nodes.iter_mut()
}
#[inline]
pub fn insert(&mut self, node: T) {
assert!(self.range.contains(&node.id()));
self.nodes.push(node);
}
#[inline]
pub fn clear(&mut self) {
self.nodes.clear();
}
#[inline]
pub fn retain<F>(&mut self, f: F)
where
F: FnMut(&T) -> bool,
{
self.nodes.retain(f);
}
#[must_use]
pub fn split(self) -> (Bucket<T, Instant>, Bucket<T, Instant>) {
let Bucket {
mut nodes,
refresh_deadline,
range,
} = self;
let middle = internal::middle(*range.end(), *range.start());
assert_ne!(*range.start(), *range.end(), "cannot split bucket");
let lower_bucket_range = *range.start()..=middle;
let upper_bucket_range = internal::next(middle)..=*range.end();
debug_assert!(*lower_bucket_range.start() <= *lower_bucket_range.end());
debug_assert!(*upper_bucket_range.start() <= *upper_bucket_range.end());
let mut upper_bucket_nodes = Vec::default();
let mut index = 0;
while index < nodes.len() {
if upper_bucket_range.contains(&nodes[index].id()) {
upper_bucket_nodes.push(nodes.remove(index));
} else {
index += 1;
}
}
let lower_bucket = Bucket {
range: lower_bucket_range,
nodes,
refresh_deadline: refresh_deadline.clone(),
};
let upper_bucket = Bucket {
range: upper_bucket_range,
nodes: upper_bucket_nodes,
refresh_deadline,
};
(lower_bucket, upper_bucket)
}
}
mod internal {
use crate::dht::node::Id;
#[must_use]
pub(super) fn rand_in_inclusive_range<R>(
range: &core::ops::RangeInclusive<Id>,
rng: &mut R,
) -> Id
where
R: rand::Rng,
{
#[must_use]
#[inline]
fn difference(first: Id, second: Id) -> Id {
let bigger: [u8; 20];
let mut smaller: [u8; 20];
if first < second {
bigger = second.0;
smaller = first.0;
} else {
bigger = first.0;
smaller = second.0;
}
smaller = twos_complement(smaller);
let (bigger, _) = overflowing_add(bigger, &smaller);
Id::new(bigger)
}
let data_bit_diff = difference(*range.end(), *range.start());
let rand_bits: [u8; 20] = randomize_up_to(<[u8; 20]>::from(data_bit_diff), rng);
let (rand_bits, _) = overflowing_add(rand_bits, &<[u8; 20]>::from(*range.start()));
Id::from(rand_bits)
}
#[must_use]
#[inline]
pub(super) const fn middle(first: Id, second: Id) -> Id {
let data = first.0;
let (data, overflow) = overflowing_add(data, &second.0);
let mut data = shift_right(data);
if overflow {
data[0] |= 0x80;
}
Id::new(data)
}
#[must_use]
#[inline]
pub(super) const fn next(id: Id) -> Id {
let (data, overflowing) = overflowing_add(
id.0,
&[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
);
assert!(!overflowing);
Id::new(data)
}
#[must_use]
#[inline]
const fn overflowing_add(mut value: [u8; 20], other: &[u8; 20]) -> ([u8; 20], bool) {
let mut carry_over: u8 = 0;
let mut idx = 19;
loop {
let (partial_val, overflow) = value[idx].overflowing_add(other[idx]);
let (final_val, carry_over_overflow) = partial_val.overflowing_add(carry_over);
value[idx] = final_val;
carry_over = if carry_over_overflow || overflow {
1
} else {
0
};
if idx == 0 {
break;
}
idx -= 1;
}
(value, carry_over == 1)
}
#[must_use]
#[inline]
const fn twos_complement(mut value: [u8; 20]) -> [u8; 20] {
let mut idx = 0;
loop {
value[idx] = !(value[idx]);
if idx == 19 {
break;
}
idx += 1;
}
let one_bit = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1];
let (value, _) = overflowing_add(value, &one_bit);
value
}
#[inline]
#[must_use]
const fn shift_right(mut value: [u8; 20]) -> [u8; 20] {
let mut add_high_bit = false;
let mut idx = 0;
loop {
let is_lower_bit_set = (value[idx] & 0x01) == 1;
value[idx] >>= 1;
if add_high_bit {
value[idx] |= 0x80;
}
add_high_bit = is_lower_bit_set;
if idx == 19 {
break;
}
idx += 1;
}
value
}
#[must_use]
fn randomize_up_to<R>(end: [u8; 20], rng: &mut R) -> [u8; 20]
where
R: rand::Rng,
{
let mut data = [0; 20];
let mut lower_than_max = false;
for idx in 0..data.len() {
data[idx] = if lower_than_max {
rng.gen()
} else {
let idx_val = end[idx];
let val = rng.gen_range(0..=idx_val);
if val < idx_val {
lower_than_max = true;
}
val
};
}
data
}
#[cfg(test)]
mod tests {
use super::*;
#[cfg(all(feature = "alloc", not(feature = "std")))]
use alloc::{format, vec};
#[cfg(feature = "std")]
use std::{format, vec};
#[test]
fn test_debug() {
let node_id = Id::max();
let debug_str = format!("{node_id:?}");
assert_eq!(debug_str, "Id(FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF)");
}
#[test]
fn test_overflowing_add() {
let bytes: [u8; 20] = [
0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
];
let (bytes, overflow) = overflowing_add(
bytes,
&[
1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20,
],
);
assert!(!overflow);
assert_eq!(
bytes,
[1, 3, 5, 7, 9, 11, 13, 15, 17, 19, 21, 23, 25, 27, 29, 31, 33, 35, 37, 39,]
);
}
#[test]
fn test_twos_complement() {
let mut bytes: [u8; 20] = [
0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
];
bytes = twos_complement(bytes);
assert_eq!(
bytes,
[
255, 254, 253, 252, 251, 250, 249, 248, 247, 246, 245, 244, 243, 242, 241, 240,
239, 238, 237, 237
]
);
}
#[test]
fn test_shift_right() {
let bytes: [u8; 20] = [
0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
];
let bytes = shift_right(bytes);
assert_eq!(
bytes,
[0, 0, 129, 1, 130, 2, 131, 3, 132, 4, 133, 5, 134, 6, 135, 7, 136, 8, 137, 9]
);
}
#[test]
fn test_id_ord() {
let mut node_ids = vec![
Id::from([0xff; 20]),
Id::from([0x00; 20]),
middle(Id::from([0xff; 20]), Id::from([0x00; 20])),
];
node_ids.sort();
assert_eq!(
node_ids,
vec![
Id::from([0x00; 20]),
middle(Id::from([0xff; 20]), Id::from([0x00; 20])),
Id::from([0xff; 20]),
]
);
}
#[test]
fn test_id_distance_ord() {
let mut node_ids = vec![
Id::from([0x00; 20]),
middle(Id::from([0xff; 20]), Id::from([0x00; 20])),
Id::from([0xff; 20]),
];
let pivot_id = middle(Id::from([0xef; 20]), Id::from([0x00; 20]));
node_ids.sort_by_key(|a| a.distance(pivot_id));
assert_eq!(
node_ids,
vec![
middle(Id::from([0xff; 20]), Id::from([0x00; 20])),
Id::from([0x00; 20]),
Id::from([0xff; 20]),
]
);
}
}
}
#[derive(Debug, Clone)]
pub struct Table<T, Instant> {
pivot: Id,
buckets: Vec<Bucket<T, Instant>>,
}
impl<T, Instant> Table<T, Instant>
where
T: Node,
Instant: crate::time::Instant,
{
#[must_use]
#[inline]
pub fn new(pivot: Id, refresh_deadline: Instant) -> Self {
Self {
pivot,
buckets: vec![Bucket::new(Id::min()..=Id::max(), refresh_deadline)],
}
}
#[must_use]
#[inline]
pub fn pivot(&self) -> Id {
self.pivot
}
#[must_use]
#[inline]
pub fn len(&self) -> usize {
self.buckets.iter().map(Bucket::len).sum()
}
#[must_use]
#[inline]
pub fn is_empty(&self) -> bool {
self.len() == 0
}
#[inline]
pub fn iter(&self) -> impl Iterator<Item = &'_ Bucket<T, Instant>> {
self.buckets.iter()
}
#[inline]
pub fn iter_mut(&mut self) -> impl Iterator<Item = &'_ mut Bucket<T, Instant>> {
self.buckets.iter_mut()
}
#[allow(clippy::missing_panics_doc)]
#[must_use]
#[inline]
pub fn find(&self, id: &Id) -> &Bucket<T, Instant> {
let bucket = self.buckets.iter().find(|b| b.range.contains(id)).unwrap();
bucket
}
#[allow(clippy::missing_panics_doc)]
#[must_use]
#[inline]
pub fn find_mut(&mut self, id: &Id) -> &mut Bucket<T, Instant> {
let bucket = self
.buckets
.iter_mut()
.find(|b| b.range.contains(id))
.unwrap();
bucket
}
#[must_use]
#[inline]
pub fn find_bucket_to_refresh(&mut self, now: &Instant) -> Option<&mut Bucket<T, Instant>> {
self.buckets
.iter_mut()
.find(|b| *b.refresh_deadline() <= *now)
}
#[allow(clippy::missing_panics_doc)]
pub fn split_last(&mut self) {
let bucket = self.buckets.pop().unwrap();
debug_assert!(bucket.range.contains(&self.pivot));
let (lower_bucket, upper_bucket) = bucket.split();
if lower_bucket.range.contains(&self.pivot) {
self.buckets.push(upper_bucket);
self.buckets.push(lower_bucket);
} else {
self.buckets.push(lower_bucket);
self.buckets.push(upper_bucket);
}
debug_assert!(self
.buckets
.last()
.map(|b| { b.range.contains(&self.pivot) })
.unwrap_or_default());
}
}
#[cfg(test)]
#[cfg(feature = "std")]
mod tests {
use super::*;
#[derive(Clone, Debug, PartialEq)]
struct MyNode {
id: Id,
}
impl Node for MyNode {
fn id(&self) -> Id {
self.id
}
}
#[test]
fn test_split_bucket() {
let now = std::time::Instant::now();
let bucket: Bucket<MyNode, _> = Bucket::new(Id::min()..=Id::max(), now);
let (first_bucket, second_bucket) = bucket.split();
assert_eq!(
first_bucket.range,
Id::min()
..=Id::from([
0x7f, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff
])
);
assert_eq!(
second_bucket.range,
Id::from([
0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
])..=Id::max()
);
}
#[test]
fn test_split_range_three() {
let now = std::time::Instant::now();
let min_id = Id::min();
let max_id = Id::from([
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x02,
]);
let bucket: Bucket<MyNode, _> = Bucket::new(min_id..=max_id, now);
let (first_bucket, second_bucket) = bucket.split();
assert_eq!(
first_bucket.range,
Id::from([
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
])
..=Id::from([
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01,
])
);
assert_eq!(
second_bucket.range,
Id::from([
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x02,
])
..=Id::from([
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02,
])
);
}
#[test]
fn test_split_range_two() {
let now = std::time::Instant::now();
let min_id = Id::from([
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x01,
]);
let max_id = Id::from([
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x02,
]);
let bucket: Bucket<MyNode, _> = Bucket::new(min_id..=max_id, now);
let (first_bucket, second_bucket) = bucket.split();
assert_eq!(
first_bucket.range,
Id::from([
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x01,
])
..=Id::from([
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01,
])
);
assert_eq!(
second_bucket.range,
Id::from([
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x02,
])
..=Id::from([
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02,
])
);
}
#[test]
fn test_split_range_two_min() {
let now = std::time::Instant::now();
let min_id = Id::min();
let max_id = Id::from([
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x01,
]);
let bucket: Bucket<MyNode, _> = Bucket::new(min_id..=max_id, now);
let (first_bucket, second_bucket) = bucket.split();
assert_eq!(first_bucket.range, Id::min()..=Id::min());
assert_eq!(
second_bucket.range,
Id::from([
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x01,
])
..=Id::from([
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01,
])
);
}
#[test]
fn test_split_range_two_max() {
let now = std::time::Instant::now();
let min_id = Id::from([
0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
0xff, 0xff, 0xff, 0xff, 0xff, 0xfe,
]);
let max_id = Id::max();
let bucket: Bucket<MyNode, _> = Bucket::new(min_id..=max_id, now);
let (first_bucket, second_bucket) = bucket.split();
assert_eq!(
first_bucket.range,
Id::from([
0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
0xff, 0xff, 0xff, 0xff, 0xff, 0xfe,
])
..=Id::from([
0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xfe,
])
);
assert_eq!(
second_bucket.range,
Id::from([
0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
])..=Id::max()
);
}
#[test]
#[should_panic(expected = "cannot split bucket")]
fn test_split_range_one() {
let now = std::time::Instant::now();
let min_id = Id::from([
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x01,
]);
let max_id = Id::from([
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x01,
]);
let bucket: Bucket<MyNode, _> = Bucket::new(min_id..=max_id, now);
let (_, _) = bucket.split();
}
#[test]
#[should_panic(expected = "cannot split bucket")]
fn test_split_range_one_min() {
let now = std::time::Instant::now();
let bucket: Bucket<MyNode, _> = Bucket::new(Id::min()..=Id::min(), now);
let (_, _) = bucket.split();
}
#[test]
#[should_panic(expected = "cannot split bucket")]
fn test_split_range_one_max() {
let now = std::time::Instant::now();
let bucket: Bucket<MyNode, _> = Bucket::new(Id::max()..=Id::max(), now);
let (_, _) = bucket.split();
}
}