use dashmap::DashMap;
use std::sync::Arc;
use std::path::Path;
use std::fs;
use std::io;
const SHARD_COUNT: usize = 256;
pub struct StringIndex {
shards: Vec<Arc<DashMap<Vec<u8>, [u8; 32]>>>,
manifest_path: Option<std::path::PathBuf>,
}
impl StringIndex {
pub fn new() -> Self {
let mut shards = Vec::with_capacity(SHARD_COUNT);
for _ in 0..SHARD_COUNT {
shards.push(Arc::new(DashMap::new()));
}
Self { shards, manifest_path: None }
}
pub fn with_persistence(path: &Path) -> io::Result<Self> {
let manifest_path = path.join("string_index.json");
let mut idx = Self::new();
idx.manifest_path = Some(manifest_path.clone());
if manifest_path.exists() {
idx.load_from_disk()?;
}
Ok(idx)
}
fn load_from_disk(&self) -> io::Result<()> {
let path = self.manifest_path.as_ref().ok_or_else(|| {
io::Error::new(io::ErrorKind::Other, "no manifest path configured")
})?;
let data = fs::read_to_string(path)?;
let entries: Vec<(Vec<u8>, [u8; 32])> = serde_json::from_str(&data)
.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
for (key, fixed_key) in entries {
self.insert(&key, fixed_key);
}
Ok(())
}
pub fn save_to_disk(&self) -> io::Result<()> {
let path = self.manifest_path.as_ref().ok_or_else(|| {
io::Error::new(io::ErrorKind::Other, "no manifest path configured")
})?;
use rayon::prelude::*;
let entries: Vec<(Vec<u8>, [u8; 32])> = self.shards
.par_iter()
.flat_map(|shard| {
shard.iter().map(|entry| {
(entry.key().clone(), *entry.value())
}).collect::<Vec<_>>()
})
.collect();
let json = serde_json::to_string(&entries)
.map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
let temp_path = path.with_extension("json.tmp");
fs::write(&temp_path, json)?;
fs::rename(temp_path, path)?;
Ok(())
}
#[inline]
pub fn insert(&self, string_key: &[u8], fixed_key: [u8; 32]) {
let shard_id = if string_key.is_empty() {
0
} else {
string_key[0] as usize
};
self.shards[shard_id].insert(string_key.to_vec(), fixed_key);
}
#[inline]
pub fn remove(&self, string_key: &[u8]) {
let shard_id = if string_key.is_empty() {
0
} else {
string_key[0] as usize
};
self.shards[shard_id].remove(string_key);
}
#[inline]
pub fn get(&self, string_key: &[u8]) -> Option<[u8; 32]> {
let shard_id = if string_key.is_empty() {
0
} else {
string_key[0] as usize
};
self.shards[shard_id].get(string_key).map(|entry| *entry.value())
}
pub fn scan_prefix(&self, prefix: &[u8]) -> Vec<(Vec<u8>, [u8; 32])> {
use rayon::prelude::*;
if prefix.is_empty() {
return self.shards
.par_iter()
.flat_map(|shard| {
shard.iter().map(|entry| {
(entry.key().clone(), *entry.value())
}).collect::<Vec<_>>()
})
.collect();
}
let first_byte = prefix[0];
let mut results: Vec<(Vec<u8>, [u8; 32])> = self.shards[first_byte as usize]
.iter()
.filter_map(|entry| {
let key = entry.key();
if key.starts_with(prefix) {
Some((key.clone(), *entry.value()))
} else {
None
}
})
.collect();
results.sort_by(|a, b| a.0.cmp(&b.0));
results
}
pub fn scan_from_reverse(&self, start: &[u8], limit: usize) -> Vec<(Vec<u8>, [u8; 32])> {
use rayon::prelude::*;
let mut results: Vec<(Vec<u8>, [u8; 32])> = self.shards
.par_iter()
.flat_map(|shard| {
shard.iter().filter_map(|entry| {
let key = entry.key();
if key.as_slice() < start {
Some((key.clone(), *entry.value()))
} else {
None
}
}).collect::<Vec<_>>()
})
.collect();
results.sort_by(|a, b| b.0.cmp(&a.0));
results.truncate(limit);
results
}
pub fn clear(&self) {
for shard in &self.shards {
shard.clear();
}
}
pub fn len(&self) -> usize {
self.shards.iter().map(|s| s.len()).sum()
}
pub fn is_empty(&self) -> bool {
self.shards.iter().all(|s| s.is_empty())
}
}
impl Default for StringIndex {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_insert_and_get() {
let idx = StringIndex::new();
let key1 = b"batch:123";
let fixed1 = [1u8; 32];
idx.insert(key1, fixed1);
assert_eq!(idx.get(key1), Some(fixed1));
assert_eq!(idx.get(b"batch:456"), None);
}
#[test]
fn test_prefix_scan() {
let idx = StringIndex::new();
idx.insert(b"batch:1", [1u8; 32]);
idx.insert(b"batch:2", [2u8; 32]);
idx.insert(b"batch:10", [10u8; 32]);
idx.insert(b"tx:1", [100u8; 32]);
let results = idx.scan_prefix(b"batch:");
assert_eq!(results.len(), 3);
assert_eq!(results[0].0, b"batch:1");
assert_eq!(results[1].0, b"batch:10");
assert_eq!(results[2].0, b"batch:2");
}
#[test]
fn test_remove() {
let idx = StringIndex::new();
let key1 = b"batch:123";
let fixed1 = [1u8; 32];
idx.insert(key1, fixed1);
assert_eq!(idx.len(), 1);
idx.remove(key1);
assert_eq!(idx.len(), 0);
assert_eq!(idx.get(key1), None);
}
#[test]
fn test_scan_from_reverse() {
let idx = StringIndex::new();
idx.insert(b"key:1", [1u8; 32]);
idx.insert(b"key:2", [2u8; 32]);
idx.insert(b"key:3", [3u8; 32]);
idx.insert(b"key:4", [4u8; 32]);
let results = idx.scan_from_reverse(b"key:3", 10);
assert_eq!(results.len(), 2);
assert_eq!(results[0].0, b"key:2");
assert_eq!(results[1].0, b"key:1");
}
}