use cid::Cid;
use parking_lot::RwLock;
use std::collections::{HashMap, HashSet};
use std::path::Path;
use crate::{PinInfo, PinMode, PinResult};
#[derive(Debug, Default, serde::Serialize, serde::Deserialize)]
struct PinState {
direct: HashMap<String, Option<String>>,
recursive: HashMap<String, Option<String>>,
indirect: HashMap<String, HashSet<String>>,
}
pub struct PinStore {
state: RwLock<PinState>,
path: Option<std::path::PathBuf>,
}
impl PinStore {
pub fn new() -> Self {
Self {
state: RwLock::new(PinState::default()),
path: None,
}
}
pub fn open(path: impl AsRef<Path>) -> PinResult<Self> {
let path = path.as_ref().to_path_buf();
let state = if path.exists() {
let data = std::fs::read_to_string(&path)?;
serde_json::from_str(&data)?
} else {
PinState::default()
};
Ok(Self {
state: RwLock::new(state),
path: Some(path),
})
}
pub fn save(&self) -> PinResult<()> {
if let Some(ref path) = self.path {
let state = self.state.read();
let data = serde_json::to_string_pretty(&*state)?;
let tmp_path = path.with_extension("tmp");
std::fs::write(&tmp_path, &data)?;
std::fs::rename(&tmp_path, path)?;
}
Ok(())
}
pub fn is_pinned(&self, cid: &Cid) -> bool {
let cid_str = cid.to_string();
let state = self.state.read();
state.direct.contains_key(&cid_str)
|| state.recursive.contains_key(&cid_str)
|| state.indirect.contains_key(&cid_str)
}
pub fn is_pinned_with_mode(&self, cid: &Cid, mode: PinMode) -> bool {
let cid_str = cid.to_string();
let state = self.state.read();
match mode {
PinMode::Direct => state.direct.contains_key(&cid_str),
PinMode::Recursive => state.recursive.contains_key(&cid_str),
PinMode::Indirect => state.indirect.contains_key(&cid_str),
}
}
pub fn add_direct(&mut self, cid: &Cid, name: Option<String>) {
let cid_str = cid.to_string();
let mut state = self.state.write();
state.direct.insert(cid_str, name);
}
pub fn remove_direct(&mut self, cid: &Cid) {
let cid_str = cid.to_string();
let mut state = self.state.write();
state.direct.remove(&cid_str);
}
pub fn add_recursive(&mut self, cid: &Cid, name: Option<String>) {
let cid_str = cid.to_string();
let mut state = self.state.write();
state.recursive.insert(cid_str, name);
}
pub fn remove_recursive(&mut self, cid: &Cid) {
let cid_str = cid.to_string();
let mut state = self.state.write();
state.recursive.remove(&cid_str);
}
pub fn add_indirect(&mut self, cid: &Cid, pinned_by: &Cid) {
let cid_str = cid.to_string();
let pinned_by_str = pinned_by.to_string();
let mut state = self.state.write();
state
.indirect
.entry(cid_str)
.or_default()
.insert(pinned_by_str);
}
pub fn remove_indirect(&mut self, cid: &Cid, pinned_by: &Cid) {
let cid_str = cid.to_string();
let pinned_by_str = pinned_by.to_string();
let mut state = self.state.write();
if let Some(refs) = state.indirect.get_mut(&cid_str) {
refs.remove(&pinned_by_str);
if refs.is_empty() {
state.indirect.remove(&cid_str);
}
}
}
pub fn get(&self, cid: &Cid) -> Option<PinInfo> {
let cid_str = cid.to_string();
let state = self.state.read();
if let Some(name) = state.direct.get(&cid_str) {
return Some(PinInfo {
cid: cid_str,
mode: PinMode::Direct,
name: name.clone(),
});
}
if let Some(name) = state.recursive.get(&cid_str) {
return Some(PinInfo {
cid: cid_str,
mode: PinMode::Recursive,
name: name.clone(),
});
}
if state.indirect.contains_key(&cid_str) {
return Some(PinInfo {
cid: cid_str,
mode: PinMode::Indirect,
name: None,
});
}
None
}
pub fn list(&self, mode: Option<PinMode>) -> Vec<PinInfo> {
let state = self.state.read();
let mut pins = Vec::new();
let include_direct = mode.is_none() || mode == Some(PinMode::Direct);
let include_recursive = mode.is_none() || mode == Some(PinMode::Recursive);
let include_indirect = mode.is_none() || mode == Some(PinMode::Indirect);
if include_direct {
for (cid, name) in &state.direct {
pins.push(PinInfo {
cid: cid.clone(),
mode: PinMode::Direct,
name: name.clone(),
});
}
}
if include_recursive {
for (cid, name) in &state.recursive {
pins.push(PinInfo {
cid: cid.clone(),
mode: PinMode::Recursive,
name: name.clone(),
});
}
}
if include_indirect {
for cid in state.indirect.keys() {
pins.push(PinInfo {
cid: cid.clone(),
mode: PinMode::Indirect,
name: None,
});
}
}
pins
}
pub fn counts(&self) -> (usize, usize, usize) {
let state = self.state.read();
(
state.direct.len(),
state.recursive.len(),
state.indirect.len(),
)
}
}
impl Default for PinStore {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use ferripfs_blockstore::create_cid_v0;
use tempfile::tempdir;
fn test_cid(data: &[u8]) -> Cid {
create_cid_v0(data).unwrap()
}
#[test]
fn test_direct_pin_operations() {
let mut store = PinStore::new();
let cid = test_cid(b"test");
assert!(!store.is_pinned(&cid));
store.add_direct(&cid, Some("test-pin".to_string()));
assert!(store.is_pinned(&cid));
assert!(store.is_pinned_with_mode(&cid, PinMode::Direct));
assert!(!store.is_pinned_with_mode(&cid, PinMode::Recursive));
let info = store.get(&cid).unwrap();
assert_eq!(info.mode, PinMode::Direct);
assert_eq!(info.name, Some("test-pin".to_string()));
store.remove_direct(&cid);
assert!(!store.is_pinned(&cid));
}
#[test]
fn test_indirect_pin_operations() {
let mut store = PinStore::new();
let cid = test_cid(b"child");
let parent = test_cid(b"parent");
store.add_indirect(&cid, &parent);
assert!(store.is_pinned(&cid));
assert!(store.is_pinned_with_mode(&cid, PinMode::Indirect));
store.remove_indirect(&cid, &parent);
assert!(!store.is_pinned(&cid));
}
#[test]
fn test_persistence() {
let dir = tempdir().unwrap();
let path = dir.path().join("pins.json");
let cid = test_cid(b"persistent");
{
let mut store = PinStore::open(&path).unwrap();
store.add_direct(&cid, Some("saved-pin".to_string()));
store.save().unwrap();
}
{
let store = PinStore::open(&path).unwrap();
assert!(store.is_pinned(&cid));
let info = store.get(&cid).unwrap();
assert_eq!(info.name, Some("saved-pin".to_string()));
}
}
#[test]
fn test_list_with_filter() {
let mut store = PinStore::new();
let cid1 = test_cid(b"direct");
let cid2 = test_cid(b"recursive");
let cid3 = test_cid(b"indirect");
let parent = test_cid(b"parent");
store.add_direct(&cid1, None);
store.add_recursive(&cid2, None);
store.add_indirect(&cid3, &parent);
assert_eq!(store.list(None).len(), 3);
assert_eq!(store.list(Some(PinMode::Direct)).len(), 1);
assert_eq!(store.list(Some(PinMode::Recursive)).len(), 1);
assert_eq!(store.list(Some(PinMode::Indirect)).len(), 1);
}
}