thread_ast_engine/tree_sitter/

mod.rs

// SPDX-FileCopyrightText: 2022 Herrington Darkholme <2883231+HerringtonDarkholme@users.noreply.github.com>
// SPDX-FileCopyrightText: 2025 Knitli Inc. <knitli@knit.li>
// SPDX-FileContributor: Adam Poulemanos <adam@knit.li>
//
// SPDX-License-Identifier: AGPL-3.0-or-later AND MIT

//! # Tree-sitter Integration and AST Backend
//!
//! Core integration layer between thread-ast-engine and the tree-sitter parsing library.
//! Provides the foundational types and functionality for parsing source code into ASTs,
//! editing trees incrementally, and bridging tree-sitter concepts with thread-ast-engine APIs.
//!
//! ## Key Components
//!
//! - [`StrDoc`] - Document type that combines source code with parsed tree-sitter trees
//! - [`LanguageExt`] - Extension trait for languages that work with tree-sitter
//! - [`TSParseError`] - Error types for tree-sitter parsing failures
//! - Tree editing and incremental parsing support
//! - Language injection support for multi-language documents
//!
//! ## Core Concepts
//!
//! ### Documents and Parsing
//!
//! [`StrDoc`] represents a parsed document containing both source code and its tree-sitter AST.
//! It handles incremental parsing when the source is modified, automatically updating the tree
//! structure while preserving unchanged portions for performance.
//!
//! ### Language Extensions
//!
//! [`LanguageExt`] extends the base [`Language`] trait with tree-sitter specific functionality:
//! - Getting tree-sitter language objects for parsing
//! - Creating AST-grep instances
//! - Handling language injections (like JavaScript in HTML)
//!
//! ### Tree Editing
//!
//! Supports incremental tree editing through tree-sitter's edit API, allowing efficient
//! updates when source code changes without full re-parsing.
//!
//! ## Example Usage
//!
//! ```rust,no_run
//! # use thread_ast_engine::tree_sitter::{StrDoc, LanguageExt};
//! # use thread_ast_engine::Language;
//! # struct Tsx;
//! # impl Language for Tsx {
//! #     fn kind_to_id(&self, _: &str) -> u16 { 0 }
//! #     fn field_to_id(&self, _: &str) -> Option<u16> { None }
//! #     fn build_pattern(&self, _: &thread_ast_engine::PatternBuilder) -> Result<thread_ast_engine::Pattern, thread_ast_engine::PatternError> { todo!() }
//! # }
//! # impl LanguageExt for Tsx {
//! #     fn get_ts_language(&self) -> thread_ast_engine::tree_sitter::TSLanguage { todo!() }
//! # }
//!
//! // Create a document from source code
//! let doc = StrDoc::new("let x = 42;", Tsx);
//!
//! // Access the parsed tree
//! let root = doc.root_node();
//! println!("Root kind: {}", root.kind());
//!
//! // Create an AST-grep instance for pattern matching
//! let ast_grep = Tsx.ast_grep("function foo() { return 42; }");
//! let root_node = ast_grep.root();
//! ```

pub mod traversal;

use crate::node::Root;

use crate::AstGrep;
#[cfg(feature = "matching")]
use crate::Matcher;
#[cfg(feature = "matching")]
use crate::replacer::Replacer;
use crate::source::{Content, Doc, Edit, SgNode};
use crate::{Language, Position, node::KindId};
use std::borrow::Cow;
use std::num::NonZero;
use thiserror::Error;
#[cfg(feature = "matching")]
use thread_utilities::RapidMap;
pub use traversal::{TsPre, Visitor};
pub use tree_sitter::Language as TSLanguage;
use tree_sitter::{InputEdit, LanguageError, Node, Parser, Point, Tree};
pub use tree_sitter::{Point as TSPoint, Range as TSRange};

/// Errors that can occur during tree-sitter parsing operations.
///
/// Tree-sitter parsing can fail for several reasons, from language compatibility
/// issues to timeout problems. These errors provide information about what
/// went wrong during the parsing process.
#[derive(Debug, Error)]
pub enum TSParseError {
    /// The language grammar is incompatible with the parser.
    ///
    /// Occurs when trying to assign a language that the parser can't handle,
    /// typically due to version mismatches between the tree-sitter library
    /// and the language grammar.
    #[error("incompatible `Language` is assigned to a `Parser`.")]
    Language(#[from] LanguageError),

    /// Tree-sitter failed to parse the input within the configured constraints.
    ///
    /// Can be caused by several conditions:
    /// * Parsing timeout exceeded (see [`Parser::set_timeout_micros`])
    /// * Cancellation flag was triggered (see [`Parser::set_cancellation_flag`])
    /// * No language was assigned to the parser (see [`Parser::set_language`])
    /// * The input was too complex or malformed for the parser to handle
    ///
    /// Tree-sitter doesn't provide detailed error information, so this covers
    /// all general parsing failures.
    #[error("general error when tree-sitter fails to parse.")]
    TreeUnavailable,
}

#[inline]
fn parse_lang(
    parse_fn: impl Fn(&mut Parser) -> Option<Tree>,
    ts_lang: &TSLanguage,
) -> Result<Tree, TSParseError> {
    let mut parser = Parser::new();
    parser.set_language(ts_lang)?;
    if let Some(tree) = parse_fn(&mut parser) {
        Ok(tree)
    } else {
        Err(TSParseError::TreeUnavailable)
    }
}

/// Document type that combines source code with its parsed tree-sitter AST.
///
/// `StrDoc` represents a complete parsed document, holding both the original
/// source code and the tree-sitter AST. It supports incremental parsing,
/// meaning when edits are made to the source, only the affected parts of
/// the tree are re-parsed for better performance.
///
/// # Type Parameters
///
/// - `L: LanguageExt` - The language implementation that provides tree-sitter integration
///
/// # Example
///
/// ```rust,no_run
/// # use thread_ast_engine::tree_sitter::StrDoc;
/// # struct JavaScript;
/// # impl thread_ast_engine::Language for JavaScript {
/// #     fn kind_to_id(&self, _: &str) -> u16 { 0 }
/// #     fn field_to_id(&self, _: &str) -> Option<u16> { None }
/// #     fn build_pattern(&self, _: &thread_ast_engine::PatternBuilder) -> Result<thread_ast_engine::Pattern, thread_ast_engine::PatternError> { todo!() }
/// # }
/// # impl thread_ast_engine::tree_sitter::LanguageExt for JavaScript {
/// #     fn get_ts_language(&self) -> thread_ast_engine::tree_sitter::TSLanguage { todo!() }
/// # }
/// let doc = StrDoc::new("const x = 42;", JavaScript);
/// let root = doc.root_node();
/// println!("AST root: {}", root.kind());
/// ```
#[derive(Clone, Debug)]
pub struct StrDoc<L: LanguageExt> {
    /// The source code text
    pub src: String,
    /// Language implementation for parsing and node operations
    pub lang: L,
    /// The parsed tree-sitter AST
    pub tree: Tree,
}

impl<L: LanguageExt> StrDoc<L> {
    pub fn try_new(src: &str, lang: L) -> Result<Self, String> {
        let src = src.to_string();
        let ts_lang = lang.get_ts_language();
        let tree =
            parse_lang(|p| p.parse(src.as_bytes(), None), &ts_lang).map_err(|e| e.to_string())?;
        Ok(Self { src, lang, tree })
    }
    pub fn new(src: &str, lang: L) -> Self {
        Self::try_new(src, lang).expect("Parser tree error")
    }
    fn parse(&self, old_tree: Option<&Tree>) -> Result<Tree, TSParseError> {
        let source = self.get_source();
        let lang = self.get_lang().get_ts_language();
        parse_lang(|p| p.parse(source.as_bytes(), old_tree), &lang)
    }
}

impl<L: LanguageExt> Doc for StrDoc<L> {
    type Source = String;
    type Lang = L;
    type Node<'r> = Node<'r>;
    fn get_lang(&self) -> &Self::Lang {
        &self.lang
    }
    fn get_source(&self) -> &Self::Source {
        &self.src
    }
    fn do_edit(&mut self, edit: &Edit<Self::Source>) -> Result<(), String> {
        let source = &mut self.src;
        perform_edit(&mut self.tree, source, edit);
        self.tree = self.parse(Some(&self.tree)).map_err(|e| e.to_string())?;
        Ok(())
    }
    fn root_node(&self) -> Node<'_> {
        self.tree.root_node()
    }
    fn get_node_text<'a>(&'a self, node: &Self::Node<'a>) -> Cow<'a, str> {
        Cow::Borrowed(
            node.utf8_text(self.src.as_bytes())
                .expect("invalid source text encoding"),
        )
    }
}

struct NodeWalker<'tree> {
    cursor: tree_sitter::TreeCursor<'tree>,
    count: usize,
}

impl<'tree> Iterator for NodeWalker<'tree> {
    type Item = Node<'tree>;
    fn next(&mut self) -> Option<Self::Item> {
        if self.count == 0 {
            return None;
        }
        let ret = Some(self.cursor.node());
        self.cursor.goto_next_sibling();
        self.count -= 1;
        ret
    }
}

impl ExactSizeIterator for NodeWalker<'_> {
    fn len(&self) -> usize {
        self.count
    }
}

impl<'r> SgNode<'r> for Node<'r> {
    fn parent(&self) -> Option<Self> {
        Node::parent(self)
    }
    fn ancestors(&self, root: Self) -> impl Iterator<Item = Self> {
        let mut ancestor = Some(root);
        let self_id = self.id();
        std::iter::from_fn(move || {
            let inner = ancestor.take()?;
            if inner.id() == self_id {
                return None;
            }
            ancestor = inner.child_with_descendant(*self);
            Some(inner)
        })
        // We must iterate up the tree to preserve backwards compatibility
        .collect::<Vec<_>>()
        .into_iter()
        .rev()
    }
    fn dfs(&self) -> impl Iterator<Item = Self> {
        TsPre::new(self)
    }
    fn child(&self, nth: usize) -> Option<Self> {
        // TODO remove cast after migrating to tree-sitter
        Node::child(self, nth)
    }
    fn children(&self) -> impl ExactSizeIterator<Item = Self> {
        let mut cursor = self.walk();
        cursor.goto_first_child();
        NodeWalker {
            cursor,
            count: self.child_count(),
        }
    }
    fn child_by_field_id(&self, field_id: u16) -> Option<Self> {
        Node::child_by_field_id(self, field_id)
    }
    fn next(&self) -> Option<Self> {
        self.next_sibling()
    }
    fn prev(&self) -> Option<Self> {
        self.prev_sibling()
    }
    fn next_all(&self) -> impl Iterator<Item = Self> {
        // if root is none, use self as fallback to return a type-stable Iterator
        let node = self.parent().unwrap_or(*self);
        let mut cursor = node.walk();
        cursor.goto_first_child_for_byte(self.start_byte());
        std::iter::from_fn(move || {
            if cursor.goto_next_sibling() {
                Some(cursor.node())
            } else {
                None
            }
        })
    }
    fn prev_all(&self) -> impl Iterator<Item = Self> {
        // if root is none, use self as fallback to return a type-stable Iterator
        let node = self.parent().unwrap_or(*self);
        let mut cursor = node.walk();
        cursor.goto_first_child_for_byte(self.start_byte());
        std::iter::from_fn(move || {
            if cursor.goto_previous_sibling() {
                Some(cursor.node())
            } else {
                None
            }
        })
    }
    fn is_named(&self) -> bool {
        Node::is_named(self)
    }
    /// N.B. it is different from `is_named` && `is_leaf`
    /// if a `Node` has no named children.
    fn is_named_leaf(&self) -> bool {
        self.named_child_count() == 0
    }
    fn is_leaf(&self) -> bool {
        self.child_count() == 0
    }
    fn kind(&self) -> Cow<'_, str> {
        Cow::Borrowed(Node::kind(self))
    }
    fn kind_id(&self) -> KindId {
        Node::kind_id(self)
    }
    fn node_id(&self) -> usize {
        self.id()
    }
    fn range(&self) -> std::ops::Range<usize> {
        self.start_byte()..self.end_byte()
    }
    fn start_pos(&self) -> Position {
        let pos = self.start_position();
        let byte = self.start_byte();
        Position::new(pos.row, pos.column, byte)
    }
    fn end_pos(&self) -> Position {
        let pos = self.end_position();
        let byte = self.end_byte();
        Position::new(pos.row, pos.column, byte)
    }
    // missing node is a tree-sitter specific concept
    fn is_missing(&self) -> bool {
        Node::is_missing(self)
    }
    fn is_error(&self) -> bool {
        Node::is_error(self)
    }

    fn field(&self, name: &str) -> Option<Self> {
        self.child_by_field_name(name)
    }
    fn field_children(&self, field_id: Option<u16>) -> impl Iterator<Item = Self> {
        let field_id = field_id.and_then(NonZero::new);
        let mut cursor = self.walk();
        cursor.goto_first_child();
        // if field_id is not found, iteration is done
        let mut done = field_id.is_none();

        std::iter::from_fn(move || {
            if done {
                return None;
            }
            while cursor.field_id() != field_id {
                if !cursor.goto_next_sibling() {
                    return None;
                }
            }
            let ret = cursor.node();
            if !cursor.goto_next_sibling() {
                done = true;
            }
            Some(ret)
        })
    }
}

pub fn perform_edit<S: ContentExt>(tree: &mut Tree, input: &mut S, edit: &Edit<S>) -> InputEdit {
    let edit = input.accept_edit(edit);
    tree.edit(&edit);
    edit
}

/// Extension trait for languages that integrate with tree-sitter parsing.
///
/// `LanguageExt` extends the base [`Language`] trait with tree-sitter specific
/// functionality. Languages implementing this trait can be used with tree-sitter
/// parsers to create ASTs, handle language injections, and work with the
/// thread-ast-engine ecosystem.
///
/// # Key Capabilities
///
/// - **AST Creation**: Convenient methods to create AST-grep instances
/// - **Tree-sitter Integration**: Access to underlying tree-sitter language objects
/// - **Language Injection**: Support for multi-language documents (e.g., JavaScript in HTML)
/// - **Parsing**: Direct parsing of source code into [`StrDoc`] instances
///
/// # Example Implementation
///
/// ```rust,ignore
/// use thread_ast_engine::tree_sitter::{LanguageExt, TSLanguage};
/// use thread_ast_engine::Language;
///
/// #[derive(Clone)]
/// struct JavaScript;
///
/// impl Language for JavaScript {
///     // ... base Language implementation
/// }
///
/// impl LanguageExt for JavaScript {
///     fn get_ts_language(&self) -> TSLanguage {
///         tree_sitter_javascript::LANGUAGE.into()
///     }
/// }
/// ```
pub trait LanguageExt: Language {
    /// Create an [`AstGrep`] instance for parsing and pattern matching.
    ///
    /// Convenience method that parses the source code and returns an [`AstGrep`]
    /// instance ready for pattern matching and tree manipulation.
    ///
    /// # Parameters
    ///
    /// - `source` - Source code to parse
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// let ast = JavaScript.ast_grep("const x = 42;");
    /// let root = ast.root();
    /// ```
    fn ast_grep<S: AsRef<str>>(&self, source: S) -> AstGrep<StrDoc<Self>> {
        AstGrep::new(source, self.clone())
    }

    /// Get the tree-sitter language object for parsing.
    ///
    /// Returns the tree-sitter language grammar that this language uses
    /// for parsing source code. This is the core integration point between
    /// thread-ast-engine and tree-sitter.
    ///
    /// # Returns
    ///
    /// The tree-sitter [`TSLanguage`] object for this language
    fn get_ts_language(&self) -> TSLanguage;

    /// List of languages that can be injected into this language.
    ///
    /// For languages that support embedding other languages (like HTML with CSS/JavaScript),
    /// returns the names of languages that can be injected. Returns `None` if this
    /// language doesn't support injections.
    ///
    /// # Returns
    ///
    /// Array of injectable language names, or `None` if no injections supported
    fn injectable_languages(&self) -> Option<&'static [&'static str]> {
        None
    }

    /// Extract language injection regions from a parsed document.
    ///
    /// Analyzes the AST to find regions where other languages are embedded.
    /// For example, finds JavaScript code blocks within HTML `<script>` tags
    /// or CSS within `<style>` tags.
    ///
    /// Returns a map where keys are language names and values are lists of
    /// byte ranges in the source where that language appears.
    ///
    /// See [tree-sitter documentation](https://tree-sitter.github.io/tree-sitter/using-parsers#multi-language-documents)
    /// for more details on language injection.
    ///
    /// # Parameters
    ///
    /// - `root` - The root node of the document to analyze
    ///
    /// # Returns
    ///
    /// Map of language names to their byte ranges in the document
    #[cfg(feature = "matching")]
    fn extract_injections<L: LanguageExt>(
        &self,
        _root: crate::Node<StrDoc<L>>,
    ) -> RapidMap<String, Vec<TSRange>> {
        RapidMap::default()
    }
}

fn position_for_offset(input: &[u8], offset: usize) -> Point {
    debug_assert!(offset <= input.len());
    let (mut row, mut col) = (0, 0);
    for c in &input[0..offset] {
        if *c as char == '\n' {
            row += 1;
            col = 0;
        } else {
            col += 1;
        }
    }
    Point::new(row, col)
}

impl<L: LanguageExt> AstGrep<StrDoc<L>> {
    pub fn new<S: AsRef<str>>(src: S, lang: L) -> Self {
        Self::str(src.as_ref(), lang)
    }

    pub fn source(&self) -> &str {
        self.doc.get_source().as_str()
    }

    pub fn generate(self) -> String {
        self.doc.src
    }
}

pub trait ContentExt: Content {
    fn accept_edit(&mut self, edit: &Edit<Self>) -> InputEdit;
}
impl ContentExt for String {
    fn accept_edit(&mut self, edit: &Edit<Self>) -> InputEdit {
        let start_byte = edit.position;
        let old_end_byte = edit.position + edit.deleted_length;
        let new_end_byte = edit.position + edit.inserted_text.len();
        let input = unsafe { self.as_mut_vec() };
        let start_position = position_for_offset(input, start_byte);
        let old_end_position = position_for_offset(input, old_end_byte);
        input.splice(start_byte..old_end_byte, edit.inserted_text.clone());
        let new_end_position = position_for_offset(input, new_end_byte);
        InputEdit {
            start_byte,
            old_end_byte,
            new_end_byte,
            start_position,
            old_end_position,
            new_end_position,
        }
    }
}

impl<L: LanguageExt> Root<StrDoc<L>> {
    pub fn str(src: &str, lang: L) -> Self {
        Self::try_new(src, lang).expect("should parse")
    }
    pub fn try_new(src: &str, lang: L) -> Result<Self, String> {
        let doc = StrDoc::try_new(src, lang)?;
        Ok(Self { doc })
    }
    pub fn get_text(&self) -> &str {
        &self.doc.src
    }
    #[cfg(feature = "matching")]
    pub fn get_injections<F: Fn(&str) -> Option<L>>(&self, get_lang: F) -> Vec<Self> {
        let root = self.root();
        let range = self.lang().extract_injections(root);
        range
            .into_iter()
            .filter_map(|(lang, ranges)| {
                let lang = get_lang(&lang)?;
                let source = self.doc.get_source();
                let mut parser = Parser::new();
                parser.set_included_ranges(&ranges).ok()?;
                parser.set_language(&lang.get_ts_language()).ok()?;
                let tree = parser.parse(source, None)?;
                Some(Self {
                    doc: StrDoc {
                        src: self.doc.src.clone(),
                        lang,
                        tree,
                    },
                })
            })
            .collect()
    }
}

pub struct DisplayContext<'r> {
    /// content for the matched node
    pub matched: Cow<'r, str>,
    /// content before the matched node
    pub leading: &'r str,
    /// content after the matched node
    pub trailing: &'r str,
    /// zero-based start line of the context
    pub start_line: usize,
}

/// these methods are only for `StrDoc`
impl<'r, L: LanguageExt> crate::Node<'r, StrDoc<L>> {
    #[doc(hidden)]
    #[must_use]
    pub fn display_context(&self, before: usize, after: usize) -> DisplayContext<'r> {
        let source = self.root.doc.get_source().as_str();
        let bytes = source.as_bytes();
        let start = self.inner.start_byte();
        let end = self.inner.end_byte();
        let (mut leading, mut trailing) = (start, end);
        let mut lines_before = before + 1;
        while leading > 0 {
            if bytes[leading - 1] == b'\n' {
                lines_before -= 1;
                if lines_before == 0 {
                    break;
                }
            }
            leading -= 1;
        }
        let mut lines_after = after + 1;
        // tree-sitter will append line ending to source so trailing can be out of bound
        trailing = trailing.min(bytes.len());
        while trailing < bytes.len() {
            if bytes[trailing] == b'\n' {
                lines_after -= 1;
                if lines_after == 0 {
                    break;
                }
            }
            trailing += 1;
        }
        // lines_before means we matched all context, offset is `before` itself
        let offset = if lines_before == 0 {
            before
        } else {
            // otherwise, there are fewer than `before` line in src, compute the actual line
            before + 1 - lines_before
        };
        DisplayContext {
            matched: self.text(),
            leading: &source[leading..start],
            trailing: &source[end..trailing],
            start_line: self.start_pos().line() - offset,
        }
    }

    #[cfg(feature = "matching")]
    pub fn replace_all<M: Matcher, R: Replacer<StrDoc<L>>>(
        &self,
        matcher: M,
        replacer: R,
    ) -> Vec<Edit<String>> {
        // TODO: support nested matches like Some(Some(1)) with pattern Some($A)
        Visitor::new(&matcher)
            .reentrant(false)
            .visit(self.clone())
            .map(|matched| matched.make_edit(&matcher, &replacer))
            .collect()
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::language::Tsx;
    use tree_sitter::Point;

    fn parse(src: &str) -> Result<Tree, TSParseError> {
        parse_lang(|p| p.parse(src, None), &Tsx.get_ts_language())
    }

    #[test]
    fn test_tree_sitter() -> Result<(), TSParseError> {
        let tree = parse("var a = 1234")?;
        let root_node = tree.root_node();
        assert_eq!(root_node.kind(), "program");
        assert_eq!(root_node.start_position().column, 0);
        assert_eq!(root_node.end_position().column, 12);
        assert_eq!(
            root_node.to_sexp(),
            "(program (variable_declaration (variable_declarator name: (identifier) value: (number))))"
        );
        Ok(())
    }

    #[test]
    fn test_object_literal() -> Result<(), TSParseError> {
        let tree = parse("{a: $X}")?;
        let root_node = tree.root_node();
        // wow this is not label. technically it is wrong but practically it is better LOL
        assert_eq!(
            root_node.to_sexp(),
            "(program (expression_statement (object (pair key: (property_identifier) value: (identifier)))))"
        );
        Ok(())
    }

    #[test]
    fn test_string() -> Result<(), TSParseError> {
        let tree = parse("'$A'")?;
        let root_node = tree.root_node();
        assert_eq!(
            root_node.to_sexp(),
            "(program (expression_statement (string (string_fragment))))"
        );
        Ok(())
    }

    #[test]
    fn test_row_col() -> Result<(), TSParseError> {
        let tree = parse("😄")?;
        let root = tree.root_node();
        assert_eq!(root.start_position(), Point::new(0, 0));
        // NOTE: Point in tree-sitter is counted in bytes instead of char
        assert_eq!(root.end_position(), Point::new(0, 4));
        Ok(())
    }

    #[test]
    fn test_edit() -> Result<(), TSParseError> {
        let mut src = "a + b".to_string();
        let mut tree = parse(&src)?;
        let _ = perform_edit(
            &mut tree,
            &mut src,
            &Edit {
                position: 1,
                deleted_length: 0,
                inserted_text: " * b".into(),
            },
        );
        let tree2 = parse_lang(|p| p.parse(&src, Some(&tree)), &Tsx.get_ts_language())?;
        assert_eq!(
            tree.root_node().to_sexp(),
            "(program (expression_statement (binary_expression left: (identifier) right: (identifier))))"
        );
        assert_eq!(
            tree2.root_node().to_sexp(),
            "(program (expression_statement (binary_expression left: (binary_expression left: (identifier) right: (identifier)) right: (identifier))))"
        );
        Ok(())
    }
}