boundgraph_slice/

ast.rs

//! Condensing source code: keep the shape, elide the bodies.
//!
//! # Why code cannot be truncated
//!
//! Cutting a source file at a byte offset destroys exactly what a coding agent
//! needs. The signatures are spread through the file, the imports are at the
//! top, and the function the agent was asked about is as likely to be at the
//! end as anywhere. A model handed the first two hundred lines of a thousand
//! line file knows less than one handed every signature in it.
//!
//! So this slicer parses, keeps the declarations, and replaces each body with a
//! marker. A four hundred line module becomes forty lines of signatures that
//! still say what the module *is*.
//!
//! # The output is still source
//!
//! Every elision marker is valid syntax in the language it replaces — a block
//! comment in Rust and JavaScript, `...` in Python. The result parses. That is
//! not decoration: a model reading something that looks like code but does not
//! parse will reproduce the malformation in its patch, and a downstream syntax
//! checker will then blame the agent.
//!
//! # Relevance, not file order
//!
//! When the cap will not hold every signature, the ones that go are chosen by
//! relevance to what the step was asked to do rather than by position. The
//! terms come from the caller — the runtime holds the
//! [`WorkOrder`](boundgraph_core::WorkOrder) and constructs the slicer with terms
//! drawn from its intent — because the [`Slicer`] contract takes raw output and
//! a cap and nothing else, and widening it so one slicer could see the work
//! order would put a model's proposal into every slicer's reach.
//!
//! Whatever survives is emitted in source order. Reordering a file to put the
//! relevant parts first would make the line numbers lie.
//!
//! # When it cannot parse
//!
//! It degrades to the log slicer rather than failing the step. An unsupported
//! language, a file that is mostly a syntax error, a template with a foreign
//! preprocessor — all of them are still *text*, and text sliced by the log
//! slicer is far more useful than a refused step. The condensate says which
//! slicer actually produced it, so nothing downstream is misled about what it
//! is reading.

use async_trait::async_trait;
use boundgraph_core::{SchemaId, SlicerId};
use serde_json::json;
use tree_sitter::{Node, Parser};

use crate::log::LogSlicer;
use crate::slicer::{Raw, Slice, SliceError, Slicer};

/// The schema an outline condensate declares.
pub const AST_SCHEMA: &str = "boundgraph.code.v1";

/// A language this slicer can parse.
///
/// Declared explicitly rather than sniffed. A slicer that guessed would
/// occasionally guess wrong and emit an outline of the wrong shape, which is
/// worse than declining — declining degrades to the log slicer, which is always
/// correct if unexciting.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[non_exhaustive]
pub enum Language {
    /// Rust.
    Rust,
    /// Python.
    Python,
    /// JavaScript, including JSX.
    JavaScript,
}

impl Language {
    /// Every language this build can parse.
    pub const SUPPORTED: &'static [Self] = &[Self::Rust, Self::Python, Self::JavaScript];

    /// The language a path or hint names, if it is one that is supported.
    pub fn of_hint(hint: &str) -> Option<Self> {
        let lower = hint.to_ascii_lowercase();
        let extension = lower.rsplit('.').next()?;
        match extension {
            "rs" => Some(Self::Rust),
            "py" | "pyi" => Some(Self::Python),
            "js" | "mjs" | "cjs" | "jsx" => Some(Self::JavaScript),
            _ => None,
        }
    }

    /// The name used in the payload.
    pub fn name(self) -> &'static str {
        match self {
            Self::Rust => "rust",
            Self::Python => "python",
            Self::JavaScript => "javascript",
        }
    }

    fn grammar(self) -> tree_sitter::Language {
        match self {
            Self::Rust => tree_sitter_rust::LANGUAGE.into(),
            Self::Python => tree_sitter_python::LANGUAGE.into(),
            Self::JavaScript => tree_sitter_javascript::LANGUAGE.into(),
        }
    }

    /// The node kinds worth keeping: the declarations that make up the shape.
    fn declarations(self) -> &'static [&'static str] {
        match self {
            Self::Rust => &[
                "use_declaration",
                "mod_item",
                "function_item",
                "struct_item",
                "enum_item",
                "trait_item",
                "impl_item",
                "type_item",
                "const_item",
                "static_item",
                "macro_definition",
            ],
            Self::Python => &[
                "import_statement",
                "import_from_statement",
                "function_definition",
                "class_definition",
                "decorated_definition",
            ],
            Self::JavaScript => &[
                "import_statement",
                "export_statement",
                "function_declaration",
                "generator_function_declaration",
                "class_declaration",
                "method_definition",
                "lexical_declaration",
            ],
        }
    }

    /// How this language spells "a body was removed here".
    ///
    /// Chosen so the result still parses. `...` is a real expression in Python;
    /// a block comment is valid wherever a statement could be in Rust and
    /// JavaScript.
    fn elision(self, lines: usize) -> String {
        match self {
            Self::Rust | Self::JavaScript => format!("{{ /* … {lines} lines … */ }}"),
            Self::Python => format!("...  # … {lines} lines …"),
        }
    }

    /// Whether a node is a body that should be replaced rather than shown.
    fn is_body(self, kind: &str) -> bool {
        match self {
            Self::Rust => kind == "block",
            Self::Python => kind == "block",
            Self::JavaScript => kind == "statement_block" || kind == "class_body",
        }
    }
}

/// One declaration kept from a file.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Item {
    /// Where it starts in the original, counting from one.
    pub line: usize,
    /// The declaration with its bodies elided.
    pub text: String,
    /// The node kind, for a reader who wants to know what was kept.
    pub kind: String,
    /// How relevant it looked to what the step was asked to do.
    relevance: usize,
}

/// Condenses source code to an outline.
///
/// # Example
///
/// ```
/// use boundgraph_slice::{AstSlicer, MemoryStore, Raw, Slicer};
///
/// let source = r#"
/// pub fn parse_header(input: &[u8]) -> Result<Header, Error> {
///     let mut cursor = 0;
///     // ... forty more lines ...
///     Ok(Header::default())
/// }
/// "#;
///
/// let raw = Raw::new(source).hinting("src/parse.rs");
/// let condensate = tokio::runtime::Builder::new_current_thread()
///     .build()?
///     .block_on(AstSlicer::new().condense(&raw, 300, &MemoryStore::new()))?;
///
/// // The signature survived; the body did not.
/// let rendered = format!("{:?}", condensate.payload);
/// assert!(rendered.contains("parse_header"));
/// # Ok::<(), Box<dyn std::error::Error>>(())
/// ```
#[derive(Debug, Clone, Default)]
pub struct AstSlicer {
    relevant: Vec<String>,
    fallback: LogSlicer,
}

impl AstSlicer {
    /// A slicer with no relevance terms, keeping declarations in file order.
    pub fn new() -> Self {
        Self {
            relevant: Vec::new(),
            fallback: LogSlicer::new(),
        }
    }

    /// Prefer declarations mentioning any of these terms.
    ///
    /// The runtime supplies them from the work order's intent — see the module
    /// documentation for why they arrive here rather than through the trait.
    #[must_use]
    pub fn relevant_to<S: Into<String>>(mut self, terms: impl IntoIterator<Item = S>) -> Self {
        self.relevant
            .extend(terms.into_iter().map(|t| t.into().to_ascii_lowercase()));
        self
    }

    /// How well a declaration matches the terms.
    fn score(&self, text: &str) -> usize {
        if self.relevant.is_empty() {
            return 0;
        }
        let lower = text.to_ascii_lowercase();
        self.relevant.iter().filter(|t| lower.contains(*t)).count()
    }

    /// Render a node with every body inside it replaced by a marker.
    fn outline_of(language: Language, node: Node<'_>, source: &str) -> String {
        let mut out = String::new();
        let mut at = node.start_byte();

        fn walk(
            language: Language,
            node: Node<'_>,
            source: &str,
            at: &mut usize,
            out: &mut String,
            depth: usize,
        ) {
            let mut cursor = node.walk();
            for child in node.children(&mut cursor) {
                if language.is_body(child.kind()) {
                    // Emit everything up to the body, then the marker, and skip
                    // the body itself — never a partial one.
                    out.push_str(&source[*at..child.start_byte()]);
                    let lines = source[child.start_byte()..child.end_byte()].lines().count();
                    out.push_str(&language.elision(lines));
                    *at = child.end_byte();
                    continue;
                }
                // Descend, so a body nested inside an `impl` is elided too, but
                // not so far that this becomes quadratic on a deep file.
                if child.child_count() > 0 && depth < 4 {
                    walk(language, child, source, at, out, depth + 1);
                }
            }
        }

        walk(language, node, source, &mut at, &mut out, 0);
        out.push_str(&source[at..node.end_byte()]);
        out.trim_end().to_string()
    }

    /// Extract the declarations of a parsed file.
    fn items(&self, language: Language, source: &str) -> Option<Vec<Item>> {
        let mut parser = Parser::new();
        parser.set_language(&language.grammar()).ok()?;
        let tree = parser.parse(source, None)?;
        let root = tree.root_node();

        let kinds = language.declarations();
        let mut items = Vec::new();
        let mut cursor = root.walk();

        for child in root.children(&mut cursor) {
            if !kinds.contains(&child.kind()) {
                continue;
            }
            let text = Self::outline_of(language, child, source);
            if text.trim().is_empty() {
                continue;
            }
            items.push(Item {
                line: child.start_position().row + 1,
                relevance: self.score(&text),
                kind: child.kind().to_string(),
                text,
            });
        }

        // Tree-sitter recovers from errors, so a tree with ERROR nodes is
        // normal for a file with one typo and useless for a file that is not
        // really source at all. Extracting nothing from an erroring tree is the
        // signal that this is the second case.
        if root.has_error() && items.is_empty() {
            return None;
        }
        Some(items)
    }

    fn payload(language: Language, kept: &[Item], all: &[Item]) -> serde_json::Value {
        let outline = kept
            .iter()
            .map(|i| i.text.as_str())
            .collect::<Vec<_>>()
            .join("\n\n");

        let mut payload = json!({
            "language": language.name(),
            "outline": outline,
            "counts": { "items": all.len(), "shown": kept.len() },
        });

        if kept.len() < all.len() {
            let shown: std::collections::BTreeSet<usize> = kept.iter().map(|i| i.line).collect();
            let dropped: Vec<&str> = all
                .iter()
                .filter(|i| !shown.contains(&i.line))
                .map(|i| i.kind.as_str())
                .collect();
            payload["dropped"] = json!(dropped.len());
            payload["dropped_kinds"] = json!(
                dropped
                    .iter()
                    .collect::<std::collections::BTreeSet<_>>()
                    .into_iter()
                    .collect::<Vec<_>>()
            );
        }
        payload
    }
}

#[async_trait]
impl Slicer for AstSlicer {
    fn id(&self) -> SlicerId {
        SlicerId::new("ast", 1)
    }

    async fn slice(&self, raw: &Raw, cap: usize) -> Result<Slice, SliceError> {
        let language = raw.hint.as_deref().and_then(Language::of_hint);
        let source = raw.text();

        // Either the language is not one this build parses, or the file is not
        // really source. Both are text, and text is what the log slicer is for.
        let Some(language) = language else {
            return self.fallback.slice(raw, cap).await;
        };
        let Some(items) = self.items(language, &source) else {
            return self.fallback.slice(raw, cap).await;
        };
        if items.is_empty() {
            return self.fallback.slice(raw, cap).await;
        }

        let counter = self.counter();
        let fits = |value: &serde_json::Value| {
            counter.count(&serde_json::to_string(value).unwrap_or_default()) <= cap
        };

        let floor = Self::payload(language, &[], &items);
        if !fits(&floor) {
            return Err(SliceError::CapTooSmall {
                cap,
                minimum: counter.count(&serde_json::to_string(&floor).unwrap_or_default()),
            });
        }

        // Choose by relevance, then by position so ties are stable.
        let mut order: Vec<usize> = (0..items.len()).collect();
        order.sort_by_key(|&i| (std::cmp::Reverse(items[i].relevance), items[i].line));

        let mut chosen: Vec<usize> = Vec::new();
        for &index in &order {
            let mut candidate = chosen.clone();
            candidate.push(index);
            // Emitted in source order: reordering a file would make the line
            // numbers lie.
            candidate.sort_by_key(|&i| items[i].line);

            let kept: Vec<Item> = candidate.iter().map(|&i| items[i].clone()).collect();
            if fits(&Self::payload(language, &kept, &items)) {
                chosen = candidate;
            }
        }

        chosen.sort_by_key(|&i| items[i].line);
        let kept: Vec<Item> = chosen.iter().map(|&i| items[i].clone()).collect();

        Ok(Slice {
            schema: SchemaId::new(AST_SCHEMA, 1),
            payload: Self::payload(language, &kept, &items),
            // Bodies are always elided, so an outline is lossy even when every
            // declaration survived. Reporting otherwise would tell a downstream
            // step it has the whole file when it has only the shape.
            lossy: true,
        })
    }
}

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

    const RUST: &str = r#"
use std::collections::BTreeMap;

pub struct Header {
    pub version: u8,
    pub length: u32,
}

pub fn parse_header(input: &[u8]) -> Result<Header, Error> {
    let mut cursor = 0;
    let version = input[cursor];
    cursor += 1;
    let length = u32::from_be_bytes(input[cursor..cursor + 4].try_into()?);
    Ok(Header { version, length })
}

pub fn unrelated_helper(x: u32) -> u32 {
    let mut total = 0;
    for i in 0..x {
        total += i;
    }
    total
}

impl Header {
    pub fn is_valid(&self) -> bool {
        self.version > 0
    }
}
"#;

    const PYTHON: &str = r#"
import os
from typing import Optional

class Parser:
    def parse_header(self, data: bytes) -> Optional[dict]:
        cursor = 0
        version = data[cursor]
        cursor += 1
        return {"version": version}

def unrelated(x):
    total = 0
    for i in range(x):
        total += i
    return total
"#;

    const JAVASCRIPT: &str = r#"
import { readFile } from "fs";

export function parseHeader(input) {
    let cursor = 0;
    const version = input[cursor];
    cursor += 1;
    return { version };
}

function unrelated(x) {
    let total = 0;
    for (let i = 0; i < x; i++) { total += i; }
    return total;
}
"#;

    fn raw_of(source: &str, hint: &str) -> Raw {
        Raw::new(source.as_bytes().to_vec()).hinting(hint)
    }

    async fn outline(source: &str, hint: &str, cap: usize) -> String {
        let slice = AstSlicer::new()
            .slice(&raw_of(source, hint), cap)
            .await
            .expect("a satisfiable cap");
        slice.payload["outline"]
            .as_str()
            .unwrap_or_default()
            .to_string()
    }

    // ── structure is kept ───────────────────────────────────────────────────

    /// The criterion: module, type, and function signatures survive.
    #[tokio::test]
    async fn rust_signatures_and_types_survive() {
        let out = outline(RUST, "src/parse.rs", 2_000).await;

        assert!(out.contains("use std::collections::BTreeMap"));
        assert!(out.contains("pub struct Header"));
        assert!(out.contains("pub fn parse_header(input: &[u8]) -> Result<Header, Error>"));
        assert!(out.contains("impl Header"));
    }

    /// The criterion: bodies are elided with a marker, never cut mid-token.
    #[tokio::test]
    async fn function_bodies_are_replaced_by_a_marker() {
        let out = outline(RUST, "src/parse.rs", 2_000).await;

        assert!(
            !out.contains("u32::from_be_bytes"),
            "the body should be gone"
        );
        assert!(out.contains("lines"), "and marked: {out}");
        // The struct's fields are its shape, not a body, so they stay.
        assert!(out.contains("pub version: u8"));
    }

    /// The criterion: the result parses. A model shown malformed code will
    /// reproduce the malformation in its patch.
    #[tokio::test]
    async fn the_outline_is_still_parseable_source() {
        for (source, hint, language) in [
            (RUST, "a.rs", Language::Rust),
            (PYTHON, "a.py", Language::Python),
            (JAVASCRIPT, "a.js", Language::JavaScript),
        ] {
            let out = outline(source, hint, 4_000).await;

            let mut parser = Parser::new();
            parser.set_language(&language.grammar()).unwrap();
            let tree = parser.parse(&out, None).expect("the outline parses");
            assert!(
                !tree.root_node().has_error(),
                "{} outline does not parse:\n{out}",
                language.name()
            );
        }
    }

    #[tokio::test]
    async fn python_and_javascript_are_outlined_too() {
        let py = outline(PYTHON, "a.py", 2_000).await;
        assert!(py.contains("class Parser"));
        assert!(py.contains("import os"));
        assert!(!py.contains("cursor = 0"), "the body should be gone: {py}");

        let js = outline(JAVASCRIPT, "a.js", 2_000).await;
        assert!(js.contains("export function parseHeader"));
        assert!(
            !js.contains("const version"),
            "the body should be gone: {js}"
        );
    }

    // ── relevance ───────────────────────────────────────────────────────────

    /// The criterion: elision is driven by relevance, not file order.
    /// `unrelated_helper` comes *before* `impl Header` in the file, so file
    /// order alone would keep the wrong one.
    #[tokio::test]
    async fn a_tight_cap_keeps_what_the_step_was_asked_about() {
        let raw = raw_of(RUST, "src/parse.rs");

        let relevant = AstSlicer::new()
            .relevant_to(["parse_header"])
            .slice(&raw, 120)
            .await
            .unwrap();
        let out = relevant.payload["outline"].as_str().unwrap();

        assert!(
            out.contains("parse_header"),
            "the relevant function must survive: {out}"
        );
        assert!(
            !out.contains("unrelated_helper"),
            "the irrelevant one should have gone first: {out}"
        );
    }

    #[tokio::test]
    async fn without_terms_the_order_is_positional_and_stable() {
        let raw = raw_of(RUST, "src/parse.rs");
        let first = AstSlicer::new().slice(&raw, 200).await.unwrap();
        for _ in 0..5 {
            assert_eq!(AstSlicer::new().slice(&raw, 200).await.unwrap(), first);
        }
    }

    /// Reordering a file to put the relevant parts first would make the line
    /// numbers lie.
    #[tokio::test]
    async fn whatever_survives_is_emitted_in_source_order() {
        let out = AstSlicer::new()
            .relevant_to(["is_valid"])
            .slice(&raw_of(RUST, "src/parse.rs"), 4_000)
            .await
            .unwrap();
        let text = out.payload["outline"].as_str().unwrap();

        let use_at = text.find("use std").expect("the import");
        let impl_at = text.find("impl Header").expect("the impl");
        assert!(use_at < impl_at, "source order, not relevance order");
    }

    // ── degrading ───────────────────────────────────────────────────────────

    /// The criterion: an unsupported language degrades predictably.
    #[tokio::test]
    async fn an_unsupported_language_falls_back_to_the_log_slicer() {
        let raw = raw_of("SELECT * FROM users WHERE id = 1;\n", "query.sql");
        let slice = AstSlicer::new().slice(&raw, 400).await.unwrap();

        assert_eq!(
            slice.schema,
            SchemaId::new(crate::log::LOG_SCHEMA, 1),
            "it should be a log condensate, and say so"
        );
        assert!(slice.payload.get("lines").is_some());
    }

    /// The criterion: unparseable input degrades rather than failing the step.
    #[tokio::test]
    async fn input_that_is_not_really_source_falls_back_rather_than_failing() {
        let raw = raw_of(
            "\u{1}\u{2} not code at all >>>>>> <<<<<< ####\nnor is this line\n",
            "broken.rs",
        );
        let slice = AstSlicer::new().slice(&raw, 400).await.unwrap();

        assert_eq!(slice.schema, SchemaId::new(crate::log::LOG_SCHEMA, 1));
    }

    /// Tree-sitter recovers from errors, so a file with one typo still yields
    /// an outline — falling back on any error at all would give up on most
    /// real files an agent is mid-way through editing.
    #[tokio::test]
    async fn a_file_with_one_syntax_error_still_produces_an_outline() {
        let source = "
pub fn good(x: u32) -> u32 { x + 1 }

pub fn broken(y: u32 -> u32 { y }

pub struct Fine { pub a: u8 }
";
        let raw = raw_of(source, "src/lib.rs");
        let slice = AstSlicer::new().slice(&raw, 2_000).await.unwrap();

        assert_eq!(slice.schema, SchemaId::new(AST_SCHEMA, 1), "not a fallback");
        let out = slice.payload["outline"].as_str().unwrap();
        assert!(out.contains("good") || out.contains("Fine"), "{out}");
    }

    #[tokio::test]
    async fn a_file_with_no_hint_falls_back() {
        let slice = AstSlicer::new()
            .slice(&Raw::new(RUST.as_bytes().to_vec()), 800)
            .await
            .unwrap();
        assert_eq!(slice.schema, SchemaId::new(crate::log::LOG_SCHEMA, 1));
    }

    // ── declared support ────────────────────────────────────────────────────

    #[test]
    fn the_supported_languages_are_declared_explicitly() {
        assert_eq!(Language::SUPPORTED.len(), 3);
        assert_eq!(Language::of_hint("src/main.rs"), Some(Language::Rust));
        assert_eq!(Language::of_hint("app/views.py"), Some(Language::Python));
        assert_eq!(
            Language::of_hint("web/index.jsx"),
            Some(Language::JavaScript)
        );
        assert_eq!(Language::of_hint("Cargo.toml"), None);
        assert_eq!(Language::of_hint("README"), None);
        assert_eq!(Language::of_hint("query.sql"), None);
    }

    #[test]
    fn every_supported_language_has_a_working_grammar() {
        for language in Language::SUPPORTED {
            let mut parser = Parser::new();
            parser
                .set_language(&language.grammar())
                .unwrap_or_else(|e| panic!("{} grammar rejected: {e}", language.name()));
        }
    }

    #[test]
    fn every_elision_marker_is_valid_in_its_own_language() {
        for language in Language::SUPPORTED {
            let marker = language.elision(42);
            assert!(marker.contains("42"), "{}", language.name());
        }
        assert!(Language::Python.elision(3).starts_with("..."));
        assert!(Language::Rust.elision(3).starts_with('{'));
    }

    // ── the contract ────────────────────────────────────────────────────────

    #[tokio::test]
    async fn the_output_never_exceeds_the_cap() {
        let store = MemoryStore::new();
        let raw = raw_of(RUST, "src/parse.rs");

        for cap in [120, 200, 400, 800, 2_000, 8_000] {
            let condensate = AstSlicer::new()
                .relevant_to(["parse_header"])
                .condense(&raw, cap, &store)
                .await
                .unwrap_or_else(|e| panic!("cap {cap} should be satisfiable: {e}"));
            assert_eq!(condensate.cap, cap);
        }
    }

    /// An outline has always dropped the bodies, so it is never lossless.
    #[tokio::test]
    async fn an_outline_is_always_lossy() {
        let slice = AstSlicer::new()
            .slice(&raw_of(RUST, "src/parse.rs"), 8_000)
            .await
            .unwrap();

        assert!(slice.lossy);
        assert_eq!(slice.payload["counts"]["items"], 5);
        assert_eq!(slice.payload["counts"]["shown"], 5);
    }

    #[tokio::test]
    async fn what_was_dropped_is_counted() {
        let slice = AstSlicer::new()
            .relevant_to(["parse_header"])
            .slice(&raw_of(RUST, "src/parse.rs"), 130)
            .await
            .unwrap();

        assert!(slice.payload["dropped"].as_u64().unwrap() > 0);
        assert!(slice.payload["dropped_kinds"].is_array());
    }
}