umi_memory/memory/
working.rs

1//! Working Memory - Session-Scoped KV Store with TTL
2//!
3//! TigerStyle: Explicit limits, TTL expiration, simulation-first testing.
4//!
5//! # Design
6//!
7//! Working memory is a bounded KV store (~1MB) for session state.
8//! Entries expire after TTL (default 1 hour).
9//!
10//! # Simulation-First
11//!
12//! Tests are written BEFORE implementation. This file starts with tests
13//! and minimal stubs. Implementation follows to make tests pass.
14
15use std::collections::HashMap;
16
17use crate::constants::{
18    WORKING_MEMORY_ENTRIES_COUNT_MAX, WORKING_MEMORY_ENTRY_SIZE_BYTES_MAX,
19    WORKING_MEMORY_SIZE_BYTES_MAX, WORKING_MEMORY_TTL_SECS_DEFAULT,
20};
21
22// =============================================================================
23// Error Types
24// =============================================================================
25
26/// Errors from working memory operations.
27#[derive(Debug, Clone, thiserror::Error)]
28pub enum WorkingMemoryError {
29    /// Entry too large
30    #[error("entry too large: {size_bytes} bytes exceeds max {max_bytes}")]
31    EntryTooLarge {
32        /// Size of the entry
33        size_bytes: usize,
34        /// Maximum allowed
35        max_bytes: usize,
36    },
37
38    /// Memory full
39    #[error("working memory full: {current_bytes}/{max_bytes} bytes")]
40    MemoryFull {
41        /// Current used bytes
42        current_bytes: usize,
43        /// Maximum allowed bytes
44        max_bytes: usize,
45    },
46
47    /// Too many entries
48    #[error("too many entries: {count} exceeds max {max_count}")]
49    TooManyEntries {
50        /// Current entry count
51        count: usize,
52        /// Maximum allowed
53        max_count: usize,
54    },
55
56    /// Key too long
57    #[error("key too long: {len} bytes exceeds max {max_len}")]
58    KeyTooLong {
59        /// Key length
60        len: usize,
61        /// Maximum allowed
62        max_len: usize,
63    },
64
65    /// TTL too long
66    #[error("TTL too long: {ttl_secs} seconds exceeds max {max_secs}")]
67    TtlTooLong {
68        /// Requested TTL
69        ttl_secs: u64,
70        /// Maximum allowed
71        max_secs: u64,
72    },
73}
74
75/// Result type for working memory operations.
76pub type WorkingMemoryResult<T> = Result<T, WorkingMemoryError>;
77
78// =============================================================================
79// Configuration
80// =============================================================================
81
82/// Configuration for working memory.
83#[derive(Debug, Clone)]
84pub struct WorkingMemoryConfig {
85    /// Maximum total size in bytes
86    pub max_bytes: usize,
87    /// Maximum number of entries
88    pub max_entries: usize,
89    /// Default TTL in milliseconds
90    pub default_ttl_ms: u64,
91    /// Maximum key length
92    pub max_key_len: usize,
93}
94
95impl Default for WorkingMemoryConfig {
96    fn default() -> Self {
97        Self {
98            max_bytes: WORKING_MEMORY_SIZE_BYTES_MAX,
99            max_entries: WORKING_MEMORY_ENTRIES_COUNT_MAX,
100            default_ttl_ms: WORKING_MEMORY_TTL_SECS_DEFAULT * 1000,
101            max_key_len: 256,
102        }
103    }
104}
105
106// =============================================================================
107// Entry Type
108// =============================================================================
109
110/// A single entry in working memory.
111#[derive(Debug, Clone)]
112struct Entry {
113    /// The value bytes
114    value: Vec<u8>,
115    /// Size in bytes (cached for quick access)
116    size_bytes: usize,
117    /// Creation timestamp (ms since epoch) - kept for debugging/stats
118    #[allow(dead_code)]
119    created_at_ms: u64,
120    /// Expiration timestamp (ms since epoch)
121    expires_at_ms: u64,
122}
123
124// =============================================================================
125// Working Memory
126// =============================================================================
127
128/// Working memory - session-scoped KV store with TTL.
129///
130/// TigerStyle:
131/// - Bounded capacity (~1MB)
132/// - TTL-based expiration
133/// - Explicit size tracking
134/// - DST-compatible via set_clock_ms()
135#[derive(Debug)]
136pub struct WorkingMemory {
137    /// Configuration
138    config: WorkingMemoryConfig,
139    /// Entries indexed by key
140    entries: HashMap<String, Entry>,
141    /// Current total size in bytes (keys + values)
142    current_bytes: usize,
143    /// Clock source for timestamps (ms since epoch)
144    clock_ms: u64,
145}
146
147impl WorkingMemory {
148    /// Create a new working memory with default configuration.
149    #[must_use]
150    pub fn new() -> Self {
151        Self::with_config(WorkingMemoryConfig::default())
152    }
153
154    /// Create a new working memory with custom configuration.
155    #[must_use]
156    pub fn with_config(config: WorkingMemoryConfig) -> Self {
157        Self {
158            config,
159            entries: HashMap::new(),
160            current_bytes: 0,
161            clock_ms: 0,
162        }
163    }
164
165    /// Set the internal clock (for DST).
166    ///
167    /// TigerStyle: Explicit time control for simulation.
168    pub fn set_clock_ms(&mut self, ms: u64) {
169        self.clock_ms = ms;
170    }
171
172    /// Get the internal clock value.
173    #[must_use]
174    pub fn clock_ms(&self) -> u64 {
175        self.clock_ms
176    }
177
178    /// Set an entry with optional TTL.
179    ///
180    /// If TTL is None, uses the default TTL from config.
181    /// If key already exists, the entry is replaced and TTL is reset.
182    ///
183    /// # Errors
184    /// Returns error if entry is too large, memory is full, or too many entries.
185    pub fn set(&mut self, key: &str, value: &[u8], ttl_ms: Option<u64>) -> WorkingMemoryResult<()> {
186        let value_len = value.len();
187        let key_len = key.len();
188        let entry_size = key_len + value_len;
189
190        // Precondition: key length
191        if key_len > self.config.max_key_len {
192            return Err(WorkingMemoryError::KeyTooLong {
193                len: key_len,
194                max_len: self.config.max_key_len,
195            });
196        }
197
198        // Precondition: entry size
199        if value_len > WORKING_MEMORY_ENTRY_SIZE_BYTES_MAX {
200            return Err(WorkingMemoryError::EntryTooLarge {
201                size_bytes: value_len,
202                max_bytes: WORKING_MEMORY_ENTRY_SIZE_BYTES_MAX,
203            });
204        }
205
206        // Calculate size delta (account for existing entry if overwriting)
207        let old_size = self
208            .entries
209            .get(key)
210            .map(|e| key_len + e.size_bytes)
211            .unwrap_or(0);
212        let projected_size = self.current_bytes - old_size + entry_size;
213
214        // Check capacity
215        if projected_size > self.config.max_bytes {
216            return Err(WorkingMemoryError::MemoryFull {
217                current_bytes: self.current_bytes,
218                max_bytes: self.config.max_bytes,
219            });
220        }
221
222        // Check entry count (only if not overwriting)
223        let is_new_key = !self.entries.contains_key(key);
224        if is_new_key && self.entries.len() >= self.config.max_entries {
225            return Err(WorkingMemoryError::TooManyEntries {
226                count: self.entries.len(),
227                max_count: self.config.max_entries,
228            });
229        }
230
231        // Calculate TTL
232        let ttl = ttl_ms.unwrap_or(self.config.default_ttl_ms);
233        let expires_at_ms = self.clock_ms.saturating_add(ttl);
234
235        // Create entry
236        let entry = Entry {
237            value: value.to_vec(),
238            size_bytes: value_len,
239            created_at_ms: self.clock_ms,
240            expires_at_ms,
241        };
242
243        self.entries.insert(key.to_string(), entry);
244        self.current_bytes = projected_size;
245
246        // Postcondition
247        assert!(
248            self.current_bytes <= self.config.max_bytes,
249            "size invariant violated"
250        );
251
252        Ok(())
253    }
254
255    /// Get an entry by key.
256    ///
257    /// Returns None if key doesn't exist or entry is expired.
258    #[must_use]
259    pub fn get(&self, key: &str) -> Option<&[u8]> {
260        self.entries.get(key).and_then(|entry| {
261            if entry.expires_at_ms > self.clock_ms {
262                Some(entry.value.as_slice())
263            } else {
264                None // Expired
265            }
266        })
267    }
268
269    /// Delete an entry by key.
270    ///
271    /// Returns true if entry existed (even if expired), false otherwise.
272    pub fn delete(&mut self, key: &str) -> bool {
273        if let Some(entry) = self.entries.remove(key) {
274            let entry_size = key.len() + entry.size_bytes;
275            self.current_bytes = self.current_bytes.saturating_sub(entry_size);
276            true
277        } else {
278            false
279        }
280    }
281
282    /// Check if a key exists (and is not expired).
283    #[must_use]
284    pub fn exists(&self, key: &str) -> bool {
285        self.entries
286            .get(key)
287            .map(|entry| entry.expires_at_ms > self.clock_ms)
288            .unwrap_or(false)
289    }
290
291    /// Remove all expired entries.
292    ///
293    /// Returns the number of entries removed.
294    pub fn cleanup_expired(&mut self) -> usize {
295        let clock = self.clock_ms;
296        let expired_keys: Vec<String> = self
297            .entries
298            .iter()
299            .filter(|(_, entry)| entry.expires_at_ms <= clock)
300            .map(|(key, _)| key.clone())
301            .collect();
302
303        let count = expired_keys.len();
304        for key in expired_keys {
305            self.delete(&key);
306        }
307
308        count
309    }
310
311    /// Get used bytes (keys + values).
312    #[must_use]
313    pub fn used_bytes(&self) -> usize {
314        self.current_bytes
315    }
316
317    /// Get available bytes.
318    #[must_use]
319    pub fn available_bytes(&self) -> usize {
320        self.config.max_bytes.saturating_sub(self.current_bytes)
321    }
322
323    /// Get entry count (including expired).
324    #[must_use]
325    pub fn entry_count(&self) -> usize {
326        self.entries.len()
327    }
328
329    /// Check if empty.
330    #[must_use]
331    pub fn is_empty(&self) -> bool {
332        self.entries.is_empty()
333    }
334
335    /// Clear all entries.
336    pub fn clear(&mut self) {
337        self.entries.clear();
338        self.current_bytes = 0;
339    }
340
341    /// Get configuration.
342    #[must_use]
343    pub fn config(&self) -> &WorkingMemoryConfig {
344        &self.config
345    }
346}
347
348impl Default for WorkingMemory {
349    fn default() -> Self {
350        Self::new()
351    }
352}
353
354// =============================================================================
355// TESTS - Written FIRST (Simulation-First)
356// =============================================================================
357
358#[cfg(test)]
359mod tests {
360    use super::*;
361
362    // =========================================================================
363    // Basic CRUD Tests
364    // =========================================================================
365
366    #[test]
367    fn test_new_working_memory() {
368        let wm = WorkingMemory::new();
369        assert_eq!(wm.used_bytes(), 0);
370        assert_eq!(wm.entry_count(), 0);
371        assert!(wm.is_empty());
372    }
373
374    #[test]
375    fn test_set_and_get() {
376        let mut wm = WorkingMemory::new();
377        wm.set("key1", b"value1", None).unwrap();
378
379        assert_eq!(wm.get("key1"), Some(b"value1".as_slice()));
380        assert!(wm.exists("key1"));
381        assert_eq!(wm.entry_count(), 1);
382    }
383
384    #[test]
385    fn test_set_overwrites() {
386        let mut wm = WorkingMemory::new();
387        wm.set("key1", b"value1", None).unwrap();
388        wm.set("key1", b"new_value", None).unwrap();
389
390        assert_eq!(wm.get("key1"), Some(b"new_value".as_slice()));
391        assert_eq!(wm.entry_count(), 1);
392    }
393
394    #[test]
395    fn test_get_nonexistent() {
396        let wm = WorkingMemory::new();
397        assert_eq!(wm.get("nonexistent"), None);
398        assert!(!wm.exists("nonexistent"));
399    }
400
401    #[test]
402    fn test_delete() {
403        let mut wm = WorkingMemory::new();
404        wm.set("key1", b"value1", None).unwrap();
405
406        assert!(wm.delete("key1"));
407        assert_eq!(wm.get("key1"), None);
408        assert!(!wm.exists("key1"));
409        assert_eq!(wm.entry_count(), 0);
410    }
411
412    #[test]
413    fn test_delete_nonexistent() {
414        let mut wm = WorkingMemory::new();
415        assert!(!wm.delete("nonexistent"));
416    }
417
418    #[test]
419    fn test_clear() {
420        let mut wm = WorkingMemory::new();
421        wm.set("key1", b"value1", None).unwrap();
422        wm.set("key2", b"value2", None).unwrap();
423
424        wm.clear();
425
426        assert!(wm.is_empty());
427        assert_eq!(wm.used_bytes(), 0);
428    }
429
430    // =========================================================================
431    // Size Tracking Tests
432    // =========================================================================
433
434    #[test]
435    fn test_size_tracking() {
436        let mut wm = WorkingMemory::new();
437
438        // key1 (4 bytes) + value1 (6 bytes) = 10 bytes
439        wm.set("key1", b"value1", None).unwrap();
440        assert_eq!(wm.used_bytes(), 10);
441
442        // key2 (4 bytes) + value2 (6 bytes) = 10 bytes, total = 20
443        wm.set("key2", b"value2", None).unwrap();
444        assert_eq!(wm.used_bytes(), 20);
445
446        // Delete key1, should free 10 bytes
447        wm.delete("key1");
448        assert_eq!(wm.used_bytes(), 10);
449    }
450
451    #[test]
452    fn test_overwrite_size_tracking() {
453        let mut wm = WorkingMemory::new();
454
455        // key1 (4) + short (5) = 9 bytes
456        wm.set("key1", b"short", None).unwrap();
457        assert_eq!(wm.used_bytes(), 9);
458
459        // key1 (4) + much_longer_value (17) = 21 bytes
460        wm.set("key1", b"much_longer_value", None).unwrap();
461        assert_eq!(wm.used_bytes(), 21);
462    }
463
464    // =========================================================================
465    // Capacity Limit Tests
466    // =========================================================================
467
468    #[test]
469    fn test_entry_too_large() {
470        let mut wm = WorkingMemory::new();
471        let large_value = vec![0u8; WORKING_MEMORY_ENTRY_SIZE_BYTES_MAX + 1];
472
473        let result = wm.set("key", &large_value, None);
474        assert!(matches!(
475            result,
476            Err(WorkingMemoryError::EntryTooLarge { .. })
477        ));
478    }
479
480    #[test]
481    fn test_memory_full() {
482        let config = WorkingMemoryConfig {
483            max_bytes: 100, // Very small for testing
484            ..Default::default()
485        };
486        let mut wm = WorkingMemory::with_config(config);
487
488        // Fill up most of the space
489        wm.set("key1", &vec![0u8; 80], None).unwrap();
490
491        // This should fail - not enough space
492        let result = wm.set("key2", &vec![0u8; 50], None);
493        assert!(matches!(result, Err(WorkingMemoryError::MemoryFull { .. })));
494    }
495
496    #[test]
497    fn test_too_many_entries() {
498        let config = WorkingMemoryConfig {
499            max_entries: 3,
500            ..Default::default()
501        };
502        let mut wm = WorkingMemory::with_config(config);
503
504        wm.set("key1", b"v", None).unwrap();
505        wm.set("key2", b"v", None).unwrap();
506        wm.set("key3", b"v", None).unwrap();
507
508        let result = wm.set("key4", b"v", None);
509        assert!(matches!(
510            result,
511            Err(WorkingMemoryError::TooManyEntries { .. })
512        ));
513    }
514
515    #[test]
516    fn test_overwrite_does_not_increase_entry_count() {
517        let config = WorkingMemoryConfig {
518            max_entries: 2,
519            ..Default::default()
520        };
521        let mut wm = WorkingMemory::with_config(config);
522
523        wm.set("key1", b"v", None).unwrap();
524        wm.set("key2", b"v", None).unwrap();
525
526        // This should succeed - overwriting existing key
527        wm.set("key1", b"new_value", None).unwrap();
528        assert_eq!(wm.entry_count(), 2);
529    }
530}
531
532// =============================================================================
533// DST Tests - Use SimClock for TTL testing
534// =============================================================================
535
536#[cfg(test)]
537mod dst_tests {
538    use super::*;
539    use crate::dst::{SimConfig, Simulation};
540
541    #[tokio::test]
542    async fn test_ttl_expiration() {
543        let sim = Simulation::new(SimConfig::with_seed(42));
544
545        sim.run(|env| async move {
546            let mut wm = WorkingMemory::new();
547            wm.set_clock_ms(env.clock.now_ms());
548
549            // Set entry with 1 second TTL
550            wm.set("key1", b"value1", Some(1000)).unwrap();
551            assert!(wm.exists("key1"));
552            assert_eq!(wm.get("key1"), Some(b"value1".as_slice()));
553
554            // Advance time by 500ms - not expired yet
555            env.clock.advance_ms(500);
556            wm.set_clock_ms(env.clock.now_ms());
557            assert!(wm.exists("key1"));
558
559            // Advance time by 600ms (total 1100ms) - now expired
560            env.clock.advance_ms(600);
561            wm.set_clock_ms(env.clock.now_ms());
562            assert!(!wm.exists("key1"));
563            assert_eq!(wm.get("key1"), None);
564
565            Ok::<(), std::convert::Infallible>(())
566        })
567        .await
568        .unwrap();
569    }
570
571    #[tokio::test]
572    async fn test_overwrite_resets_ttl() {
573        let sim = Simulation::new(SimConfig::with_seed(42));
574
575        sim.run(|env| async move {
576            let mut wm = WorkingMemory::new();
577            wm.set_clock_ms(env.clock.now_ms());
578
579            // Set entry with 1 second TTL
580            wm.set("key1", b"value1", Some(1000)).unwrap();
581
582            // Advance time by 800ms
583            env.clock.advance_ms(800);
584            wm.set_clock_ms(env.clock.now_ms());
585
586            // Overwrite - should reset TTL
587            wm.set("key1", b"new_value", Some(1000)).unwrap();
588
589            // Advance time by 800ms (total 1600ms from start, but only 800ms from overwrite)
590            env.clock.advance_ms(800);
591            wm.set_clock_ms(env.clock.now_ms());
592
593            // Should still exist because TTL was reset
594            assert!(wm.exists("key1"));
595
596            // Advance another 300ms (1100ms from overwrite) - now expired
597            env.clock.advance_ms(300);
598            wm.set_clock_ms(env.clock.now_ms());
599            assert!(!wm.exists("key1"));
600
601            Ok::<(), std::convert::Infallible>(())
602        })
603        .await
604        .unwrap();
605    }
606
607    #[tokio::test]
608    async fn test_cleanup_expired() {
609        let sim = Simulation::new(SimConfig::with_seed(42));
610
611        sim.run(|env| async move {
612            let mut wm = WorkingMemory::new();
613            wm.set_clock_ms(env.clock.now_ms());
614
615            // Set entries with different TTLs
616            wm.set("short", b"v", Some(500)).unwrap(); // 500ms TTL
617            wm.set("medium", b"v", Some(1000)).unwrap(); // 1s TTL
618            wm.set("long", b"v", Some(2000)).unwrap(); // 2s TTL
619
620            assert_eq!(wm.entry_count(), 3);
621
622            // Advance 600ms - "short" expired
623            env.clock.advance_ms(600);
624            wm.set_clock_ms(env.clock.now_ms());
625
626            let removed = wm.cleanup_expired();
627            assert_eq!(removed, 1);
628            assert_eq!(wm.entry_count(), 2);
629
630            // Advance another 500ms (1100ms total) - "medium" expired
631            env.clock.advance_ms(500);
632            wm.set_clock_ms(env.clock.now_ms());
633
634            let removed = wm.cleanup_expired();
635            assert_eq!(removed, 1);
636            assert_eq!(wm.entry_count(), 1);
637
638            // "long" should still exist
639            assert!(wm.exists("long"));
640
641            Ok::<(), std::convert::Infallible>(())
642        })
643        .await
644        .unwrap();
645    }
646
647    #[tokio::test]
648    async fn test_cleanup_frees_memory() {
649        let sim = Simulation::new(SimConfig::with_seed(42));
650
651        sim.run(|env| async move {
652            let mut wm = WorkingMemory::new();
653            wm.set_clock_ms(env.clock.now_ms());
654
655            // Set entry
656            wm.set("key1", b"value1", Some(500)).unwrap();
657            let used_before = wm.used_bytes();
658            assert!(used_before > 0);
659
660            // Expire it
661            env.clock.advance_ms(600);
662            wm.set_clock_ms(env.clock.now_ms());
663
664            wm.cleanup_expired();
665            assert_eq!(wm.used_bytes(), 0);
666
667            Ok::<(), std::convert::Infallible>(())
668        })
669        .await
670        .unwrap();
671    }
672
673    #[tokio::test]
674    async fn test_default_ttl() {
675        let sim = Simulation::new(SimConfig::with_seed(42));
676
677        sim.run(|env| async move {
678            let config = WorkingMemoryConfig {
679                default_ttl_ms: 1000, // 1 second default
680                ..Default::default()
681            };
682            let mut wm = WorkingMemory::with_config(config);
683            wm.set_clock_ms(env.clock.now_ms());
684
685            // Set without explicit TTL - should use default
686            wm.set("key1", b"value1", None).unwrap();
687
688            // After 900ms - still exists
689            env.clock.advance_ms(900);
690            wm.set_clock_ms(env.clock.now_ms());
691            assert!(wm.exists("key1"));
692
693            // After 200ms more (1100ms total) - expired
694            env.clock.advance_ms(200);
695            wm.set_clock_ms(env.clock.now_ms());
696            assert!(!wm.exists("key1"));
697
698            Ok::<(), std::convert::Infallible>(())
699        })
700        .await
701        .unwrap();
702    }
703
704    #[tokio::test]
705    async fn test_multiple_entries_different_ttls() {
706        let sim = Simulation::new(SimConfig::with_seed(42));
707
708        sim.run(|env| async move {
709            let mut wm = WorkingMemory::new();
710            wm.set_clock_ms(env.clock.now_ms());
711
712            // Set entries at different times with same TTL
713            wm.set("first", b"v", Some(1000)).unwrap();
714
715            env.clock.advance_ms(300);
716            wm.set_clock_ms(env.clock.now_ms());
717            wm.set("second", b"v", Some(1000)).unwrap();
718
719            env.clock.advance_ms(300);
720            wm.set_clock_ms(env.clock.now_ms());
721            wm.set("third", b"v", Some(1000)).unwrap();
722
723            // Now at 600ms from start
724            // first expires at 1000ms, second at 1300ms, third at 1600ms
725
726            // At 800ms - all exist
727            env.clock.advance_ms(200);
728            wm.set_clock_ms(env.clock.now_ms());
729            assert!(wm.exists("first"));
730            assert!(wm.exists("second"));
731            assert!(wm.exists("third"));
732
733            // At 1100ms - first expired
734            env.clock.advance_ms(300);
735            wm.set_clock_ms(env.clock.now_ms());
736            assert!(!wm.exists("first"));
737            assert!(wm.exists("second"));
738            assert!(wm.exists("third"));
739
740            Ok::<(), std::convert::Infallible>(())
741        })
742        .await
743        .unwrap();
744    }
745}
746
747// =============================================================================
748// Property-Based Tests
749// =============================================================================
750
751#[cfg(test)]
752mod property_tests {
753    use super::*;
754    use crate::dst::{
755        DeterministicRng, PropertyTest, PropertyTestable, SimClock, TimeAdvanceConfig,
756    };
757
758    /// Operations that can be performed on WorkingMemory
759    #[derive(Debug, Clone)]
760    enum WorkingMemoryOp {
761        Set {
762            key: String,
763            value_len: usize,
764            ttl_ms: u64,
765        },
766        Get {
767            key: String,
768        },
769        Delete {
770            key: String,
771        },
772        Cleanup,
773    }
774
775    /// Wrapper for property testing
776    struct WorkingMemoryWrapper {
777        inner: WorkingMemory,
778        /// Track keys we've set (for generating realistic ops)
779        known_keys: Vec<String>,
780    }
781
782    impl PropertyTestable for WorkingMemoryWrapper {
783        type Operation = WorkingMemoryOp;
784
785        fn generate_operation(&self, rng: &mut DeterministicRng) -> Self::Operation {
786            let op_type = rng.next_usize(0, 3); // 0-3 inclusive = 4 options
787
788            match op_type {
789                0 => {
790                    // Set - either new key or existing
791                    let key = if !self.known_keys.is_empty() && rng.next_bool(0.3) {
792                        // Use existing key (next_usize is inclusive, so use len-1)
793                        let idx = rng.next_usize(0, self.known_keys.len() - 1);
794                        self.known_keys[idx].clone()
795                    } else {
796                        // New key
797                        format!("key_{}", rng.next_usize(0, 999))
798                    };
799                    let value_len = rng.next_usize(1, 1000);
800                    let ttl_ms = rng.next_usize(100, 10000) as u64;
801                    WorkingMemoryOp::Set {
802                        key,
803                        value_len,
804                        ttl_ms,
805                    }
806                }
807                1 => {
808                    // Get - prefer existing keys
809                    let key = if !self.known_keys.is_empty() && rng.next_bool(0.7) {
810                        let idx = rng.next_usize(0, self.known_keys.len() - 1);
811                        self.known_keys[idx].clone()
812                    } else {
813                        format!("key_{}", rng.next_usize(0, 999))
814                    };
815                    WorkingMemoryOp::Get { key }
816                }
817                2 => {
818                    // Delete - prefer existing keys
819                    let key = if !self.known_keys.is_empty() && rng.next_bool(0.5) {
820                        let idx = rng.next_usize(0, self.known_keys.len() - 1);
821                        self.known_keys[idx].clone()
822                    } else {
823                        format!("key_{}", rng.next_usize(0, 999))
824                    };
825                    WorkingMemoryOp::Delete { key }
826                }
827                _ => WorkingMemoryOp::Cleanup,
828            }
829        }
830
831        fn apply_operation(&mut self, op: &Self::Operation, clock: &SimClock) {
832            self.inner.set_clock_ms(clock.now_ms());
833
834            match op {
835                WorkingMemoryOp::Set {
836                    key,
837                    value_len,
838                    ttl_ms,
839                } => {
840                    let value = vec![0u8; *value_len];
841                    if self.inner.set(key, &value, Some(*ttl_ms)).is_ok() {
842                        if !self.known_keys.contains(key) {
843                            self.known_keys.push(key.clone());
844                        }
845                    }
846                }
847                WorkingMemoryOp::Get { key } => {
848                    let _ = self.inner.get(key);
849                }
850                WorkingMemoryOp::Delete { key } => {
851                    if self.inner.delete(key) {
852                        self.known_keys.retain(|k| k != key);
853                    }
854                }
855                WorkingMemoryOp::Cleanup => {
856                    self.inner.cleanup_expired();
857                    // Update known_keys based on what still exists
858                    self.known_keys.retain(|k| self.inner.exists(k));
859                }
860            }
861        }
862
863        fn check_invariants(&self) -> Result<(), String> {
864            // Invariant 1: used_bytes <= max_bytes
865            if self.inner.used_bytes() > self.inner.config().max_bytes {
866                return Err(format!(
867                    "used_bytes {} exceeds max {}",
868                    self.inner.used_bytes(),
869                    self.inner.config().max_bytes
870                ));
871            }
872
873            // Invariant 2: entry_count <= max_entries
874            if self.inner.entry_count() > self.inner.config().max_entries {
875                return Err(format!(
876                    "entry_count {} exceeds max {}",
877                    self.inner.entry_count(),
878                    self.inner.config().max_entries
879                ));
880            }
881
882            // Invariant 3: if empty, used_bytes should be 0
883            if self.inner.is_empty() && self.inner.used_bytes() != 0 {
884                return Err(format!(
885                    "is_empty() but used_bytes is {}",
886                    self.inner.used_bytes()
887                ));
888            }
889
890            Ok(())
891        }
892
893        fn describe_state(&self) -> String {
894            format!(
895                "WorkingMemory {{ entries: {}, bytes: {}/{}, known_keys: {} }}",
896                self.inner.entry_count(),
897                self.inner.used_bytes(),
898                self.inner.config().max_bytes,
899                self.known_keys.len()
900            )
901        }
902    }
903
904    #[test]
905    fn test_property_invariants() {
906        let wm = WorkingMemoryWrapper {
907            inner: WorkingMemory::new(),
908            known_keys: Vec::new(),
909        };
910
911        PropertyTest::new(42)
912            .with_max_operations(500)
913            .with_time_advance(TimeAdvanceConfig::random(0, 5000, 0.3))
914            .run_and_assert(wm);
915    }
916
917    #[test]
918    fn test_property_invariants_small_capacity() {
919        let config = WorkingMemoryConfig {
920            max_bytes: 10_000, // 10KB
921            max_entries: 50,
922            ..Default::default()
923        };
924        let wm = WorkingMemoryWrapper {
925            inner: WorkingMemory::with_config(config),
926            known_keys: Vec::new(),
927        };
928
929        PropertyTest::new(12345)
930            .with_max_operations(1000)
931            .with_time_advance(TimeAdvanceConfig::random(0, 2000, 0.5))
932            .run_and_assert(wm);
933    }
934
935    #[test]
936    fn test_property_multi_seed() {
937        for seed in [0, 1, 42, 12345, 99999] {
938            let wm = WorkingMemoryWrapper {
939                inner: WorkingMemory::new(),
940                known_keys: Vec::new(),
941            };
942
943            PropertyTest::new(seed)
944                .with_max_operations(200)
945                .with_time_advance(TimeAdvanceConfig::random(0, 1000, 0.4))
946                .run_and_assert(wm);
947        }
948    }
949}
umi_memory/memory/working.rs

umi_memory/memory/
working.rs
