1use std::collections::HashMap;
4use std::sync::atomic::{AtomicU64, Ordering};
5use std::time::{SystemTime, UNIX_EPOCH};
6
7use serde::{Deserialize, Serialize};
8
9use crate::types::{AgentId, MemoryId, Timestamp};
10
11#[derive(Debug, Clone, Serialize, Deserialize)]
13pub struct Version {
14 pub version_id: u64,
15 pub memory_id: MemoryId,
16 pub agent_id: AgentId,
17 pub timestamp: Timestamp,
18 pub data_hash: u64,
19}
20
21#[derive(Debug)]
23pub struct VersionStore {
24 versions: HashMap<MemoryId, Vec<Version>>,
25 next_version: AtomicU64,
26}
27
28impl Default for VersionStore {
29 fn default() -> Self {
30 Self {
31 versions: HashMap::new(),
32 next_version: AtomicU64::new(1),
33 }
34 }
35}
36
37fn now_micros() -> Timestamp {
38 SystemTime::now()
39 .duration_since(UNIX_EPOCH)
40 .unwrap_or_default()
41 .as_micros() as Timestamp
42}
43
44impl VersionStore {
45 pub fn new() -> Self {
46 Self::default()
47 }
48
49 pub fn record_write(&mut self, memory_id: MemoryId, agent_id: AgentId, data_hash: u64) -> u64 {
51 let vid = self.next_version.fetch_add(1, Ordering::Relaxed);
52 let version = Version {
53 version_id: vid,
54 memory_id,
55 agent_id,
56 timestamp: now_micros(),
57 data_hash,
58 };
59 self.versions.entry(memory_id).or_default().push(version);
60 vid
61 }
62
63 pub fn get_latest(&self, memory_id: MemoryId) -> Option<&Version> {
65 self.versions.get(&memory_id).and_then(|v| v.last())
66 }
67
68 pub fn get_history(&self, memory_id: MemoryId) -> Vec<&Version> {
70 self.versions
71 .get(&memory_id)
72 .map(|v| v.iter().collect())
73 .unwrap_or_default()
74 }
75
76 pub fn get_version_at(&self, memory_id: MemoryId, timestamp: Timestamp) -> Option<&Version> {
78 self.versions
79 .get(&memory_id)
80 .and_then(|v| v.iter().rev().find(|ver| ver.timestamp <= timestamp))
81 }
82}
83
84#[cfg(test)]
85mod tests {
86 use super::*;
87 use uuid::Uuid;
88
89 #[test]
90 fn record_and_get_latest() {
91 let mut store = VersionStore::new();
92 let mid = Uuid::new_v4();
93 let aid = Uuid::new_v4();
94 store.record_write(mid, aid, 111);
95 store.record_write(mid, aid, 222);
96 assert_eq!(store.get_latest(mid).unwrap().data_hash, 222);
97 }
98
99 #[test]
100 fn version_ids_increment() {
101 let mut store = VersionStore::new();
102 let mid = Uuid::new_v4();
103 let aid = Uuid::new_v4();
104 let v1 = store.record_write(mid, aid, 1);
105 let v2 = store.record_write(mid, aid, 2);
106 assert_eq!(v2, v1 + 1);
107 }
108
109 #[test]
110 fn get_history() {
111 let mut store = VersionStore::new();
112 let mid = Uuid::new_v4();
113 let aid = Uuid::new_v4();
114 store.record_write(mid, aid, 10);
115 store.record_write(mid, aid, 20);
116 store.record_write(mid, aid, 30);
117 assert_eq!(store.get_history(mid).len(), 3);
118 }
119
120 #[test]
121 fn get_version_at() {
122 let mut store = VersionStore::new();
123 let mid = Uuid::new_v4();
124 let aid = Uuid::new_v4();
125 store.record_write(mid, aid, 1);
126 let ver = store.get_version_at(mid, u64::MAX);
128 assert!(ver.is_some());
129 }
130
131 #[test]
132 fn empty_history() {
133 let store = VersionStore::new();
134 assert!(store.get_latest(Uuid::new_v4()).is_none());
135 assert!(store.get_history(Uuid::new_v4()).is_empty());
136 }
137}