use std::collections::HashMap;
use std::fs;
use std::path::PathBuf;
use std::time::{Duration, SystemTime};
use serde::{Deserialize, Serialize};
use super::file_memory::{FileMemoryError, FileMemoryStore};
const LEDGER: &str = ".memory.json";
#[derive(Debug, Clone, Copy)]
pub struct ForgetConfig {
pub ttl: Duration,
pub min_importance: f64,
pub weak_grace: Duration,
}
impl Default for ForgetConfig {
fn default() -> Self {
Self {
ttl: Duration::from_secs(30 * 24 * 3600), min_importance: 0.5,
weak_grace: Duration::from_secs(30 * 24 * 3600), }
}
}
impl ForgetConfig {
pub fn new() -> Self {
Self::default()
}
pub fn with_ttl(mut self, ttl: Duration) -> Self {
self.ttl = ttl;
self
}
pub fn with_min_importance(mut self, min_importance: f64) -> Self {
self.min_importance = min_importance;
self
}
pub fn with_weak_grace(mut self, weak_grace: Duration) -> Self {
self.weak_grace = weak_grace;
self
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct MemoryMeta {
importance: f64,
created_at: u64,
last_access_at: u64,
}
impl MemoryMeta {
fn new(importance: f64, now: SystemTime) -> Self {
Self {
importance,
created_at: epoch_ms(now),
last_access_at: epoch_ms(now),
}
}
}
fn epoch_ms(t: SystemTime) -> u64 {
t.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64
}
fn from_epoch_ms(ms: u64) -> SystemTime {
std::time::UNIX_EPOCH + Duration::from_millis(ms)
}
pub struct ForgettingMemory {
files: FileMemoryStore,
config: ForgetConfig,
meta: HashMap<String, MemoryMeta>,
}
impl std::fmt::Debug for ForgettingMemory {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ForgettingMemory")
.field("root", &self.files.root())
.field("config", &self.config)
.field("live", &self.meta.len())
.finish_non_exhaustive()
}
}
impl ForgettingMemory {
pub fn new(files: FileMemoryStore, config: ForgetConfig) -> Result<Self, FileMemoryError> {
let meta = Self::load_ledger(files.root())?;
Ok(Self {
files,
config,
meta,
})
}
pub fn config(&self) -> ForgetConfig {
self.config
}
fn ledger_path(root: &std::path::Path) -> PathBuf {
root.join(LEDGER)
}
fn load_ledger(root: &std::path::Path) -> Result<HashMap<String, MemoryMeta>, FileMemoryError> {
let path = Self::ledger_path(root);
if !path.exists() {
return Ok(HashMap::new());
}
let raw = fs::read_to_string(&path)?;
serde_json::from_str(&raw).map_err(|e| FileMemoryError::UnsafeName {
name: LEDGER.to_string(),
reason: format!("malformed ledger: {e}"),
})
}
fn persist_ledger(&self) -> Result<(), FileMemoryError> {
let path = Self::ledger_path(self.files.root());
let raw =
serde_json::to_string_pretty(&self.meta).map_err(|e| FileMemoryError::UnsafeName {
name: LEDGER.to_string(),
reason: format!("serialize ledger: {e}"),
})?;
fs::write(path, raw)?;
Ok(())
}
pub fn remember(
&mut self,
name: &str,
content: &str,
importance: f64,
now: SystemTime,
) -> Result<(), FileMemoryError> {
if !(0.0..=1.0).contains(&importance) {
return Err(FileMemoryError::UnsafeName {
name: name.to_string(),
reason: format!("importance {importance} outside [0, 1]"),
});
}
self.files.write(name, content)?;
self.meta
.insert(name.to_string(), MemoryMeta::new(importance, now));
self.persist_ledger()?;
Ok(())
}
pub fn recall(&mut self, name: &str, now: SystemTime) -> Result<String, FileMemoryError> {
let content = self.files.view(name)?;
if let Some(meta) = self.meta.get_mut(name) {
meta.last_access_at = epoch_ms(now);
self.persist_ledger()?;
}
Ok(content)
}
pub fn importance(&self, name: &str) -> Option<f64> {
self.meta.get(name).map(|m| m.importance)
}
pub fn forget(&mut self, name: &str) -> Result<(), FileMemoryError> {
self.files.delete(name)?;
self.meta.remove(name);
self.persist_ledger()?;
Ok(())
}
fn idle(&self, name: &str, now: SystemTime) -> Option<Duration> {
let meta = self.meta.get(name)?;
let last = from_epoch_ms(meta.last_access_at);
now.duration_since(last).ok()
}
pub fn should_forget(&self, name: &str, now: SystemTime) -> bool {
let Some(idle) = self.idle(name, now) else {
return true; };
if idle >= self.config.ttl {
return true;
}
let weak = self
.meta
.get(name)
.is_none_or(|m| m.importance < self.config.min_importance);
weak && idle >= self.config.weak_grace
}
pub fn consolidate(&mut self, now: SystemTime) -> Result<usize, FileMemoryError> {
let doomed: Vec<String> = self
.meta
.keys()
.filter(|name| self.should_forget(name, now))
.cloned()
.collect();
let count = doomed.len();
for name in doomed {
let _ = self.files.delete(&name); self.meta.remove(&name);
}
if count > 0 {
self.persist_ledger()?;
}
Ok(count)
}
pub fn live_at(&self, now: SystemTime) -> Vec<String> {
let mut live: Vec<String> = self
.meta
.keys()
.filter(|name| !self.should_forget(name, now))
.cloned()
.collect();
live.sort();
live
}
}
#[cfg(test)]
mod tests {
use super::*;
const T0: std::time::SystemTime = std::time::UNIX_EPOCH;
fn store_and_config(
ttl_secs: u64,
min_imp: f64,
grace_secs: u64,
) -> (tempfile::TempDir, FileMemoryStore, ForgetConfig) {
let dir = tempfile::tempdir().unwrap();
let files = FileMemoryStore::new(dir.path()).unwrap();
let config = ForgetConfig::new()
.with_ttl(Duration::from_secs(ttl_secs))
.with_min_importance(min_imp)
.with_weak_grace(Duration::from_secs(grace_secs));
(dir, files, config)
}
#[test]
fn remember_recall_persists_content() {
let (_d, files, cfg) = store_and_config(100, 0.4, 100);
let mut m = ForgettingMemory::new(files, cfg).unwrap();
m.remember("site", "the docs live under /docs", 0.9, T0)
.unwrap();
assert_eq!(m.recall("site", T0).unwrap(), "the docs live under /docs");
assert_eq!(m.importance("site"), Some(0.9));
}
#[test]
fn active_memory_survives_before_ttl() {
let (_d, files, cfg) = store_and_config(100, 0.4, 100);
let mut m = ForgettingMemory::new(files, cfg).unwrap();
m.remember("a", "x", 0.9, T0).unwrap();
let later = T0 + Duration::from_secs(90); assert!(!m.should_forget("a", later));
assert_eq!(m.consolidate(later).unwrap(), 0);
}
#[test]
fn idle_memory_forgotten_at_ttl() {
let (_d, files, cfg) = store_and_config(100, 0.4, 100);
let mut m = ForgettingMemory::new(files, cfg).unwrap();
m.remember("a", "x", 0.9, T0).unwrap();
let past = T0 + Duration::from_secs(101);
assert!(m.should_forget("a", past));
assert_eq!(m.consolidate(past).unwrap(), 1);
assert!(m.live_at(past).is_empty());
assert!(matches!(
m.files.view("a"),
Err(FileMemoryError::NotFound(_))
));
}
#[test]
fn recall_refreshes_lifespan() {
let (_d, files, cfg) = store_and_config(100, 0.4, 100);
let mut m = ForgettingMemory::new(files, cfg).unwrap();
m.remember("a", "x", 0.9, T0).unwrap();
let _ = m.recall("a", T0 + Duration::from_secs(80)).unwrap();
let later = T0 + Duration::from_secs(150);
assert!(!m.should_forget("a", later));
}
#[test]
fn weak_memory_pruned_within_ttl_after_grace() {
let (_d, files, cfg) = store_and_config(1000, 0.5, 50);
let mut m = ForgettingMemory::new(files, cfg).unwrap();
m.remember("weak", "trivia", 0.1, T0).unwrap();
let later = T0 + Duration::from_secs(60);
assert!(m.should_forget("weak", later));
assert_eq!(m.consolidate(later).unwrap(), 1);
}
#[test]
fn importance_out_of_range_rejected() {
let (_d, files, cfg) = store_and_config(100, 0.4, 100);
let mut m = ForgettingMemory::new(files, cfg).unwrap();
assert!(m.remember("a", "x", 1.5, T0).is_err());
assert!(m.remember("b", "x", -0.1, T0).is_err());
}
#[test]
fn ledger_survives_restart() {
let (dir, files, cfg) = store_and_config(1000, 0.5, 1000);
{
let mut m = ForgettingMemory::new(files, cfg).unwrap();
m.remember("a", "content", 0.8, T0).unwrap();
}
let files2 = FileMemoryStore::new(dir.path()).unwrap();
let m2 = ForgettingMemory::new(files2, cfg).unwrap();
assert!(m2.live_at(T0).contains(&"a".to_string()));
assert_eq!(m2.importance("a"), Some(0.8));
}
}