zincite-query 0.4.0

//! Direct CST navigation and source projections for the fixed query pipeline.
use std::ops::Range;

use serde_json::{Value, json};
use zincite_syntax::{NodeKind, SyntaxElement, SyntaxNode, TokenKind, is_utf8_boundary};

use crate::language::{Stage, StageKind};
use crate::{Input, ItemKind, QueryError, QueryResult, Work, identifier_identity};

/// One original CST node, with its original-file byte coordinates.
#[derive(Clone, Debug)]
pub struct SelectedNode<'a> {
    pub node: &'a SyntaxNode,
    pub file: &'a str,
    pub range: Range<usize>,
}

#[derive(Debug)]
pub struct NodeSelection<'a> {
    pub(crate) input: &'a Input,
    pub(crate) nodes: Vec<SelectedNode<'a>>,
}

impl<'a> NodeSelection<'a> {
    pub fn input(&self) -> &'a Input {
        self.input
    }
    pub fn nodes(&self) -> &[SelectedNode<'a>] {
        &self.nodes
    }
}

/// A validated JSON emitter retaining the native result until rendering.
#[derive(Debug)]
pub struct JsonResult<'a> {
    result: Box<QueryResult<'a>>,
}

impl<'a> JsonResult<'a> {
    pub fn result(&self) -> &QueryResult<'a> {
        &self.result
    }
    pub(crate) fn render(&self) -> Vec<u8> {
        // Evaluation checks text conversion while the immutable input is borrowed.
        let value = json_value(&self.result).expect("JSON text was validated during evaluation");
        let mut bytes = serde_json::to_vec(&value).expect("inspection values serialize");
        bytes.push(b'\n');
        bytes
    }
}

pub(crate) fn type_error(range: &Range<usize>, message: &str) -> QueryError {
    QueryError::query(range.clone(), message)
}

pub(crate) fn require_nodes(
    result: &QueryResult<'_>,
    range: &Range<usize>,
) -> Result<(), QueryError> {
    if matches!(result, QueryResult::Selection(_) | QueryResult::Nodes(_)) {
        Ok(())
    } else {
        Err(type_error(range, "stage requires selected nodes"))
    }
}

pub(crate) fn into_nodes<'a>(
    result: QueryResult<'a>,
    range: &Range<usize>,
    work: &mut Work,
) -> Result<NodeSelection<'a>, QueryError> {
    match result {
        QueryResult::Nodes(nodes) => Ok(nodes),
        QueryResult::Selection(selection) => {
            work.produce(selection.items.len(), range)?;
            Ok(NodeSelection {
                input: selection.input,
                nodes: selection
                    .items
                    .into_iter()
                    .map(|item| SelectedNode {
                        node: item.node,
                        file: selection.input.file(),
                        range: item.range,
                    })
                    .collect(),
            })
        }
        _ => Err(type_error(range, "stage requires selected nodes")),
    }
}

pub(crate) fn navigate<'a>(
    result: QueryResult<'a>,
    stage: &Stage,
    work: &mut Work,
    limit: usize,
) -> Result<QueryResult<'a>, QueryError> {
    let selection = into_nodes(result, &stage.range, work)?;
    let input = selection.input;
    if let StageKind::Range(range) = &stage.kind
        && (range.end > input.bytes.len()
            || !is_utf8_boundary(&input.bytes, range.start)
            || !is_utf8_boundary(&input.bytes, range.end))
    {
        return Err(type_error(
            &stage.range,
            "range bounds must be valid original-input byte boundaries",
        ));
    }
    let mut nodes = Vec::new();
    for selected in selection.nodes {
        let mut pending = vec![selected.node];
        while let Some(node) = pending.pop() {
            work.visit(&stage.range)?;
            let original = input.original_range(node.range());
            let retain = match &stage.kind {
                StageKind::Expressions => is_expression(node.kind()),
                StageKind::Children => !std::ptr::eq(node, selected.node),
                StageKind::Subtree => true,
                StageKind::Range(range) => {
                    range.start < range.end
                        && range.start <= original.start
                        && original.end <= range.end
                }
                _ => unreachable!(),
            };
            if retain {
                check_collection(nodes.len(), limit, &stage.range)?;
                work.produce(1, &stage.range)?;
                nodes.push(SelectedNode {
                    node,
                    file: input.file(),
                    range: original,
                });
            }
            let descend = match stage.kind {
                StageKind::Expressions | StageKind::Subtree => true,
                StageKind::Children => std::ptr::eq(node, selected.node),
                _ => false,
            };
            if descend {
                work.spend(node.children().len(), &stage.range)?;
                for child in node.children().iter().rev() {
                    if let SyntaxElement::Node(child) = child {
                        check_collection(pending.len(), limit, &stage.range)?;
                        pending.push(child);
                    }
                }
            }
        }
    }
    Ok(QueryResult::Nodes(NodeSelection { input, nodes }))
}

pub(crate) fn check_collection(
    length: usize,
    limit: usize,
    range: &Range<usize>,
) -> Result<(), QueryError> {
    if length >= limit {
        Err(type_error(range, "query collection limit exceeded"))
    } else {
        Ok(())
    }
}

pub(crate) fn project<'a>(
    result: QueryResult<'a>,
    stage: &Stage,
    work: &mut Work,
    limit: usize,
) -> Result<QueryResult<'a>, QueryError> {
    let selection = into_nodes(result, &stage.range, work)?;
    let mut values = Vec::new();
    for selected in selection.nodes {
        work.visit(&stage.range)?;
        work.spend(selected.node.children().len(), &stage.range)?;
        if matches!(stage.kind, StageKind::Text | StageKind::AnnotationNames) {
            work.spend(selected.range.len(), &stage.range)?;
        }
        let names = match stage.kind {
            StageKind::Names => direct_names(selection.input.syntax(), selected.node),
            StageKind::CallNames => call_name(selection.input, selected.node)
                .into_iter()
                .collect(),
            StageKind::AnnotationNames => annotation_name(selection.input, selected.node)
                .into_iter()
                .collect(),
            StageKind::Text => vec![original_text(selection.input, &selected.range)?],
            _ => unreachable!(),
        };
        for name in names {
            check_collection(values.len(), limit, &stage.range)?;
            if !matches!(stage.kind, StageKind::Text) {
                work.spend(name.len(), &stage.range)?;
            }
            work.produce(1, &stage.range)?;
            values.push(name.to_owned());
        }
    }
    Ok(QueryResult::Strings(values))
}

pub(crate) fn direct_names<'a>(
    input: &'a zincite_syntax::ParsedFile,
    node: &SyntaxNode,
) -> Vec<&'a str> {
    node.children()
        .iter()
        .filter_map(|child| {
            let SyntaxElement::Token(index) = child else {
                return None;
            };
            let token = &input.tokens()[*index];
            matches!(
                token.kind,
                TokenKind::Identifier | TokenKind::QuotedIdentifier
            )
            .then(|| identifier_identity(&input.source()[token.range.clone()]))
        })
        .collect()
}

pub(crate) fn call_name<'a>(input: &'a Input, node: &SyntaxNode) -> Option<&'a str> {
    if !matches!(
        node.kind(),
        NodeKind::CallExpression | NodeKind::GeneratorCallExpression
    ) {
        return None;
    }
    direct_names(input.syntax(), node).first().copied().or_else(|| {
        node.children().iter().any(|child| matches!(child, SyntaxElement::Token(index) if input.syntax().tokens()[*index].kind == TokenKind::Anonymous)).then_some("_")
    })
}

fn annotation_name<'a>(input: &'a Input, node: &SyntaxNode) -> Option<&'a str> {
    if node.kind() != NodeKind::Annotation {
        return None;
    }
    let mut head = node.child_nodes().next()?;
    while matches!(
        head.kind(),
        NodeKind::AnnotatedExpression
            | NodeKind::ArrayAccessExpression
            | NodeKind::FieldAccessExpression
            | NodeKind::ParenthesizedExpression
    ) {
        head = head.child_nodes().next()?;
    }
    call_name(input, head).or_else(|| direct_names(input.syntax(), head).first().copied()).or_else(|| {
        head.children().iter().any(|child| matches!(child, SyntaxElement::Token(index) if input.syntax().tokens()[*index].kind == TokenKind::Output)).then_some("output")
    })
}

pub(crate) fn original_text<'a>(
    input: &'a Input,
    range: &Range<usize>,
) -> Result<&'a str, QueryError> {
    std::str::from_utf8(&input.bytes[range.clone()]).map_err(|_| QueryError::input(range.clone(), "text and JSON inspection require UTF-8 source text; source emission preserves opaque comment bytes"))
}

pub(crate) fn json<'a>(
    result: QueryResult<'a>,
    work: &mut Work,
    range: &Range<usize>,
) -> Result<QueryResult<'a>, QueryError> {
    validate_json(&result, work, range)?;
    Ok(QueryResult::Json(JsonResult {
        result: Box::new(result),
    }))
}

pub(crate) fn validate_json(
    result: &QueryResult<'_>,
    work: &mut Work,
    range: &Range<usize>,
) -> Result<(), QueryError> {
    match result {
        QueryResult::Selection(selection) => {
            for item in &selection.items {
                work.spend(item.range.len().max(1), range)?;
            }
        }
        QueryResult::Nodes(selection) => {
            for node in &selection.nodes {
                work.spend(node.range.len().max(1), range)?;
            }
        }
        QueryResult::Strings(values) => {
            for value in values {
                work.spend(value.len().max(1), range)?;
            }
        }
        QueryResult::Tally(values) => {
            for (value, _) in values {
                work.spend(value.len().max(1), range)?;
            }
        }
        QueryResult::Literals(selection) => {
            crate::structured::validate_json(selection, work, range)?
        }
        QueryResult::Count(_) => {}
        _ => return Err(type_error(range, "json requires an inspection result")),
    }
    match result {
        QueryResult::Selection(selection) => {
            for item in &selection.items {
                original_text(selection.input, &item.range)?;
                work.spend(selection.input.file().len(), range)?;
            }
        }
        QueryResult::Nodes(selection) => {
            for node in &selection.nodes {
                original_text(selection.input, &node.range)?;
                work.spend(selection.input.file().len(), range)?;
            }
        }
        _ => {}
    }
    Ok(())
}

fn node_json(input: &Input, node: &SyntaxNode, range: &Range<usize>) -> Result<Value, QueryError> {
    Ok(json!({
        "kind": kind_name(node.kind()),
        "names": direct_names(input.syntax(), node),
        "call_names": call_name(input, node).into_iter().collect::<Vec<_>>(),
        "annotation_names": annotation_name(input, node).into_iter().collect::<Vec<_>>(),
        "file": input.file(),
        "range": {"start": range.start, "end": range.end},
        "text": original_text(input, range)?,
    }))
}

pub(crate) fn json_value(result: &QueryResult<'_>) -> Result<Value, QueryError> {
    match result {
        QueryResult::Selection(selection) => selection
            .items
            .iter()
            .map(|item| node_json(selection.input, item.node, &item.range))
            .collect::<Result<Vec<_>, _>>()
            .map(Value::Array),
        QueryResult::Nodes(selection) => selection
            .nodes
            .iter()
            .map(|node| node_json(selection.input, node.node, &node.range))
            .collect::<Result<Vec<_>, _>>()
            .map(Value::Array),
        QueryResult::Literals(selection) => crate::structured::json_value(selection),
        QueryResult::Strings(values) => Ok(json!(values)),
        QueryResult::Count(count) => Ok(json!(count)),
        QueryResult::Tally(values) => Ok(tally_json(values)),
        _ => unreachable!("only validated inspection results reach the JSON renderer"),
    }
}

pub(crate) fn kind_name(kind: NodeKind) -> String {
    let written = format!("{kind:?}");
    let mut name = String::new();
    for (index, character) in written.chars().enumerate() {
        if character.is_ascii_uppercase() && index != 0 {
            name.push('_');
        }
        name.push(character.to_ascii_lowercase());
    }
    name
}

pub(crate) fn known_kind(name: &str) -> bool {
    name == "reference"
        || ItemKind::parse(name).is_some()
        || SYNTAX_KINDS.iter().any(|kind| kind_name(*kind) == name)
}

fn is_expression(kind: NodeKind) -> bool {
    matches!(
        kind,
        NodeKind::Expression
            | NodeKind::InterpolatedString
            | NodeKind::UnaryExpression
            | NodeKind::BinaryExpression
            | NodeKind::ParenthesizedExpression
            | NodeKind::CallExpression
            | NodeKind::GeneratorCallExpression
            | NodeKind::ConditionalExpression
            | NodeKind::LetExpression
            | NodeKind::SetComprehension
            | NodeKind::ArrayComprehension
            | NodeKind::IndexedArrayComprehension
            | NodeKind::SetLiteral
            | NodeKind::ArrayLiteral
            | NodeKind::MatrixLiteral
            | NodeKind::TupleLiteral
            | NodeKind::RecordLiteral
            | NodeKind::ArrayAccessExpression
            | NodeKind::FieldAccessExpression
            | NodeKind::AnnotatedExpression
            | NodeKind::RangeExpression
    )
}

const SYNTAX_KINDS: &[NodeKind] = &[
    NodeKind::Root,
    NodeKind::EnumDeclaration,
    NodeKind::TypeAlias,
    NodeKind::EnumDefinition,
    NodeKind::EnumCases,
    NodeKind::EnumCase,
    NodeKind::EnumConstructor,
    NodeKind::TypeInstConcatenation,
    NodeKind::Declaration,
    NodeKind::FunctionDeclaration,
    NodeKind::PredicateDeclaration,
    NodeKind::AnnotationCapture,
    NodeKind::TestDeclaration,
    NodeKind::AnnotationDeclaration,
    NodeKind::ParameterList,
    NodeKind::Parameter,
    NodeKind::TypeInstVariable,
    NodeKind::ScalarType,
    NodeKind::TupleType,
    NodeKind::RecordType,
    NodeKind::RecordField,
    NodeKind::TupleLiteral,
    NodeKind::RecordLiteral,
    NodeKind::RecordLiteralField,
    NodeKind::FieldAccessExpression,
    NodeKind::DomainType,
    NodeKind::SetType,
    NodeKind::SetCardinality,
    NodeKind::ArrayType,
    NodeKind::ArrayIndexBinding,
    NodeKind::ListType,
    NodeKind::Assignment,
    NodeKind::Constraint,
    NodeKind::Include,
    NodeKind::Output,
    NodeKind::Solve,
    NodeKind::SolveMinimize,
    NodeKind::SolveMaximize,
    NodeKind::Expression,
    NodeKind::InterpolatedString,
    NodeKind::UnaryExpression,
    NodeKind::BinaryExpression,
    NodeKind::ParenthesizedExpression,
    NodeKind::CallExpression,
    NodeKind::GeneratorCallExpression,
    NodeKind::ConditionalExpression,
    NodeKind::ConditionalBranch,
    NodeKind::ElseBranch,
    NodeKind::LetExpression,
    NodeKind::LetBlock,
    NodeKind::SetComprehension,
    NodeKind::ArrayComprehension,
    NodeKind::IndexedArrayComprehension,
    NodeKind::GeneratorList,
    NodeKind::Generator,
    NodeKind::WhereFilter,
    NodeKind::SetLiteral,
    NodeKind::ArrayLiteral,
    NodeKind::IndexedArrayEntry,
    NodeKind::IndexTuple,
    NodeKind::MatrixLiteral,
    NodeKind::MatrixColumnIndices,
    NodeKind::MatrixRow,
    NodeKind::ArrayAccessExpression,
    NodeKind::NamedArgument,
    NodeKind::AnnotatedExpression,
    NodeKind::Annotation,
    NodeKind::RangeExpression,
    NodeKind::Error,
];

/// Admit owned inspection JSON before constructing its containers and strings.
pub(crate) fn admit_json(
    result: &QueryResult<'_>,
    work: &mut Work,
    range: &Range<usize>,
) -> Result<(), QueryError> {
    work.produce(1, range)?;
    let mut node =
        |input: &Input, node: &SyntaxNode, source: &Range<usize>| -> Result<(), QueryError> {
            // Includes the temporary direct-name vector and all row metadata fields.
            work.produce(25 + 2 * node.children().len(), range)?;
            work.spend(160 + input.file().len() + source.len(), range)?;
            for name in direct_names(input.syntax(), node)
                .into_iter()
                .chain(call_name(input, node))
                .chain(annotation_name(input, node))
            {
                work.spend(name.len(), range)?;
            }
            Ok(())
        };
    match result {
        QueryResult::Selection(s) => {
            for item in &s.items {
                node(s.input, item.node, &item.range)?;
            }
        }
        QueryResult::Nodes(s) => {
            for item in &s.nodes {
                node(s.input, item.node, &item.range)?;
            }
        }
        QueryResult::Strings(values) => {
            work.produce(values.len(), range)?;
            for value in values {
                work.spend(value.len(), range)?;
            }
        }
        QueryResult::Count(_) => {}
        QueryResult::Tally(values) => {
            work.produce(2 * values.len(), range)?;
            for (key, _) in values {
                work.spend(key.len(), range)?;
            }
        }
        QueryResult::Literals(s) => crate::structured::admit_json(s, work, range)?,
        _ => unreachable!("read-only inspection leaves"),
    }
    Ok(())
}

pub(crate) fn tally_json(values: &[(String, usize)]) -> Value {
    Value::Object(
        values
            .iter()
            .map(|(key, count)| (key.clone(), json!(count)))
            .collect(),
    )
}