Skip to main content

zincite_query/
inspection.rs

1//! Direct CST navigation and source projections for the fixed query pipeline.
2use std::ops::Range;
3
4use serde_json::{Value, json};
5use zincite_syntax::{NodeKind, SyntaxElement, SyntaxNode, TokenKind, is_utf8_boundary};
6
7use crate::language::{Stage, StageKind};
8use crate::{Input, ItemKind, QueryError, QueryResult, Work, identifier_identity};
9
10/// One original CST node, with its original-file byte coordinates.
11#[derive(Clone, Debug)]
12pub struct SelectedNode<'a> {
13    pub node: &'a SyntaxNode,
14    pub file: &'a str,
15    pub range: Range<usize>,
16}
17
18#[derive(Debug)]
19pub struct NodeSelection<'a> {
20    pub(crate) input: &'a Input,
21    pub(crate) nodes: Vec<SelectedNode<'a>>,
22}
23
24impl<'a> NodeSelection<'a> {
25    pub fn input(&self) -> &'a Input {
26        self.input
27    }
28    pub fn nodes(&self) -> &[SelectedNode<'a>] {
29        &self.nodes
30    }
31}
32
33/// A validated JSON emitter retaining the native result until rendering.
34#[derive(Debug)]
35pub struct JsonResult<'a> {
36    result: Box<QueryResult<'a>>,
37}
38
39impl<'a> JsonResult<'a> {
40    pub fn result(&self) -> &QueryResult<'a> {
41        &self.result
42    }
43    pub(crate) fn render(&self) -> Vec<u8> {
44        // Evaluation checks text conversion while the immutable input is borrowed.
45        let value = json_value(&self.result).expect("JSON text was validated during evaluation");
46        let mut bytes = serde_json::to_vec(&value).expect("inspection values serialize");
47        bytes.push(b'\n');
48        bytes
49    }
50}
51
52pub(crate) fn type_error(range: &Range<usize>, message: &str) -> QueryError {
53    QueryError::query(range.clone(), message)
54}
55
56pub(crate) fn require_nodes(
57    result: &QueryResult<'_>,
58    range: &Range<usize>,
59) -> Result<(), QueryError> {
60    if matches!(result, QueryResult::Selection(_) | QueryResult::Nodes(_)) {
61        Ok(())
62    } else {
63        Err(type_error(range, "stage requires selected nodes"))
64    }
65}
66
67pub(crate) fn into_nodes<'a>(
68    result: QueryResult<'a>,
69    range: &Range<usize>,
70    work: &mut Work,
71) -> Result<NodeSelection<'a>, QueryError> {
72    match result {
73        QueryResult::Nodes(nodes) => Ok(nodes),
74        QueryResult::Selection(selection) => {
75            work.produce(selection.items.len(), range)?;
76            Ok(NodeSelection {
77                input: selection.input,
78                nodes: selection
79                    .items
80                    .into_iter()
81                    .map(|item| SelectedNode {
82                        node: item.node,
83                        file: selection.input.file(),
84                        range: item.range,
85                    })
86                    .collect(),
87            })
88        }
89        _ => Err(type_error(range, "stage requires selected nodes")),
90    }
91}
92
93pub(crate) fn navigate<'a>(
94    result: QueryResult<'a>,
95    stage: &Stage,
96    work: &mut Work,
97    limit: usize,
98) -> Result<QueryResult<'a>, QueryError> {
99    let selection = into_nodes(result, &stage.range, work)?;
100    let input = selection.input;
101    if let StageKind::Range(range) = &stage.kind
102        && (range.end > input.bytes.len()
103            || !is_utf8_boundary(&input.bytes, range.start)
104            || !is_utf8_boundary(&input.bytes, range.end))
105    {
106        return Err(type_error(
107            &stage.range,
108            "range bounds must be valid original-input byte boundaries",
109        ));
110    }
111    let mut nodes = Vec::new();
112    for selected in selection.nodes {
113        let mut pending = vec![selected.node];
114        while let Some(node) = pending.pop() {
115            work.visit(&stage.range)?;
116            let original = input.original_range(node.range());
117            let retain = match &stage.kind {
118                StageKind::Expressions => is_expression(node.kind()),
119                StageKind::Children => !std::ptr::eq(node, selected.node),
120                StageKind::Subtree => true,
121                StageKind::Range(range) => {
122                    range.start < range.end
123                        && range.start <= original.start
124                        && original.end <= range.end
125                }
126                _ => unreachable!(),
127            };
128            if retain {
129                check_collection(nodes.len(), limit, &stage.range)?;
130                work.produce(1, &stage.range)?;
131                nodes.push(SelectedNode {
132                    node,
133                    file: input.file(),
134                    range: original,
135                });
136            }
137            let descend = match stage.kind {
138                StageKind::Expressions | StageKind::Subtree => true,
139                StageKind::Children => std::ptr::eq(node, selected.node),
140                _ => false,
141            };
142            if descend {
143                work.spend(node.children().len(), &stage.range)?;
144                for child in node.children().iter().rev() {
145                    if let SyntaxElement::Node(child) = child {
146                        check_collection(pending.len(), limit, &stage.range)?;
147                        pending.push(child);
148                    }
149                }
150            }
151        }
152    }
153    Ok(QueryResult::Nodes(NodeSelection { input, nodes }))
154}
155
156pub(crate) fn check_collection(
157    length: usize,
158    limit: usize,
159    range: &Range<usize>,
160) -> Result<(), QueryError> {
161    if length >= limit {
162        Err(type_error(range, "query collection limit exceeded"))
163    } else {
164        Ok(())
165    }
166}
167
168pub(crate) fn project<'a>(
169    result: QueryResult<'a>,
170    stage: &Stage,
171    work: &mut Work,
172    limit: usize,
173) -> Result<QueryResult<'a>, QueryError> {
174    let selection = into_nodes(result, &stage.range, work)?;
175    let mut values = Vec::new();
176    for selected in selection.nodes {
177        work.visit(&stage.range)?;
178        work.spend(selected.node.children().len(), &stage.range)?;
179        if matches!(stage.kind, StageKind::Text | StageKind::AnnotationNames) {
180            work.spend(selected.range.len(), &stage.range)?;
181        }
182        let names = match stage.kind {
183            StageKind::Names => direct_names(selection.input.syntax(), selected.node),
184            StageKind::CallNames => call_name(selection.input, selected.node)
185                .into_iter()
186                .collect(),
187            StageKind::AnnotationNames => annotation_name(selection.input, selected.node)
188                .into_iter()
189                .collect(),
190            StageKind::Text => vec![original_text(selection.input, &selected.range)?],
191            _ => unreachable!(),
192        };
193        for name in names {
194            check_collection(values.len(), limit, &stage.range)?;
195            if !matches!(stage.kind, StageKind::Text) {
196                work.spend(name.len(), &stage.range)?;
197            }
198            work.produce(1, &stage.range)?;
199            values.push(name.to_owned());
200        }
201    }
202    Ok(QueryResult::Strings(values))
203}
204
205pub(crate) fn direct_names<'a>(
206    input: &'a zincite_syntax::ParsedFile,
207    node: &SyntaxNode,
208) -> Vec<&'a str> {
209    node.children()
210        .iter()
211        .filter_map(|child| {
212            let SyntaxElement::Token(index) = child else {
213                return None;
214            };
215            let token = &input.tokens()[*index];
216            matches!(
217                token.kind,
218                TokenKind::Identifier | TokenKind::QuotedIdentifier
219            )
220            .then(|| identifier_identity(&input.source()[token.range.clone()]))
221        })
222        .collect()
223}
224
225pub(crate) fn call_name<'a>(input: &'a Input, node: &SyntaxNode) -> Option<&'a str> {
226    if !matches!(
227        node.kind(),
228        NodeKind::CallExpression | NodeKind::GeneratorCallExpression
229    ) {
230        return None;
231    }
232    direct_names(input.syntax(), node).first().copied().or_else(|| {
233        node.children().iter().any(|child| matches!(child, SyntaxElement::Token(index) if input.syntax().tokens()[*index].kind == TokenKind::Anonymous)).then_some("_")
234    })
235}
236
237fn annotation_name<'a>(input: &'a Input, node: &SyntaxNode) -> Option<&'a str> {
238    if node.kind() != NodeKind::Annotation {
239        return None;
240    }
241    let mut head = node.child_nodes().next()?;
242    while matches!(
243        head.kind(),
244        NodeKind::AnnotatedExpression
245            | NodeKind::ArrayAccessExpression
246            | NodeKind::FieldAccessExpression
247            | NodeKind::ParenthesizedExpression
248    ) {
249        head = head.child_nodes().next()?;
250    }
251    call_name(input, head).or_else(|| direct_names(input.syntax(), head).first().copied()).or_else(|| {
252        head.children().iter().any(|child| matches!(child, SyntaxElement::Token(index) if input.syntax().tokens()[*index].kind == TokenKind::Output)).then_some("output")
253    })
254}
255
256pub(crate) fn original_text<'a>(
257    input: &'a Input,
258    range: &Range<usize>,
259) -> Result<&'a str, QueryError> {
260    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"))
261}
262
263pub(crate) fn json<'a>(
264    result: QueryResult<'a>,
265    work: &mut Work,
266    range: &Range<usize>,
267) -> Result<QueryResult<'a>, QueryError> {
268    validate_json(&result, work, range)?;
269    Ok(QueryResult::Json(JsonResult {
270        result: Box::new(result),
271    }))
272}
273
274pub(crate) fn validate_json(
275    result: &QueryResult<'_>,
276    work: &mut Work,
277    range: &Range<usize>,
278) -> Result<(), QueryError> {
279    match result {
280        QueryResult::Selection(selection) => {
281            for item in &selection.items {
282                work.spend(item.range.len().max(1), range)?;
283            }
284        }
285        QueryResult::Nodes(selection) => {
286            for node in &selection.nodes {
287                work.spend(node.range.len().max(1), range)?;
288            }
289        }
290        QueryResult::Strings(values) => {
291            for value in values {
292                work.spend(value.len().max(1), range)?;
293            }
294        }
295        QueryResult::Tally(values) => {
296            for (value, _) in values {
297                work.spend(value.len().max(1), range)?;
298            }
299        }
300        QueryResult::Literals(selection) => {
301            crate::structured::validate_json(selection, work, range)?
302        }
303        QueryResult::Count(_) => {}
304        _ => return Err(type_error(range, "json requires an inspection result")),
305    }
306    match result {
307        QueryResult::Selection(selection) => {
308            for item in &selection.items {
309                original_text(selection.input, &item.range)?;
310                work.spend(selection.input.file().len(), range)?;
311            }
312        }
313        QueryResult::Nodes(selection) => {
314            for node in &selection.nodes {
315                original_text(selection.input, &node.range)?;
316                work.spend(selection.input.file().len(), range)?;
317            }
318        }
319        _ => {}
320    }
321    Ok(())
322}
323
324fn node_json(input: &Input, node: &SyntaxNode, range: &Range<usize>) -> Result<Value, QueryError> {
325    Ok(json!({
326        "kind": kind_name(node.kind()),
327        "names": direct_names(input.syntax(), node),
328        "call_names": call_name(input, node).into_iter().collect::<Vec<_>>(),
329        "annotation_names": annotation_name(input, node).into_iter().collect::<Vec<_>>(),
330        "file": input.file(),
331        "range": {"start": range.start, "end": range.end},
332        "text": original_text(input, range)?,
333    }))
334}
335
336pub(crate) fn json_value(result: &QueryResult<'_>) -> Result<Value, QueryError> {
337    match result {
338        QueryResult::Selection(selection) => selection
339            .items
340            .iter()
341            .map(|item| node_json(selection.input, item.node, &item.range))
342            .collect::<Result<Vec<_>, _>>()
343            .map(Value::Array),
344        QueryResult::Nodes(selection) => selection
345            .nodes
346            .iter()
347            .map(|node| node_json(selection.input, node.node, &node.range))
348            .collect::<Result<Vec<_>, _>>()
349            .map(Value::Array),
350        QueryResult::Literals(selection) => crate::structured::json_value(selection),
351        QueryResult::Strings(values) => Ok(json!(values)),
352        QueryResult::Count(count) => Ok(json!(count)),
353        QueryResult::Tally(values) => Ok(tally_json(values)),
354        _ => unreachable!("only validated inspection results reach the JSON renderer"),
355    }
356}
357
358pub(crate) fn kind_name(kind: NodeKind) -> String {
359    let written = format!("{kind:?}");
360    let mut name = String::new();
361    for (index, character) in written.chars().enumerate() {
362        if character.is_ascii_uppercase() && index != 0 {
363            name.push('_');
364        }
365        name.push(character.to_ascii_lowercase());
366    }
367    name
368}
369
370pub(crate) fn known_kind(name: &str) -> bool {
371    name == "reference"
372        || ItemKind::parse(name).is_some()
373        || SYNTAX_KINDS.iter().any(|kind| kind_name(*kind) == name)
374}
375
376fn is_expression(kind: NodeKind) -> bool {
377    matches!(
378        kind,
379        NodeKind::Expression
380            | NodeKind::InterpolatedString
381            | NodeKind::UnaryExpression
382            | NodeKind::BinaryExpression
383            | NodeKind::ParenthesizedExpression
384            | NodeKind::CallExpression
385            | NodeKind::GeneratorCallExpression
386            | NodeKind::ConditionalExpression
387            | NodeKind::LetExpression
388            | NodeKind::SetComprehension
389            | NodeKind::ArrayComprehension
390            | NodeKind::IndexedArrayComprehension
391            | NodeKind::SetLiteral
392            | NodeKind::ArrayLiteral
393            | NodeKind::MatrixLiteral
394            | NodeKind::TupleLiteral
395            | NodeKind::RecordLiteral
396            | NodeKind::ArrayAccessExpression
397            | NodeKind::FieldAccessExpression
398            | NodeKind::AnnotatedExpression
399            | NodeKind::RangeExpression
400    )
401}
402
403const SYNTAX_KINDS: &[NodeKind] = &[
404    NodeKind::Root,
405    NodeKind::EnumDeclaration,
406    NodeKind::TypeAlias,
407    NodeKind::EnumDefinition,
408    NodeKind::EnumCases,
409    NodeKind::EnumCase,
410    NodeKind::EnumConstructor,
411    NodeKind::TypeInstConcatenation,
412    NodeKind::Declaration,
413    NodeKind::FunctionDeclaration,
414    NodeKind::PredicateDeclaration,
415    NodeKind::AnnotationCapture,
416    NodeKind::TestDeclaration,
417    NodeKind::AnnotationDeclaration,
418    NodeKind::ParameterList,
419    NodeKind::Parameter,
420    NodeKind::TypeInstVariable,
421    NodeKind::ScalarType,
422    NodeKind::TupleType,
423    NodeKind::RecordType,
424    NodeKind::RecordField,
425    NodeKind::TupleLiteral,
426    NodeKind::RecordLiteral,
427    NodeKind::RecordLiteralField,
428    NodeKind::FieldAccessExpression,
429    NodeKind::DomainType,
430    NodeKind::SetType,
431    NodeKind::SetCardinality,
432    NodeKind::ArrayType,
433    NodeKind::ArrayIndexBinding,
434    NodeKind::ListType,
435    NodeKind::Assignment,
436    NodeKind::Constraint,
437    NodeKind::Include,
438    NodeKind::Output,
439    NodeKind::Solve,
440    NodeKind::SolveMinimize,
441    NodeKind::SolveMaximize,
442    NodeKind::Expression,
443    NodeKind::InterpolatedString,
444    NodeKind::UnaryExpression,
445    NodeKind::BinaryExpression,
446    NodeKind::ParenthesizedExpression,
447    NodeKind::CallExpression,
448    NodeKind::GeneratorCallExpression,
449    NodeKind::ConditionalExpression,
450    NodeKind::ConditionalBranch,
451    NodeKind::ElseBranch,
452    NodeKind::LetExpression,
453    NodeKind::LetBlock,
454    NodeKind::SetComprehension,
455    NodeKind::ArrayComprehension,
456    NodeKind::IndexedArrayComprehension,
457    NodeKind::GeneratorList,
458    NodeKind::Generator,
459    NodeKind::WhereFilter,
460    NodeKind::SetLiteral,
461    NodeKind::ArrayLiteral,
462    NodeKind::IndexedArrayEntry,
463    NodeKind::IndexTuple,
464    NodeKind::MatrixLiteral,
465    NodeKind::MatrixColumnIndices,
466    NodeKind::MatrixRow,
467    NodeKind::ArrayAccessExpression,
468    NodeKind::NamedArgument,
469    NodeKind::AnnotatedExpression,
470    NodeKind::Annotation,
471    NodeKind::RangeExpression,
472    NodeKind::Error,
473];
474
475/// Admit owned inspection JSON before constructing its containers and strings.
476pub(crate) fn admit_json(
477    result: &QueryResult<'_>,
478    work: &mut Work,
479    range: &Range<usize>,
480) -> Result<(), QueryError> {
481    work.produce(1, range)?;
482    let mut node =
483        |input: &Input, node: &SyntaxNode, source: &Range<usize>| -> Result<(), QueryError> {
484            // Includes the temporary direct-name vector and all row metadata fields.
485            work.produce(25 + 2 * node.children().len(), range)?;
486            work.spend(160 + input.file().len() + source.len(), range)?;
487            for name in direct_names(input.syntax(), node)
488                .into_iter()
489                .chain(call_name(input, node))
490                .chain(annotation_name(input, node))
491            {
492                work.spend(name.len(), range)?;
493            }
494            Ok(())
495        };
496    match result {
497        QueryResult::Selection(s) => {
498            for item in &s.items {
499                node(s.input, item.node, &item.range)?;
500            }
501        }
502        QueryResult::Nodes(s) => {
503            for item in &s.nodes {
504                node(s.input, item.node, &item.range)?;
505            }
506        }
507        QueryResult::Strings(values) => {
508            work.produce(values.len(), range)?;
509            for value in values {
510                work.spend(value.len(), range)?;
511            }
512        }
513        QueryResult::Count(_) => {}
514        QueryResult::Tally(values) => {
515            work.produce(2 * values.len(), range)?;
516            for (key, _) in values {
517                work.spend(key.len(), range)?;
518            }
519        }
520        QueryResult::Literals(s) => crate::structured::admit_json(s, work, range)?,
521        _ => unreachable!("read-only inspection leaves"),
522    }
523    Ok(())
524}
525
526pub(crate) fn tally_json(values: &[(String, usize)]) -> Value {
527    Value::Object(
528        values
529            .iter()
530            .map(|(key, count)| (key.clone(), json!(count)))
531            .collect(),
532    )
533}