magi-code 0.77.1

//! Shared bounded source reading and cooperative parsing for embedded AST tools.
use super::*;
use ast_grep_core::{
    Language, Pattern,
    matcher::{PatternBuilder, PatternError},
    tree_sitter::{LanguageExt, StrDoc, TSLanguage},
};
use ast_grep_language::SupportLang;
use std::{borrow::Cow, ops::ControlFlow};
use tree_sitter::{ParseOptions, Parser, Tree};

pub(super) const MAX_FILE_BYTES: usize = 256 * 1024;
pub(super) const MAX_TREE_NODES: usize = 16_384;
const MAX_TREE_DEPTH: usize = 64;
const MAX_PATTERN_NODES: usize = 256;
pub(super) const MAX_FILES: usize = 10_000;
pub(super) const MAX_SCAN_BYTES: usize = 128 * 1024 * 1024;

#[derive(Clone)]
pub(super) struct EmbeddedLanguage {
    pub language: SupportLang,
    pub cancellation: AgentCancellation,
    pub started: Instant,
    pub timeout: Duration,
}

impl EmbeddedLanguage {
    pub fn stopped(&self) -> bool {
        self.cancellation.is_canceled() || self.started.elapsed() >= self.timeout
    }

    pub fn parse(&self, src: String, max_nodes: usize) -> anyhow::Result<StrDoc<Self>> {
        self.cancellation.check()?;
        anyhow::ensure!(!self.stopped(), "ast-grep deadline exceeded");
        let mut parser = Parser::new();
        parser.set_language(&self.get_ts_language())?;
        let mut progress = |_: &tree_sitter::ParseState| {
            if self.stopped() {
                ControlFlow::Break(())
            } else {
                ControlFlow::Continue(())
            }
        };
        let tree = parser.parse_with_options(
            &mut |offset, _| &src.as_bytes()[offset..],
            None,
            Some(ParseOptions::new().progress_callback(&mut progress)),
        );
        self.cancellation.check()?;
        let tree = tree.ok_or_else(|| anyhow::anyhow!("ast-grep deadline exceeded"))?;
        self.check_tree(&tree, max_nodes)?;
        Ok(StrDoc {
            src,
            lang: self.clone(),
            tree,
        })
    }

    // Bound recursion/captures inside individual library calls before extraction.
    fn check_tree(&self, tree: &Tree, max_nodes: usize) -> anyhow::Result<()> {
        let mut cursor = tree.walk();
        let mut depth = 0;
        let mut nodes = 0;
        loop {
            self.cancellation.check()?;
            anyhow::ensure!(!self.stopped(), "ast-grep deadline exceeded");
            nodes += 1;
            anyhow::ensure!(
                nodes <= max_nodes && depth <= MAX_TREE_DEPTH,
                "ast-grep syntax tree exceeds node/depth budget"
            );
            if cursor.goto_first_child() {
                depth += 1;
                continue;
            }
            loop {
                if cursor.goto_next_sibling() {
                    break;
                }
                if !cursor.goto_parent() {
                    return Ok(());
                }
                depth -= 1;
            }
        }
    }
}

impl Language for EmbeddedLanguage {
    fn pre_process_pattern<'q>(&self, query: &'q str) -> Cow<'q, str> {
        self.language.pre_process_pattern(query)
    }
    fn meta_var_char(&self) -> char {
        self.language.meta_var_char()
    }
    fn expando_char(&self) -> char {
        self.language.expando_char()
    }
    fn kind_to_id(&self, kind: &str) -> u16 {
        self.language.kind_to_id(kind)
    }
    fn field_to_id(&self, field: &str) -> Option<u16> {
        self.language.field_to_id(field)
    }
    fn build_pattern(&self, builder: &PatternBuilder) -> Result<Pattern, PatternError> {
        builder.build(|src| {
            self.parse(src.to_owned(), MAX_PATTERN_NODES)
                .map_err(|e| e.to_string())
        })
    }
}
impl LanguageExt for EmbeddedLanguage {
    fn get_ts_language(&self) -> TSLanguage {
        self.language.get_ts_language()
    }
}

// Whole UTF-8 files only. Nonblocking/no-follow open rejects raced FIFOs and
// symlinks; ordinary filesystem I/O itself is not preemptible.
pub(super) fn read_source(path: &Path, scanned: &mut usize) -> Option<String> {
    use std::{fs::OpenOptions, io::Read};
    if !path.symlink_metadata().ok()?.file_type().is_file() {
        return None;
    }
    let mut options = OpenOptions::new();
    options.read(true);
    #[cfg(unix)]
    {
        use std::os::unix::fs::OpenOptionsExt;
        options.custom_flags(libc::O_NONBLOCK | libc::O_NOFOLLOW);
    }
    let file = options.open(path).ok()?;
    if !file.metadata().ok()?.is_file() {
        return None;
    }
    let remaining = MAX_SCAN_BYTES.saturating_sub(*scanned);
    let mut bytes = Vec::new();
    let read = file
        .take((MAX_FILE_BYTES + 1).min(remaining) as u64)
        .read_to_end(&mut bytes);
    *scanned += bytes.len();
    read.ok()?;
    if bytes.len() > MAX_FILE_BYTES || bytes.len() == remaining {
        return None;
    }
    String::from_utf8(bytes).ok()
}