Added some features

Author: anaslimem

src/core/state_machine.rs

@@ -55,6 +55,19 @@ impl StateMachine {
         let id = entry.id;
         let timestamp = entry.created_at;
 
+        // If this ID already exists, remove its previous temporal index entry first.
+        if let Some(previous) = self.memories.get(&id) {
+            let old_ts = previous.created_at;
+            let mut remove_old_bucket = false;
+            if let Some(ids) = self.temporal_index.get_mut(&old_ts) {
+                ids.retain(|&mid| mid != id);
+                remove_old_bucket = ids.is_empty();
+            }
+            if remove_old_bucket {
+                self.temporal_index.remove(&old_ts);
+            }
+        }
+
         self.memories.insert(id, entry);
 
         // Add to temporal index
@@ -170,6 +183,13 @@ impl StateMachine {
         self.memories.len()
     }
 
+    /// Get all memories in deterministic ID order.
+    pub fn all_memories(&self) -> Vec<&MemoryEntry> {
+        let mut entries: Vec<&MemoryEntry> = self.memories.values().collect();
+        entries.sort_by_key(|e| e.id);
+        entries
+    }
+
     pub fn is_empty(&self) -> bool {
         self.memories.is_empty()
     }
@@ -262,6 +282,38 @@ mod tests {
         assert_eq!(range[1].id, MemoryId(3));
     }
 
+    #[test]
+    fn test_insert_update_replaces_old_temporal_timestamp() {
+        let mut sm = StateMachine::new();
+        sm.insert_memory(create_test_entry(1, "default", 1000))
+            .unwrap();
+        sm.insert_memory(create_test_entry(1, "default", 3000))
+            .unwrap();
+
+        assert_eq!(sm.len(), 1);
+        assert!(sm.get_memories_in_time_range(1000, 1000).is_empty());
+        let new_range = sm.get_memories_in_time_range(3000, 3000);
+        assert_eq!(new_range.len(), 1);
+        assert_eq!(new_range[0].id, MemoryId(1));
+    }
+
+    #[test]
+    fn test_all_memories_deterministic_order() {
+        let mut sm = StateMachine::new();
+        sm.insert_memory(create_test_entry(3, "default", 1000))
+            .unwrap();
+        sm.insert_memory(create_test_entry(1, "default", 1000))
+            .unwrap();
+        sm.insert_memory(create_test_entry(2, "default", 1000))
+            .unwrap();
+
+        let all = sm.all_memories();
+        assert_eq!(all.len(), 3);
+        assert_eq!(all[0].id, MemoryId(1));
+        assert_eq!(all[1].id, MemoryId(2));
+        assert_eq!(all[2].id, MemoryId(3));
+    }
+
     #[test]
     fn test_namespace_filtering() {
         let mut sm = StateMachine::new();

src/engine.rs

@@ -204,7 +204,7 @@ impl Engine {
         let mut unique_ids = std::collections::HashSet::new();
         let mut candidates: Vec<(MemoryId, u64, f32)> = self
             .state_machine
-            .get_memories_in_time_range(0, u64::MAX)
+            .all_memories()
             .into_iter()
             .filter(|entry| unique_ids.insert(entry.id))
             .map(|entry| (entry.id, entry.created_at, entry.importance))
@@ -240,7 +240,7 @@ impl Engine {
         let entries = self.state_machine.len();
         let bytes = self
             .state_machine
-            .get_memories_in_time_range(0, u64::MAX)
+            .all_memories()
             .into_iter()
             .map(Self::estimate_memory_bytes)
             .sum();

src/index/temporal.rs

@@ -77,7 +77,7 @@ impl TemporalIndex {
     pub fn get_earliest_timestamp(state_machine: &StateMachine) -> Option<u64> {
         // Find first memory in the system (brute force for now)
         let mut earliest = u64::MAX;
-        for entry in state_machine.get_memories_in_time_range(0, u64::MAX) {
+        for entry in state_machine.all_memories() {
             if entry.created_at < earliest {
                 earliest = entry.created_at;
             }
@@ -94,7 +94,7 @@ impl TemporalIndex {
     pub fn get_latest_timestamp(state_machine: &StateMachine) -> Option<u64> {
         // Find last memory in the system
         let mut latest: Option<u64> = None;
-        for entry in state_machine.get_memories_in_time_range(0, u64::MAX) {
+        for entry in state_machine.all_memories() {
             latest = Some(match latest {
                 Some(current) => current.max(entry.created_at),
                 None => entry.created_at,
@@ -144,7 +144,7 @@ impl TemporalIndex {
     pub fn get_oldest(state_machine: &StateMachine, count: usize) -> Result<Vec<MemoryId>> {
         // Get all memories and sort by timestamp
         let mut all = Vec::new();
-        for entry in state_machine.get_memories_in_time_range(0, u64::MAX) {
+        for entry in state_machine.all_memories() {
             all.push((entry.id, entry.created_at));
         }
 
@@ -165,7 +165,7 @@ impl TemporalIndex {
     pub fn get_newest(state_machine: &StateMachine, count: usize) -> Result<Vec<MemoryId>> {
         // Get all memories and sort by timestamp
         let mut all = Vec::new();
-        for entry in state_machine.get_memories_in_time_range(0, u64::MAX) {
+        for entry in state_machine.all_memories() {
             all.push((entry.id, entry.created_at));
         }
 

src/store.rs

@@ -110,11 +110,7 @@ impl MnemosStore {
     pub fn rebuild_vector_index(&mut self) -> Result<usize> {
         self.indexes = IndexLayer::new(self.indexes.vector.dimension());
         let mut indexed = 0usize;
-        for entry in self
-            .engine
-            .get_state_machine()
-            .get_memories_in_time_range(0, u64::MAX)
-        {
+        for entry in self.engine.get_state_machine().all_memories() {
             if let Some(embedding) = entry.embedding.clone() {
                 self.indexes.vector_index_mut().index(entry.id, embedding)?;
                 indexed += 1;