Skip to main content

lean_ctx/core/
entropy.rs

1use std::collections::{HashMap, HashSet};
2use std::hash::{Hash, Hasher};
3use std::io::Write;
4
5use flate2::Compression;
6use flate2::write::GzEncoder;
7
8use super::tokens::{count_tokens, encode_tokens};
9
10const BPE_ENTROPY_THRESHOLD: f64 = 1.0;
11
12/// Result of entropy-based compression: output text, token counts, and techniques used.
13#[derive(Debug)]
14pub struct EntropyResult {
15    pub output: String,
16    pub original_tokens: usize,
17    pub compressed_tokens: usize,
18    pub techniques: Vec<String>,
19}
20
21impl EntropyResult {
22    /// Returns the percentage of tokens saved by compression.
23    pub fn savings_percent(&self) -> f64 {
24        if self.original_tokens == 0 {
25            return 0.0;
26        }
27        let saved = self.original_tokens.saturating_sub(self.compressed_tokens);
28        (saved as f64 / self.original_tokens as f64) * 100.0
29    }
30}
31
32/// Computes Shannon entropy (bits) over character frequencies in the text.
33pub fn shannon_entropy(text: &str) -> f64 {
34    if text.is_empty() {
35        return 0.0;
36    }
37    let mut freq: HashMap<char, usize> = HashMap::new();
38    let total = text.chars().count();
39
40    for c in text.chars() {
41        *freq.entry(c).or_default() += 1;
42    }
43
44    freq.values().fold(0.0_f64, |acc, &count| {
45        let p = count as f64 / total as f64;
46        acc - p * p.log2()
47    })
48}
49
50/// Shannon entropy over already-encoded BPE token IDs (o200k_base).
51pub fn token_entropy_from_ids(tokens: &[u32]) -> f64 {
52    if tokens.is_empty() {
53        return 0.0;
54    }
55    let total = tokens.len();
56    let mut freq: HashMap<u32, usize> = HashMap::new();
57    for &t in tokens {
58        *freq.entry(t).or_default() += 1;
59    }
60    // Sum in a canonical (sorted-by-count) order: f64 addition is not
61    // associative, so folding over HashMap iteration order — random per instance
62    // — would make the entropy, and thus `density`/`entropy` mode output,
63    // non-deterministic across processes and defeat prompt caching (#498).
64    let mut counts: Vec<usize> = freq.into_values().collect();
65    counts.sort_unstable();
66    counts.iter().fold(0.0_f64, |acc, &count| {
67        let p = count as f64 / total as f64;
68        acc - p * p.log2()
69    })
70}
71
72/// Shannon entropy over BPE token IDs (o200k_base).
73/// More LLM-relevant than character entropy since LLMs process BPE tokens.
74pub fn token_entropy(text: &str) -> f64 {
75    let tokens = encode_tokens(text);
76    token_entropy_from_ids(&tokens)
77}
78
79/// Normalized Shannon entropy over encoded token IDs: H(X) / log₂(n), n = unique token count.
80pub fn normalized_token_entropy_from_ids(tokens: &[u32]) -> f64 {
81    if tokens.is_empty() {
82        return 0.0;
83    }
84    let total = tokens.len();
85    let mut freq: HashMap<u32, usize> = HashMap::new();
86    for &t in tokens {
87        *freq.entry(t).or_default() += 1;
88    }
89    let n_unique = freq.len();
90    if n_unique <= 1 {
91        return 0.0;
92    }
93    // Canonical summation order for determinism (#498); see `token_entropy_from_ids`.
94    let mut counts: Vec<usize> = freq.into_values().collect();
95    counts.sort_unstable();
96    let h = counts.iter().fold(0.0_f64, |acc, &count| {
97        let p = count as f64 / total as f64;
98        acc - p * p.log2()
99    });
100    let h_max = (n_unique as f64).log2();
101    h / h_max
102}
103
104/// Normalized Shannon entropy: H(X) / log₂(n) where n = number of unique symbols.
105/// Returns a value in [0, 1] where 0 = perfectly predictable, 1 = maximum entropy.
106/// This makes thresholds comparable across different alphabet sizes.
107pub fn normalized_token_entropy(text: &str) -> f64 {
108    let tokens = encode_tokens(text);
109    normalized_token_entropy_from_ids(&tokens)
110}
111
112/// Computes word-set Jaccard similarity between two strings (0.0–1.0).
113pub fn jaccard_similarity(a: &str, b: &str) -> f64 {
114    let set_a: HashSet<&str> = a.split_whitespace().collect();
115    let set_b: HashSet<&str> = b.split_whitespace().collect();
116
117    let intersection = set_a.intersection(&set_b).count();
118    let union = set_a.union(&set_b).count();
119
120    if union == 0 {
121        return 0.0;
122    }
123    intersection as f64 / union as f64
124}
125
126/// N-gram Jaccard similarity — preserves word order (unlike word-set Jaccard).
127pub fn ngram_jaccard(a: &str, b: &str, n: usize) -> f64 {
128    let set_a = ngram_set(a, n);
129    let set_b = ngram_set(b, n);
130
131    let intersection = set_a.intersection(&set_b).count();
132    let union = set_a.union(&set_b).count();
133
134    if union == 0 {
135        return 0.0;
136    }
137    intersection as f64 / union as f64
138}
139
140fn ngram_set(text: &str, n: usize) -> HashSet<Vec<String>> {
141    let words: Vec<&str> = text.split_whitespace().collect();
142    if words.len() < n {
143        let mut set = HashSet::new();
144        if !words.is_empty() {
145            set.insert(words.iter().map(std::string::ToString::to_string).collect());
146        }
147        return set;
148    }
149    words
150        .windows(n)
151        .map(|w| w.iter().map(std::string::ToString::to_string).collect())
152        .collect()
153}
154
155/// Minhash signature for approximate Jaccard via LSH.
156/// Uses k independent hash functions (polynomial hashing with different seeds).
157pub fn minhash_signature(text: &str, n: usize, k: usize) -> Vec<u64> {
158    let ngrams = ngram_set(text, n);
159    if ngrams.is_empty() {
160        return vec![u64::MAX; k];
161    }
162    let mut signature = vec![u64::MAX; k];
163    for ngram in &ngrams {
164        for (i, min) in signature.iter_mut().enumerate() {
165            let h = hash_with_seed(ngram, i as u64);
166            if h < *min {
167                *min = h;
168            }
169        }
170    }
171    signature
172}
173
174/// Approximate Jaccard from two minhash signatures.
175pub fn minhash_similarity(sig_a: &[u64], sig_b: &[u64]) -> f64 {
176    if sig_a.len() != sig_b.len() || sig_a.is_empty() {
177        return 0.0;
178    }
179    let matches = sig_a
180        .iter()
181        .zip(sig_b.iter())
182        .filter(|(a, b)| a == b)
183        .count();
184    matches as f64 / sig_a.len() as f64
185}
186
187fn hash_with_seed<T: Hash>(value: &T, seed: u64) -> u64 {
188    let mut hasher = std::collections::hash_map::DefaultHasher::new();
189    seed.hash(&mut hasher);
190    value.hash(&mut hasher);
191    hasher.finish()
192}
193
194/// Kolmogorov complexity proxy: K(x) ≈ len(gzip(x)) / len(x).
195/// Lower values = more compressible = more redundant.
196pub fn kolmogorov_proxy(content: &str) -> f64 {
197    if content.is_empty() {
198        return 1.0;
199    }
200    let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
201    encoder.write_all(content.as_bytes()).ok();
202    let compressed = encoder.finish().unwrap_or_default();
203    compressed.len() as f64 / content.len() as f64
204}
205
206/// Classification of content compressibility based on Kolmogorov proxy (gzip ratio).
207#[derive(Debug, Clone, Copy, PartialEq, Eq)]
208pub enum CompressibilityClass {
209    High,
210    Medium,
211    Low,
212}
213
214impl CompressibilityClass {
215    /// Returns a human-readable label with the Kolmogorov threshold range.
216    pub fn label(&self) -> &'static str {
217        match self {
218            Self::High => "high (K<0.3)",
219            Self::Medium => "medium (0.3≤K<0.6)",
220            Self::Low => "low (K≥0.6)",
221        }
222    }
223}
224
225/// Classify how compressible content is based on gzip ratio.
226pub fn compressibility_class(content: &str) -> CompressibilityClass {
227    let k = kolmogorov_proxy(content);
228    if k < 0.3 {
229        CompressibilityClass::High
230    } else if k < 0.6 {
231        CompressibilityClass::Medium
232    } else {
233        CompressibilityClass::Low
234    }
235}
236
237/// Compresses content by removing low-entropy lines and deduplicating patterns.
238pub fn entropy_compress(content: &str) -> EntropyResult {
239    entropy_compress_with_thresholds(content, BPE_ENTROPY_THRESHOLD, 0.7, &[])
240}
241
242/// Entropy compression with the opportunistic semantic redundancy filter
243/// (#544) pinned OFF. The regular path uses the shared embedding engine
244/// whenever it happens to be loaded, so its output depends on runtime state.
245/// Benchmarks and the scorecard (#211) need run-to-run and machine-to-machine
246/// reproducibility, so they must go through this entry point.
247pub fn entropy_compress_deterministic(content: &str) -> EntropyResult {
248    entropy_compress_inner(content, BPE_ENTROPY_THRESHOLD, 0.7, &[], false, &[])
249}
250
251/// Entropy compression with file-type-adaptive thresholds and event emission.
252/// `force_keep` lines (explicit `protect` tokens, #709) survive verbatim; pass
253/// `&[]` to keep the pre-protect behaviour byte-identical (#498).
254pub fn entropy_compress_adaptive(
255    content: &str,
256    path: &str,
257    force_keep: &[String],
258) -> EntropyResult {
259    let thresholds = super::adaptive_thresholds::adaptive_thresholds(path, content);
260    let before_lines = content.lines().count() as u32;
261    let result = entropy_compress_with_thresholds(
262        content,
263        thresholds.bpe_entropy,
264        thresholds.jaccard,
265        force_keep,
266    );
267    let after_lines = result.output.lines().count() as u32;
268
269    if before_lines != after_lines {
270        super::events::emit(super::events::EventKind::Compression {
271            path: path.to_string(),
272            before_lines,
273            after_lines,
274            strategy: "entropy_adaptive".to_string(),
275            kept_line_count: after_lines,
276            removed_line_count: before_lines.saturating_sub(after_lines),
277        });
278    }
279
280    result
281}
282
283/// Like [`entropy_compress_adaptive`] but overrides the learned BPE-entropy
284/// threshold (e.g. from the aggressiveness knob) while keeping the file-adaptive
285/// jaccard. Pure function of its inputs (#498). Higher `bpe_entropy` drops more
286/// low-information lines.
287pub fn entropy_compress_with_threshold(
288    content: &str,
289    path: &str,
290    bpe_entropy: f64,
291    force_keep: &[String],
292) -> EntropyResult {
293    let thresholds = super::adaptive_thresholds::adaptive_thresholds(path, content);
294    entropy_compress_with_thresholds(content, bpe_entropy, thresholds.jaccard, force_keep)
295}
296
297/// Task-conditioned entropy compression: lines that would normally be dropped
298/// for low entropy are kept if they contain task-relevant keywords.  This is
299/// the Information Bottleneck proxy: we compress away only what is neither
300/// surprising (high H) *nor* task-relevant (mentions goal concepts).
301/// Falls back to pure entropy when `task_keywords` is empty.
302pub fn entropy_compress_task_conditioned(
303    content: &str,
304    path: &str,
305    task_keywords: &[String],
306    force_keep: &[String],
307) -> EntropyResult {
308    let thresholds = super::adaptive_thresholds::adaptive_thresholds(path, content);
309    let before_lines = content.lines().count() as u32;
310    let result = entropy_compress_with_task(
311        content,
312        thresholds.bpe_entropy,
313        thresholds.jaccard,
314        task_keywords,
315        force_keep,
316    );
317    let after_lines = result.output.lines().count() as u32;
318    if before_lines != after_lines {
319        super::events::emit(super::events::EventKind::Compression {
320            path: path.to_string(),
321            before_lines,
322            after_lines,
323            strategy: "entropy_task_conditioned".to_string(),
324            kept_line_count: after_lines,
325            removed_line_count: before_lines.saturating_sub(after_lines),
326        });
327    }
328    result
329}
330
331/// Real line embedder for the semantic redundancy filter (#544): uses the
332/// shared neural engine only when it is ALREADY loaded (never blocks a read
333/// on a model load) and only for files small enough that per-line embedding
334/// stays within the read-path latency budget. Embeddings are cached by line
335/// hash so re-reads cost nothing.
336#[cfg(feature = "embeddings")]
337fn line_embedder(line_count: usize) -> impl Fn(&str) -> Option<Vec<f32>> {
338    use std::collections::HashMap;
339    use std::sync::Mutex;
340
341    const MAX_LINES_FOR_SEMANTIC: usize = 400;
342    const CACHE_CAP: usize = 8192;
343    static LINE_EMBED_CACHE: Mutex<Option<HashMap<u64, Vec<f32>>>> = Mutex::new(None);
344
345    let engine = if line_count <= MAX_LINES_FOR_SEMANTIC {
346        crate::tools::ctx_knowledge::embeddings::embedding_engine_nonblocking()
347    } else {
348        None
349    };
350
351    move |line: &str| {
352        let engine = engine?;
353        if line.len() < 8 {
354            return None;
355        }
356        let mut hasher = std::collections::hash_map::DefaultHasher::new();
357        std::hash::Hash::hash(&line, &mut hasher);
358        let key = std::hash::Hasher::finish(&hasher);
359
360        if let Ok(mut guard) = LINE_EMBED_CACHE.lock()
361            && let Some(hit) = guard.get_or_insert_with(HashMap::new).get(&key)
362        {
363            return Some(hit.clone());
364        }
365        let emb = engine.embed(line).ok()?;
366        if let Ok(mut guard) = LINE_EMBED_CACHE.lock() {
367            let map = guard.get_or_insert_with(HashMap::new);
368            if map.len() >= CACHE_CAP {
369                map.clear();
370            }
371            map.insert(key, emb.clone());
372        }
373        Some(emb)
374    }
375}
376
377#[cfg(not(feature = "embeddings"))]
378fn line_embedder(_line_count: usize) -> impl Fn(&str) -> Option<Vec<f32>> {
379    |_: &str| None
380}
381
382fn entropy_compress_with_task(
383    content: &str,
384    entropy_threshold: f64,
385    jaccard_threshold: f64,
386    task_keywords: &[String],
387    force_keep: &[String],
388) -> EntropyResult {
389    entropy_compress_inner(
390        content,
391        entropy_threshold,
392        jaccard_threshold,
393        task_keywords,
394        true,
395        force_keep,
396    )
397}
398
399fn entropy_compress_inner(
400    content: &str,
401    entropy_threshold: f64,
402    jaccard_threshold: f64,
403    task_keywords: &[String],
404    semantic: bool,
405    force_keep: &[String],
406) -> EntropyResult {
407    let original_tokens = count_tokens(content);
408    let mut lines: Vec<&str> = content.lines().collect();
409    let mut techniques = Vec::new();
410
411    let kw_lower: Vec<String> = task_keywords.iter().map(|k| k.to_lowercase()).collect();
412    let original_count = lines.len();
413    let mut task_rescued = 0usize;
414    // Semantic redundancy (#544): when the embedding engine is already
415    // loaded, kept lines that embed near-identically to earlier kept lines
416    // are dropped (MMR). Without the engine the closure returns None and the
417    // decision path is byte-identical to the Zipf-only filter.
418    // `semantic=false` (deterministic contract) requests an embedder above
419    // the size cutoff, which never resolves an engine.
420    let embed = line_embedder(if semantic { original_count } else { usize::MAX });
421    let mut scoring_ctx = super::surprise::ScoringCtx::new();
422    lines.retain(|line| {
423        let trimmed = line.trim();
424        // Explicit protect tokens (#709) win over every lossy heuristic: a line
425        // containing one is force-kept verbatim before any threshold is applied.
426        if super::protect::line_is_protected(line, force_keep) {
427            return true;
428        }
429        if super::surprise::should_keep_line_semantic(
430            trimmed,
431            entropy_threshold,
432            &embed,
433            &mut scoring_ctx,
434        ) {
435            return true;
436        }
437        // Task-conditioned rescue: keep low-entropy lines that mention task keywords.
438        if !kw_lower.is_empty() {
439            let lower = trimmed.to_lowercase();
440            if kw_lower.iter().any(|kw| lower.contains(kw.as_str())) {
441                task_rescued += 1;
442                return true;
443            }
444        }
445        false
446    });
447    let removed = original_count - lines.len();
448    if removed > 0 || task_rescued > 0 {
449        let mut msg = format!("⊘ {removed} low-entropy lines (BPE H<{entropy_threshold:.2})");
450        if task_rescued > 0 {
451            msg.push_str(&format!(" [+{task_rescued} task-rescued]"));
452        }
453        techniques.push(msg);
454    }
455
456    let blocks = extract_blocks(&lines);
457    let groups = find_pattern_groups(&blocks, jaccard_threshold);
458    let mut dedup_count = 0;
459    for group in &groups {
460        if group.len() > 1 {
461            dedup_count += group.len() - 1;
462        }
463    }
464    if dedup_count > 0 {
465        techniques.push(format!("⊘ {dedup_count} duplicate patterns (J≥0.7)"));
466    }
467
468    let mut result: Vec<String> = Vec::new();
469    let mut skip_indices: std::collections::HashSet<usize> = std::collections::HashSet::new();
470    for group in &groups {
471        if group.len() > 1 {
472            for &idx in &group[1..] {
473                // Protected lines (#709) are a hard keep: never dedup them away,
474                // not just exempt them from the entropy threshold above.
475                if !super::protect::line_is_protected(lines[idx], force_keep) {
476                    skip_indices.insert(idx);
477                }
478            }
479        }
480    }
481    for (i, line) in lines.iter().enumerate() {
482        if !skip_indices.contains(&i) {
483            result.push(line.to_string());
484        }
485    }
486
487    let mut collapsed = Vec::new();
488    let mut blank_count = 0;
489    for line in &result {
490        if line.trim().is_empty() {
491            blank_count += 1;
492            if blank_count <= 1 {
493                collapsed.push(line.clone());
494            }
495        } else {
496            blank_count = 0;
497            collapsed.push(line.clone());
498        }
499    }
500    let output = collapsed.join("\n");
501    let compressed_tokens = count_tokens(&output);
502
503    // Safeguard: BPE re-tokenization of a line subset can, for tiny adversarial
504    // inputs, exceed the original token count (e.g. a dropped trailing newline
505    // merges differently). Never inflate — fall back to the original verbatim.
506    let final_output = if compressed_tokens > original_tokens {
507        content.to_string()
508    } else {
509        output
510    };
511    let final_tokens = if compressed_tokens > original_tokens {
512        original_tokens
513    } else {
514        compressed_tokens
515    };
516
517    EntropyResult {
518        output: final_output,
519        original_tokens,
520        compressed_tokens: final_tokens,
521        techniques,
522    }
523}
524
525fn entropy_compress_with_thresholds(
526    content: &str,
527    entropy_threshold: f64,
528    jaccard_threshold: f64,
529    force_keep: &[String],
530) -> EntropyResult {
531    entropy_compress_with_task(
532        content,
533        entropy_threshold,
534        jaccard_threshold,
535        &[],
536        force_keep,
537    )
538}
539
540/// Budget-based compression to a target density (SDE principle: aim for a
541/// *target* information density instead of maximum compression).
542///
543/// Keeps the highest-entropy lines — in original order — until the BPE token
544/// budget `original_tokens * target` is exhausted. Deterministic: ties break
545/// on line index. `target` is clamped to [0.05, 1.0]; the top-scored line is
546/// always kept so output is never empty for non-empty input.
547pub fn entropy_compress_to_density(content: &str, target: f64) -> EntropyResult {
548    let target = target.clamp(0.05, 1.0);
549    let original_tokens = count_tokens(content);
550    if content.is_empty() || original_tokens == 0 {
551        return EntropyResult {
552            output: String::new(),
553            original_tokens,
554            compressed_tokens: 0,
555            techniques: vec![format!("density target={target:.2} (empty input)")],
556        };
557    }
558
559    #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
560    let budget = ((original_tokens as f64) * target).ceil() as usize;
561
562    let lines: Vec<&str> = content.lines().collect();
563    let mut scored: Vec<(usize, f64, usize)> = lines
564        .iter()
565        .enumerate()
566        .map(|(i, l)| {
567            let trimmed = l.trim();
568            // +1 approximates the newline token; keeps the per-line sum close
569            // to the whole-file count so the budget is meaningful.
570            let toks = count_tokens(trimmed).max(1);
571            (i, token_entropy(trimmed), toks)
572        })
573        .collect();
574    scored.sort_by(|a, b| {
575        b.1.partial_cmp(&a.1)
576            .unwrap_or(std::cmp::Ordering::Equal)
577            .then(a.0.cmp(&b.0))
578    });
579
580    let mut keep = vec![false; lines.len()];
581    let mut used = 0usize;
582    let mut kept_count = 0usize;
583    // The single most informative line is always kept, even if it alone
584    // exceeds the budget — a density target that drops the highest-signal
585    // line would be self-defeating. Everything else fills greedily.
586    if let Some(&(idx, _, toks)) = scored.first() {
587        keep[idx] = true;
588        used += toks;
589        kept_count += 1;
590    }
591    for &(idx, _h, toks) in scored.iter().skip(1) {
592        if used + toks > budget {
593            // Greedy knapsack: skip lines that overshoot, smaller ones may fit.
594            continue;
595        }
596        keep[idx] = true;
597        used += toks;
598        kept_count += 1;
599    }
600
601    let mut out_lines: Vec<&str> = Vec::with_capacity(kept_count);
602    for (i, line) in lines.iter().enumerate() {
603        if keep[i] {
604            out_lines.push(line);
605        }
606    }
607    let output = out_lines.join("\n");
608    let compressed_tokens = count_tokens(&output);
609
610    let dropped = lines.len() - kept_count;
611    EntropyResult {
612        output,
613        original_tokens,
614        compressed_tokens,
615        techniques: vec![format!(
616            "density target={target:.2} budget={budget} tok, kept {kept_count}/{} lines (⊘ {dropped})",
617            lines.len()
618        )],
619    }
620}
621
622/// Per-line entropy statistics for a block of content.
623#[derive(Debug)]
624pub struct EntropyAnalysis {
625    pub avg_entropy: f64,
626    pub low_entropy_count: usize,
627    pub high_entropy_count: usize,
628    pub total_lines: usize,
629}
630
631/// Analyzes per-line BPE token entropy, counting low/high entropy lines.
632pub fn analyze_entropy(content: &str) -> EntropyAnalysis {
633    let lines: Vec<&str> = content.lines().collect();
634    let total = lines.len();
635    let mut sum = 0.0;
636    let mut low = 0;
637    let mut high = 0;
638    let mut counted = 0;
639
640    for line in &lines {
641        let trimmed = line.trim();
642        if trimmed.is_empty() {
643            continue;
644        }
645        let h = token_entropy(trimmed);
646        sum += h;
647        counted += 1;
648        if h < BPE_ENTROPY_THRESHOLD {
649            low += 1;
650        }
651        if h > 3.0 {
652            high += 1;
653        }
654    }
655
656    EntropyAnalysis {
657        avg_entropy: if counted > 0 {
658            sum / counted as f64
659        } else {
660            0.0
661        },
662        low_entropy_count: low,
663        high_entropy_count: high,
664        total_lines: total,
665    }
666}
667
668struct Block {
669    content: String,
670}
671
672fn extract_blocks(lines: &[&str]) -> Vec<Block> {
673    let mut blocks = Vec::new();
674    let mut current = String::new();
675
676    for line in lines {
677        let trimmed = line.trim();
678        if trimmed.is_empty() && !current.is_empty() {
679            blocks.push(Block {
680                content: current.clone(),
681            });
682            current.clear();
683        } else if !trimmed.is_empty() {
684            current.push_str(trimmed);
685            current.push('\n');
686        }
687    }
688
689    if !current.is_empty() {
690        blocks.push(Block { content: current });
691    }
692
693    blocks
694}
695
696fn find_pattern_groups(blocks: &[Block], threshold: f64) -> Vec<Vec<usize>> {
697    // Exact n-gram Jaccard, but with precomputed n-gram sets per block to avoid
698    // rebuilding allocations per pair. Includes a size-ratio impossibility check
699    // (max possible Jaccard is |A|/|B| for |A|<=|B|).
700    let sets: Vec<HashSet<Vec<String>>> = blocks.iter().map(|b| ngram_set(&b.content, 2)).collect();
701    let sizes: Vec<usize> = sets.iter().map(std::collections::HashSet::len).collect();
702
703    let mut groups: Vec<Vec<usize>> = Vec::new();
704    let mut assigned: HashSet<usize> = HashSet::new();
705
706    for i in 0..blocks.len() {
707        if assigned.contains(&i) {
708            continue;
709        }
710        let mut group = vec![i];
711        for j in (i + 1)..blocks.len() {
712            if assigned.contains(&j) {
713                continue;
714            }
715            let size_i = sizes[i];
716            let size_j = sizes[j];
717            let min_sz = size_i.min(size_j);
718            let max_sz = size_i.max(size_j);
719            if max_sz > 0 && (min_sz as f64) < (threshold * max_sz as f64) {
720                continue;
721            }
722            let inter = sets[i].intersection(&sets[j]).count();
723            let union = size_i + size_j - inter;
724            if union > 0 && (inter as f64 / union as f64) >= threshold {
725                group.push(j);
726                assigned.insert(j);
727            }
728        }
729        if group.len() > 1 {
730            assigned.insert(i);
731        }
732        groups.push(group);
733    }
734
735    groups
736}
737
738#[cfg(test)]
739mod tests {
740    use super::*;
741
742    #[test]
743    fn protect_force_keeps_lines_in_entropy() {
744        // A block of identical low-information lines that entropy + dedup would
745        // normally collapse to a single survivor (or drop entirely).
746        let mut content = String::new();
747        for _ in 0..15 {
748            content.push_str("boilerplate noise line\n");
749        }
750        let baseline =
751            entropy_compress_inner(&content, BPE_ENTROPY_THRESHOLD, 0.7, &[], false, &[]);
752        let protected = entropy_compress_inner(
753            &content,
754            BPE_ENTROPY_THRESHOLD,
755            0.7,
756            &[],
757            false,
758            &["boilerplate".to_string()],
759        );
760        let baseline_hits = baseline.output.matches("boilerplate").count();
761        let protected_hits = protected.output.matches("boilerplate").count();
762        // Every protected line survives verbatim — a hard keep over both the
763        // entropy threshold and the dedup pass.
764        assert_eq!(
765            protected_hits, 15,
766            "all protected lines must survive: {}",
767            protected.output
768        );
769        assert!(
770            protected_hits >= baseline_hits,
771            "protect must keep at least as many lines as the baseline"
772        );
773    }
774
775    #[test]
776    fn empty_force_keep_is_byte_identical() {
777        // The protect feature must not perturb the unprotected path (#498).
778        let content = "fn a() {}\n// note\nlet x = 1;\nlet x = 1;\n// note\n";
779        let a = entropy_compress_inner(content, BPE_ENTROPY_THRESHOLD, 0.7, &[], false, &[]);
780        let b = entropy_compress_deterministic(content);
781        assert_eq!(a.output, b.output);
782    }
783
784    #[test]
785    fn shannon_entropy_empty_is_zero() {
786        assert_eq!(shannon_entropy(""), 0.0);
787    }
788
789    #[test]
790    fn shannon_entropy_single_char() {
791        assert_eq!(shannon_entropy("aaaa"), 0.0);
792    }
793
794    #[test]
795    fn shannon_entropy_high_for_varied_text() {
796        let varied = "abcdefghijklmnopqrstuvwxyz0123456789";
797        let uniform = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa";
798        assert!(
799            shannon_entropy(varied) > shannon_entropy(uniform),
800            "varied text should have higher entropy"
801        );
802    }
803
804    #[test]
805    fn jaccard_identical_is_one() {
806        let sim = jaccard_similarity("hello world", "hello world");
807        assert!((sim - 1.0).abs() < f64::EPSILON);
808    }
809
810    #[test]
811    fn jaccard_disjoint_is_zero() {
812        let sim = jaccard_similarity("abc", "xyz");
813        assert_eq!(sim, 0.0);
814    }
815
816    #[test]
817    fn jaccard_partial_overlap() {
818        let sim = jaccard_similarity("hello world", "hello rust");
819        assert!(sim > 0.0 && sim < 1.0);
820    }
821
822    #[test]
823    fn entropy_compress_produces_output() {
824        let content = "fn main() {\n    println!(\"hello\");\n}\n\n// comment\n// another comment\n\nfn helper() {\n    let x = 42;\n}\n";
825        let result = entropy_compress(content);
826        assert!(!result.output.is_empty(), "should produce non-empty output");
827        assert!(result.compressed_tokens <= result.original_tokens);
828    }
829
830    #[test]
831    fn entropy_result_savings() {
832        let r = EntropyResult {
833            output: "short".to_string(),
834            original_tokens: 100,
835            compressed_tokens: 60,
836            techniques: vec!["test".to_string()],
837        };
838        assert!((r.savings_percent() - 40.0).abs() < 0.1);
839    }
840
841    #[test]
842    fn entropy_result_zero_original() {
843        let r = EntropyResult {
844            output: String::new(),
845            original_tokens: 0,
846            compressed_tokens: 0,
847            techniques: vec![],
848        };
849        assert_eq!(r.savings_percent(), 0.0);
850    }
851
852    #[test]
853    fn token_entropy_empty_is_zero() {
854        assert_eq!(token_entropy(""), 0.0);
855    }
856
857    #[test]
858    fn token_entropy_single_repeated_token() {
859        assert_eq!(token_entropy("}"), 0.0);
860    }
861
862    #[test]
863    fn token_entropy_higher_for_diverse_code() {
864        let diverse = "let result = compute_something(x, y, z);";
865        let repetitive = "aaaa aaaa aaaa aaaa";
866        assert!(
867            token_entropy(diverse) > token_entropy(repetitive),
868            "diverse code should have higher BPE token entropy"
869        );
870    }
871
872    #[test]
873    fn token_entropy_vs_char_entropy_differ() {
874        let code = "fn main() { println!(\"hello world\"); }";
875        let te = token_entropy(code);
876        let ce = shannon_entropy(code);
877        assert!(te != ce, "BPE and char entropy should differ for code");
878    }
879
880    #[test]
881    fn ngram_jaccard_preserves_order() {
882        let a = "a b c d";
883        let b = "d c b a";
884        let word_j = jaccard_similarity(a, b);
885        let ngram_j = ngram_jaccard(a, b, 2);
886        assert!(
887            ngram_j < word_j,
888            "reordered text should have lower bigram Jaccard ({ngram_j}) than word Jaccard ({word_j})"
889        );
890    }
891
892    #[test]
893    fn ngram_jaccard_identical_is_one() {
894        let text = "fn main() { println!(\"hello\"); }";
895        let j = ngram_jaccard(text, text, 2);
896        assert!((j - 1.0).abs() < f64::EPSILON);
897    }
898
899    #[test]
900    fn ngram_jaccard_disjoint_is_zero() {
901        let j = ngram_jaccard("alpha beta gamma", "delta epsilon zeta", 2);
902        assert_eq!(j, 0.0);
903    }
904
905    #[test]
906    fn minhash_approximates_jaccard() {
907        let a = "fn main() { let x = 1; let y = 2; let z = x + y; println!(z); }";
908        let b = "fn main() { let x = 1; let y = 2; let z = x + y; return z; }";
909        let exact = ngram_jaccard(a, b, 2);
910        let sig_a = minhash_signature(a, 2, 128);
911        let sig_b = minhash_signature(b, 2, 128);
912        let approx = minhash_similarity(&sig_a, &sig_b);
913        assert!(
914            (exact - approx).abs() < 0.2,
915            "minhash ({approx}) should approximate exact ({exact}) within 0.2"
916        );
917    }
918
919    #[test]
920    fn minhash_empty_text() {
921        let sig = minhash_signature("", 2, 64);
922        assert!(sig.iter().all(|&v| v == u64::MAX));
923    }
924
925    #[test]
926    fn kolmogorov_empty_is_one() {
927        assert_eq!(kolmogorov_proxy(""), 1.0);
928    }
929
930    #[test]
931    fn kolmogorov_repetitive_is_low() {
932        let repetitive = "aaa\n".repeat(1000);
933        let k = kolmogorov_proxy(&repetitive);
934        assert!(
935            k < 0.1,
936            "highly repetitive text should compress well: K={k}"
937        );
938    }
939
940    #[test]
941    fn kolmogorov_diverse_is_higher() {
942        let repetitive = "aaa\n".repeat(500);
943        let diverse = (0..500)
944            .map(|i| format!("line_{i}_unique_content_{}", i * 17 % 97))
945            .collect::<Vec<_>>()
946            .join("\n");
947        assert!(
948            kolmogorov_proxy(&diverse) > kolmogorov_proxy(&repetitive),
949            "diverse content should have higher K than repetitive"
950        );
951    }
952
953    #[test]
954    fn compressibility_class_repetitive_is_high() {
955        let text = "use std::io;\n".repeat(200);
956        assert_eq!(compressibility_class(&text), CompressibilityClass::High);
957    }
958
959    #[test]
960    fn kolmogorov_diverse_higher_than_repetitive() {
961        let rep = "test\n".repeat(500);
962        let diverse = (0..500)
963            .map(|i| format!("unique_line_{i}_xk{}", i * 31 % 1000))
964            .collect::<Vec<_>>()
965            .join("\n");
966        assert!(
967            kolmogorov_proxy(&diverse) > kolmogorov_proxy(&rep),
968            "diverse content should have higher K"
969        );
970    }
971
972    fn density_fixture() -> String {
973        (0..120)
974            .map(|i| {
975                if i % 3 == 0 {
976                    format!(
977                        "fn compute_value_{i}(input: &str, flags: u32) -> Result<Output, Error> {{"
978                    )
979                } else if i % 3 == 1 {
980                    format!("    let intermediate_{i} = transform(input, flags ^ {i});")
981                } else {
982                    "}".to_string()
983                }
984            })
985            .collect::<Vec<_>>()
986            .join("\n")
987    }
988
989    #[test]
990    fn density_respects_token_budget() {
991        let content = density_fixture();
992        let orig = count_tokens(&content);
993        for target in [0.3, 0.5, 0.7] {
994            let r = entropy_compress_to_density(&content, target);
995            let actual = r.compressed_tokens as f64 / orig as f64;
996            assert!(
997                actual <= target + 0.10,
998                "target {target}: actual density {actual:.2} exceeds budget"
999            );
1000            assert!(!r.output.is_empty());
1001        }
1002    }
1003
1004    #[test]
1005    fn density_is_deterministic() {
1006        let content = density_fixture();
1007        let a = entropy_compress_to_density(&content, 0.4);
1008        let b = entropy_compress_to_density(&content, 0.4);
1009        assert_eq!(a.output, b.output);
1010        assert_eq!(a.compressed_tokens, b.compressed_tokens);
1011    }
1012
1013    #[test]
1014    fn density_target_one_keeps_everything() {
1015        let content = density_fixture();
1016        let r = entropy_compress_to_density(&content, 1.0);
1017        assert_eq!(r.output, content);
1018    }
1019
1020    #[test]
1021    fn density_clamps_out_of_range_target() {
1022        let content = density_fixture();
1023        let low = entropy_compress_to_density(&content, 0.0);
1024        assert!(!low.output.is_empty(), "clamped to 0.05, never empty");
1025        let high = entropy_compress_to_density(&content, 5.0);
1026        assert_eq!(high.output, content, "clamped to 1.0 keeps all");
1027    }
1028
1029    #[test]
1030    fn density_empty_input() {
1031        let r = entropy_compress_to_density("", 0.5);
1032        assert_eq!(r.compressed_tokens, 0);
1033        assert!(r.output.is_empty());
1034    }
1035
1036    #[test]
1037    fn density_prefers_high_entropy_lines() {
1038        let content =
1039            "}\n}\n}\nlet complex_result = compute_unique_hash(seed, nonce, payload);\n}\n}";
1040        let r = entropy_compress_to_density(content, 0.6);
1041        assert!(
1042            r.output.contains("compute_unique_hash"),
1043            "high-entropy line must survive: {}",
1044            r.output
1045        );
1046    }
1047}