use dashmap::DashMap;
use rustc_hash::{FxHashMap, FxHashSet};
use std::collections::BTreeMap;
use std::sync::RwLock;
pub use rustc_hash::FxBuildHasher;
pub type Int64Map<V> = FxHashMap<i64, V>;
pub type UInt64Map<V> = FxHashMap<u64, V>;
pub type UsizeMap<V> = FxHashMap<usize, V>;
pub type Int64Set = FxHashSet<i64>;
pub type UInt64Set = FxHashSet<u64>;
pub type UsizeSet = FxHashSet<usize>;
pub type ConcurrentInt64Map<V> = DashMap<i64, V, FxBuildHasher>;
pub type ConcurrentUInt64Map<V> = DashMap<u64, V, FxBuildHasher>;
pub type ConcurrentUsizeMap<V> = DashMap<usize, V, FxBuildHasher>;
pub type OrderedInt64Map<V> = RwLock<BTreeMap<i64, V>>;
pub fn new_int64_map<V>() -> Int64Map<V> {
FxHashMap::default()
}
pub fn new_int64_map_with_capacity<V>(capacity: usize) -> Int64Map<V> {
FxHashMap::with_capacity_and_hasher(capacity, FxBuildHasher)
}
pub fn new_uint64_map<V>() -> UInt64Map<V> {
FxHashMap::default()
}
pub fn new_uint64_map_with_capacity<V>(capacity: usize) -> UInt64Map<V> {
FxHashMap::with_capacity_and_hasher(capacity, FxBuildHasher)
}
pub fn new_usize_map<V>() -> UsizeMap<V> {
FxHashMap::default()
}
pub fn new_usize_map_with_capacity<V>(capacity: usize) -> UsizeMap<V> {
FxHashMap::with_capacity_and_hasher(capacity, FxBuildHasher)
}
pub fn new_concurrent_int64_map<V>() -> ConcurrentInt64Map<V> {
DashMap::with_hasher(FxBuildHasher)
}
pub fn new_concurrent_int64_map_with_capacity<V>(capacity: usize) -> ConcurrentInt64Map<V> {
DashMap::with_capacity_and_hasher(capacity, FxBuildHasher)
}
pub fn new_concurrent_uint64_map<V>() -> ConcurrentUInt64Map<V> {
DashMap::with_hasher(FxBuildHasher)
}
pub fn new_concurrent_uint64_map_with_capacity<V>(capacity: usize) -> ConcurrentUInt64Map<V> {
DashMap::with_capacity_and_hasher(capacity, FxBuildHasher)
}
pub fn new_concurrent_usize_map<V>() -> ConcurrentUsizeMap<V> {
DashMap::with_hasher(FxBuildHasher)
}
pub fn new_concurrent_usize_map_with_capacity<V>(capacity: usize) -> ConcurrentUsizeMap<V> {
DashMap::with_capacity_and_hasher(capacity, FxBuildHasher)
}
pub fn new_ordered_int64_map<V>() -> OrderedInt64Map<V> {
RwLock::new(BTreeMap::new())
}
pub fn new_int64_set() -> Int64Set {
FxHashSet::default()
}
pub fn new_int64_set_with_capacity(capacity: usize) -> Int64Set {
FxHashSet::with_capacity_and_hasher(capacity, FxBuildHasher)
}
pub fn new_uint64_set() -> UInt64Set {
FxHashSet::default()
}
pub fn new_uint64_set_with_capacity(capacity: usize) -> UInt64Set {
FxHashSet::with_capacity_and_hasher(capacity, FxBuildHasher)
}
pub fn new_usize_set() -> UsizeSet {
FxHashSet::default()
}
pub fn new_usize_set_with_capacity(capacity: usize) -> UsizeSet {
FxHashSet::with_capacity_and_hasher(capacity, FxBuildHasher)
}
struct Segment<V> {
data: std::sync::RwLock<Int64Map<V>>,
}
pub struct SegmentInt64Map<V> {
segments: Box<[Segment<V>]>,
segment_mask: usize,
}
impl<V> SegmentInt64Map<V> {
pub fn new() -> Self {
Self::with_segments(16, 64)
}
pub fn with_segments(segment_count: usize, capacity_per_segment: usize) -> Self {
let segment_count = segment_count.next_power_of_two();
let segment_mask = segment_count - 1;
let segments: Vec<Segment<V>> = (0..segment_count)
.map(|_| Segment {
data: std::sync::RwLock::new(FxHashMap::with_capacity_and_hasher(
capacity_per_segment,
FxBuildHasher,
)),
})
.collect();
Self {
segments: segments.into_boxed_slice(),
segment_mask,
}
}
#[inline]
fn segment_index(&self, key: i64) -> usize {
let h = (key as u64).wrapping_mul(0x9E3779B97F4A7C15);
((h >> 32) as usize) & self.segment_mask
}
#[inline]
pub fn get(&self, key: &i64) -> Option<V>
where
V: Clone,
{
let idx = self.segment_index(*key);
let guard = self.segments[idx].data.read().unwrap();
guard.get(key).cloned()
}
#[inline]
pub fn contains_key(&self, key: &i64) -> bool {
let idx = self.segment_index(*key);
let guard = self.segments[idx].data.read().unwrap();
guard.contains_key(key)
}
#[inline]
pub fn insert(&self, key: i64, value: V) {
let idx = self.segment_index(key);
let mut guard = self.segments[idx].data.write().unwrap();
guard.insert(key, value);
}
#[inline]
pub fn insert_if_not_exists(&self, key: i64, value: V) -> (Option<V>, bool)
where
V: Clone,
{
let idx = self.segment_index(key);
let mut guard = self.segments[idx].data.write().unwrap();
if let Some(existing) = guard.get(&key) {
(Some(existing.clone()), false)
} else {
guard.insert(key, value);
(None, true)
}
}
#[inline]
pub fn remove(&self, key: &i64) -> Option<V> {
let idx = self.segment_index(*key);
let mut guard = self.segments[idx].data.write().unwrap();
guard.remove(key)
}
pub fn len(&self) -> usize {
self.segments
.iter()
.map(|s| s.data.read().unwrap().len())
.sum()
}
pub fn is_empty(&self) -> bool {
self.segments
.iter()
.all(|s| s.data.read().unwrap().is_empty())
}
pub fn clear(&self) {
for segment in self.segments.iter() {
segment.data.write().unwrap().clear();
}
}
pub fn iter_collect(&self) -> Vec<(i64, V)>
where
V: Clone,
{
let mut result = Vec::new();
for segment in self.segments.iter() {
let guard = segment.data.read().unwrap();
for (k, v) in guard.iter() {
result.push((*k, v.clone()));
}
}
result
}
}
impl<V> Default for SegmentInt64Map<V> {
fn default() -> Self {
Self::new()
}
}
pub fn new_segment_int64_map<V>() -> SegmentInt64Map<V> {
SegmentInt64Map::new()
}
pub fn new_segment_int64_map_with_segments<V>(
segment_count: usize,
capacity_per_segment: usize,
) -> SegmentInt64Map<V> {
SegmentInt64Map::with_segments(segment_count, capacity_per_segment)
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Arc;
use std::thread;
#[test]
fn test_int64_map_basic() {
let mut map: Int64Map<String> = new_int64_map();
map.insert(1, "one".to_string());
map.insert(2, "two".to_string());
map.insert(3, "three".to_string());
assert_eq!(map.get(&1), Some(&"one".to_string()));
assert_eq!(map.get(&2), Some(&"two".to_string()));
assert_eq!(map.get(&3), Some(&"three".to_string()));
assert_eq!(map.get(&4), None);
assert_eq!(map.remove(&2), Some("two".to_string()));
assert_eq!(map.get(&2), None);
assert_eq!(map.len(), 2);
}
#[test]
fn test_int64_map_with_capacity() {
let map: Int64Map<i32> = new_int64_map_with_capacity(100);
assert!(map.capacity() >= 100);
}
#[test]
fn test_int64_map_negative_keys() {
let mut map: Int64Map<String> = new_int64_map();
map.insert(-1, "negative one".to_string());
map.insert(i64::MIN, "min".to_string());
map.insert(i64::MAX, "max".to_string());
assert_eq!(map.get(&-1), Some(&"negative one".to_string()));
assert_eq!(map.get(&i64::MIN), Some(&"min".to_string()));
assert_eq!(map.get(&i64::MAX), Some(&"max".to_string()));
}
#[test]
fn test_uint64_map_basic() {
let mut map: UInt64Map<String> = new_uint64_map();
map.insert(0, "zero".to_string());
map.insert(u64::MAX, "max".to_string());
assert_eq!(map.get(&0), Some(&"zero".to_string()));
assert_eq!(map.get(&u64::MAX), Some(&"max".to_string()));
}
#[test]
fn test_usize_map_basic() {
let mut map: UsizeMap<String> = new_usize_map();
map.insert(0, "zero".to_string());
map.insert(100, "hundred".to_string());
assert_eq!(map.get(&0), Some(&"zero".to_string()));
assert_eq!(map.get(&100), Some(&"hundred".to_string()));
}
#[test]
fn test_int64_set_basic() {
let mut set = new_int64_set();
set.insert(1);
set.insert(2);
set.insert(3);
assert!(set.contains(&1));
assert!(set.contains(&2));
assert!(set.contains(&3));
assert!(!set.contains(&4));
set.remove(&2);
assert!(!set.contains(&2));
}
#[test]
fn test_concurrent_int64_map_basic() {
let map: ConcurrentInt64Map<String> = new_concurrent_int64_map();
map.insert(1, "one".to_string());
map.insert(2, "two".to_string());
assert_eq!(map.get(&1).map(|v| v.clone()), Some("one".to_string()));
assert_eq!(map.get(&2).map(|v| v.clone()), Some("two".to_string()));
assert!(map.get(&3).is_none());
map.remove(&1);
assert!(map.get(&1).is_none());
}
#[test]
fn test_concurrent_int64_map_multithreaded() {
let map: Arc<ConcurrentInt64Map<i64>> = Arc::new(new_concurrent_int64_map());
let num_threads: i64 = 8;
let ops_per_thread: i64 = 1000;
let handles: Vec<_> = (0..num_threads)
.map(|t| {
let map: Arc<ConcurrentInt64Map<i64>> = Arc::clone(&map);
thread::spawn(move || {
let base = t * ops_per_thread;
for i in 0..ops_per_thread {
map.insert(base + i, (base + i) * 2);
}
for i in 0..ops_per_thread {
let key = base + i;
let expected = key * 2;
assert_eq!(map.get(&key).map(|v| *v), Some(expected));
}
})
})
.collect();
for handle in handles {
handle.join().unwrap();
}
assert_eq!(map.len(), (num_threads * ops_per_thread) as usize);
}
#[test]
fn test_concurrent_int64_map_with_capacity() {
let map: ConcurrentInt64Map<i32> = new_concurrent_int64_map_with_capacity(100);
assert!(map.capacity() >= 100);
}
#[test]
fn test_concurrent_map_update() {
let map: ConcurrentInt64Map<i64> = new_concurrent_int64_map();
map.insert(1, 100);
assert_eq!(*map.get(&1).unwrap(), 100);
map.insert(1, 200);
assert_eq!(*map.get(&1).unwrap(), 200);
map.entry(2).or_insert(300);
assert_eq!(*map.get(&2).unwrap(), 300);
}
#[test]
fn test_concurrent_map_iter() {
let map: ConcurrentInt64Map<i64> = new_concurrent_int64_map();
for i in 0..10 {
map.insert(i, i * 10);
}
let sum: i64 = map.iter().map(|r| *r.value()).sum();
assert_eq!(sum, 450); }
#[test]
fn test_map_type_inference() {
let mut int_map: Int64Map<i32> = new_int64_map();
int_map.insert(1, 42);
let mut string_map: Int64Map<String> = new_int64_map();
string_map.insert(1, "hello".to_string());
let mut vec_map: Int64Map<Vec<u8>> = new_int64_map();
vec_map.insert(1, vec![1, 2, 3]);
assert_eq!(int_map.get(&1), Some(&42));
assert_eq!(string_map.get(&1), Some(&"hello".to_string()));
assert_eq!(vec_map.get(&1), Some(&vec![1, 2, 3]));
}
#[test]
fn test_set_operations() {
let mut set1 = new_int64_set();
let mut set2 = new_int64_set();
for i in 0..5 {
set1.insert(i);
}
for i in 3..8 {
set2.insert(i);
}
let intersection: Int64Set = set1.intersection(&set2).copied().collect();
assert_eq!(intersection.len(), 2);
let union: Int64Set = set1.union(&set2).copied().collect();
assert_eq!(union.len(), 8); }
}