use std::collections::HashMap;
use super::block::{MemoryBlock, MemoryBlockId, MemoryBlockType};
use crate::constants::{
CORE_MEMORY_BLOCK_SIZE_BYTES_MAX, CORE_MEMORY_SIZE_BYTES_MAX, CORE_MEMORY_SIZE_BYTES_MIN,
};
#[derive(Debug, Clone, thiserror::Error)]
pub enum CoreMemoryError {
#[error("core memory full: {current_bytes}/{max_bytes} bytes, need {requested_bytes}")]
Full {
current_bytes: usize,
max_bytes: usize,
requested_bytes: usize,
},
#[error("block not found: {block_type}")]
BlockNotFound {
block_type: String,
},
#[error("block too large: {size_bytes} bytes exceeds max {max_bytes}")]
BlockTooLarge {
size_bytes: usize,
max_bytes: usize,
},
#[error("too many blocks: {count} exceeds max {max_count}")]
TooManyBlocks {
count: usize,
max_count: usize,
},
}
pub type CoreMemoryResult<T> = Result<T, CoreMemoryError>;
#[derive(Debug, Clone)]
pub struct CoreMemoryConfig {
pub max_bytes: usize,
}
impl CoreMemoryConfig {
#[must_use]
pub fn new(max_bytes: usize) -> Self {
assert!(
max_bytes >= CORE_MEMORY_SIZE_BYTES_MIN,
"max_bytes {} below minimum {}",
max_bytes,
CORE_MEMORY_SIZE_BYTES_MIN
);
assert!(
max_bytes <= CORE_MEMORY_SIZE_BYTES_MAX,
"max_bytes {} exceeds maximum {}",
max_bytes,
CORE_MEMORY_SIZE_BYTES_MAX
);
Self { max_bytes }
}
}
impl Default for CoreMemoryConfig {
fn default() -> Self {
Self {
max_bytes: CORE_MEMORY_SIZE_BYTES_MAX,
}
}
}
#[derive(Debug)]
pub struct CoreMemory {
config: CoreMemoryConfig,
blocks_by_type: HashMap<MemoryBlockType, MemoryBlock>,
current_bytes: usize,
clock_ms: u64,
}
impl CoreMemory {
#[must_use]
pub fn new() -> Self {
Self::with_config(CoreMemoryConfig::default())
}
#[must_use]
pub fn with_config(config: CoreMemoryConfig) -> Self {
Self {
config,
blocks_by_type: HashMap::new(),
current_bytes: 0,
clock_ms: 0,
}
}
pub fn set_clock_ms(&mut self, ms: u64) {
self.clock_ms = ms;
}
#[must_use]
pub fn clock_ms(&self) -> u64 {
self.clock_ms
}
pub fn set_block(
&mut self,
block_type: MemoryBlockType,
content: impl Into<String>,
) -> CoreMemoryResult<MemoryBlockId> {
let content = content.into();
let new_size = content.len();
if new_size > CORE_MEMORY_BLOCK_SIZE_BYTES_MAX {
return Err(CoreMemoryError::BlockTooLarge {
size_bytes: new_size,
max_bytes: CORE_MEMORY_BLOCK_SIZE_BYTES_MAX,
});
}
let old_size = self
.blocks_by_type
.get(&block_type)
.map(|b| b.size_bytes())
.unwrap_or(0);
let projected_size = self.current_bytes - old_size + new_size;
if projected_size > self.config.max_bytes {
return Err(CoreMemoryError::Full {
current_bytes: self.current_bytes,
max_bytes: self.config.max_bytes,
requested_bytes: new_size,
});
}
let block = MemoryBlock::new(block_type, content, self.clock_ms);
let id = block.id();
self.blocks_by_type.insert(block_type, block);
self.current_bytes = projected_size;
assert!(
self.current_bytes <= self.config.max_bytes,
"size invariant violated"
);
Ok(id)
}
pub fn set_block_with_label(
&mut self,
block_type: MemoryBlockType,
label: impl Into<String>,
content: impl Into<String>,
) -> CoreMemoryResult<MemoryBlockId> {
let label = label.into();
let content = content.into();
let new_size = content.len();
if new_size > CORE_MEMORY_BLOCK_SIZE_BYTES_MAX {
return Err(CoreMemoryError::BlockTooLarge {
size_bytes: new_size,
max_bytes: CORE_MEMORY_BLOCK_SIZE_BYTES_MAX,
});
}
let old_size = self
.blocks_by_type
.get(&block_type)
.map(|b| b.size_bytes())
.unwrap_or(0);
let projected_size = self.current_bytes - old_size + new_size;
if projected_size > self.config.max_bytes {
return Err(CoreMemoryError::Full {
current_bytes: self.current_bytes,
max_bytes: self.config.max_bytes,
requested_bytes: new_size,
});
}
let block = MemoryBlock::with_label(block_type, label, content, self.clock_ms);
let id = block.id();
self.blocks_by_type.insert(block_type, block);
self.current_bytes = projected_size;
Ok(id)
}
#[must_use]
pub fn get_block(&self, block_type: MemoryBlockType) -> Option<&MemoryBlock> {
self.blocks_by_type.get(&block_type)
}
#[must_use]
pub fn get_content(&self, block_type: MemoryBlockType) -> Option<&str> {
self.blocks_by_type.get(&block_type).map(|b| b.content())
}
#[must_use]
pub fn has_block(&self, block_type: MemoryBlockType) -> bool {
self.blocks_by_type.contains_key(&block_type)
}
pub fn remove_block(&mut self, block_type: MemoryBlockType) -> CoreMemoryResult<MemoryBlock> {
match self.blocks_by_type.remove(&block_type) {
Some(block) => {
self.current_bytes -= block.size_bytes();
assert!(
self.current_bytes <= self.config.max_bytes,
"size invariant violated after removal"
);
Ok(block)
}
None => Err(CoreMemoryError::BlockNotFound {
block_type: block_type.to_string(),
}),
}
}
pub fn clear(&mut self) {
self.blocks_by_type.clear();
self.current_bytes = 0;
assert_eq!(self.current_bytes, 0, "size must be zero after clear");
}
#[must_use]
pub fn block_count(&self) -> usize {
self.blocks_by_type.len()
}
#[must_use]
pub fn used_bytes(&self) -> usize {
self.current_bytes
}
#[must_use]
pub fn available_bytes(&self) -> usize {
self.config.max_bytes.saturating_sub(self.current_bytes)
}
#[must_use]
pub fn max_bytes(&self) -> usize {
self.config.max_bytes
}
#[must_use]
pub fn utilization(&self) -> f64 {
if self.config.max_bytes == 0 {
return 0.0;
}
self.current_bytes as f64 / self.config.max_bytes as f64
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.blocks_by_type.is_empty()
}
pub fn blocks_ordered(&self) -> impl Iterator<Item = &MemoryBlock> {
MemoryBlockType::all_ordered()
.iter()
.filter_map(|bt| self.blocks_by_type.get(bt))
}
#[must_use]
pub fn render(&self) -> String {
let mut output = String::with_capacity(self.current_bytes + 256);
output.push_str("<core_memory>\n");
for block in self.blocks_ordered() {
output.push_str(&block.render());
output.push('\n');
}
output.push_str("</core_memory>");
output
}
#[must_use]
pub fn config(&self) -> &CoreMemoryConfig {
&self.config
}
}
impl Default for CoreMemory {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_core_memory_new() {
let core = CoreMemory::new();
assert_eq!(core.used_bytes(), 0);
assert_eq!(core.max_bytes(), CORE_MEMORY_SIZE_BYTES_MAX);
assert!(core.is_empty());
}
#[test]
fn test_core_memory_with_config() {
let config = CoreMemoryConfig::new(16 * 1024);
let core = CoreMemory::with_config(config);
assert_eq!(core.max_bytes(), 16 * 1024);
}
#[test]
fn test_set_block() {
let mut core = CoreMemory::new();
let id = core.set_block(MemoryBlockType::System, "Hello").unwrap();
assert!(!id.as_uuid().is_nil());
assert!(core.has_block(MemoryBlockType::System));
assert_eq!(core.get_content(MemoryBlockType::System), Some("Hello"));
assert_eq!(core.used_bytes(), 5);
}
#[test]
fn test_set_block_replaces() {
let mut core = CoreMemory::new();
core.set_block(MemoryBlockType::System, "Hello").unwrap();
assert_eq!(core.used_bytes(), 5);
core.set_block(MemoryBlockType::System, "Hi").unwrap();
assert_eq!(core.used_bytes(), 2);
assert_eq!(core.get_content(MemoryBlockType::System), Some("Hi"));
}
#[test]
fn test_set_block_with_label() {
let mut core = CoreMemory::new();
core.set_block_with_label(MemoryBlockType::Facts, "prefs", "Likes cats")
.unwrap();
let block = core.get_block(MemoryBlockType::Facts).unwrap();
assert_eq!(block.label(), Some("prefs"));
assert_eq!(block.content(), "Likes cats");
}
#[test]
fn test_remove_block() {
let mut core = CoreMemory::new();
core.set_block(MemoryBlockType::System, "Hello").unwrap();
let removed = core.remove_block(MemoryBlockType::System).unwrap();
assert_eq!(removed.content(), "Hello");
assert!(!core.has_block(MemoryBlockType::System));
assert_eq!(core.used_bytes(), 0);
}
#[test]
fn test_remove_block_not_found() {
let mut core = CoreMemory::new();
let result = core.remove_block(MemoryBlockType::System);
assert!(matches!(result, Err(CoreMemoryError::BlockNotFound { .. })));
}
#[test]
fn test_clear() {
let mut core = CoreMemory::new();
core.set_block(MemoryBlockType::System, "Hello").unwrap();
core.set_block(MemoryBlockType::Human, "World").unwrap();
core.clear();
assert!(core.is_empty());
assert_eq!(core.used_bytes(), 0);
}
#[test]
fn test_capacity_limit() {
let config = CoreMemoryConfig::new(CORE_MEMORY_SIZE_BYTES_MIN); let mut core = CoreMemory::with_config(config);
let content = "x".repeat(CORE_MEMORY_SIZE_BYTES_MIN - 100);
core.set_block(MemoryBlockType::System, content).unwrap();
let result = core.set_block(MemoryBlockType::Human, "x".repeat(200));
assert!(matches!(result, Err(CoreMemoryError::Full { .. })));
}
#[test]
fn test_block_too_large() {
let mut core = CoreMemory::new();
let content = "x".repeat(CORE_MEMORY_BLOCK_SIZE_BYTES_MAX + 1);
let result = core.set_block(MemoryBlockType::System, content);
assert!(matches!(result, Err(CoreMemoryError::BlockTooLarge { .. })));
}
#[test]
fn test_utilization() {
let mut core = CoreMemory::new();
assert_eq!(core.utilization(), 0.0);
let content = "x".repeat(CORE_MEMORY_SIZE_BYTES_MAX / 2);
core.set_block(MemoryBlockType::System, content).unwrap();
let util = core.utilization();
assert!(util > 0.49 && util < 0.51);
}
#[test]
fn test_render_empty() {
let core = CoreMemory::new();
let rendered = core.render();
assert_eq!(rendered, "<core_memory>\n</core_memory>");
}
#[test]
fn test_render_with_blocks() {
let mut core = CoreMemory::new();
core.set_block(MemoryBlockType::System, "Be helpful.")
.unwrap();
core.set_block(MemoryBlockType::Human, "User: Alice")
.unwrap();
let rendered = core.render();
assert!(rendered.starts_with("<core_memory>"));
assert!(rendered.ends_with("</core_memory>"));
assert!(rendered.contains("Be helpful."));
assert!(rendered.contains("User: Alice"));
let sys_pos = rendered.find("system").unwrap();
let human_pos = rendered.find("human").unwrap();
assert!(sys_pos < human_pos);
}
#[test]
fn test_render_order() {
let mut core = CoreMemory::new();
core.set_block(MemoryBlockType::Scratch, "5").unwrap();
core.set_block(MemoryBlockType::Goals, "4").unwrap();
core.set_block(MemoryBlockType::Facts, "3").unwrap();
core.set_block(MemoryBlockType::Human, "2").unwrap();
core.set_block(MemoryBlockType::Persona, "1").unwrap();
core.set_block(MemoryBlockType::System, "0").unwrap();
let rendered = core.render();
let positions: Vec<usize> = ["system", "persona", "human", "facts", "goals", "scratch"]
.iter()
.map(|s| rendered.find(s).unwrap())
.collect();
for i in 1..positions.len() {
assert!(
positions[i] > positions[i - 1],
"render order should be by priority"
);
}
}
#[test]
fn test_clock_ms() {
let mut core = CoreMemory::new();
assert_eq!(core.clock_ms(), 0);
core.set_clock_ms(5000);
assert_eq!(core.clock_ms(), 5000);
core.set_block(MemoryBlockType::System, "Test").unwrap();
let block = core.get_block(MemoryBlockType::System).unwrap();
assert_eq!(block.created_at_ms(), 5000);
}
#[test]
fn test_blocks_ordered_iterator() {
let mut core = CoreMemory::new();
core.set_block(MemoryBlockType::Scratch, "scratch").unwrap();
core.set_block(MemoryBlockType::System, "system").unwrap();
let blocks: Vec<_> = core.blocks_ordered().collect();
assert_eq!(blocks.len(), 2);
assert_eq!(blocks[0].block_type(), MemoryBlockType::System);
assert_eq!(blocks[1].block_type(), MemoryBlockType::Scratch);
}
#[test]
#[should_panic(expected = "max_bytes")]
fn test_config_below_minimum() {
let _ = CoreMemoryConfig::new(100);
}
#[test]
#[should_panic(expected = "max_bytes")]
fn test_config_above_maximum() {
let _ = CoreMemoryConfig::new(CORE_MEMORY_SIZE_BYTES_MAX + 1);
}
}
#[cfg(test)]
mod dst_tests {
use super::*;
use crate::dst::{SimConfig, Simulation};
#[tokio::test]
async fn test_core_memory_with_sim_clock() {
let sim = Simulation::new(SimConfig::with_seed(42));
sim.run(|env| async move {
let mut core = CoreMemory::new();
core.set_clock_ms(env.clock.now_ms());
core.set_block(MemoryBlockType::System, "Initial").unwrap();
let block = core.get_block(MemoryBlockType::System).unwrap();
assert_eq!(block.created_at_ms(), 0);
env.clock.advance_ms(1000);
core.set_clock_ms(env.clock.now_ms());
core.set_block(MemoryBlockType::System, "Updated").unwrap();
let block = core.get_block(MemoryBlockType::System).unwrap();
assert_eq!(block.created_at_ms(), 1000);
Ok::<(), std::convert::Infallible>(())
})
.await
.unwrap();
}
#[tokio::test]
async fn test_core_memory_determinism() {
let mut results1 = Vec::new();
let mut results2 = Vec::new();
let sim1 = Simulation::new(SimConfig::with_seed(12345));
sim1.run(|mut env| async move {
let mut core = CoreMemory::new();
for _ in 0..5 {
env.clock.advance_ms(100);
core.set_clock_ms(env.clock.now_ms());
let content = format!("block_{}", env.rng.next_usize(0, 1000));
core.set_block(MemoryBlockType::Scratch, &content).unwrap();
results1.push(content);
}
Ok::<(), std::convert::Infallible>(())
})
.await
.unwrap();
let sim2 = Simulation::new(SimConfig::with_seed(12345));
sim2.run(|mut env| async move {
let mut core = CoreMemory::new();
for _ in 0..5 {
env.clock.advance_ms(100);
core.set_clock_ms(env.clock.now_ms());
let content = format!("block_{}", env.rng.next_usize(0, 1000));
core.set_block(MemoryBlockType::Scratch, &content).unwrap();
results2.push(content);
}
Ok::<(), std::convert::Infallible>(())
})
.await
.unwrap();
}
#[tokio::test]
async fn test_core_memory_time_tracking() {
let sim = Simulation::new(SimConfig::with_seed(42));
sim.run(|env| async move {
let mut core = CoreMemory::new();
let mut timestamps = Vec::new();
for i in 0..3 {
env.clock.advance_ms(500);
core.set_clock_ms(env.clock.now_ms());
let content = format!("Block {}", i);
let block_type = match i {
0 => MemoryBlockType::System,
1 => MemoryBlockType::Human,
_ => MemoryBlockType::Facts,
};
core.set_block(block_type, content).unwrap();
timestamps.push(env.clock.now_ms());
}
assert_eq!(timestamps, vec![500, 1000, 1500]);
assert_eq!(
core.get_block(MemoryBlockType::System)
.unwrap()
.created_at_ms(),
500
);
assert_eq!(
core.get_block(MemoryBlockType::Human)
.unwrap()
.created_at_ms(),
1000
);
assert_eq!(
core.get_block(MemoryBlockType::Facts)
.unwrap()
.created_at_ms(),
1500
);
Ok::<(), std::convert::Infallible>(())
})
.await
.unwrap();
}
#[tokio::test]
async fn test_core_memory_capacity_under_simulation() {
let sim = Simulation::new(SimConfig::with_seed(42));
sim.run(|env| async move {
let config = CoreMemoryConfig::new(CORE_MEMORY_SIZE_BYTES_MIN);
let mut core = CoreMemory::with_config(config);
core.set_clock_ms(env.clock.now_ms());
let mut total_added = 0;
let block_types = [
MemoryBlockType::System,
MemoryBlockType::Persona,
MemoryBlockType::Human,
];
for block_type in &block_types {
let size = 1000; let content = "x".repeat(size);
match core.set_block(*block_type, content) {
Ok(_) => total_added += size,
Err(CoreMemoryError::Full { .. }) => break,
Err(e) => panic!("Unexpected error: {:?}", e),
}
env.clock.advance_ms(100);
core.set_clock_ms(env.clock.now_ms());
}
assert!(total_added >= 3000);
assert!(core.used_bytes() <= CORE_MEMORY_SIZE_BYTES_MIN);
Ok::<(), std::convert::Infallible>(())
})
.await
.unwrap();
}
#[tokio::test]
async fn test_render_deterministic() {
let sim = Simulation::new(SimConfig::with_seed(42));
sim.run(|env| async move {
let mut core = CoreMemory::new();
core.set_clock_ms(env.clock.now_ms());
core.set_block(MemoryBlockType::System, "System prompt")
.unwrap();
core.set_block(MemoryBlockType::Human, "User info").unwrap();
core.set_block(MemoryBlockType::Facts, "Key facts").unwrap();
let rendered = core.render();
assert!(rendered.starts_with("<core_memory>"));
assert!(rendered.ends_with("</core_memory>"));
let sys_pos = rendered.find("type=\"system\"").unwrap();
let human_pos = rendered.find("type=\"human\"").unwrap();
let facts_pos = rendered.find("type=\"facts\"").unwrap();
assert!(sys_pos < human_pos);
assert!(human_pos < facts_pos);
Ok::<(), std::convert::Infallible>(())
})
.await
.unwrap();
}
}
#[cfg(test)]
mod edge_case_tests {
use super::*;
#[test]
fn test_unicode_content_size() {
let mut core = CoreMemory::new();
core.set_block(MemoryBlockType::System, "こんにちは")
.unwrap();
assert_eq!(core.used_bytes(), 15);
assert_eq!(
core.get_content(MemoryBlockType::System)
.unwrap()
.chars()
.count(),
5
);
}
#[test]
fn test_empty_string_content() {
let mut core = CoreMemory::new();
core.set_block(MemoryBlockType::System, "").unwrap();
assert_eq!(core.used_bytes(), 0);
assert_eq!(core.get_content(MemoryBlockType::System), Some(""));
assert!(core.has_block(MemoryBlockType::System));
}
#[test]
fn test_empty_label() {
let mut core = CoreMemory::new();
core.set_block_with_label(MemoryBlockType::Facts, "", "content")
.unwrap();
let block = core.get_block(MemoryBlockType::Facts).unwrap();
assert_eq!(block.label(), Some(""));
}
#[test]
fn test_max_length_label() {
use crate::constants::CORE_MEMORY_BLOCK_LABEL_BYTES_MAX;
let mut core = CoreMemory::new();
let max_label = "x".repeat(CORE_MEMORY_BLOCK_LABEL_BYTES_MAX);
core.set_block_with_label(MemoryBlockType::Facts, &max_label, "content")
.unwrap();
let block = core.get_block(MemoryBlockType::Facts).unwrap();
assert_eq!(
block.label().unwrap().len(),
CORE_MEMORY_BLOCK_LABEL_BYTES_MAX
);
}
#[test]
fn test_whitespace_content() {
let mut core = CoreMemory::new();
core.set_block(MemoryBlockType::Scratch, " \n\t ")
.unwrap();
assert_eq!(core.used_bytes(), 7);
}
}