meta_language/

query.rs

use std::collections::HashMap;
use std::fmt;

mod parse;

use crate::link_network::{Link, LinkId, LinkNetwork, LinkType};
use parse::{QueryParser, tokenize};

/// Structural query over links.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct LinkQuery {
    link_type: Option<LinkType>,
    term: Option<String>,
    language: Option<String>,
    named: Option<bool>,
    pattern: Option<QueryPattern>,
    pattern_source: Option<String>,
    predicates: Vec<QueryPredicate>,
}

impl LinkQuery {
    /// Creates an empty query that matches every link.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Creates a query restricted to a link type.
    #[must_use]
    pub fn by_type(link_type: LinkType) -> Self {
        Self::new().with_link_type(link_type)
    }

    /// Parses a tree-sitter-query-like S-expression query.
    ///
    /// The structural engine binds captures and leaves predicate evaluation to
    /// a caller-provided [`QueryPredicateHost`].
    pub fn from_sexpression(source: &str) -> Result<Self, QueryParseError> {
        let tokens = tokenize(source)?;
        let mut parser = QueryParser::new(tokens);
        let (pattern, predicates) = parser.parse()?;
        Ok(Self {
            pattern: Some(pattern),
            pattern_source: Some(source.to_string()),
            predicates,
            ..Self::default()
        })
    }

    /// Restricts matches to a link type.
    #[must_use]
    pub const fn with_link_type(mut self, link_type: LinkType) -> Self {
        self.link_type = Some(link_type);
        self
    }

    /// Restricts matches to a term.
    #[must_use]
    pub fn with_term(mut self, term: impl Into<String>) -> Self {
        self.term = Some(term.into());
        self
    }

    /// Restricts matches to a language label.
    #[must_use]
    pub fn with_language(mut self, language: impl Into<String>) -> Self {
        self.language = Some(language.into());
        self
    }

    /// Restricts matches by the named flag.
    #[must_use]
    pub const fn with_named(mut self, named: bool) -> Self {
        self.named = Some(named);
        self
    }

    pub(crate) fn matches_in_network(
        &self,
        network: &LinkNetwork,
        link: &Link,
        predicate_host: &impl QueryPredicateHost,
    ) -> Vec<QueryMatch> {
        if !self.matches_metadata(link) {
            return Vec::new();
        }

        let captures = self.pattern.as_ref().map_or_else(
            || vec![QueryCaptures::default()],
            |pattern| pattern.match_root(network, link.id()),
        );

        captures
            .into_iter()
            .filter(|captures| {
                self.predicates
                    .iter()
                    .all(|predicate| predicate_host.evaluate(predicate, captures, network))
            })
            .map(|captures| QueryMatch::new(link.id(), captures))
            .collect()
    }

    fn matches_metadata(&self, link: &Link) -> bool {
        let metadata = link.metadata();
        self.link_type
            .is_none_or(|link_type| metadata.link_type() == Some(link_type))
            && self
                .term
                .as_deref()
                .is_none_or(|term| metadata.term() == Some(term))
            && self
                .language
                .as_deref()
                .is_none_or(|language| metadata.language() == Some(language))
            && self.named.is_none_or(|named| metadata.is_named() == named)
    }

    pub(crate) const fn link_type_filter(&self) -> Option<LinkType> {
        self.link_type
    }

    pub(crate) fn term_filter(&self) -> Option<&str> {
        self.term.as_deref()
    }

    pub(crate) fn language_filter(&self) -> Option<&str> {
        self.language.as_deref()
    }

    pub(crate) const fn named_filter(&self) -> Option<bool> {
        self.named
    }

    pub(crate) fn pattern_source(&self) -> Option<&str> {
        self.pattern_source.as_deref()
    }
}

/// A structural query match and its capture bindings.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct QueryMatch {
    link_id: LinkId,
    captures: QueryCaptures,
}

impl QueryMatch {
    const fn new(link_id: LinkId, captures: QueryCaptures) -> Self {
        Self { link_id, captures }
    }

    /// Link selected by the query root pattern.
    #[must_use]
    pub const fn link_id(&self) -> LinkId {
        self.link_id
    }

    /// Captures bound while matching this result.
    #[must_use]
    pub const fn captures(&self) -> &QueryCaptures {
        &self.captures
    }
}

/// Ordered capture bindings for a query match.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct QueryCaptures {
    values: Vec<QueryCapture>,
}

impl QueryCaptures {
    fn with_capture(mut self, name: &str, link_id: LinkId) -> Self {
        self.values.push(QueryCapture {
            name: name.to_string(),
            link_id,
        });
        self
    }

    /// Returns the first link bound to a capture name.
    #[must_use]
    pub fn first(&self, name: &str) -> Option<LinkId> {
        self.values
            .iter()
            .find(|capture| capture.name == name)
            .map(|capture| capture.link_id)
    }

    /// Iterates capture bindings in match order.
    pub fn iter(&self) -> impl Iterator<Item = &QueryCapture> {
        self.values.iter()
    }
}

/// One capture binding.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct QueryCapture {
    name: String,
    link_id: LinkId,
}

impl QueryCapture {
    /// Capture name without the leading `@`.
    #[must_use]
    pub fn name(&self) -> &str {
        &self.name
    }

    /// Link bound to this capture.
    #[must_use]
    pub const fn link_id(&self) -> LinkId {
        self.link_id
    }
}

/// Host hook for evaluating text, regex, semantic, or other predicates.
pub trait QueryPredicateHost {
    /// Returns whether a predicate accepts the current capture set.
    fn evaluate(
        &self,
        predicate: &QueryPredicate,
        captures: &QueryCaptures,
        network: &LinkNetwork,
    ) -> bool;
}

/// Predicate host that rejects every predicate: the host of a query run
/// without one, as the JavaScript `rejectPredicateHost`.
#[derive(Clone, Copy, Debug, Default)]
pub struct RejectPredicateHost;

impl QueryPredicateHost for RejectPredicateHost {
    fn evaluate(
        &self,
        _predicate: &QueryPredicate,
        _captures: &QueryCaptures,
        _network: &LinkNetwork,
    ) -> bool {
        false
    }
}

/// A query for the concept link spelled `term`, as the JavaScript
/// `queryByConceptTerm`.
#[must_use]
pub fn query_by_concept_term(term: &str) -> LinkQuery {
    LinkQuery::by_type(LinkType::Concept).with_term(term)
}

/// Structural child and field lookups for one network state, as the
/// JavaScript `QueryIndex`: built once, so matching does not rescan every
/// link for every node.
///
/// A tree stored bottom-up, every child referencing its parent first, and a
/// tree stored top-down, a `Syntax` link referencing its children tokens
/// included, hold the same children.
#[derive(Clone, Debug)]
pub struct QueryIndex<'a> {
    network: &'a LinkNetwork,
    children: HashMap<LinkId, Vec<LinkId>>,
    fields: HashMap<LinkId, Vec<(LinkId, LinkId)>>,
}

impl<'a> QueryIndex<'a> {
    /// Indexes the structural children and fields of `network`.
    #[must_use]
    pub fn new(network: &'a LinkNetwork) -> Self {
        let link_type = |id: &LinkId| {
            network
                .link(*id)
                .and_then(|link| link.metadata().link_type())
        };
        let top_down = network.links().any(|link| {
            link.metadata().link_type() == Some(LinkType::Syntax)
                && link
                    .references()
                    .iter()
                    .any(|reference| link_type(reference) == Some(LinkType::Token))
        });
        let mut children: HashMap<LinkId, Vec<LinkId>> = HashMap::new();
        let mut fields: HashMap<LinkId, Vec<(LinkId, LinkId)>> = HashMap::new();
        for link in network.links() {
            let references = link.references();
            let Some(first) = references.first() else {
                continue;
            };
            let kind = link.metadata().link_type();
            if kind == Some(LinkType::Field) {
                if let [parent, label, child] = references {
                    fields.entry(*parent).or_default().push((*label, *child));
                }
            } else if top_down {
                if kind == Some(LinkType::Syntax) {
                    for reference in references {
                        if !matches!(
                            link_type(reference),
                            Some(LinkType::Field | LinkType::Trivia)
                        ) {
                            children.entry(link.id()).or_default().push(*reference);
                        }
                    }
                }
            } else if kind != Some(LinkType::Trivia) {
                children.entry(*first).or_default().push(link.id());
            }
        }
        Self {
            network,
            children,
            fields,
        }
    }

    /// The link `id` names.
    #[must_use]
    pub fn link(&self, id: LinkId) -> Option<&'a Link> {
        self.network.link(id)
    }

    /// The structural children of `parent`, in network order, without field
    /// and trivia links.
    #[must_use]
    pub fn structural_children(&self, parent: LinkId) -> &[LinkId] {
        self.children.get(&parent).map_or(&[], Vec::as_slice)
    }

    /// The children `parent` holds under field `label`.
    #[must_use]
    pub fn field_targets(&self, parent: LinkId, label: &str) -> Vec<LinkId> {
        self.fields
            .get(&parent)
            .into_iter()
            .flatten()
            .filter(|(label_id, _)| {
                self.network
                    .link(*label_id)
                    .and_then(|link| link.metadata().term())
                    == Some(label)
            })
            .map(|(_, child)| *child)
            .collect()
    }
}

/// Predicate expression parsed from an S-expression query.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct QueryPredicate {
    name: String,
    arguments: Vec<QueryPredicateArgument>,
}

impl QueryPredicate {
    /// Predicate name without the leading `#`.
    #[must_use]
    pub fn name(&self) -> &str {
        &self.name
    }

    /// Predicate arguments in source order.
    #[must_use]
    pub fn arguments(&self) -> &[QueryPredicateArgument] {
        &self.arguments
    }
}

/// Predicate argument parsed from an S-expression query.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum QueryPredicateArgument {
    /// Capture reference such as `@name`.
    Capture(String),
    /// Host literal such as `"main"` or an unquoted atom.
    Literal(String),
}

impl QueryPredicateArgument {
    /// Returns the referenced capture name, if this argument is a capture.
    #[must_use]
    pub fn capture_name(&self) -> Option<&str> {
        match self {
            Self::Capture(name) => Some(name),
            Self::Literal(_) => None,
        }
    }

    /// Returns the literal value, if this argument is a literal.
    #[must_use]
    pub fn literal(&self) -> Option<&str> {
        match self {
            Self::Capture(_) => None,
            Self::Literal(value) => Some(value),
        }
    }
}

/// Error returned while parsing an S-expression query.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct QueryParseError {
    message: String,
}

impl QueryParseError {
    pub(crate) fn new(message: impl Into<String>) -> Self {
        Self {
            message: message.into(),
        }
    }
}

impl fmt::Display for QueryParseError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(&self.message)
    }
}

impl std::error::Error for QueryParseError {}

#[derive(Clone, Debug, PartialEq, Eq)]
struct QueryPattern {
    root: QueryNodePattern,
    capture: Option<String>,
}

impl QueryPattern {
    fn match_root(&self, network: &LinkNetwork, link_id: LinkId) -> Vec<QueryCaptures> {
        let captures = self
            .root
            .matches(network, link_id, QueryCaptures::default());
        captures
            .into_iter()
            .map(|captures| {
                if let Some(capture) = &self.capture {
                    captures.with_capture(capture, link_id)
                } else {
                    captures
                }
            })
            .collect()
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
struct QueryNodePattern {
    kind: QueryNodeKind,
    children: Vec<QueryChildPattern>,
}

impl QueryNodePattern {
    fn matches(
        &self,
        network: &LinkNetwork,
        link_id: LinkId,
        captures: QueryCaptures,
    ) -> Vec<QueryCaptures> {
        let Some(link) = network.link(link_id) else {
            return Vec::new();
        };
        if !self.kind.matches(link) || self.has_negated_field(network, link_id) {
            return Vec::new();
        }

        let structural_children = structural_children(network, link_id);
        let context = MatchContext {
            network,
            parent: link_id,
            structural_children: &structural_children,
        };
        match_child_patterns(&context, &self.children, 0, false, captures)
            .into_iter()
            .map(|(_child_index, captures)| captures)
            .collect()
    }

    fn has_negated_field(&self, network: &LinkNetwork, link_id: LinkId) -> bool {
        self.children.iter().any(|child| {
            if let QueryChildPattern::NegatedField(label) = child {
                has_field(network, link_id, label)
            } else {
                false
            }
        })
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
enum QueryNodeKind {
    Exact(String),
    Wildcard,
}

impl QueryNodeKind {
    fn matches(&self, link: &Link) -> bool {
        match self {
            Self::Exact(kind) => link.metadata().term() == Some(kind.as_str()),
            Self::Wildcard => true,
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
enum QueryChildPattern {
    Anchor,
    NegatedField(String),
    Pattern(QueryChildExpression),
}

#[derive(Clone, Debug, PartialEq, Eq)]
struct QueryChildExpression {
    field: Option<String>,
    expression: QueryExpression,
    capture: Option<String>,
    quantifier: QueryQuantifier,
}

#[derive(Clone, Debug, PartialEq, Eq)]
enum QueryExpression {
    Node(QueryNodePattern),
    Alternation(Vec<QueryNodePattern>),
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum QueryQuantifier {
    One,
    ZeroOrOne,
    ZeroOrMore,
    OneOrMore,
}

struct MatchContext<'a> {
    network: &'a LinkNetwork,
    parent: LinkId,
    structural_children: &'a [LinkId],
}

fn match_child_patterns(
    context: &MatchContext<'_>,
    patterns: &[QueryChildPattern],
    child_index: usize,
    anchored: bool,
    captures: QueryCaptures,
) -> Vec<(usize, QueryCaptures)> {
    let Some((pattern, remaining)) = patterns.split_first() else {
        return vec![(child_index, captures)];
    };

    match pattern {
        QueryChildPattern::Anchor => {
            if remaining.is_empty() || remaining.iter().all(is_negated_field_pattern) {
                if child_index == context.structural_children.len() {
                    match_child_patterns(context, remaining, child_index, true, captures)
                } else {
                    Vec::new()
                }
            } else {
                match_child_patterns(context, remaining, child_index, true, captures)
            }
        }
        QueryChildPattern::NegatedField(_) => {
            match_child_patterns(context, remaining, child_index, anchored, captures)
        }
        QueryChildPattern::Pattern(expression) => match_quantified_expression(
            context,
            expression,
            remaining,
            child_index,
            anchored,
            &captures,
        ),
    }
}

const fn is_negated_field_pattern(pattern: &QueryChildPattern) -> bool {
    matches!(pattern, QueryChildPattern::NegatedField(_))
}

fn match_quantified_expression(
    context: &MatchContext<'_>,
    expression: &QueryChildExpression,
    remaining: &[QueryChildPattern],
    child_index: usize,
    anchored: bool,
    captures: &QueryCaptures,
) -> Vec<(usize, QueryCaptures)> {
    match expression.quantifier {
        QueryQuantifier::One => match_one_then_continue(
            context,
            expression,
            remaining,
            child_index,
            anchored,
            captures,
        ),
        QueryQuantifier::ZeroOrOne => {
            let mut results =
                match_child_patterns(context, remaining, child_index, anchored, captures.clone());
            results.extend(match_one_then_continue(
                context,
                expression,
                remaining,
                child_index,
                anchored,
                captures,
            ));
            results
        }
        QueryQuantifier::ZeroOrMore => {
            let mut results =
                match_child_patterns(context, remaining, child_index, anchored, captures.clone());
            results.extend(match_repeated_expression(
                context,
                expression,
                remaining,
                child_index,
                anchored,
                captures,
            ));
            results
        }
        QueryQuantifier::OneOrMore => match_repeated_expression(
            context,
            expression,
            remaining,
            child_index,
            anchored,
            captures,
        ),
    }
}

fn match_repeated_expression(
    context: &MatchContext<'_>,
    expression: &QueryChildExpression,
    remaining: &[QueryChildPattern],
    child_index: usize,
    anchored: bool,
    captures: &QueryCaptures,
) -> Vec<(usize, QueryCaptures)> {
    let mut results = Vec::new();
    for (next_index, next_captures) in
        match_expression_at_positions(context, expression, child_index, anchored, captures)
    {
        results.extend(match_child_patterns(
            context,
            remaining,
            next_index,
            false,
            next_captures.clone(),
        ));
        results.extend(match_repeated_expression(
            context,
            expression,
            remaining,
            next_index,
            true,
            &next_captures,
        ));
    }
    results
}

fn match_one_then_continue(
    context: &MatchContext<'_>,
    expression: &QueryChildExpression,
    remaining: &[QueryChildPattern],
    child_index: usize,
    anchored: bool,
    captures: &QueryCaptures,
) -> Vec<(usize, QueryCaptures)> {
    match_expression_at_positions(context, expression, child_index, anchored, captures)
        .into_iter()
        .flat_map(|(next_index, captures)| {
            match_child_patterns(context, remaining, next_index, false, captures)
        })
        .collect()
}

fn match_expression_at_positions(
    context: &MatchContext<'_>,
    expression: &QueryChildExpression,
    child_index: usize,
    anchored: bool,
    captures: &QueryCaptures,
) -> Vec<(usize, QueryCaptures)> {
    let positions: Box<dyn Iterator<Item = usize>> = if anchored {
        Box::new(std::iter::once(child_index))
    } else {
        Box::new(child_index..context.structural_children.len())
    };

    positions
        .filter(|position| *position < context.structural_children.len())
        .flat_map(|position| {
            let child = context.structural_children[position];
            match_expression(
                context.network,
                context.parent,
                child,
                expression,
                captures.clone(),
            )
            .into_iter()
            .map(move |captures| (position + 1, captures))
        })
        .collect()
}

fn match_expression(
    network: &LinkNetwork,
    parent: LinkId,
    child: LinkId,
    expression: &QueryChildExpression,
    captures: QueryCaptures,
) -> Vec<QueryCaptures> {
    if let Some(field) = &expression.field {
        let field_targets = field_targets(network, parent, field);
        if !field_targets.contains(&child) {
            return Vec::new();
        }
    }

    let matches = match &expression.expression {
        QueryExpression::Node(node) => node.matches(network, child, captures),
        QueryExpression::Alternation(alternatives) => alternatives
            .iter()
            .flat_map(|alternative| alternative.matches(network, child, captures.clone()))
            .collect(),
    };

    matches
        .into_iter()
        .map(|captures| {
            if let Some(capture) = &expression.capture {
                captures.with_capture(capture, child)
            } else {
                captures
            }
        })
        .collect()
}

fn structural_children(network: &LinkNetwork, parent: LinkId) -> Vec<LinkId> {
    network
        .links()
        .filter(|link| link.references().first().copied() == Some(parent))
        .filter(|link| {
            !matches!(
                link.metadata().link_type(),
                Some(LinkType::Field | LinkType::Trivia)
            )
        })
        .map(Link::id)
        .collect()
}

fn has_field(network: &LinkNetwork, parent: LinkId, label: &str) -> bool {
    !field_targets(network, parent, label).is_empty()
}

fn field_targets(network: &LinkNetwork, parent: LinkId, label: &str) -> Vec<LinkId> {
    network
        .links()
        .filter(|link| link.metadata().link_type() == Some(LinkType::Field))
        .filter_map(|link| {
            let references = link.references();
            let [field_parent, field_label, field_child] = references else {
                return None;
            };
            if *field_parent != parent {
                return None;
            }
            let label_link = network.link(*field_label)?;
            (label_link.metadata().term() == Some(label)).then_some(*field_child)
        })
        .collect()
}