amql_engine/

compression.rs

//! Pure in-memory compression measurement for benchmarking.
//!
//! Computes the ratio of raw source bytes to serialized annotation bytes.
//! No filesystem or subprocess access — works in WASM, CLI, MCP, and Rust.

use crate::store::Annotation;
use crate::types::RelativePath;
use serde::{Deserialize, Serialize};

/// A source file entry for in-memory benchmarking.
#[cfg_attr(feature = "ts", derive(ts_rs::TS))]
#[cfg_attr(feature = "flow", derive(flowjs_rs::Flow))]
#[cfg_attr(feature = "ts", ts(export))]
#[cfg_attr(feature = "flow", flow(export))]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SourceEntry {
    /// Relative file path.
    pub file: RelativePath,
    /// Full source text.
    pub source: String,
}

/// Compression measurement result from in-memory data.
#[cfg_attr(feature = "ts", derive(ts_rs::TS))]
#[cfg_attr(feature = "flow", derive(flowjs_rs::Flow))]
#[cfg_attr(feature = "ts", ts(export))]
#[cfg_attr(feature = "flow", flow(export))]
#[derive(Debug, Clone, Serialize)]
pub struct CompressionResult {
    /// Total source bytes across all files.
    pub source_bytes: usize,
    /// Estimated source tokens (bytes / 4).
    pub source_tokens: usize,
    /// Serialized annotation JSON bytes.
    pub annotation_bytes: usize,
    /// Estimated annotation tokens (bytes / 4).
    pub annotation_tokens: usize,
    /// Number of files that have at least one annotation.
    pub covered_files: usize,
    /// Total number of files.
    pub total_files: usize,
    /// File coverage ratio (covered / total).
    pub file_coverage: f64,
    /// Compression ratio: source_bytes / annotation_bytes. `None` when annotation_bytes = 0.
    pub ratio: Option<f64>,
    /// Per-file breakdown.
    pub files: Vec<FileCompression>,
}

/// Per-file compression detail.
#[cfg_attr(feature = "ts", derive(ts_rs::TS))]
#[cfg_attr(feature = "flow", derive(flowjs_rs::Flow))]
#[cfg_attr(feature = "ts", ts(export))]
#[cfg_attr(feature = "flow", flow(export))]
#[derive(Debug, Clone, Serialize)]
pub struct FileCompression {
    /// Relative file path.
    pub file: RelativePath,
    /// Source bytes for this file.
    pub source_bytes: usize,
    /// Annotation count for this file.
    pub annotation_count: usize,
    /// Annotation JSON bytes for this file.
    pub annotation_bytes: usize,
    /// Per-file ratio. `None` when annotation_bytes = 0.
    pub ratio: Option<f64>,
}

/// Measure compression from pre-loaded source files and annotations.
///
/// Pure function — no filesystem, no subprocesses. Computes source bytes vs
/// annotation JSON bytes per file and in aggregate. Returns the compression
/// ratio (higher = more compression) and per-file breakdown.
pub fn measure_compression(
    sources: &[SourceEntry],
    annotations: &[Annotation],
) -> CompressionResult {
    let total_files = sources.len();
    let source_bytes: usize = sources.iter().map(|s| s.source.len()).sum();
    let source_tokens = source_bytes / 4;

    // Group annotations by file
    let mut by_file: rustc_hash::FxHashMap<&str, Vec<&Annotation>> =
        rustc_hash::FxHashMap::default();
    for ann in annotations {
        by_file.entry(ann.file.as_ref()).or_default().push(ann);
    }

    let mut files = Vec::with_capacity(total_files);
    let mut covered_files = 0usize;
    let mut total_ann_bytes = 0usize;

    for entry in sources {
        let file_anns = by_file.get(AsRef::<str>::as_ref(&entry.file));
        let ann_count =
            file_anns.map_or(0, |v| v.iter().map(|a| count_annotations(a)).sum::<usize>());
        let ann_bytes = file_anns.map_or(0, |v| {
            serde_json::to_string(v).map(|s| s.len()).unwrap_or(0)
        });
        if ann_count > 0 {
            covered_files += 1;
        }
        total_ann_bytes += ann_bytes;

        files.push(FileCompression {
            file: entry.file.clone(),
            source_bytes: entry.source.len(),
            annotation_count: ann_count,
            annotation_bytes: ann_bytes,
            ratio: if ann_bytes > 0 {
                Some(entry.source.len() as f64 / ann_bytes as f64)
            } else {
                None
            },
        });
    }

    let file_coverage = if total_files > 0 {
        covered_files as f64 / total_files as f64
    } else {
        0.0
    };

    CompressionResult {
        source_bytes,
        source_tokens,
        annotation_bytes: total_ann_bytes,
        annotation_tokens: total_ann_bytes / 4,
        covered_files,
        total_files,
        file_coverage,
        ratio: if total_ann_bytes > 0 {
            Some(source_bytes as f64 / total_ann_bytes as f64)
        } else {
            None
        },
        files,
    }
}

/// Count annotations recursively (including nested children).
fn count_annotations(ann: &Annotation) -> usize {
    1 + ann.children.iter().map(count_annotations).sum::<usize>()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::store::Annotation;
    use crate::types::{Binding, RelativePath, TagName};

    fn make_annotation(file: &str, tag: &str, binding: &str) -> Annotation {
        Annotation {
            tag: TagName::from(tag),
            binding: Binding::from(binding),
            file: RelativePath::from(file),
            attrs: rustc_hash::FxHashMap::default(),
            children: Vec::new(),
        }
    }

    #[test]
    fn empty_input() {
        // Act
        let result = measure_compression(&[], &[]);

        // Assert
        assert_eq!(result.total_files, 0, "no files");
        assert_eq!(result.source_bytes, 0, "no source bytes");
        assert!(result.ratio.is_none(), "no ratio when no annotations");
    }

    #[test]
    fn basic_compression() {
        // Arrange
        let sources = vec![SourceEntry {
            file: RelativePath::from("src/main.ts"),
            source: "function hello() {\n  console.log('hello');\n}\n".to_string(),
        }];
        let annotations = vec![make_annotation("src/main.ts", "function", "hello")];

        // Act
        let result = measure_compression(&sources, &annotations);

        // Assert
        assert_eq!(result.total_files, 1, "one file");
        assert_eq!(result.covered_files, 1, "one covered");
        assert!(result.ratio.is_some(), "has ratio");
        assert!(result.ratio.unwrap() > 0.0, "positive ratio");
        assert_eq!(result.files.len(), 1, "one file detail");
        assert_eq!(result.files[0].annotation_count, 1, "one annotation");
    }

    #[test]
    fn uncovered_file() {
        // Arrange
        let sources = vec![
            SourceEntry {
                file: RelativePath::from("src/a.ts"),
                source: "export const a = 1;".to_string(),
            },
            SourceEntry {
                file: RelativePath::from("src/b.ts"),
                source: "export const b = 2;".to_string(),
            },
        ];
        let annotations = vec![make_annotation("src/a.ts", "const", "a")];

        // Act
        let result = measure_compression(&sources, &annotations);

        // Assert
        assert_eq!(result.total_files, 2, "two files");
        assert_eq!(result.covered_files, 1, "one covered");
        assert_eq!(result.file_coverage, 0.5, "50% coverage");
        assert!(
            result.files[1].ratio.is_none(),
            "uncovered file has no ratio"
        );
    }
}