use crate::error::Error;
use crate::node::{Address, Center, Node};
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Bucket {
nodes: Vec<Node>,
limit: usize,
}
impl Bucket {
pub fn new(limit: usize) -> Self {
Bucket {
nodes: Vec::new(),
limit,
}
}
pub fn sort(&mut self) {
self.nodes.sort();
}
pub fn try_add(&mut self, node: Node) -> Result<(), Error> {
if self.len() == self.limit {
return Err(Error::Full);
} else {
self.nodes.push(node);
self.dedup();
return Ok(());
}
}
pub fn add(&mut self, node: Node) {
if self.len() < self.limit {
self.nodes.push(node);
self.sort();
self.dedup();
} else {
if let Some(first) = self.nodes.first_mut() {
if !first.is_reachable() {
*first = node;
}
self.sort();
self.dedup();
}
}
}
pub fn split(self, center: &Center, ul: u8) -> (Self, Self) {
let mut near = Bucket::new(self.limit);
let mut far = Bucket::new(self.limit);
for i in self.nodes {
let index = (i.address.clone() ^ center.public.clone())[0];
if index < (ul / 2) {
near.add(i);
} else {
far.add(i);
}
}
(near, far)
}
pub fn find(&self, search: &Address) -> Option<&Node> {
let index = self.nodes.iter().position(|e| &e.address == search);
match index {
Some(i) => self.nodes.get(i),
None => None,
}
}
pub fn find_mut(&mut self, search: &Address) -> Option<&mut Node> {
let index = self.nodes.iter().position(|e| &e.address == search);
match index {
Some(i) => self.nodes.get_mut(i),
None => None,
}
}
pub fn get(&self, limit: usize) -> Vec<&Node> {
let mut targets = Vec::new();
for i in &self.nodes {
targets.push(i);
}
targets.sort();
targets.truncate(limit);
return targets;
}
pub fn remove(&mut self, target: &Address) -> Result<(), Error> {
let index = self.nodes.iter().position(|e| &e.address == target);
match index {
Some(i) => {
self.nodes.remove(i);
Ok(())
}
None => Err(Error::Unknown),
}
}
pub fn capacity(&self) -> usize {
self.limit
}
pub fn dedup(&mut self) {
self.sort_by_address();
self.nodes.dedup_by(|a, b| a.address == b.address);
}
pub fn len(&self) -> usize {
self.nodes.len()
}
fn sort_by_address(&mut self) {
self.nodes
.sort_by(|a, b| a.address.partial_cmp(&b.address).unwrap());
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::node::Address;
use sodiumoxide::crypto::box_::curve25519xsalsa20poly1305::SecretKey;
#[test]
fn test_leaf_add() {
let mut bucket = gen_bucket(20);
let node = gen_node("test");
let ret = bucket.try_add(node);
assert_eq!(ret.is_err(), false);
}
#[test]
fn test_leaf_add_error() {
let mut leaf = gen_bucket(1);
let node = gen_node("test");
leaf.try_add(node).unwrap();
let node = gen_node("test2");
let ret = leaf.try_add(node);
assert_eq!(ret.is_err(), true);
}
#[test]
fn test_bucket_add_error_lim() {
let mut bucket = gen_bucket(1);
let node = gen_node("test");
bucket.try_add(node).unwrap();
let node = gen_node("test2");
let ret = bucket.try_add(node);
assert_eq!(ret.is_err(), true);
}
#[test]
fn test_bucket_add_disregard() {
let mut bucket = gen_bucket(1);
let node = gen_node("test");
bucket.add(node);
let node = gen_node("test2");
bucket.add(node);
assert_eq!(bucket.len(), 1);
}
#[test]
fn test_bucket_dedup() {
let mut bucket = gen_bucket(10);
let node = gen_node("abc");
bucket.add(node);
let node = gen_node("abc");
bucket.add(node);
bucket.dedup();
assert_eq!(bucket.len(), 1);
}
#[test]
fn test_bucket_split_root() {
let mut root = Bucket::new(20);
root.add(gen_node("first"));
root.add(gen_node("second"));
root.add(gen_node("another"));
let center = gen_center();
let (near, far) = root.split(¢er, 255);
assert_eq!(near.len(), 2);
assert_eq!(far.len(), 1);
}
#[test]
fn test_bucket_get() {
let mut root = Bucket::new(20);
root.add(gen_node("first"));
root.add(gen_node("second"));
root.add(gen_node("another"));
let targets = root.get(2);
assert_eq!(targets.len(), 2);
}
#[test]
fn test_bucket_remove() {
let mut root = Bucket::new(20);
root.add(gen_node("first"));
root.add(gen_node("second"));
let target = gen_node("first").address;
root.remove(&target).unwrap();
assert_eq!(root.len(), 1);
}
#[test]
fn test_bucket_remove_empty() {
let mut root = Bucket::new(20);
let target = gen_node("first").address;
assert_eq!(root.remove(&target).is_err(), true);
}
fn gen_bucket(l: usize) -> Bucket {
Bucket::new(l)
}
fn gen_node(s: &str) -> Node {
Node::new(Address::generate(s), None)
}
fn gen_center() -> Center {
let mut b = [0; 32];
b[0] = 42;
let s = SecretKey::from_slice(&b).unwrap();
Center::new(s, String::from(""), 8080)
}
}