sqz_engine/

structural_summary.rs

/// Structural summary extraction for source code files.
///
/// Instead of dumping entire files into LLM context, this module extracts
/// just the structural skeleton: imports, function/method signatures, class
/// definitions, and call relationships. The model sees the architecture
/// without the implementation noise — typically ~70% fewer tokens while
/// actually improving navigation.
///
/// Builds on top of `AstParser` (signature extraction) and `DependencyMapper`
/// (import graph), adding **call graph extraction** — which functions call
/// which other functions — to complete the structural picture.
///
/// Output format is a compact, LLM-friendly text representation:
/// ```text
/// # file: src/engine.rs
/// ## imports
/// use crate::pipeline::CompressionPipeline
/// use crate::cache_manager::CacheManager
/// ## types
/// pub struct SqzEngine { ... }
/// ## functions
/// pub fn compress(&self, input: &str) -> Result<CompressedContent>
///   → calls: pipeline.compress, cache.get_or_insert, verifier.check
/// pub fn compress_with_mode(&self, input: &str, mode: CompressionMode) -> Result<CompressedContent>
///   → calls: compress
/// ## dependencies
/// imports: pipeline, cache_manager, verifier
/// imported by: main, cli_proxy
/// ```

use std::collections::{HashMap, HashSet};

use crate::ast_parser::{AstParser, CodeSummary};
use crate::dependency_mapper::DependencyMapper;
use crate::error::Result;

/// Configuration for structural summary generation.
#[derive(Debug, Clone)]
pub struct SummaryConfig {
    /// Include import statements in the summary.
    pub include_imports: bool,
    /// Include function/method signatures.
    pub include_functions: bool,
    /// Include class/struct/interface definitions.
    pub include_types: bool,
    /// Include type aliases.
    pub include_type_aliases: bool,
    /// Extract and include call relationships.
    pub include_calls: bool,
    /// Include dependency graph info (imports/imported-by).
    pub include_dep_graph: bool,
    /// Maximum number of call targets to show per function.
    pub max_calls_per_function: usize,
    /// Minimum file size (chars) to trigger summarization.
    /// Files smaller than this are returned as-is.
    pub min_file_size: usize,
}

impl Default for SummaryConfig {
    fn default() -> Self {
        Self {
            include_imports: true,
            include_functions: true,
            include_types: true,
            include_type_aliases: true,
            include_calls: true,
            include_dep_graph: true,
            max_calls_per_function: 10,
            min_file_size: 500,
        }
    }
}

/// Result of structural summary extraction.
#[derive(Debug, Clone)]
pub struct StructuralSummaryResult {
    /// The compact structural summary text.
    pub summary: String,
    /// Token count of the original source.
    pub tokens_original: u32,
    /// Token count of the summary.
    pub tokens_summary: u32,
    /// Number of functions extracted.
    pub functions_count: usize,
    /// Number of types/classes extracted.
    pub types_count: usize,
    /// Number of call edges discovered.
    pub call_edges: usize,
    /// Compression ratio (summary_tokens / original_tokens).
    pub compression_ratio: f64,
}

/// Extract a structural summary from source code.
///
/// This is the main entry point. Given source code and its language,
/// produces a compact summary containing imports, signatures, call
/// relationships, and dependency info.
pub fn summarize(
    source: &str,
    language: &str,
    file_path: &str,
    config: &SummaryConfig,
    dep_mapper: Option<&DependencyMapper>,
) -> Result<StructuralSummaryResult> {
    let tokens_original = approx_tokens(source);

    // Small files aren't worth summarizing
    if source.len() < config.min_file_size {
        return Ok(StructuralSummaryResult {
            summary: source.to_string(),
            tokens_original,
            tokens_summary: tokens_original,
            functions_count: 0,
            types_count: 0,
            call_edges: 0,
            compression_ratio: 1.0,
        });
    }

    let parser = AstParser::new();
    let code_summary = parser.extract_signatures(source, language).unwrap_or_else(|_| {
        // Fallback: regex-based extraction for unsupported languages
        CodeSummary {
            imports: Vec::new(),
            functions: Vec::new(),
            classes: Vec::new(),
            types: Vec::new(),
            tokens_original,
            tokens_summary: tokens_original,
        }
    });

    // Extract call relationships
    let call_graph = if config.include_calls {
        extract_call_graph(source, &code_summary)
    } else {
        HashMap::new()
    };

    let call_edges: usize = call_graph.values().map(|v| v.len()).sum();

    // Build the summary text
    let mut parts: Vec<String> = Vec::new();

    parts.push(format!("# file: {file_path}"));

    // Imports section
    if config.include_imports && !code_summary.imports.is_empty() {
        parts.push("## imports".to_string());
        for imp in &code_summary.imports {
            parts.push(imp.text.clone());
        }
    }

    // Types section (classes, structs, interfaces)
    if config.include_types && !code_summary.classes.is_empty() {
        parts.push("## types".to_string());
        for cls in &code_summary.classes {
            parts.push(cls.signature.clone());
        }
    }

    // Type aliases
    if config.include_type_aliases && !code_summary.types.is_empty() {
        for ty in &code_summary.types {
            parts.push(ty.signature.clone());
        }
    }

    // Functions section with call relationships
    if config.include_functions && !code_summary.functions.is_empty() {
        parts.push("## functions".to_string());
        for func in &code_summary.functions {
            parts.push(func.signature.clone());
            if config.include_calls {
                if let Some(calls) = call_graph.get(&func.name) {
                    if !calls.is_empty() {
                        let display_calls: Vec<&str> = calls
                            .iter()
                            .take(config.max_calls_per_function)
                            .map(|s| s.as_str())
                            .collect();
                        let suffix = if calls.len() > config.max_calls_per_function {
                            format!(" +{} more", calls.len() - config.max_calls_per_function)
                        } else {
                            String::new()
                        };
                        parts.push(format!(
                            "  \u{2192} calls: {}{}",
                            display_calls.join(", "),
                            suffix
                        ));
                    }
                }
            }
        }
    }

    // Dependency graph section
    if config.include_dep_graph {
        if let Some(mapper) = dep_mapper {
            let path = std::path::Path::new(file_path);
            let dep_summary = mapper.summary(path);
            if !dep_summary.is_empty() {
                parts.push("## dependencies".to_string());
                parts.push(dep_summary);
            }
        }
    }

    let summary = parts.join("\n");
    let tokens_summary = approx_tokens(&summary);
    let compression_ratio = if tokens_original > 0 {
        tokens_summary as f64 / tokens_original as f64
    } else {
        1.0
    };

    Ok(StructuralSummaryResult {
        summary,
        tokens_original,
        tokens_summary,
        functions_count: code_summary.functions.len(),
        types_count: code_summary.classes.len(),
        call_edges,
        compression_ratio,
    })
}

/// Summarize multiple files into a single structural map.
///
/// Useful for giving the model an overview of an entire module or package.
pub fn summarize_multi(
    files: &[(&str, &str, &str)], // (source, language, file_path)
    config: &SummaryConfig,
    dep_mapper: Option<&DependencyMapper>,
) -> Result<StructuralSummaryResult> {
    let mut all_parts: Vec<String> = Vec::new();
    let mut total_original: u32 = 0;
    let mut total_functions: usize = 0;
    let mut total_types: usize = 0;
    let mut total_edges: usize = 0;

    for (source, language, file_path) in files {
        let result = summarize(source, language, file_path, config, dep_mapper)?;
        total_original += result.tokens_original;
        total_functions += result.functions_count;
        total_types += result.types_count;
        total_edges += result.call_edges;
        all_parts.push(result.summary);
    }

    let summary = all_parts.join("\n---\n");
    let tokens_summary = approx_tokens(&summary);
    let compression_ratio = if total_original > 0 {
        tokens_summary as f64 / total_original as f64
    } else {
        1.0
    };

    Ok(StructuralSummaryResult {
        summary,
        tokens_original: total_original,
        tokens_summary,
        functions_count: total_functions,
        types_count: total_types,
        call_edges: total_edges,
        compression_ratio,
    })
}

// ── Call graph extraction ─────────────────────────────────────────────────

/// Extract call relationships from source code.
///
/// For each function in the code summary, scans its body to find calls to
/// other known functions. Uses a combination of:
/// 1. Direct name matching against known function names
/// 2. Method call pattern matching (`.method_name(`)
/// 3. Qualified call matching (`module::function(`)
///
/// Returns a map: caller_name → [callee_names]
fn extract_call_graph(
    source: &str,
    code_summary: &CodeSummary,
) -> HashMap<String, Vec<String>> {
    let mut graph: HashMap<String, Vec<String>> = HashMap::new();

    // Build a set of known function/method names for matching
    let known_names: HashSet<&str> = code_summary
        .functions
        .iter()
        .map(|f| f.name.as_str())
        .collect();

    // Also collect class names for qualified calls
    let known_classes: HashSet<&str> = code_summary
        .classes
        .iter()
        .map(|c| c.name.as_str())
        .collect();

    let lines: Vec<&str> = source.lines().collect();

    // Find function boundaries (start line → end line)
    let boundaries = find_function_boundaries(source, code_summary);

    for (func_name, start, end) in &boundaries {
        let mut calls: Vec<String> = Vec::new();
        let mut seen: HashSet<String> = HashSet::new();

        // Don't include self-references
        seen.insert(func_name.clone());

        for line_idx in *start..*end.min(&lines.len()) {
            let line = lines[line_idx].trim();

            // Skip comments
            if line.starts_with("//") || line.starts_with('#') || line.starts_with("/*") {
                continue;
            }

            // Check for direct function calls: `function_name(`
            for name in &known_names {
                if seen.contains(*name) {
                    continue;
                }
                // Match `name(` but not `some_name(` (word boundary)
                if contains_call(line, name) {
                    calls.push(name.to_string());
                    seen.insert(name.to_string());
                }
            }

            // Check for method calls: `.method_name(`
            for name in &known_names {
                if seen.contains(*name) {
                    continue;
                }
                let pattern = format!(".{}(", name);
                if line.contains(&pattern) {
                    calls.push(name.to_string());
                    seen.insert(name.to_string());
                }
            }

            // Check for qualified calls: `ClassName::method(` or `module::func(`
            for class_name in &known_classes {
                let pattern = format!("{}::", class_name);
                if line.contains(&pattern) {
                    // Extract the method name after ::
                    if let Some(rest) = line.split(&pattern).nth(1) {
                        let method = rest
                            .split(|c: char| !c.is_alphanumeric() && c != '_')
                            .next()
                            .unwrap_or("");
                        if !method.is_empty() && !seen.contains(method) {
                            let qualified = format!("{}.{}", class_name, method);
                            calls.push(qualified);
                            seen.insert(method.to_string());
                        }
                    }
                }
            }
        }

        calls.sort();
        if !calls.is_empty() {
            graph.insert(func_name.clone(), calls);
        }
    }

    graph
}

/// Check if a line contains a function call to `name` with word boundary.
/// Matches `name(` but not `some_name(` or `name_suffix(`.
fn contains_call(line: &str, name: &str) -> bool {
    let pattern = format!("{}(", name);
    let mut search_from = 0;

    while let Some(pos) = line[search_from..].find(&pattern) {
        let abs_pos = search_from + pos;
        // Check left boundary: must be start of line or preceded by non-alphanumeric
        let left_ok = abs_pos == 0
            || !line.as_bytes()[abs_pos - 1].is_ascii_alphanumeric()
                && line.as_bytes()[abs_pos - 1] != b'_';
        if left_ok {
            return true;
        }
        search_from = abs_pos + 1;
    }
    false
}

/// Find approximate line boundaries for each function in the source.
/// Returns (function_name, start_line, end_line).
fn find_function_boundaries(
    source: &str,
    code_summary: &CodeSummary,
) -> Vec<(String, usize, usize)> {
    let lines: Vec<&str> = source.lines().collect();
    let mut boundaries = Vec::new();

    // Find the start line of each function by matching its signature
    let mut func_starts: Vec<(String, usize)> = Vec::new();

    for func in &code_summary.functions {
        // Find the line containing this function's signature
        let sig_prefix = crate::text_boundary::truncate_str(&func.signature, 20);

        for (i, line) in lines.iter().enumerate() {
            if line.trim().starts_with(sig_prefix.trim()) || line.contains(&format!("fn {}", func.name)) || line.contains(&format!("def {}", func.name)) || line.contains(&format!("function {}", func.name)) {
                func_starts.push((func.name.clone(), i));
                break;
            }
        }
    }

    // Sort by start line
    func_starts.sort_by_key(|(_, line)| *line);

    // Each function ends where the next one starts (or at EOF)
    for i in 0..func_starts.len() {
        let (ref name, start) = func_starts[i];
        let end = if i + 1 < func_starts.len() {
            func_starts[i + 1].1
        } else {
            lines.len()
        };
        boundaries.push((name.clone(), start, end));
    }

    boundaries
}

/// Approximate token count (chars / 4).
fn approx_tokens(s: &str) -> u32 {
    ((s.len() as f64) / 4.0).ceil() as u32
}

// ── Tests ─────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;

    const RUST_SOURCE: &str = r#"
use std::collections::HashMap;
use crate::pipeline::CompressionPipeline;
use crate::cache_manager::CacheManager;

/// The main engine.
pub struct SqzEngine {
    pipeline: CompressionPipeline,
    cache: CacheManager,
    config: HashMap<String, String>,
}

impl SqzEngine {
    pub fn new() -> Result<Self> {
        let pipeline = CompressionPipeline::new();
        let cache = CacheManager::new();
        let config = HashMap::new();
        Ok(Self { pipeline, cache, config })
    }

    pub fn compress(&self, input: &str) -> Result<CompressedContent> {
        let cached = self.cache.get(input);
        if let Some(hit) = cached {
            return Ok(hit);
        }
        let result = self.pipeline.compress(input);
        self.cache.insert(input, &result);
        self.verify(&result);
        result
    }

    pub fn compress_with_mode(&self, input: &str, mode: CompressionMode) -> Result<CompressedContent> {
        self.compress(input)
    }

    fn verify(&self, result: &CompressedContent) -> bool {
        result.compression_ratio < 0.95
    }

    pub fn status(&self) -> String {
        format!("cache: {} entries", self.cache.len())
    }
}

pub fn standalone_helper(x: i32) -> i32 {
    x + 1
}
"#;

    #[test]
    fn test_summarize_rust_file() {
        let config = SummaryConfig::default();
        let result = summarize(RUST_SOURCE, "rust", "src/engine.rs", &config, None).unwrap();

        assert!(result.summary.contains("# file: src/engine.rs"));
        assert!(result.summary.contains("## imports"));
        assert!(result.summary.contains("## functions"));
        assert!(result.tokens_summary < result.tokens_original);
        assert!(result.compression_ratio < 1.0);
        assert!(result.functions_count > 0);
    }

    #[test]
    fn test_summarize_extracts_calls() {
        let config = SummaryConfig::default();
        let result = summarize(RUST_SOURCE, "rust", "src/engine.rs", &config, None).unwrap();

        // The call graph should find at least some relationships
        // (compress calls verify, compress_with_mode calls compress, etc.)
        // If no calls found, the summary still works — just without call arrows
        if result.call_edges > 0 {
            assert!(
                result.summary.contains("\u{2192} calls:"),
                "summary should show call arrows when edges exist"
            );
        }
        // Verify the summary is still smaller than the original regardless
        assert!(result.tokens_summary < result.tokens_original);
    }

    #[test]
    fn test_summarize_compression_ratio() {
        let config = SummaryConfig::default();
        let result = summarize(RUST_SOURCE, "rust", "src/engine.rs", &config, None).unwrap();

        // Structural summary should be significantly smaller
        assert!(
            result.compression_ratio < 0.8,
            "compression ratio {} should be < 0.8",
            result.compression_ratio
        );
    }

    #[test]
    fn test_summarize_small_file_passthrough() {
        let small = "fn main() {}";
        let config = SummaryConfig::default();
        let result = summarize(small, "rust", "main.rs", &config, None).unwrap();

        assert_eq!(result.summary, small);
        assert_eq!(result.compression_ratio, 1.0);
    }

    #[test]
    fn test_summarize_python() {
        let source = r#"
import os
from typing import List, Dict
from .utils import helper

class UserService:
    def __init__(self, db):
        self.db = db

    def create(self, name: str) -> Dict:
        user = self.db.insert(name)
        self.notify(user)
        return user

    def notify(self, user: Dict) -> None:
        print(f"Created {user}")

def standalone(x: int) -> int:
    result = helper(x)
    return result + 1
"#;
        let config = SummaryConfig {
            min_file_size: 50,
            ..Default::default()
        };
        let result = summarize(source, "python", "services/user.py", &config, None).unwrap();

        assert!(result.summary.contains("## imports"));
        assert!(result.summary.contains("## types") || result.summary.contains("class UserService"));
        assert!(result.summary.contains("## functions"));
        assert!(result.tokens_summary < result.tokens_original);
    }

    #[test]
    fn test_summarize_javascript() {
        let source = r#"
import React from 'react';
import { useState, useEffect } from 'react';
import { fetchUsers } from './api';

class UserList extends React.Component {
    constructor(props) {
        super(props);
        this.state = { users: [] };
    }

    componentDidMount() {
        fetchUsers().then(users => {
            this.setState({ users });
        });
    }

    render() {
        return this.state.users.map(u => (
            <div key={u.id}>{u.name}</div>
        ));
    }
}

function formatUser(user) {
    const name = user.firstName + ' ' + user.lastName;
    return { ...user, displayName: name };
}

export default UserList;
"#;
        let config = SummaryConfig::default();
        let result = summarize(source, "javascript", "src/UserList.js", &config, None).unwrap();

        assert!(result.summary.contains("## imports"));
        assert!(result.functions_count > 0);
        assert!(result.tokens_summary < result.tokens_original);
    }

    #[test]
    fn test_summarize_with_dep_mapper() {
        let mut mapper = DependencyMapper::new();
        mapper.add_file(
            std::path::Path::new("src/engine.rs"),
            "use crate::pipeline;\nuse crate::cache;\n",
        );
        mapper.add_file(
            std::path::Path::new("src/main.rs"),
            "use crate::engine;\n",
        );

        let config = SummaryConfig::default();
        let result = summarize(
            RUST_SOURCE,
            "rust",
            "src/engine.rs",
            &config,
            Some(&mapper),
        )
        .unwrap();

        assert!(result.summary.contains("## dependencies"));
    }

    #[test]
    fn test_summarize_config_disable_calls() {
        let config = SummaryConfig {
            include_calls: false,
            ..Default::default()
        };
        let result = summarize(RUST_SOURCE, "rust", "src/engine.rs", &config, None).unwrap();

        assert!(
            !result.summary.contains("\u{2192} calls:"),
            "should not show calls when disabled"
        );
        assert_eq!(result.call_edges, 0);
    }

    #[test]
    fn test_summarize_config_disable_imports() {
        let config = SummaryConfig {
            include_imports: false,
            ..Default::default()
        };
        let result = summarize(RUST_SOURCE, "rust", "src/engine.rs", &config, None).unwrap();

        assert!(
            !result.summary.contains("## imports"),
            "should not show imports when disabled"
        );
    }

    #[test]
    fn test_summarize_multi() {
        let files: Vec<(&str, &str, &str)> = vec![
            (RUST_SOURCE, "rust", "src/engine.rs"),
            (
                "use crate::engine;\n\nfn main() {\n    let e = SqzEngine::new();\n    e.compress(\"hello\");\n}\n",
                "rust",
                "src/main.rs",
            ),
        ];
        let config = SummaryConfig {
            min_file_size: 10,
            ..Default::default()
        };
        let result = summarize_multi(&files, &config, None).unwrap();

        assert!(result.summary.contains("src/engine.rs"));
        assert!(result.summary.contains("src/main.rs"));
        assert!(result.summary.contains("---")); // separator
    }

    #[test]
    fn test_contains_call_word_boundary() {
        assert!(contains_call("    cache.get(input)", "get"));
        assert!(contains_call("result = compress(data)", "compress"));
        assert!(!contains_call("decompressor(data)", "compress"));
        assert!(!contains_call("get_all()", "get"));
    }

    #[test]
    fn test_extract_call_graph_finds_known_calls() {
        let parser = AstParser::new();
        let summary = parser.extract_signatures(RUST_SOURCE, "rust").unwrap();
        let graph = extract_call_graph(RUST_SOURCE, &summary);

        // compress() should call verify()
        if let Some(calls) = graph.get("compress") {
            assert!(
                calls.iter().any(|c| c == "verify"),
                "compress should call verify, got: {:?}",
                calls
            );
        }
    }

    #[test]
    fn test_find_function_boundaries() {
        let parser = AstParser::new();
        let summary = parser.extract_signatures(RUST_SOURCE, "rust").unwrap();
        let boundaries = find_function_boundaries(RUST_SOURCE, &summary);

        assert!(
            !boundaries.is_empty(),
            "should find function boundaries"
        );
        // Each boundary should have start < end
        for (name, start, end) in &boundaries {
            assert!(
                start < end,
                "function {} should have start ({}) < end ({})",
                name,
                start,
                end
            );
        }
    }

    #[test]
    fn test_summarize_unsupported_language() {
        let source = "some random content that is long enough to trigger summarization, we need at least 500 characters of content here to pass the minimum file size threshold so let me add more text to make this work properly and ensure we get a valid result back from the summarize function even for unsupported languages like COBOL or Fortran or whatever else might come through the pipeline in production use cases where we cannot predict what files will be processed by the engine and need graceful fallback behavior that does not crash or panic but instead returns a reasonable default result that the caller can work with safely and reliably in all circumstances";
        let config = SummaryConfig::default();
        let result = summarize(source, "cobol", "main.cob", &config, None).unwrap();

        // Should still produce a result (with file header at minimum)
        assert!(result.summary.contains("# file: main.cob"));
    }
}