mirror of https://github.com/rapiz1/rathole.git
115 lines
2.9 KiB
Rust
115 lines
2.9 KiB
Rust
use std::borrow::Borrow;
|
|
use std::collections::HashMap;
|
|
use std::hash::{Hash, Hasher};
|
|
|
|
struct RawItem<K1, K2, V>(*mut (K1, K2, V));
|
|
unsafe impl<K1, K2, V> Send for RawItem<K1, K2, V> {}
|
|
unsafe impl<K1, K2, V> Sync for RawItem<K1, K2, V> {}
|
|
|
|
pub struct MultiMap<K1, K2, V> {
|
|
map1: HashMap<Key<K1>, RawItem<K1, K2, V>>,
|
|
map2: HashMap<Key<K2>, RawItem<K1, K2, V>>,
|
|
}
|
|
|
|
struct Key<T>(*const T);
|
|
|
|
unsafe impl<T> Send for Key<T> {}
|
|
unsafe impl<T> Sync for Key<T> {}
|
|
|
|
impl<T> Borrow<T> for Key<T> {
|
|
fn borrow(&self) -> &T {
|
|
unsafe { &*self.0 }
|
|
}
|
|
}
|
|
|
|
impl<T: Hash> Hash for Key<T> {
|
|
fn hash<H: Hasher>(&self, state: &mut H) {
|
|
(self.borrow() as &T).hash(state)
|
|
}
|
|
}
|
|
|
|
impl<T: PartialEq> PartialEq for Key<T> {
|
|
fn eq(&self, other: &Self) -> bool {
|
|
(self.borrow() as &T).eq(other.borrow())
|
|
}
|
|
}
|
|
|
|
impl<T: Eq> Eq for Key<T> {}
|
|
|
|
impl<K1, K2, V> MultiMap<K1, K2, V> {
|
|
pub fn new() -> Self {
|
|
MultiMap {
|
|
map1: HashMap::new(),
|
|
map2: HashMap::new(),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[allow(dead_code)]
|
|
impl<K1, K2, V> MultiMap<K1, K2, V>
|
|
where
|
|
K1: Hash + Eq + Send,
|
|
K2: Hash + Eq + Send,
|
|
V: Send,
|
|
{
|
|
pub fn insert(&mut self, k1: K1, k2: K2, v: V) -> Result<(), (K1, K2, V)> {
|
|
if self.map1.contains_key(&k1) || self.map2.contains_key(&k2) {
|
|
return Err((k1, k2, v));
|
|
}
|
|
let item = Box::new((k1, k2, v));
|
|
let k1 = Key(&item.0);
|
|
let k2 = Key(&item.1);
|
|
let item = Box::into_raw(item);
|
|
self.map1.insert(k1, RawItem(item));
|
|
self.map2.insert(k2, RawItem(item));
|
|
Ok(())
|
|
}
|
|
|
|
pub fn get1(&self, k1: &K1) -> Option<&V> {
|
|
let item = self.map1.get(k1)?;
|
|
let item = unsafe { &*item.0 };
|
|
Some(&item.2)
|
|
}
|
|
|
|
pub fn get1_mut(&mut self, k1: &K1) -> Option<&mut V> {
|
|
let item = self.map1.get(k1)?;
|
|
let item = unsafe { &mut *item.0 };
|
|
Some(&mut item.2)
|
|
}
|
|
|
|
pub fn get2(&self, k2: &K2) -> Option<&V> {
|
|
let item = self.map2.get(k2)?;
|
|
let item = unsafe { &*item.0 };
|
|
Some(&item.2)
|
|
}
|
|
|
|
pub fn get_mut2(&mut self, k2: &K2) -> Option<&mut V> {
|
|
let item = self.map2.get(k2)?;
|
|
let item = unsafe { &mut *item.0 };
|
|
Some(&mut item.2)
|
|
}
|
|
|
|
pub fn remove1(&mut self, k1: &K1) -> Option<V> {
|
|
let item = self.map1.remove(k1)?;
|
|
let item = unsafe { Box::from_raw(item.0) };
|
|
self.map2.remove(&item.1);
|
|
Some(item.2)
|
|
}
|
|
|
|
pub fn remove2(&mut self, k2: &K2) -> Option<V> {
|
|
let item = self.map2.remove(k2)?;
|
|
let item = unsafe { Box::from_raw(item.0) };
|
|
self.map1.remove(&item.0);
|
|
Some(item.2)
|
|
}
|
|
}
|
|
|
|
impl<K1, K2, V> Drop for MultiMap<K1, K2, V> {
|
|
fn drop(&mut self) {
|
|
self.map1.clear();
|
|
self.map2
|
|
.drain()
|
|
.for_each(|(_, item)| drop(unsafe { Box::from_raw(item.0) }));
|
|
}
|
|
}
|