Implement a Concurrent Hash Map in Rust (Simplified DashMap)

Concurrent data structures are essential for building performant multi-threaded applications. This challenge asks you to implement a simplified version of a concurrent hash map, similar in concept to dashmap, which allows safe and efficient concurrent access to data from multiple threads. You will need to manage shared mutable state carefully to avoid race conditions and ensure data integrity.

Problem Description

Your task is to create a concurrent hash map data structure in Rust. This data structure should allow multiple threads to insert, retrieve, and remove key-value pairs concurrently without data corruption or deadlocks. For simplicity, we will focus on a basic implementation that uses a fixed number of segments (or buckets) to partition the hash map, allowing for more fine-grained locking.

Key Requirements:

Concurrent Access: The hash map must support insert, get, and remove operations that can be called from multiple threads simultaneously.
Segmentation/Bucketing: The hash map should be divided into a predefined number of segments. Each segment will have its own lock, allowing operations on different segments to proceed in parallel.
Thread Safety: All operations must be thread-safe. This means using appropriate synchronization primitives (e.g., Mutex or RwLock) to protect shared data.
Basic API: Implement the following methods:
- new(num_segments: usize) -> Self: Creates a new concurrent hash map with the specified number of segments.
- insert(&self, key: K, value: V) -> Option<V>: Inserts a key-value pair. If the key already exists, it updates the value and returns the old value.
- get(&self, key: &K) -> Option<&V>: Retrieves a reference to the value associated with the key.
- remove(&self, key: &K) -> Option<V>: Removes the key-value pair and returns the removed value.
Generics: The hash map should be generic over key type K and value type V. The key type K must implement Eq + Hash + Clone.

Expected Behavior:

When multiple threads attempt to insert, get, or remove elements, the operations should complete without panics or data loss.
Operations targeting different segments should not block each other unnecessarily.
insert should correctly handle both new keys and existing keys.
get should return Some(&V) if the key exists and None otherwise.
remove should return Some(V) if the key existed and None otherwise.

Edge Cases:

Empty Map: Operations on an empty map should behave correctly.
Key Collisions: Keys that hash to the same segment should be handled correctly within that segment's data structure (e.g., using a Vec or another HashMap per segment).
Zero Segments: The new function should ideally handle or panic for num_segments == 0.

Examples

Example 1: Basic Insertion and Retrieval

use std::sync::Arc;
use std::thread;

// Assume your ConcurrentHashMap is defined as `ConcurrentHashMap<String, i32>`
let map = Arc::new(ConcurrentHashMap::new(4)); // 4 segments

let mut handles = vec![];

for i in 0..10 {
    let map_clone = Arc::clone(&map);
    handles.push(thread::spawn(move || {
        map_clone.insert(format!("key_{}", i), i * 10);
    }));
}

for handle in handles {
    handle.join().unwrap();
}

assert_eq!(map.get(&"key_3".to_string()), Some(&30));
assert_eq!(map.get(&"key_7".to_string()), Some(&70));
assert_eq!(map.get(&"key_10".to_string()), None);

Explanation: Ten threads are spawned, each inserting a unique key-value pair into the concurrent hash map. After all insertions, we verify that a few keys can be retrieved correctly.

Example 2: Concurrent Updates and Removals

use std::sync::Arc;
use std::thread;

let map = Arc::new(ConcurrentHashMap::new(2)); // 2 segments

// Initial insertions
map.insert("apple".to_string(), 1);
map.insert("banana".to_string(), 2);

let mut handles = vec![];

// Thread 1: Updates "apple" and inserts "cherry"
handles.push(thread::spawn({
    let map_clone = Arc::clone(&map);
    move || {
        map_clone.insert("apple".to_string(), 10);
        map_clone.insert("cherry".to_string(), 3);
    }
}));

// Thread 2: Removes "banana" and tries to get "apple"
handles.push(thread::spawn({
    let map_clone = Arc::clone(&map);
    move || {
        let removed = map_clone.remove(&"banana".to_string());
        assert_eq!(removed, Some(2));
        let apple_val = map_clone.get(&"apple".to_string());
        // We can't assert the exact value here due to potential race with Thread 1
        // but we expect it to be Some.
        assert!(apple_val.is_some());
    }
}));

for handle in handles {
    handle.join().unwrap();
}

assert_eq!(map.get(&"apple".to_string()), Some(&10));
assert_eq!(map.get(&"banana".to_string()), None);
assert_eq!(map.get(&"cherry".to_string()), Some(&3));

Explanation: Two threads operate concurrently. One updates an existing value and inserts a new one. The other removes an existing entry and attempts to retrieve a value that might be concurrently updated. The final assertions confirm the state of the map after these concurrent operations.

Example 3: Handling Key Collisions (Conceptual)

This example is conceptual as the exact output depends on the hashing function and segment distribution.

Input data: Keys that are known to hash to the same segment. Consider keys: "a", "q", "A" (if case-insensitive hashing is used, these might collide).

// Assuming a hash function where "a" and "q" hash to the same segment.
let map = ConcurrentHashMap::new(1); // Single segment for simplicity in demonstration

map.insert("a".to_string(), 100);
map.insert("q".to_string(), 200);

assert_eq!(map.get(&"a".to_string()), Some(&100));
assert_eq!(map.get(&"q".to_string()), Some(&200));

let old_val = map.insert("a".to_string(), 150);
assert_eq!(old_val, Some(100));
assert_eq!(map.get(&"a".to_string()), Some(&150));

let removed_val = map.remove(&"q".to_string());
assert_eq!(removed_val, Some(200));
assert_eq!(map.get(&"q".to_string()), None);

Explanation: When using a single segment (or if multiple keys hash to the same segment), the internal data structure of that segment (e.g., a std::collections::HashMap) must correctly handle multiple keys. This example shows that even with potential collisions within a segment, individual keys can be inserted, retrieved, updated, and removed correctly.

Constraints

The number of segments (num_segments) will be between 1 and 256, inclusive.
Key types K must implement Eq, Hash, and Clone.
Value types V can be any type.
The total number of operations in testing scenarios will not exceed 10,000.
The solution should aim for good performance in concurrent scenarios, meaning that the critical sections protected by locks should be as small as possible.

Notes

You will need to decide how to store the data within each segment. A std::collections::HashMap per segment is a common and reasonable choice.
Consider using std::sync::RwLock for segments that allow multiple readers concurrently, which can improve performance for read-heavy workloads. However, std::sync::Mutex is also a valid and simpler starting point.
The get method should return a reference (&V). This implies careful lifetime management, especially with concurrent access. The reference returned must be valid as long as the entry exists in the map and the lock protecting it is held (or the reference is conceptually tied to the segment's lock).
You'll need to use Arc to share the ConcurrentHashMap instance across threads.
The Hash trait from std::hash will be crucial for determining which segment a key belongs to. You can use hashbrown::hash_map::DefaultHasher or a similar hashing mechanism.
The calculation of the segment index from a key's hash can be done using the modulo operator: hash % num_segments.
The remove operation should also return the owned value, not a reference.
Consider how to handle the lifetime of references returned by get in a thread-safe manner. For this simplified challenge, you might consider returning a reference that is tied to the lifetime of the ConcurrentHashMap itself, assuming the map is not modified in a way that invalidates the reference externally (which is guaranteed by the get method's signature usually). A more robust solution might involve internal reference counting or other techniques, but for this problem, a direct reference is likely sufficient.