use std::borrow::Borrow;
use crate::adnl;
use crate::proto;
use crate::util::*;
pub struct Buckets {
local_id: [u8; 32],
buckets: Box<[FastDashMap<adnl::NodeIdShort, proto::dht::NodeOwned>; 256]>,
}
impl Buckets {
pub fn new(local_id: &adnl::NodeIdShort) -> Self {
Self {
local_id: *local_id.as_slice(),
buckets: Box::new([(); 256].map(|_| Default::default())),
}
}
pub fn iter(&self) -> std::slice::Iter<FastDashMap<adnl::NodeIdShort, proto::dht::NodeOwned>> {
self.buckets.iter()
}
pub fn insert(&self, peer_id: &adnl::NodeIdShort, peer: proto::dht::NodeOwned) {
use dashmap::mapref::entry::Entry;
let affinity = get_affinity(&self.local_id, peer_id.borrow());
match self.buckets[affinity as usize].entry(*peer_id) {
Entry::Occupied(mut entry) => {
if entry.get().version < peer.version {
entry.insert(peer);
}
}
Entry::Vacant(entry) => {
entry.insert(peer);
}
}
}
pub fn find<T>(&self, peer_id: T, k: u32) -> proto::dht::NodesOwned
where
T: Borrow<[u8; 32]>,
{
let key1 = &self.local_id;
let key2: &[u8; 32] = peer_id.borrow();
let mut nodes = Vec::new();
'outer: for i in 0..32 {
let mut distance = i as u8 * 8;
let mut diff = key1[i] ^ key2[i];
while diff != 0 {
let equal_bits = diff.leading_zeros() as u8; distance += equal_bits;
let bucket = &self.buckets[distance as usize];
for item in bucket.iter() {
nodes.push(item.value().clone());
if nodes.len() >= k as usize {
break 'outer;
}
}
if equal_bits < 7 {
diff <<= equal_bits + 1;
distance = distance.saturating_add(1);
} else {
continue 'outer;
}
}
}
proto::dht::NodesOwned { nodes }
}
}
impl<'a> IntoIterator for &'a Buckets {
type Item = &'a FastDashMap<adnl::NodeIdShort, proto::dht::NodeOwned>;
type IntoIter = std::slice::Iter<'a, FastDashMap<adnl::NodeIdShort, proto::dht::NodeOwned>>;
fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}
pub fn get_affinity(key1: &[u8; 32], key2: &[u8; 32]) -> u8 {
for i in 0..32 {
let diff = key1[i] ^ key2[i];
if diff != 0 {
return (i * 8 + diff.leading_zeros() as usize) as u8;
}
}
255
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn same_affinity() {
assert_eq!(get_affinity(&[0xaa; 32], &[0xaa; 32]), 255);
}
}