use std::collections::HashMap;
use sha2::{Sha256, Digest};
use serde::{Deserialize, Serialize};
use anyhow::Result;
#[derive(Debug)]
pub struct ContentDeduplicator {
hash_to_storage: HashMap<String, String>,
ref_counts: HashMap<String, u32>,
storage_to_hash: HashMap<String, String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DedupInfo {
pub is_reference: bool,
pub original_storage_id: Option<String>,
pub hash: String,
pub ref_count: u32,
}
impl ContentDeduplicator {
pub fn new() -> Self {
Self {
hash_to_storage: HashMap::new(),
ref_counts: HashMap::new(),
storage_to_hash: HashMap::new(),
}
}
pub fn calculate_hash(data: &[u8]) -> String {
let mut hasher = Sha256::new();
hasher.update(data);
format!("{:x}", hasher.finalize())
}
pub fn check_duplicate(&mut self, hash: &str) -> Option<String> {
if let Some(storage_id) = self.hash_to_storage.get(hash) {
*self.ref_counts.entry(storage_id.clone()).or_insert(0) += 1;
Some(storage_id.clone())
} else {
None
}
}
pub fn register_file(&mut self, hash: String, storage_id: String) {
self.hash_to_storage.insert(hash.clone(), storage_id.clone());
self.storage_to_hash.insert(storage_id.clone(), hash);
self.ref_counts.insert(storage_id, 1);
}
pub fn remove_reference(&mut self, storage_id: &str) -> bool {
if let Some(count) = self.ref_counts.get_mut(storage_id) {
*count -= 1;
if *count == 0 {
self.ref_counts.remove(storage_id);
if let Some(hash) = self.storage_to_hash.remove(storage_id) {
self.hash_to_storage.remove(&hash);
}
true } else {
false }
} else {
true }
}
pub fn remove_hash_reference(&mut self, hash: &str) -> bool {
if let Some(storage_id) = self.hash_to_storage.get(hash) {
let storage_id = storage_id.clone(); self.remove_reference(&storage_id)
} else {
true }
}
pub fn add_hash_reference(&mut self, hash: &str, storage_id: &str) {
if let Some(existing_storage_id) = self.hash_to_storage.get(hash) {
if existing_storage_id == storage_id {
*self.ref_counts.entry(storage_id.to_string()).or_insert(0) += 1;
}
} else {
self.hash_to_storage.insert(hash.to_string(), storage_id.to_string());
self.storage_to_hash.insert(storage_id.to_string(), hash.to_string());
*self.ref_counts.entry(storage_id.to_string()).or_insert(0) += 1;
}
}
pub fn get_dedup_info(&self, storage_id: &str) -> Option<DedupInfo> {
if let Some(hash) = self.storage_to_hash.get(storage_id) {
let ref_count = self.ref_counts.get(storage_id).copied().unwrap_or(0);
Some(DedupInfo {
is_reference: ref_count > 1,
original_storage_id: None, hash: hash.clone(),
ref_count,
})
} else {
None
}
}
pub fn get_reference_info(&self, hash: &str) -> Option<DedupInfo> {
if let Some(storage_id) = self.hash_to_storage.get(hash) {
let ref_count = self.ref_counts.get(storage_id).copied().unwrap_or(0);
Some(DedupInfo {
is_reference: true,
original_storage_id: Some(storage_id.clone()),
hash: hash.to_string(),
ref_count,
})
} else {
None
}
}
pub fn get_stats(&self) -> DedupStats {
let total_files = self.ref_counts.values().sum::<u32>();
let unique_files = self.ref_counts.len() as u32;
let duplicate_files = total_files.saturating_sub(unique_files);
DedupStats {
total_files,
unique_files,
duplicate_files,
dedup_ratio: if total_files > 0 {
duplicate_files as f32 / total_files as f32
} else {
0.0
},
}
}
pub fn rebuild_from_index(&mut self, entries: Vec<(String, String, u32)>) -> Result<()> {
self.hash_to_storage.clear();
self.ref_counts.clear();
self.storage_to_hash.clear();
for (storage_id, hash, ref_count) in entries {
self.hash_to_storage.insert(hash.clone(), storage_id.clone());
self.storage_to_hash.insert(storage_id.clone(), hash);
self.ref_counts.insert(storage_id, ref_count);
}
Ok(())
}
}
impl Default for ContentDeduplicator {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone)]
pub struct DedupStats {
pub total_files: u32,
pub unique_files: u32,
pub duplicate_files: u32,
pub dedup_ratio: f32,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_deduplicator_basic() {
let mut dedup = ContentDeduplicator::new();
let hash1 = "abc123".to_string();
assert_eq!(dedup.check_duplicate(&hash1), None);
dedup.register_file(hash1.clone(), "storage1".to_string());
assert_eq!(dedup.check_duplicate(&hash1), Some("storage1".to_string()));
let info = dedup.get_dedup_info("storage1").unwrap();
assert_eq!(info.ref_count, 2); }
#[test]
fn test_hash_calculation() {
let data = b"Hello, World!";
let hash = ContentDeduplicator::calculate_hash(data);
assert!(!hash.is_empty());
assert_eq!(hash.len(), 64); }
#[test]
fn test_remove_reference() {
let mut dedup = ContentDeduplicator::new();
dedup.register_file("hash1".to_string(), "storage1".to_string());
dedup.check_duplicate("hash1");
assert!(!dedup.remove_reference("storage1"));
assert!(dedup.remove_reference("storage1"));
}
#[test]
fn test_remove_reference_by_hash() {
let mut dedup = ContentDeduplicator::new();
dedup.register_file("hash1".to_string(), "storage1".to_string());
dedup.check_duplicate("hash1");
assert!(!dedup.remove_hash_reference("hash1"));
assert!(dedup.remove_hash_reference("hash1"));
}
#[test]
fn test_add_reference_by_hash() {
let mut dedup = ContentDeduplicator::new();
dedup.register_file("hash1".to_string(), "storage1".to_string());
dedup.add_hash_reference("hash1", "storage1");
let info = dedup.get_dedup_info("storage1").unwrap();
assert_eq!(info.ref_count, 2);
dedup.add_hash_reference("hash2", "storage2");
assert_eq!(dedup.hash_to_storage.get("hash2"), Some(&"storage2".to_string()));
}
}