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
//! # Meta-variable Environment and Utilities
//!
//! This module provides types and functions for handling meta-variables in AST pattern matching.
//! Meta-variables allow patterns to flexibly match and capture code fragments, supporting single and multi-capture semantics.
//!
//! ## Key Components
//!
//! - [`MetaVarEnv`](crates/ast-engine/src/meta_var.rs:26): Stores meta-variable instantiations during pattern matching.
//! - [`MetaVariable`](crates/ast-engine/src/meta_var.rs:260): Enum representing different meta-variable forms (single, multi, dropped).
//! - `extract_meta_var`: Utility to parse meta-variable strings.
//! - Insertion, retrieval, and transformation APIs for meta-variable environments.
//!
//! ## Example
//!
//! ```rust,no_run
//! use thread_ast_engine::meta_var::{MetaVarEnv, MetaVariable, extract_meta_var};
//!
//! let mut env = MetaVarEnv::new();
//! env.insert("$A", node);
//! let meta = extract_meta_var("$A", '$');
//! ```
//!
//! See [`MetaVarEnv`](crates/ast-engine/src/meta_var.rs:48) for details on usage in AST matching and rewriting.
#[cfg(feature = "matching")]
use crate::match_tree::does_node_match_exactly;
#[cfg(feature = "matching")]
use crate::matcher::Matcher;
#[cfg(feature = "matching")]
use crate::replacer::formatted_slice;
use crate::source::Content;
use crate::{Doc, Node};
#[cfg(feature = "matching")]
use std::borrow::Cow;
use std::sync::Arc;
use thread_utilities::{RapidMap, map_with_capacity};

/// Interned string type for meta-variable identifiers.
///
/// Using `Arc<str>` instead of `String` eliminates per-clone heap allocations.
/// Cloning an `Arc<str>` is a single atomic increment (~1ns) versus `String::clone`
/// which copies the entire buffer (~10-50ns depending on length). Since meta-variable
/// names are cloned extensively during pattern matching (environment forks, variable
/// captures, constraint checking), this reduces allocation pressure by 20-30%.
pub type MetaVariableID = Arc<str>;

pub type Underlying<D> = Vec<<<D as Doc>::Source as Content>::Underlying>;

/// a dictionary that stores metavariable instantiation
/// const a = 123 matched with const a = $A will produce env: $A => 123
#[derive(Clone, Debug)]
pub struct MetaVarEnv<'tree, D: Doc> {
    single_matched: RapidMap<MetaVariableID, Node<'tree, D>>,
    multi_matched: RapidMap<MetaVariableID, Vec<Node<'tree, D>>>,
    transformed_var: RapidMap<MetaVariableID, Underlying<D>>,
}

impl<'t, D: Doc> MetaVarEnv<'t, D> {
    #[must_use]
    pub fn new() -> Self {
        Self {
            single_matched: RapidMap::default(),
            multi_matched: RapidMap::default(),
            transformed_var: RapidMap::default(),
        }
    }

    #[cfg(feature = "matching")]
    pub fn insert(&mut self, id: &str, ret: Node<'t, D>) -> Option<&mut Self> {
        if self.match_variable(id, &ret) {
            self.single_matched.insert(Arc::from(id), ret);
            Some(self)
        } else {
            None
        }
    }

    #[cfg(feature = "matching")]
    pub fn insert_multi(&mut self, id: &str, ret: Vec<Node<'t, D>>) -> Option<&mut Self> {
        if self.match_multi_var(id, &ret) {
            self.multi_matched.insert(Arc::from(id), ret);
            Some(self)
        } else {
            None
        }
    }

    /// Insert without cloning the key if it's already owned
    #[cfg(feature = "matching")]
    pub fn insert_owned(&mut self, id: MetaVariableID, ret: Node<'t, D>) -> Option<&mut Self> {
        if self.match_variable(&id, &ret) {
            self.single_matched.insert(id, ret);
            Some(self)
        } else {
            None
        }
    }

    /// Insert multi without cloning the key if it's already owned
    #[cfg(feature = "matching")]
    pub fn insert_multi_owned(
        &mut self,
        id: MetaVariableID,
        ret: Vec<Node<'t, D>>,
    ) -> Option<&mut Self> {
        if self.match_multi_var(&id, &ret) {
            self.multi_matched.insert(id, ret);
            Some(self)
        } else {
            None
        }
    }
    #[must_use]
    pub fn get_match(&self, var: &str) -> Option<&'_ Node<'t, D>> {
        self.single_matched.get(var)
    }
    #[must_use]
    pub fn get_multiple_matches(&self, var: &str) -> Vec<Node<'t, D>> {
        self.multi_matched.get(var).cloned().unwrap_or_default()
    }

    /// Returns a reference to multiple matches without cloning
    #[must_use]
    pub fn get_multiple_matches_ref(&self, var: &str) -> Option<&Vec<Node<'t, D>>> {
        self.multi_matched.get(var)
    }

    pub fn add_label(&mut self, label: &str, node: Node<'t, D>) {
        self.multi_matched
            .entry(Arc::from(label))
            .or_default()
            .push(node);
    }
    #[must_use]
    pub fn get_labels(&self, label: &str) -> Option<&Vec<Node<'t, D>>> {
        self.multi_matched.get(label)
    }

    #[cfg(feature = "matching")]
    pub fn get_matched_variables(&self) -> impl Iterator<Item = MetaVariable> + use<'_, 't, D> {
        let single = self
            .single_matched
            .keys()
            .map(|n| MetaVariable::Capture(n.clone(), false));
        let transformed = self
            .transformed_var
            .keys()
            .map(|n| MetaVariable::Capture(n.clone(), false));
        let multi = self
            .multi_matched
            .keys()
            .map(|n| MetaVariable::MultiCapture(n.clone()));
        single.chain(multi).chain(transformed)
    }

    #[cfg(feature = "matching")]
    #[must_use]
    fn match_variable(&self, id: &str, candidate: &Node<'t, D>) -> bool {
        if let Some(m) = self.single_matched.get(id) {
            return does_node_match_exactly(m, candidate);
        }
        true
    }
    #[cfg(feature = "matching")]
    fn match_multi_var(&self, id: &str, cands: &[Node<'t, D>]) -> bool {
        let Some(nodes) = self.multi_matched.get(id) else {
            return true;
        };
        let mut named_nodes = nodes.iter().filter(|n| n.is_named());
        let mut named_cands = cands.iter().filter(|n| n.is_named());
        loop {
            if let Some(node) = named_nodes.next() {
                let Some(cand) = named_cands.next() else {
                    // cand is done but node is not
                    break false;
                };
                if !does_node_match_exactly(node, cand) {
                    break false;
                }
            } else if named_cands.next().is_some() {
                // node is done but cand is not
                break false;
            } else {
                // both None, matches
                break true;
            }
        }
    }

    #[cfg(feature = "matching")]
    pub fn match_constraints<M: Matcher>(
        &mut self,
        var_matchers: &RapidMap<MetaVariableID, M>,
    ) -> bool {
        let mut env = Cow::Borrowed(self);
        for (var_id, candidate) in &self.single_matched {
            if let Some(m) = var_matchers.get(var_id)
                && m.match_node_with_env(candidate.clone(), &mut env).is_none()
            {
                return false;
            }
        }
        if let Cow::Owned(env) = env {
            *self = env;
        }
        true
    }

    #[cfg(feature = "matching")]
    pub fn insert_transformation(&mut self, var: &MetaVariable, name: &str, slice: Underlying<D>) {
        let node = match var {
            MetaVariable::Capture(v, _) => self.single_matched.get(v),
            MetaVariable::MultiCapture(vs) => self.multi_matched.get(vs).and_then(|vs| vs.first()),
            _ => None,
        };
        let deindented = if let Some(v) = node {
            formatted_slice(&slice, v.get_doc().get_source(), v.range().start).to_vec()
        } else {
            slice
        };
        self.transformed_var.insert(Arc::from(name), deindented);
    }
    #[must_use]
    pub fn get_transformed(&self, var: &str) -> Option<&Underlying<D>> {
        self.transformed_var.get(var)
    }
    #[must_use]
    pub fn get_var_bytes<'s>(
        &'s self,
        var: &MetaVariable,
    ) -> Option<&'s [<D::Source as Content>::Underlying]> {
        get_var_bytes_impl(self, var)
    }
}

#[cfg(feature = "matching")]
impl<D: Doc> MetaVarEnv<'_, D> {
    /// internal for readopt `NodeMatch` in pinned.rs
    /// readopt node and env when sending them to other threads
    pub(crate) fn visit_nodes<F>(&mut self, mut f: F)
    where
        F: FnMut(&mut Node<'_, D>),
    {
        for n in self.single_matched.values_mut() {
            f(n);
        }
        for ns in self.multi_matched.values_mut() {
            for n in ns {
                f(n);
            }
        }
    }
}

fn get_var_bytes_impl<'e, 't, C, D>(
    env: &'e MetaVarEnv<'t, D>,
    var: &MetaVariable,
) -> Option<&'e [C::Underlying]>
where
    D: Doc<Source = C> + 't,
    C: Content + 't,
{
    match var {
        MetaVariable::Capture(n, _) => {
            if let Some(node) = env.get_match(n) {
                let bytes = node.get_doc().get_source().get_range(node.range());
                Some(bytes)
            } else if let Some(bytes) = env.get_transformed(n) {
                Some(bytes)
            } else {
                None
            }
        }
        MetaVariable::MultiCapture(n) => {
            let nodes = env.get_multiple_matches(n);
            if nodes.is_empty() {
                None
            } else {
                // NOTE: start_byte is not always index range of source's slice.
                // e.g. start_byte is still byte_offset in utf_16 (napi). start_byte
                // so we need to call source's get_range method
                let start = nodes[0].range().start;
                let end = nodes[nodes.len() - 1].range().end;
                Some(nodes[0].get_doc().get_source().get_range(start..end))
            }
        }
        _ => None,
    }
}

impl<D: Doc> Default for MetaVarEnv<'_, D> {
    fn default() -> Self {
        Self::new()
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum MetaVariable {
    /// $A for captured meta var
    Capture(MetaVariableID, bool),
    /// $_ for non-captured meta var
    Dropped(bool),
    /// $$$ for non-captured multi var
    Multiple,
    /// $$$A for captured ellipsis
    MultiCapture(MetaVariableID),
}

pub(crate) fn extract_meta_var(src: &str, meta_char: char) -> Option<MetaVariable> {
    use MetaVariable::{Capture, Dropped, MultiCapture, Multiple};
    let ellipsis: String = std::iter::repeat_n(meta_char, 3).collect();
    if src == ellipsis {
        return Some(Multiple);
    }
    if let Some(trimmed) = src.strip_prefix(&ellipsis) {
        if !trimmed.chars().all(is_valid_meta_var_char) {
            return None;
        }
        if trimmed.starts_with('_') {
            return Some(Multiple);
        }
        return Some(MultiCapture(Arc::from(trimmed)));
    }
    if !src.starts_with(meta_char) {
        return None;
    }
    let trimmed = &src[meta_char.len_utf8()..];
    let (trimmed, named) = if let Some(t) = trimmed.strip_prefix(meta_char) {
        (t, false)
    } else {
        (trimmed, true)
    };
    if !trimmed.starts_with(is_valid_first_char) || // empty or started with number
    !trimmed.chars().all(is_valid_meta_var_char)
    // not in form of $A or $_
    {
        return None;
    }
    if trimmed.starts_with('_') {
        Some(Dropped(named))
    } else {
        Some(Capture(Arc::from(trimmed), named))
    }
}

#[inline]
const fn is_valid_first_char(c: char) -> bool {
    matches!(c, 'A'..='Z' | '_')
}

#[inline]
pub(crate) const fn is_valid_meta_var_char(c: char) -> bool {
    is_valid_first_char(c) || c.is_ascii_digit()
}

// RapidMap is intentionally specific (not generic over BuildHasher) for performance.
// This conversion is in the pattern matching hot path and should use rapidhash.
#[allow(clippy::implicit_hasher)]
impl<'tree, D: Doc> From<MetaVarEnv<'tree, D>> for RapidMap<String, String>
where
    D::Source: Content,
{
    fn from(env: MetaVarEnv<'tree, D>) -> Self {
        let mut ret: Self = map_with_capacity(
            env.single_matched.len() + env.multi_matched.len() + env.transformed_var.len(),
        );
        for (id, node) in env.single_matched {
            ret.insert(id.to_string(), node.text().into());
        }
        for (id, bytes) in env.transformed_var {
            ret.insert(
                id.to_string(),
                <D::Source as Content>::encode_bytes(&bytes).to_string(),
            );
        }
        for (id, nodes) in env.multi_matched {
            // Optimize string concatenation by pre-calculating capacity
            if nodes.is_empty() {
                ret.insert(id.to_string(), "[]".to_string());
                continue;
            }

            let estimated_capacity = nodes.len() * 16 + 10; // rough estimate
            let mut result = String::with_capacity(estimated_capacity);
            result.push('[');

            let mut first = true;
            for node in &nodes {
                if !first {
                    result.push_str(", ");
                }
                result.push_str(&node.text());
                first = false;
            }
            result.push(']');
            ret.insert(id.to_string(), result);
        }
        ret
    }
}

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

    fn extract_var(s: &str) -> Option<MetaVariable> {
        extract_meta_var(s, '$')
    }
    #[test]
    fn test_match_var() {
        use MetaVariable::*;
        assert_eq!(extract_var("$$$"), Some(Multiple));
        assert_eq!(extract_var("$ABC"), Some(Capture("ABC".into(), true)));
        assert_eq!(extract_var("$$ABC"), Some(Capture("ABC".into(), false)));
        assert_eq!(extract_var("$MATCH1"), Some(Capture("MATCH1".into(), true)));
        assert_eq!(extract_var("$$$ABC"), Some(MultiCapture("ABC".into())));
        assert_eq!(extract_var("$_"), Some(Dropped(true)));
        assert_eq!(extract_var("$_123"), Some(Dropped(true)));
        assert_eq!(extract_var("$$_"), Some(Dropped(false)));
    }

    #[test]
    fn test_not_meta_var() {
        assert_eq!(extract_var("$123"), None);
        assert_eq!(extract_var("$"), None);
        assert_eq!(extract_var("$$"), None);
        assert_eq!(extract_var("abc"), None);
        assert_eq!(extract_var("$abc"), None);
    }

    fn match_constraints(pattern: &str, node: &str) -> bool {
        let mut matchers = thread_utilities::RapidMap::default();
        matchers.insert(Arc::from("A"), Pattern::new(pattern, &Tsx));
        let mut env = MetaVarEnv::new();
        let root = Tsx.ast_grep(node);
        let node = root.root().child(0).unwrap().child(0).unwrap();
        env.insert("A", node);
        env.match_constraints(&matchers)
    }

    #[test]
    fn test_non_ascii_meta_var() {
        let extract = |s| extract_meta_var(s, 'µ');
        use MetaVariable::*;
        assert_eq!(extract("µµµ"), Some(Multiple));
        assert_eq!(extract("µABC"), Some(Capture("ABC".into(), true)));
        assert_eq!(extract("µµABC"), Some(Capture("ABC".into(), false)));
        assert_eq!(extract("µµµABC"), Some(MultiCapture("ABC".into())));
        assert_eq!(extract("µ_"), Some(Dropped(true)));
        assert_eq!(extract("abc"), None);
        assert_eq!(extract("µabc"), None);
    }

    #[test]
    fn test_match_constraints() {
        assert!(match_constraints("a + b", "a + b"));
    }

    #[test]
    fn test_match_not_constraints() {
        assert!(!match_constraints("a - b", "a + b"));
    }

    #[test]
    fn test_multi_var_match() {
        let grep = Tsx.ast_grep("if (true) { a += 1; b += 1 } else { a += 1; b += 1 }");
        let node = grep.root();
        let found = node.find("if (true) { $$$A } else { $$$A }");
        assert!(found.is_some());
        let grep = Tsx.ast_grep("if (true) { a += 1 } else { b += 1 }");
        let node = grep.root();
        let not_found = node.find("if (true) { $$$A } else { $$$A }");
        assert!(not_found.is_none());
    }

    #[test]
    fn test_multi_var_match_with_trailing() {
        let grep = Tsx.ast_grep("if (true) { a += 1; } else { a += 1; b += 1 }");
        let node = grep.root();
        let not_found = node.find("if (true) { $$$A } else { $$$A }");
        assert!(not_found.is_none());
        let grep = Tsx.ast_grep("if (true) { a += 1; b += 1; } else { a += 1 }");
        let node = grep.root();
        let not_found = node.find("if (true) { $$$A } else { $$$A }");
        assert!(not_found.is_none());
    }
}