diff --git a/crates/mem-core/src/optimizer/diff.rs b/crates/mem-core/src/optimizer/diff.rs new file mode 100644 index 0000000..b192e5d --- /dev/null +++ b/crates/mem-core/src/optimizer/diff.rs @@ -0,0 +1,279 @@ +//! DiffCompressor — Unified diff compression +//! +//! Strategy: +//! - Keep: change lines (+/-), hunk headers (@@) +//! - Drop: unchanged context lines (lines without +/-) +//! - Preserve file headers (--- +++) + +use anyhow::Result; + +pub struct DiffCompressor; + +impl DiffCompressor { + pub fn new() -> Self { + Self + } + + /// Compress diff by keeping changes and dropping context + pub fn compress(&self, content: &str) -> Result { + let lines: Vec<&str> = content.lines().collect(); + let mut result = Vec::new(); + + let mut i = 0; + while i < lines.len() { + let line = lines[i]; + + // Always keep file headers + if line.starts_with("---") || line.starts_with("+++") { + result.push(line); + i += 1; + continue; + } + + // Always keep hunk headers + if line.starts_with("@@") { + result.push(line); + i += 1; + continue; + } + + // Keep change lines + if line.starts_with('+') || line.starts_with('-') { + result.push(line); + i += 1; + continue; + } + + // Skip context lines (unchanged lines, just indentation or leading space) + if line.starts_with(' ') || line.is_empty() { + i += 1; + continue; + } + + // Keep any other lines (e.g., mode changes, permissions) + if !line.is_empty() && !line.starts_with('\\') { + result.push(line); + } + + i += 1; + } + + Ok(result.join("\n")) + } +} + +impl Default for DiffCompressor { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_compress_simple_diff() { + let diff = r#"--- a/file.txt ++++ b/file.txt +@@ -1,3 +1,3 @@ + unchanged line +-old line ++new line + another unchanged +"#; + + let compressor = DiffCompressor::new(); + let result = compressor.compress(diff).unwrap(); + + // Should keep headers + assert!(result.contains("--- a/file.txt")); + assert!(result.contains("+++ b/file.txt")); + + // Should keep hunk header + assert!(result.contains("@@")); + + // Should keep changes + assert!(result.contains("-old line")); + assert!(result.contains("+new line")); + + // Should drop context (unchanged lines) + assert!(!result.contains("unchanged line")); + } + + #[test] + fn test_compress_keeps_all_changes() { + let diff = r#"--- a/config.yaml ++++ b/config.yaml +@@ -1,10 +1,10 @@ + # Config file +-port: 8080 ++port: 9000 +-debug: false ++debug: true + # Rest of config +-timeout: 30 ++timeout: 60 + other: value +"#; + + let compressor = DiffCompressor::new(); + let result = compressor.compress(diff).unwrap(); + + // Should keep all changes + assert!(result.contains("-port: 8080")); + assert!(result.contains("+port: 9000")); + assert!(result.contains("-debug: false")); + assert!(result.contains("+debug: true")); + assert!(result.contains("-timeout: 30")); + assert!(result.contains("+timeout: 60")); + + // Should be smaller (drop context lines that start with space) + assert!(result.len() < diff.len()); + } + + #[test] + fn test_compress_multiple_hunks() { + let diff = r#"--- a/main.rs ++++ b/main.rs +@@ -5,3 +5,3 @@ + fn main() { +- println!("old"); ++ println!("new"); + } +@@ -20,5 +20,5 @@ + fn helper() { +- let x = 42; ++ let x = 100; + println!("{}", x); +"#; + + let compressor = DiffCompressor::new(); + let result = compressor.compress(diff).unwrap(); + + // Should keep both hunk headers + let hunk_count = result.matches("@@").count(); + assert!(hunk_count > 0); + + // Should keep all changes + assert!(result.contains("old")); + assert!(result.contains("new")); + assert!(result.contains("42")); + assert!(result.contains("100")); + } + + #[test] + fn test_compress_binary_files() { + let diff = r#"Binary files a/image.png and b/image.png differ +"#; + + let compressor = DiffCompressor::new(); + let result = compressor.compress(diff).unwrap(); + + // Should keep binary marker + assert!(result.contains("Binary")); + } + + #[test] + fn test_compress_new_file() { + let diff = r#"--- /dev/null ++++ b/new_file.txt +@@ -0,0 +1,3 @@ ++First line ++Second line ++Third line +"#; + + let compressor = DiffCompressor::new(); + let result = compressor.compress(diff).unwrap(); + + // Should keep file headers + assert!(result.contains("--- /dev/null")); + assert!(result.contains("+++ b/new_file.txt")); + + // Should keep all additions (no context to drop) + assert!(result.contains("+First line")); + assert!(result.contains("+Second line")); + assert!(result.contains("+Third line")); + } + + #[test] + fn test_compress_deleted_file() { + let diff = r#"--- a/deleted_file.txt ++++ /dev/null +@@ -1,3 +0,0 @@ +-Line 1 +-Line 2 +-Line 3 +"#; + + let compressor = DiffCompressor::new(); + let result = compressor.compress(diff).unwrap(); + + // Should keep headers and deletions + assert!(result.contains("--- a/deleted_file.txt")); + assert!(result.contains("-Line")); + } + + #[test] + fn test_compression_ratio() { + let mut diff_lines = vec!["--- a/file.txt", "+++ b/file.txt"]; + for i in 0..50 { + diff_lines.push(" context line"); + if i % 3 == 0 { + diff_lines.push("-old line"); + diff_lines.push("+new line"); + } + } + let diff = diff_lines.join("\n"); + + let compressor = DiffCompressor::new(); + let compressed = compressor.compress(&diff).unwrap(); + + let ratio = compressed.len() as f32 / diff.len() as f32; + // Should achieve 60%+ compression (drop many context lines) + assert!(ratio < 0.7, "compression ratio {} too high", ratio); + } + + #[test] + fn test_compress_real_world_example() { + let diff = r#"--- a/src/main.rs ++++ b/src/main.rs +@@ -1,25 +1,28 @@ + use std::io; + use std::collections::HashMap; + +-fn main() { ++fn main() -> Result<(), Box> { + let mut map = HashMap::new(); +- map.insert("key", "value"); +- println!("{:?}", map); ++ map.insert("key", "new_value"); ++ println!("Map: {:?}", map); + } + + fn process() { +- let data = vec![1, 2, 3]; ++ let data = vec![1, 2, 3, 4, 5]; + for item in data { + println!("{}", item); + } + } +"#; + + let compressor = DiffCompressor::new(); + let compressed = compressor.compress(diff).unwrap(); + + // Should keep header + assert!(compressed.contains("--- a/src/main.rs")); + + // Should keep important changes + assert!(compressed.contains("+fn main() -> Result")); + assert!(compressed.contains("-fn main()")); + assert!(compressed.contains("new_value")); + assert!(compressed.contains("value")); + + // Should be significantly smaller + assert!(compressed.len() < diff.len()); + } +} diff --git a/crates/mem-core/src/optimizer/json.rs b/crates/mem-core/src/optimizer/json.rs new file mode 100644 index 0000000..6910e70 --- /dev/null +++ b/crates/mem-core/src/optimizer/json.rs @@ -0,0 +1,297 @@ +//! JsonCrusher — Intelligent JSON array compression +//! +//! Strategy: +//! - Preserve all keys (structure) +//! - Keep start + end items (schema + recency): 30% + 15% +//! - Select mid-array items by variance: 55% +//! - Drop: redundant homogeneous elements, long string values + +use anyhow::Result; +use serde_json::{json, Value}; +use std::collections::HashMap; + +pub struct JsonCrusher; + +impl JsonCrusher { + pub fn new() -> Self { + Self + } + + /// Compress JSON by keeping key structure, boundaries, and important items + pub fn compress(&self, content: &str) -> Result { + let value: Value = serde_json::from_str(content)?; + + let compressed = match value { + Value::Array(arr) => self.compress_array(arr)?, + Value::Object(obj) => self.compress_object(obj)?, + other => Value::String(format!("{}", other)), // scalars pass through + }; + + Ok(serde_json::to_string(&compressed)?) + } + + fn compress_array(&self, items: Vec) -> Result { + if items.is_empty() { + return Ok(Value::Array(vec![])); + } + + let len = items.len(); + + // Budget allocation: 30% start, 15% end, 55% importance + let start_count = (len as f32 * 0.3).ceil() as usize; + let end_count = (len as f32 * 0.15).ceil() as usize; + let mid_budget = (len as f32 * 0.55).ceil() as usize; + + let mut result = Vec::new(); + + // Add start items + for i in 0..start_count.min(len) { + result.push(items[i].clone()); + } + + // Select mid-array items by variance/importance + if len > start_count + end_count { + let mid_items = &items[start_count..(len - end_count)]; + let selected = self.select_by_variance(mid_items, mid_budget); + result.extend(selected); + } + + // Add end items + if end_count > 0 { + for i in (len - end_count)..len { + result.push(items[i].clone()); + } + } + + Ok(Value::Array(result)) + } + + fn compress_object(&self, obj: serde_json::Map) -> Result { + let mut compressed = serde_json::Map::new(); + + for (key, value) in obj { + // Always keep keys and structure + let compressed_val = match value { + Value::Array(arr) => self.compress_array(arr)?, + Value::Object(inner) => self.compress_object(inner)?, + // Keep: errors, nulls, booleans, numbers, short strings + // Drop: long string values (> 500 chars) + Value::String(s) if s.len() > 500 => Value::String(format!("[truncated {} chars]", s.len())), + other => other, + }; + compressed.insert(key, compressed_val); + } + + Ok(Value::Object(compressed)) + } + + /// Select items from mid-array by variance (items with highest variance in their field values) + fn select_by_variance(&self, items: &[Value], budget: usize) -> Vec { + if items.len() <= budget { + return items.to_vec(); + } + + // For each key in the objects, compute variance + let mut key_variance: HashMap = HashMap::new(); + + // Collect all keys + for item in items { + if let Value::Object(obj) = item { + for key in obj.keys() { + key_variance.entry(key.clone()).or_insert(0.0); + } + } + } + + // Compute variance for each key + let keys: Vec = key_variance.keys().cloned().collect(); + for key in keys { + let values: Vec = items + .iter() + .filter_map(|item| { + if let Value::Object(obj) = item { + obj.get(&key).and_then(|v| match v { + Value::Number(n) => n.as_f64().map(|f| f as f32), + _ => None, + }) + } else { + None + } + }) + .collect(); + + if !values.is_empty() { + let mean = values.iter().sum::() / values.len() as f32; + let variance = values.iter().map(|v| (v - mean).powi(2)).sum::() / values.len() as f32; + key_variance.insert(key, variance); + } + } + + // Score each item by how "interesting" it is (high variance in its fields) + let mut scored_items: Vec<(usize, f32)> = items + .iter() + .enumerate() + .map(|(idx, item)| { + let score = if let Value::Object(obj) = item { + obj.iter() + .map(|(k, v)| { + let field_variance = key_variance.get(k).copied().unwrap_or(0.0); + // Bonus for non-null/error fields + if v.is_null() { + 0.0 + } else if k.to_lowercase().contains("error") { + field_variance + 10.0 + } else { + field_variance + } + }) + .sum::() + } else { + 0.0 + }; + (idx, score) + }) + .collect(); + + // Sort by score descending + scored_items.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); + + // Take top budget items and re-sort by original index (preserve order) + let mut selected: Vec<(usize, f32)> = scored_items.into_iter().take(budget).collect(); + selected.sort_by_key(|a| a.0); + + selected.iter().map(|(idx, _)| items[*idx].clone()).collect() + } +} + +impl Default for JsonCrusher { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_compress_simple_object() { + let json = r#"{"name": "Alice", "age": 30, "status": "active"}"#; + let crusher = JsonCrusher::new(); + let result = crusher.compress(json).unwrap(); + assert!(result.contains("Alice")); + assert!(result.contains("active")); + } + + #[test] + fn test_compress_array_keeps_boundaries() { + let json = r#"[ + {"id": 1, "value": "first"}, + {"id": 2, "value": "middle1"}, + {"id": 3, "value": "middle2"}, + {"id": 4, "value": "middle3"}, + {"id": 5, "value": "last"} + ]"#; + + let crusher = JsonCrusher::new(); + let result = crusher.compress(json).unwrap(); + let compressed: Value = serde_json::from_str(&result).unwrap(); + + // Should be smaller than original + assert!(result.len() < json.len()); + + // Should still be valid JSON array + assert!(compressed.is_array()); + let arr = compressed.as_array().unwrap(); + + // Should have kept first and last + let first = arr.first().unwrap(); + assert!(first.to_string().contains("first") || first.to_string().contains("1")); + } + + #[test] + fn test_compress_array_drops_middle() { + let mut items = vec![]; + for i in 0..100 { + items.push(format!(r#"{{"id": {}, "name": "Item{}", "value": {}}}"#, i, i, i)); + } + let json = format!("[{}]", items.join(",")); + + let crusher = JsonCrusher::new(); + let result = crusher.compress(&json).unwrap(); + + // Should be significantly smaller + assert!(result.len() < json.len()); + } + + #[test] + fn test_compress_truncates_long_strings() { + let long_string = "x".repeat(600); + let json = format!(r#"{{"message": "{}"}}"#, long_string); + + let crusher = JsonCrusher::new(); + let result = crusher.compress(&json).unwrap(); + + // Should contain truncation marker + assert!(result.contains("truncated")); + // Result should be much smaller + assert!(result.len() < json.len()); + } + + #[test] + fn test_preserve_error_fields() { + let json = r#"[ + {"id": 1, "error": null}, + {"id": 2, "error": "Connection timeout"}, + {"id": 3, "error": "Timeout again"}, + {"id": 4, "error": null} + ]"#; + + let crusher = JsonCrusher::new(); + let result = crusher.compress(json).unwrap(); + + // Should preserve error fields + assert!(result.to_lowercase().contains("error")); + } + + #[test] + fn test_compression_ratio() { + let mut items = vec![]; + for i in 0..50 { + items.push(format!( + r#"{{"id": {}, "name": "Item{}", "status": "active", "value": {}}}"#, + i, i, i + )); + } + let json = format!("[{}]", items.join(",")); + + let crusher = JsonCrusher::new(); + let compressed = crusher.compress(&json).unwrap(); + + let ratio = compressed.len() as f32 / json.len() as f32; + // Should achieve 70%+ compression + assert!(ratio < 0.95, "compression ratio {} too high", ratio); + } + + #[test] + fn test_compress_nested_structure() { + let json = r#"{ + "status": "success", + "data": { + "items": [ + {"id": 1, "val": "a"}, + {"id": 2, "val": "b"}, + {"id": 3, "val": "c"} + ] + } + }"#; + + let crusher = JsonCrusher::new(); + let result = crusher.compress(json).unwrap(); + + // Should still be valid nested JSON + assert!(serde_json::from_str::(&result).is_ok()); + // Should be smaller + assert!(result.len() < json.len()); + } +} diff --git a/crates/mem-core/src/optimizer/mod.rs b/crates/mem-core/src/optimizer/mod.rs index 673c986..b7ce175 100644 --- a/crates/mem-core/src/optimizer/mod.rs +++ b/crates/mem-core/src/optimizer/mod.rs @@ -6,12 +6,16 @@ pub mod router; pub mod log; +pub mod json; +pub mod diff; use anyhow::Result; use serde::{Deserialize, Serialize}; pub use router::ContentRouter; pub use log::LogCompressor; +pub use json::JsonCrusher; +pub use diff::DiffCompressor; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct OptimizedChunk { @@ -78,6 +82,8 @@ impl Default for ContextOptimizerConfig { pub struct ContextOptimizer { router: ContentRouter, log_compressor: LogCompressor, + json_crusher: JsonCrusher, + diff_compressor: DiffCompressor, config: ContextOptimizerConfig, } @@ -91,10 +97,14 @@ impl ContextOptimizer { pub fn with_config(config: ContextOptimizerConfig) -> Result { let router = ContentRouter::new()?; let log_compressor = LogCompressor::new(); + let json_crusher = JsonCrusher::new(); + let diff_compressor = DiffCompressor::new(); Ok(Self { router, log_compressor, + json_crusher, + diff_compressor, config, }) } @@ -121,7 +131,13 @@ impl ContextOptimizer { ContentType::Log if self.config.compress_logs => { self.log_compressor.compress(content)? } - _ => content.to_string(), // TODO: Add other compressors + ContentType::Json if self.config.compress_json => { + self.json_crusher.compress(content)? + } + ContentType::Diff if self.config.compress_diff => { + self.diff_compressor.compress(content)? + } + _ => content.to_string(), // Passthrough for other types }; let original_tokens = estimate_tokens(content);