thread-ast-engine 0.1.1

// 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

use super::kind::{KindMatcher, kind_utils};
use super::matcher::Matcher;
pub use super::types::{MatchStrictness, Pattern, PatternBuilder, PatternError, PatternNode};
use crate::language::Language;
use crate::match_tree::{match_end_non_recursive, match_node_non_recursive};
use crate::meta_var::{MetaVarEnv, MetaVariable};
use crate::source::SgNode;
use crate::{Doc, Node, Root};

use bit_set::BitSet;
use std::borrow::Cow;
use thread_utilities::RapidSet;

impl PatternBuilder<'_> {
    pub fn build<D, F>(&self, parse: F) -> Result<Pattern, PatternError>
    where
        F: FnOnce(&str) -> Result<D, String>,
        D: Doc,
    {
        let doc = parse(&self.src).map_err(PatternError::Parse)?;
        let root = Root::doc(doc);
        if let Some(selector) = self.selector {
            self.contextual(&root, selector)
        } else {
            self.single(&root)
        }
    }
    fn single<D: Doc>(&self, root: &Root<D>) -> Result<Pattern, PatternError> {
        let goal = root.root();
        if goal.children().len() == 0 {
            return Err(PatternError::NoContent(self.src.to_string()));
        }
        if !is_single_node(&goal.inner) {
            return Err(PatternError::MultipleNode(self.src.to_string()));
        }
        let node = Pattern::single_matcher(root);
        Ok(Pattern::from(node))
    }

    fn contextual<D: Doc>(&self, root: &Root<D>, selector: &str) -> Result<Pattern, PatternError> {
        let goal = root.root();
        let kind_matcher = KindMatcher::try_new(selector, root.lang())?;
        let Some(node) = goal.find(&kind_matcher) else {
            return Err(PatternError::NoSelectorInContext {
                context: self.src.to_string(),
                selector: selector.into(),
            });
        };
        Ok(Pattern {
            root_kind: Some(node.kind_id()),
            node: convert_node_to_pattern(node.get_node()),
            strictness: MatchStrictness::Smart,
        })
    }
}

impl PatternNode {
    // for skipping trivial nodes in goal after ellipsis
    #[must_use]
    pub const fn is_trivial(&self) -> bool {
        match self {
            Self::Terminal { is_named, .. } => !*is_named,
            _ => false,
        }
    }

    #[inline]
    #[must_use]
    pub fn fixed_string(&self) -> Cow<'_, str> {
        match &self {
            Self::Terminal { text, .. } => Cow::Borrowed(text),
            Self::MetaVar { .. } => Cow::Borrowed(""),
            Self::Internal { children, .. } => children.iter().map(|n| n.fixed_string()).fold(
                Cow::Borrowed(""),
                |longest, curr| {
                    if longest.len() >= curr.len() {
                        longest
                    } else {
                        curr
                    }
                },
            ),
        }
    }
}
impl<'r, D: Doc> From<Node<'r, D>> for PatternNode {
    fn from(node: Node<'r, D>) -> Self {
        convert_node_to_pattern(&node)
    }
}

impl<'r, D: Doc> From<Node<'r, D>> for Pattern {
    fn from(node: Node<'r, D>) -> Self {
        Self {
            node: convert_node_to_pattern(&node),
            root_kind: None,
            strictness: MatchStrictness::Smart,
        }
    }
}

fn convert_node_to_pattern<D: Doc>(node: &Node<'_, D>) -> PatternNode {
    if let Some(meta_var) = extract_var_from_node(node) {
        PatternNode::MetaVar { meta_var }
    } else if node.is_leaf() {
        PatternNode::Terminal {
            text: node.text().to_string(),
            is_named: node.is_named(),
            kind_id: node.kind_id(),
        }
    } else {
        // Pre-allocate vector with estimated capacity to reduce allocations
        let child_count = node.children().count();
        let mut children = Vec::with_capacity(child_count);

        for child in node.children() {
            if !child.is_missing() {
                children.push(PatternNode::from(child));
            }
        }

        PatternNode::Internal {
            kind_id: node.kind_id(),
            children,
        }
    }
}

fn extract_var_from_node<D: Doc>(goal: &Node<'_, D>) -> Option<MetaVariable> {
    let key = goal.text();
    goal.lang().extract_meta_var(&key)
}

#[inline]
fn is_single_node<'r, N: SgNode<'r>>(n: &N) -> bool {
    match n.children().len() {
        1 => true,
        2 => {
            let c = n.child(1).expect("second child must exist");
            // some language will have weird empty syntax node at the end
            // see golang's `$A = 0` pattern test case
            c.is_missing() || c.kind().is_empty()
        }
        _ => false,
    }
}
impl Pattern {
    #[must_use]
    pub const fn has_error(&self) -> bool {
        let kind = match &self.node {
            PatternNode::Terminal { kind_id, .. } | PatternNode::Internal { kind_id, .. } => {
                *kind_id
            }
            PatternNode::MetaVar { .. } => match self.root_kind {
                Some(k) => k,
                None => return false,
            },
        };
        kind_utils::is_error_kind(kind)
    }

    #[must_use]
    pub fn fixed_string(&self) -> Cow<'_, str> {
        self.node.fixed_string()
    }

    /// Get all defined variables in the pattern.
    /// Used for validating rules and report undefined variables.
    #[must_use]
    pub fn defined_vars(&self) -> RapidSet<&str> {
        let mut vars = RapidSet::default();
        collect_vars(&self.node, &mut vars);
        vars
    }
}

fn meta_var_name(meta_var: &MetaVariable) -> Option<&str> {
    use MetaVariable as MV;
    match meta_var {
        MV::Capture(name, _) | MV::MultiCapture(name) => Some(name),
        MV::Dropped(_) | MV::Multiple => None,
    }
}

fn collect_vars<'p>(p: &'p PatternNode, vars: &mut RapidSet<&'p str>) {
    match p {
        PatternNode::MetaVar { meta_var, .. } => {
            if let Some(name) = meta_var_name(meta_var) {
                vars.insert(name);
            }
        }
        PatternNode::Terminal { .. } => {
            // collect nothing for terminal nodes!
        }
        PatternNode::Internal { children, .. } => {
            for c in children {
                collect_vars(c, vars);
            }
        }
    }
}

impl Pattern {
    pub fn try_new<L: Language>(src: &str, lang: &L) -> Result<Self, PatternError> {
        let processed = lang.pre_process_pattern(src);
        let builder = PatternBuilder {
            selector: None,
            src: processed,
        };
        lang.build_pattern(&builder)
    }

    pub fn new<L: Language>(src: &str, lang: &L) -> Self {
        Self::try_new(src, lang).unwrap()
    }

    #[must_use]
    pub const fn with_strictness(mut self, strictness: MatchStrictness) -> Self {
        self.strictness = strictness;
        self
    }

    pub fn contextual<L: Language>(
        context: &str,
        selector: &str,
        lang: &L,
    ) -> Result<Self, PatternError> {
        let processed = lang.pre_process_pattern(context);
        let builder = PatternBuilder {
            selector: Some(selector),
            src: processed,
        };
        lang.build_pattern(&builder)
    }
    fn single_matcher<D: Doc>(root: &Root<D>) -> Node<'_, D> {
        // debug_assert!(matches!(self.style, PatternStyle::Single));
        let node = root.root();
        let mut inner = node.inner;
        while is_single_node(&inner) {
            inner = inner.child(0).unwrap();
        }
        Node { inner, root }
    }
}

impl Matcher for Pattern {
    fn match_node_with_env<'tree, D: Doc>(
        &self,
        node: Node<'tree, D>,
        env: &mut Cow<MetaVarEnv<'tree, D>>,
    ) -> Option<Node<'tree, D>> {
        if let Some(k) = self.root_kind
            && node.kind_id() != k
        {
            return None;
        }
        // do not pollute the env if pattern does not match
        let mut may_write = Cow::Borrowed(env.as_ref());
        let node = match_node_non_recursive(self, node, &mut may_write)?;
        if let Cow::Owned(map) = may_write {
            // only change env when pattern matches
            *env = Cow::Owned(map);
        }
        Some(node)
    }

    fn potential_kinds(&self) -> Option<bit_set::BitSet> {
        let kind = match self.node {
            PatternNode::Terminal { kind_id, .. } => kind_id,
            PatternNode::MetaVar { .. } => self.root_kind?,
            PatternNode::Internal { kind_id, .. } => {
                if kind_utils::is_error_kind(kind_id) {
                    // error can match any kind
                    return None;
                }
                kind_id
            }
        };

        let mut kinds = BitSet::new();
        kinds.insert(kind.into());
        Some(kinds)
    }

    fn get_match_len<D: Doc>(&self, node: Node<'_, D>) -> Option<usize> {
        let start = node.range().start;
        let end = match_end_non_recursive(self, &node)?;
        Some(end - start)
    }
}
impl std::fmt::Debug for PatternNode {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::MetaVar { meta_var, .. } => write!(f, "{meta_var:?}"),
            Self::Terminal { text, .. } => write!(f, "{text}"),
            Self::Internal { children, .. } => write!(f, "{children:?}"),
        }
    }
}

impl std::fmt::Debug for Pattern {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{:?}", self.node)
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::language::Tsx;
    use crate::matcher::MatcherExt;
    use crate::meta_var::MetaVarEnv;
    use crate::tree_sitter::StrDoc;
    use thread_utilities::RapidMap;

    fn pattern_node(s: &str) -> Root<StrDoc<Tsx>> {
        Root::str(s, Tsx)
    }

    fn test_match(s1: &str, s2: &str) {
        let pattern = Pattern::new(s1, &Tsx);
        let cand = pattern_node(s2);
        let cand = cand.root();
        assert!(
            pattern.find_node(cand.clone()).is_some(),
            "goal: {:?}, candidate: {}",
            pattern,
            cand.get_inner_node().to_sexp(),
        );
    }
    fn test_non_match(s1: &str, s2: &str) {
        let pattern = Pattern::new(s1, &Tsx);
        let cand = pattern_node(s2);
        let cand = cand.root();
        assert!(
            pattern.find_node(cand.clone()).is_none(),
            "goal: {:?}, candidate: {}",
            pattern,
            cand.get_inner_node().to_sexp(),
        );
    }

    #[test]
    fn test_meta_variable() {
        test_match("const a = $VALUE", "const a = 123");
        test_match("const $VARIABLE = $VALUE", "const a = 123");
        test_match("const $VARIABLE = $VALUE", "const a = 123");
    }

    #[test]
    fn test_whitespace() {
        test_match("function t() { }", "function t() {}");
        test_match("function t() {}", "function t() {  }");
    }

    fn match_env(goal_str: &str, cand: &str) -> RapidMap<String, String> {
        let pattern = Pattern::new(goal_str, &Tsx);
        let cand = pattern_node(cand);
        let cand = cand.root();
        let nm = pattern.find_node(cand).unwrap();
        RapidMap::from(nm.get_env().clone())
    }

    #[test]
    fn test_meta_variable_env() {
        let env = match_env("const a = $VALUE", "const a = 123");
        assert_eq!(env["VALUE"], "123");
    }

    #[test]
    fn test_pattern_should_not_pollute_env() {
        // gh issue #1164
        let pattern = Pattern::new("const $A = 114", &Tsx);
        let cand = pattern_node("const a = 514");
        let cand = cand.root().child(0).unwrap();
        let map = MetaVarEnv::new();
        let mut env = Cow::Borrowed(&map);
        let nm = pattern.match_node_with_env(cand, &mut env);
        assert!(nm.is_none());
        assert!(env.get_match("A").is_none());
        assert!(map.get_match("A").is_none());
    }

    #[test]
    fn test_match_non_atomic() {
        let env = match_env("const a = $VALUE", "const a = 5 + 3");
        assert_eq!(env["VALUE"], "5 + 3");
    }

    #[test]
    fn test_class_assignment() {
        test_match("class $C { $MEMBER = $VAL}", "class A {a = 123}");
        test_non_match("class $C { $MEMBER = $VAL; b = 123; }", "class A {a = 123}");
        // test_match("a = 123", "class A {a = 123}");
        test_non_match("a = 123", "class B {b = 123}");
    }

    #[test]
    fn test_return() {
        test_match("$A($B)", "return test(123)");
    }

    #[test]
    fn test_contextual_pattern() {
        let pattern = Pattern::contextual("class A { $F = $I }", "public_field_definition", &Tsx)
            .expect("test");
        let cand = pattern_node("class B { b = 123 }");
        assert!(pattern.find_node(cand.root()).is_some());
        let cand = pattern_node("let b = 123");
        assert!(pattern.find_node(cand.root()).is_none());
    }

    #[test]
    fn test_contextual_match_with_env() {
        let pattern = Pattern::contextual("class A { $F = $I }", "public_field_definition", &Tsx)
            .expect("test");
        let cand = pattern_node("class B { b = 123 }");
        let nm = pattern.find_node(cand.root()).expect("test");
        let env = nm.get_env();
        let env = RapidMap::from(env.clone());
        assert_eq!(env["F"], "b");
        assert_eq!(env["I"], "123");
    }

    #[test]
    fn test_contextual_unmatch_with_env() {
        let pattern = Pattern::contextual("class A { $F = $I }", "public_field_definition", &Tsx)
            .expect("test");
        let cand = pattern_node("let b = 123");
        let nm = pattern.find_node(cand.root());
        assert!(nm.is_none());
    }

    fn get_kind(kind_str: &str) -> usize {
        Tsx.kind_to_id(kind_str).into()
    }

    #[test]
    fn test_pattern_potential_kinds() {
        let pattern = Pattern::new("const a = 1", &Tsx);
        let kind = get_kind("lexical_declaration");
        let kinds = pattern.potential_kinds().expect("should have kinds");
        assert_eq!(kinds.len(), 1);
        assert!(kinds.contains(kind));
    }

    #[test]
    fn test_pattern_with_non_root_meta_var() {
        let pattern = Pattern::new("const $A = $B", &Tsx);
        let kind = get_kind("lexical_declaration");
        let kinds = pattern.potential_kinds().expect("should have kinds");
        assert_eq!(kinds.len(), 1);
        assert!(kinds.contains(kind));
    }

    #[test]
    fn test_bare_wildcard() {
        let pattern = Pattern::new("$A", &Tsx);
        // wildcard should match anything, so kinds should be None
        assert!(pattern.potential_kinds().is_none());
    }

    #[test]
    fn test_contextual_potential_kinds() {
        let pattern = Pattern::contextual("class A { $F = $I }", "public_field_definition", &Tsx)
            .expect("test");
        let kind = get_kind("public_field_definition");
        let kinds = pattern.potential_kinds().expect("should have kinds");
        assert_eq!(kinds.len(), 1);
        assert!(kinds.contains(kind));
    }

    #[test]
    fn test_contextual_wildcard() {
        let pattern =
            Pattern::contextual("class A { $F }", "property_identifier", &Tsx).expect("test");
        let kind = get_kind("property_identifier");
        let kinds = pattern.potential_kinds().expect("should have kinds");
        assert_eq!(kinds.len(), 1);
        assert!(kinds.contains(kind));
    }

    #[test]
    #[ignore = "multi-node patterns not yet implemented"]
    fn test_multi_node_pattern() {
        let pattern = Pattern::new("a;b;c;", &Tsx);
        let kinds = pattern.potential_kinds().expect("should have kinds");
        assert_eq!(kinds.len(), 1);
        test_match("a;b;c", "a;b;c;");
    }

    #[test]
    #[ignore = "multi-node patterns not yet implemented"]
    fn test_multi_node_meta_var() {
        let env = match_env("a;$B;c", "a;b;c");
        assert_eq!(env["B"], "b");
        let env = match_env("a;$B;c", "a;1+2+3;c");
        assert_eq!(env["B"], "1+2+3");
    }

    #[test]
    #[ignore = "struct layout is compiler and platform specific"]
    fn test_pattern_size() {
        assert_eq!(std::mem::size_of::<Pattern>(), 40);
    }

    #[test]
    fn test_error_kind() {
        let ret = Pattern::contextual("a", "property_identifier", &Tsx);
        assert!(ret.is_err());
        let ret = Pattern::new("123+", &Tsx);
        assert!(ret.has_error());
    }

    #[test]
    fn test_bare_wildcard_in_context() {
        let pattern =
            Pattern::contextual("class A { $F }", "property_identifier", &Tsx).expect("test");
        let cand = pattern_node("let b = 123");
        // it should not match
        assert!(pattern.find_node(cand.root()).is_none());
    }

    #[test]
    fn test_pattern_fixed_string() {
        let pattern = Pattern::new("class A { $F }", &Tsx);
        assert_eq!(pattern.fixed_string(), "class");
        let pattern =
            Pattern::contextual("class A { $F }", "property_identifier", &Tsx).expect("test");
        assert!(pattern.fixed_string().is_empty());
    }

    #[test]
    fn test_pattern_error() {
        let pattern = Pattern::try_new("", &Tsx);
        assert!(matches!(pattern, Err(PatternError::NoContent(_))));
        let pattern = Pattern::try_new("12  3344", &Tsx);
        assert!(matches!(pattern, Err(PatternError::MultipleNode(_))));
    }

    #[test]
    fn test_debug_pattern() {
        let pattern = Pattern::new("var $A = 1", &Tsx);
        assert_eq!(
            format!("{pattern:?}"),
            "[var, [Capture(\"A\", true), =, 1]]"
        );
    }

    fn defined_vars(s: &str) -> Vec<String> {
        let pattern = Pattern::new(s, &Tsx);
        let mut vars: Vec<_> = pattern
            .defined_vars()
            .into_iter()
            .map(String::from)
            .collect();
        vars.sort();
        vars
    }

    #[test]
    fn test_extract_meta_var_from_pattern() {
        let vars = defined_vars("var $A = 1");
        assert_eq!(vars, ["A"]);
    }

    #[test]
    fn test_extract_complex_meta_var() {
        let vars = defined_vars("function $FUNC($$$ARGS): $RET { $$$BODY }");
        assert_eq!(vars, ["ARGS", "BODY", "FUNC", "RET"]);
    }

    #[test]
    fn test_extract_duplicate_meta_var() {
        let vars = defined_vars("var $A = $A");
        assert_eq!(vars, ["A"]);
    }

    #[test]
    fn test_contextual_pattern_vars() {
        let pattern =
            Pattern::contextual("<div ref={$A}/>", "jsx_attribute", &Tsx).expect("correct");
        assert_eq!(pattern.defined_vars(), std::iter::once("A").collect());
    }

    #[test]
    fn test_gh_1087() {
        test_match("($P) => $F($P)", "(x) => bar(x)");
    }
}