use std::{
collections::hash_map::RandomState,
convert::TryFrom,
default::Default,
hash::{BuildHasher, Hash, Hasher},
iter::FromIterator,
marker::PhantomData,
};
use hashbrown::HashMap;
use thiserror::Error;
use crate::{anchor::Anchor, ResourceIterator, ResourceMutIterator};
#[derive(Debug, Error, PartialEq, Eq, Clone, Copy)]
pub enum Error {
#[error("configured resource capacity reached")]
CapacityLimitReached,
#[error("resource not found")]
ResourceNotFound,
}
type Result<T, E = Error> = std::result::Result<T, E>;
#[derive(Debug, Clone)]
pub struct Builder<R, B>
where
B: BuildHasher,
{
resources: Option<Vec<R>>,
hasher: B,
}
impl<R> Default for Builder<R, RandomState> {
fn default() -> Self {
Self {
hasher: RandomState::default(),
resources: None,
}
}
}
impl<R, B> Builder<R, B>
where
B: BuildHasher,
{
pub fn build<K: Hash>(self, capacity: u16) -> AnchorHash<K, R, B> {
let mut anchor = Anchor::new(capacity, 0);
let mut resources = HashMap::new();
if let Some(res) = self.resources {
for r in res {
let bucket = anchor
.add_bucket()
.expect("number of resources cannot exceed capacity");
resources.insert(bucket, r);
}
}
AnchorHash {
anchor,
hasher: self.hasher,
resources,
_key_type: PhantomData::default(),
}
}
pub fn with_hasher(builder: B) -> Self {
Self {
hasher: builder,
resources: None,
}
}
pub fn with_resources(self, resources: impl IntoIterator<Item = R>) -> Self {
Self {
resources: Some(resources.into_iter().collect()),
..self
}
}
}
impl<R, K> FromIterator<R> for AnchorHash<K, R, RandomState>
where
K: Hash,
{
fn from_iter<T: IntoIterator<Item = R>>(iter: T) -> Self {
let resources = iter.into_iter().collect::<Vec<_>>();
let n = u16::try_from(resources.len()).expect("too many resources");
Builder::default().with_resources(resources).build(n)
}
}
#[derive(Debug)]
pub struct AnchorHash<K, R, B>
where
K: Hash,
B: BuildHasher,
{
anchor: Anchor,
hasher: B,
resources: HashMap<u16, R>,
_key_type: PhantomData<K>,
}
impl<K, R, B> Clone for AnchorHash<K, R, B>
where
K: Hash,
B: BuildHasher + Clone,
R: Clone,
{
fn clone(&self) -> Self {
Self {
anchor: self.anchor.clone(),
hasher: self.hasher.clone(),
resources: self.resources.clone(),
_key_type: PhantomData::default(),
}
}
}
impl<K, R, B> AnchorHash<K, R, B>
where
K: Hash,
B: BuildHasher,
R: PartialEq,
{
pub fn get_resource(&self, key: K) -> Option<&R> {
let mut hasher = self.hasher.build_hasher();
key.hash(&mut hasher);
let key = hasher.finish();
let b = self.anchor.get_bucket(key as u32);
self.resources.get(&b)
}
pub fn add_resource(&mut self, resource: R) -> Result<()> {
let b = self
.anchor
.add_bucket()
.ok_or(Error::CapacityLimitReached)?;
assert!(self.resources.insert(b, resource).is_none());
Ok(())
}
pub fn remove_resource(&mut self, resource: &R) -> Result<()> {
let b = self
.resources
.iter()
.find(|(_k, r)| *r == resource)
.map(|(&k, _r)| k)
.ok_or(Error::ResourceNotFound)?;
self.resources.remove(&b);
self.anchor.remove_bucket(b);
Ok(())
}
pub fn resources(&self) -> ResourceIterator<'_, R> {
self.resources.values().into()
}
pub fn resources_mut(&mut self) -> ResourceMutIterator<'_, R> {
self.resources.values_mut().into()
}
}
#[cfg(test)]
mod tests {
use hashbrown::HashSet;
use super::*;
#[derive(Debug, PartialEq)]
struct BackendServer {
id: usize,
}
#[test]
fn test_build_empty_add_resource() {
let mut a: AnchorHash<usize, _, _> = Builder::default()
.with_resources(Vec::<usize>::new())
.build(10);
assert_eq!(a.resources().len(), 0);
assert!(a.get_resource(42).is_none());
a.add_resource(24).expect("should add new resource");
assert_eq!(a.get_resource(42), Some(&24));
a.remove_resource(&24).expect("should remove resource");
assert!(a.get_resource(42).is_none());
}
#[test]
fn test_build_with_resources() {
let servers = vec!["A", "B", "C", "D"];
let a: AnchorHash<usize, _, _> =
Builder::default().with_resources(servers.clone()).build(10);
let working = a.anchor.working_buckets();
assert_eq!(working.len(), servers.len());
assert_eq!(a.resources.len(), servers.len());
assert_eq!(a.resources().len(), servers.len());
for bucket in a.resources.keys() {
assert!(working.contains(bucket));
}
let values = a.resources.values().cloned().collect::<HashSet<_>>();
assert_eq!(values, servers.into_iter().collect::<HashSet<_>>());
}
#[test]
fn test_anchorhash_borrowed_resources() {
let server_a = BackendServer { id: 1 };
let server_b = BackendServer { id: 2 };
let mut a = Builder::default().build(20);
a.add_resource(&server_a).unwrap();
a.add_resource(&server_b).unwrap();
let got = a.get_resource("a key").expect("should return a resource");
assert!(got == &&server_a || got == &&server_b);
a.remove_resource(&&server_b)
.expect("removing existing resource should succeed");
let got = a.get_resource("a key").expect("should return a resource");
assert_eq!(got, &&server_a);
}
#[test]
fn test_anchorhash_owned_resources() {
let server_a = BackendServer { id: 1 };
let server_b = BackendServer { id: 2 };
let mut a = Builder::default().build(20);
a.add_resource(&server_a).unwrap();
a.add_resource(&server_b).unwrap();
let got = a.get_resource("a key").expect("should return a resource");
assert!(got == &&server_a || got == &&server_b);
a.remove_resource(&&server_b)
.expect("removing existing resource should succeed");
let got = a.get_resource("a key").expect("should return a resource");
assert_eq!(got, &&server_a);
}
#[test]
fn test_full() {
let mut a: AnchorHash<usize, _, _> = Builder::default().build(2);
a.add_resource(1).unwrap();
a.add_resource(2).unwrap();
let err = a
.add_resource(3)
.expect_err("should not allow 3rd resource for capacity == 2");
assert_eq!(err, Error::CapacityLimitReached);
}
#[test]
fn test_remove_not_found() {
let mut a: AnchorHash<usize, _, _> = Builder::default().build(2);
a.add_resource(1).unwrap();
a.add_resource(2).unwrap();
let err = a
.remove_resource(&3)
.expect_err("should not allow removing non-existant resource");
assert_eq!(err, Error::ResourceNotFound);
}
#[test]
fn test_cloneable() {
let mut a: AnchorHash<usize, _, _> = Builder::default().build(2);
a.add_resource(1).unwrap();
a.add_resource(2).unwrap();
let b = a.clone();
let got_a = a.resources().cloned().collect::<HashSet<_>>();
let got_b = b.resources().cloned().collect::<HashSet<_>>();
assert_eq!(got_a, got_b);
for i in 0..100 {
assert_eq!(a.get_resource(i), b.get_resource(i));
}
}
}