sqz_engine/
dict_compressor.rs

1/// Dictionary-based JSON Compression.
2///
3/// Exploits cross-document redundancy by maintaining a dictionary of common
4/// JSON field names, value patterns, and structural elements. Each new JSON
5/// document is compressed relative to this shared dictionary, achieving
6/// better ratios than compressing in isolation.
7///
8/// This is a pure-Rust implementation that avoids adding the zstd crate
9/// (which would add ~200KB to the binary). Instead, it uses a dictionary
10/// substitution approach optimized for JSON payloads.
11
12use std::collections::HashMap;
13
14use crate::error::Result;
15
16/// A dictionary entry mapping a common pattern to a short code.
17#[derive(Debug, Clone)]
18struct DictEntry {
19    /// The original pattern (e.g., a field name or common value).
20    pattern: String,
21    /// The short replacement code.
22    code: String,
23    /// How many times this pattern has been seen.
24    #[allow(dead_code)]
25    frequency: u32,
26}
27
28/// Configuration for dictionary compression.
29#[derive(Debug, Clone)]
30pub struct DictConfig {
31    /// Maximum number of dictionary entries.
32    /// Default: 256
33    pub max_entries: usize,
34    /// Minimum pattern length (chars) to consider for dictionary.
35    /// Default: 4
36    pub min_pattern_length: usize,
37    /// Minimum frequency before a pattern is added to the dictionary.
38    /// Default: 2
39    pub min_frequency: u32,
40}
41
42impl Default for DictConfig {
43    fn default() -> Self {
44        Self {
45            max_entries: 256,
46            min_pattern_length: 4,
47            min_frequency: 2,
48        }
49    }
50}
51
52/// Result of dictionary compression.
53#[derive(Debug, Clone)]
54pub struct DictCompressResult {
55    /// The compressed text.
56    pub data: String,
57    /// Number of substitutions made.
58    pub substitutions: usize,
59    /// Bytes saved.
60    pub bytes_saved: usize,
61    /// Whether the dictionary header was included.
62    pub has_dict_header: bool,
63}
64
65/// Dictionary-based JSON compressor.
66///
67/// Maintains a session-level dictionary of common JSON patterns and
68/// substitutes them with short codes during compression.
69pub struct DictCompressor {
70    config: DictConfig,
71    /// Pattern → DictEntry mapping.
72    entries: HashMap<String, DictEntry>,
73    /// Frequency counter for candidate patterns.
74    candidates: HashMap<String, u32>,
75    /// Next code index for new entries.
76    next_code: u16,
77    /// Built-in common JSON field names.
78    builtin_entries: Vec<DictEntry>,
79}
80
81impl DictCompressor {
82    pub fn new() -> Self {
83        Self::with_config(DictConfig::default())
84    }
85
86    pub fn with_config(config: DictConfig) -> Self {
87        let builtin = build_builtin_dictionary();
88        let mut compressor = Self {
89            config,
90            entries: HashMap::new(),
91            candidates: HashMap::new(),
92            next_code: 0,
93            builtin_entries: builtin,
94        };
95        compressor.load_builtins();
96        compressor
97    }
98
99    /// Load built-in dictionary entries.
100    fn load_builtins(&mut self) {
101        for entry in &self.builtin_entries.clone() {
102            self.entries.insert(entry.pattern.clone(), entry.clone());
103        }
104        self.next_code = self.entries.len() as u16;
105    }
106
107    /// Observe a JSON string to learn common patterns.
108    pub fn observe(&mut self, json_str: &str) {
109        // Extract field names from JSON
110        let fields = extract_json_fields(json_str);
111        for field in fields {
112            if field.len() >= self.config.min_pattern_length {
113                let count = self.candidates.entry(field.clone()).or_insert(0);
114                *count += 1;
115
116                // Promote to dictionary if frequency threshold met
117                if *count >= self.config.min_frequency
118                    && !self.entries.contains_key(&field)
119                    && self.entries.len() < self.config.max_entries
120                {
121                    let code = format!("~{:X}", self.next_code);
122                    self.next_code += 1;
123                    self.entries.insert(
124                        field.clone(),
125                        DictEntry {
126                            pattern: field,
127                            code,
128                            frequency: *count,
129                        },
130                    );
131                }
132            }
133        }
134    }
135
136    /// Compress a JSON string using the dictionary.
137    ///
138    /// Substitutes known field names and patterns with short codes.
139    /// Prepends a dictionary header if substitutions were made.
140    pub fn compress(&self, json_str: &str) -> Result<DictCompressResult> {
141        if self.entries.is_empty() {
142            return Ok(DictCompressResult {
143                data: json_str.to_string(),
144                substitutions: 0,
145                bytes_saved: 0,
146                has_dict_header: false,
147            });
148        }
149
150        let mut result = json_str.to_string();
151        let mut substitutions = 0usize;
152        let mut bytes_saved = 0usize;
153        let mut used_entries: Vec<&DictEntry> = Vec::new();
154
155        // Sort entries by pattern length (longest first) to avoid partial matches
156        let mut sorted_entries: Vec<&DictEntry> = self.entries.values().collect();
157        sorted_entries.sort_by(|a, b| {
158            b.pattern
159                .len()
160                .cmp(&a.pattern.len())
161                .then_with(|| a.pattern.cmp(&b.pattern))
162        });
163
164        for entry in &sorted_entries {
165            // Only substitute in JSON key positions: "field_name":
166            let search = format!("\"{}\":", entry.pattern);
167            let replace = format!("\"{}\":", entry.code);
168
169            if result.contains(&search) {
170                let count = result.matches(&search).count();
171                result = result.replace(&search, &replace);
172                substitutions += count;
173                bytes_saved += count * (search.len() - replace.len());
174                used_entries.push(entry);
175            }
176        }
177
178        // Only add header if we actually made substitutions
179        let has_dict_header = !used_entries.is_empty();
180        if has_dict_header {
181            let header = format_dict_header(&used_entries);
182            result = format!("{header}\n{result}");
183        }
184
185        Ok(DictCompressResult {
186            data: result,
187            substitutions,
188            bytes_saved,
189            has_dict_header,
190        })
191    }
192
193    /// Get the current dictionary size.
194    pub fn dict_size(&self) -> usize {
195        self.entries.len()
196    }
197
198    /// Reset the dictionary (e.g., on session reset).
199    pub fn reset(&mut self) {
200        self.entries.clear();
201        self.candidates.clear();
202        self.next_code = 0;
203        self.load_builtins();
204    }
205}
206
207impl Default for DictCompressor {
208    fn default() -> Self {
209        Self::new()
210    }
211}
212
213// ── Built-in dictionary ───────────────────────────────────────────────────────
214
215/// Build the default dictionary of common JSON field names.
216fn build_builtin_dictionary() -> Vec<DictEntry> {
217    let common_fields = [
218        // API response fields
219        ("status_code", "~s"),
220        ("message", "~m"),
221        ("timestamp", "~ts"),
222        ("created_at", "~ca"),
223        ("updated_at", "~ua"),
224        ("deleted_at", "~da"),
225        ("description", "~dc"),
226        ("metadata", "~md"),
227        ("properties", "~pp"),
228        ("attributes", "~at"),
229        ("parameters", "~pm"),
230        ("configuration", "~cf"),
231        ("environment", "~ev"),
232        ("dependencies", "~dp"),
233        ("permissions", "~pr"),
234        ("resources", "~rs"),
235        ("namespace", "~ns"),
236        ("annotations", "~an"),
237        // Error fields
238        ("error_code", "~ec"),
239        ("error_message", "~em"),
240        ("stack_trace", "~st"),
241        ("exception", "~ex"),
242        // Pagination
243        ("page_size", "~ps"),
244        ("next_token", "~nt"),
245        ("total_count", "~tc"),
246        // AWS/Cloud common
247        ("resource_type", "~rt"),
248        ("account_id", "~ai"),
249        ("region", "~rg"),
250    ];
251
252    common_fields
253        .iter()
254        .enumerate()
255        .map(|(_, &(pattern, code))| DictEntry {
256            pattern: pattern.to_string(),
257            code: code.to_string(),
258            frequency: 100, // built-in entries have high base frequency
259        })
260        .collect()
261}
262
263/// Extract field names from a JSON string (simple regex-free parser).
264fn extract_json_fields(json_str: &str) -> Vec<String> {
265    let mut fields = Vec::new();
266    let bytes = json_str.as_bytes();
267    let len = bytes.len();
268    let mut i = 0;
269
270    while i < len {
271        // Look for "field_name":
272        if bytes[i] == b'"' {
273            let start = i + 1;
274            i += 1;
275            // Find closing quote
276            while i < len && bytes[i] != b'"' {
277                if bytes[i] == b'\\' {
278                    i += 1; // skip escaped char
279                }
280                i += 1;
281            }
282            if i < len {
283                let end = i;
284                i += 1;
285                // Check if followed by ':'
286                // Skip whitespace
287                while i < len && (bytes[i] == b' ' || bytes[i] == b'\t' || bytes[i] == b'\n' || bytes[i] == b'\r') {
288                    i += 1;
289                }
290                if i < len && bytes[i] == b':' {
291                    if let Ok(field) = std::str::from_utf8(&bytes[start..end]) {
292                        fields.push(field.to_string());
293                    }
294                }
295            }
296        } else {
297            i += 1;
298        }
299    }
300
301    fields
302}
303
304/// Format the dictionary header for compressed output.
305fn format_dict_header(entries: &[&DictEntry]) -> String {
306    let mut lines = vec!["§dict§".to_string()];
307    for entry in entries {
308        lines.push(format!("{}={}", entry.code, entry.pattern));
309    }
310    lines.push("§/dict§".to_string());
311    lines.join("\n")
312}
313
314// ── Tests ─────────────────────────────────────────────────────────────────────
315
316#[cfg(test)]
317mod tests {
318    use super::*;
319
320    #[test]
321    fn test_new_has_builtin_entries() {
322        let comp = DictCompressor::new();
323        assert!(comp.dict_size() > 0);
324    }
325
326    #[test]
327    fn test_compress_with_known_fields() {
328        let comp = DictCompressor::new();
329        let json = r#"{"status_code":200,"message":"ok","timestamp":"2024-01-01"}"#;
330        let result = comp.compress(json).unwrap();
331        assert!(result.substitutions > 0);
332        assert!(result.bytes_saved > 0);
333        assert!(result.has_dict_header);
334        assert!(result.data.contains("§dict§"));
335    }
336
337    #[test]
338    fn test_compress_no_known_fields() {
339        let mut comp = DictCompressor::new();
340        comp.entries.clear(); // Remove builtins for this test
341        let json = r#"{"x":1,"y":2}"#;
342        let result = comp.compress(json).unwrap();
343        assert_eq!(result.substitutions, 0);
344        assert_eq!(result.bytes_saved, 0);
345        assert!(!result.has_dict_header);
346        assert_eq!(result.data, json);
347    }
348
349    #[test]
350    fn test_observe_learns_patterns() {
351        let mut comp = DictCompressor::new();
352        let initial_size = comp.dict_size();
353
354        // Observe the same JSON twice to meet min_frequency
355        let json = r#"{"custom_field_name":1,"another_long_field":2}"#;
356        comp.observe(json);
357        comp.observe(json);
358
359        assert!(comp.dict_size() > initial_size);
360    }
361
362    #[test]
363    fn test_observe_ignores_short_fields() {
364        let mut comp = DictCompressor::new();
365        let initial_size = comp.dict_size();
366        let json = r#"{"id":1,"x":2}"#;
367        comp.observe(json);
368        comp.observe(json);
369        // Short fields should not be added
370        assert_eq!(comp.dict_size(), initial_size);
371    }
372
373    #[test]
374    fn test_extract_json_fields() {
375        let json = r#"{"name":"Alice","age":30,"nested":{"key":"val"}}"#;
376        let fields = extract_json_fields(json);
377        assert!(fields.contains(&"name".to_string()));
378        assert!(fields.contains(&"age".to_string()));
379        assert!(fields.contains(&"nested".to_string()));
380        assert!(fields.contains(&"key".to_string()));
381    }
382
383    #[test]
384    fn test_extract_json_fields_with_escaped_quotes() {
385        let json = r#"{"field_with_\"quotes\"":"value","normal":1}"#;
386        let fields = extract_json_fields(json);
387        assert!(fields.contains(&"normal".to_string()));
388    }
389
390    #[test]
391    fn test_reset_clears_learned_entries() {
392        let mut comp = DictCompressor::new();
393        let json = r#"{"custom_field_name":1}"#;
394        comp.observe(json);
395        comp.observe(json);
396        let size_after_learn = comp.dict_size();
397
398        comp.reset();
399        // Should be back to just builtins
400        assert!(comp.dict_size() <= size_after_learn);
401        assert!(comp.candidates.is_empty());
402    }
403
404    #[test]
405    fn test_dict_header_format() {
406        let entry = DictEntry {
407            pattern: "status_code".to_string(),
408            code: "~s".to_string(),
409            frequency: 10,
410        };
411        let entries = vec![&entry];
412        let header = format_dict_header(&entries);
413        assert!(header.starts_with("§dict§"));
414        assert!(header.ends_with("§/dict§"));
415        assert!(header.contains("~s=status_code"));
416    }
417
418    // ── Property tests ────────────────────────────────────────────────────────
419
420    use proptest::prelude::*;
421
422    proptest! {
423        /// Compression never increases the logical content size
424        /// (excluding the dictionary header).
425        #[test]
426        fn prop_compression_does_not_increase_content(
427            field_name in "[a-z_]{5,15}",
428            value in "[a-z0-9]{1,20}",
429            repeat in 2usize..=5usize,
430        ) {
431            let mut comp = DictCompressor::new();
432            comp.entries.clear(); // Start fresh
433
434            let entry = format!("\"{}\":\"{}\"", field_name, value);
435            let json = format!("{{{}}}", std::iter::repeat(entry.as_str())
436                .take(repeat)
437                .collect::<Vec<_>>()
438                .join(","));
439
440            // Observe enough times to learn the pattern
441            for _ in 0..3 {
442                comp.observe(&json);
443            }
444
445            let result = comp.compress(&json).unwrap();
446
447            // If substitutions were made, the content portion should be shorter
448            if result.substitutions > 0 {
449                prop_assert!(
450                    result.bytes_saved > 0,
451                    "substitutions made but no bytes saved"
452                );
453            }
454        }
455    }
456}
sqz_engine/dict_compressor.rs

sqz_engine/
dict_compressor.rs
