Merge pull request #133 from tree-sitter/easier-match-ops

CHANGELOG.md

···
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## v0.10.1 -- 2023-05-15
+
+### Library
+
+#### Added
+
+- The `ast::File::try_visit_matches` method can be used to execute the stanza query matching without executing the stanza bodies.
+
 ## v0.10.0 -- 2023-05-10
 
 ### DSL

Cargo.toml

···
 [package]
 name = "tree-sitter-graph"
-version = "0.10.0"
+version = "0.10.1"
 description = "Construct graphs from parsed source code"
 homepage = "https://github.com/tree-sitter/tree-sitter-graph/"
 repository = "https://github.com/tree-sitter/tree-sitter-graph/"

src/ast.rs

···
 use tree_sitter::Language;
 use tree_sitter::Query;
 
+use crate::parser::Range;
 use crate::Identifier;
 use crate::Location;
 
···
     pub full_match_stanza_capture_index: usize,
     /// Capture index of the full match in the file query
     pub full_match_file_capture_index: usize,
-    pub location: Location,
+    pub range: Range,
 }
 
 /// A statement that can appear in a graph DSL stanza

src/checker.rs

···
         if !unused_captures.is_empty() {
             return Err(CheckError::UnusedCaptures(
                 unused_captures.join(" "),
-                self.location,
+                self.range.start,
             ));
         }
 

src/execution.rs

···
 
 use thiserror::Error;
 use tree_sitter::CaptureQuantifier;
+use tree_sitter::Node;
 use tree_sitter::QueryMatch;
 use tree_sitter::Tree;
 
 use crate::ast::File;
+use crate::ast::Stanza;
 use crate::execution::error::ExecutionError;
 use crate::functions::Functions;
 use crate::graph::Graph;
 use crate::graph::Value;
 use crate::variables::Globals;
 use crate::Identifier;
+use crate::Location;
 
 pub(crate) mod error;
 mod lazy;
···
 
         Ok(())
     }
+
+    pub fn try_visit_matches<'tree, E, F>(
+        &self,
+        tree: &'tree Tree,
+        source: &'tree str,
+        lazy: bool,
+        mut visit: F,
+    ) -> Result<(), E>
+    where
+        F: FnMut(Match<'_, 'tree>) -> Result<(), E>,
+    {
+        if lazy {
+            let file_query = self.query.as_ref().expect("missing file query");
+            self.try_visit_matches_lazy(tree, source, |stanza, mat| {
+                let named_captures = stanza
+                    .query
+                    .capture_names()
+                    .iter()
+                    .map(|name| {
+                        let index = file_query
+                            .capture_index_for_name(name)
+                            .expect("missing index for capture");
+                        let quantifier =
+                            file_query.capture_quantifiers(mat.pattern_index)[index as usize];
+                        (name, quantifier, index)
+                    })
+                    .filter(|c| c.2 != stanza.full_match_file_capture_index as u32)
+                    .collect();
+                visit(Match {
+                    mat,
+                    full_capture_index: stanza.full_match_file_capture_index as u32,
+                    named_captures,
+                    query_location: stanza.range.start,
+                })
+            })
+        } else {
+            self.try_visit_matches_strict(tree, source, |stanza, mat| {
+                let named_captures = stanza
+                    .query
+                    .capture_names()
+                    .iter()
+                    .map(|name| {
+                        let index = stanza
+                            .query
+                            .capture_index_for_name(name)
+                            .expect("missing index for capture");
+                        let quantifier = stanza.query.capture_quantifiers(0)[index as usize];
+                        (name, quantifier, index)
+                    })
+                    .filter(|c| c.2 != stanza.full_match_stanza_capture_index as u32)
+                    .collect();
+                visit(Match {
+                    mat,
+                    full_capture_index: stanza.full_match_stanza_capture_index as u32,
+                    named_captures,
+                    query_location: stanza.range.start,
+                })
+            })
+        }
+    }
+}
+
+impl Stanza {
+    pub fn try_visit_matches<'tree, E, F>(
+        &self,
+        tree: &'tree Tree,
+        source: &'tree str,
+        mut visit: F,
+    ) -> Result<(), E>
+    where
+        F: FnMut(Match<'_, 'tree>) -> Result<(), E>,
+    {
+        self.try_visit_matches_strict(tree, source, |mat| {
+            let named_captures = self
+                .query
+                .capture_names()
+                .iter()
+                .map(|name| {
+                    let index = self
+                        .query
+                        .capture_index_for_name(name)
+                        .expect("missing index for capture");
+                    let quantifier = self.query.capture_quantifiers(0)[index as usize];
+                    (name, quantifier, index)
+                })
+                .filter(|c| c.2 != self.full_match_stanza_capture_index as u32)
+                .collect();
+            visit(Match {
+                mat,
+                full_capture_index: self.full_match_stanza_capture_index as u32,
+                named_captures,
+                query_location: self.range.start,
+            })
+        })
+    }
+}
+
+pub struct Match<'a, 'tree> {
+    mat: QueryMatch<'a, 'tree>,
+    full_capture_index: u32,
+    named_captures: Vec<(&'a String, CaptureQuantifier, u32)>,
+    query_location: Location,
+}
+
+impl<'a, 'tree> Match<'a, 'tree> {
+    /// Return the top-level matched node.
+    pub fn full_capture(&self) -> Node<'tree> {
+        self.mat
+            .nodes_for_capture_index(self.full_capture_index)
+            .next()
+            .expect("missing full capture")
+    }
+
+    /// Return the matched nodes for a named capture.
+    pub fn named_captures<'s: 'a + 'tree>(
+        &'s self,
+    ) -> impl Iterator<
+        Item = (
+            &String,
+            CaptureQuantifier,
+            impl Iterator<Item = Node<'tree>> + 's,
+        ),
+    > {
+        self.named_captures
+            .iter()
+            .map(move |c| (c.0, c.1, self.mat.nodes_for_capture_index(c.2)))
+    }
+
+    /// Return the matched nodes for a named capture.
+    pub fn named_capture<'s: 'a + 'tree>(
+        &'s self,
+        name: &str,
+    ) -> Option<(CaptureQuantifier, impl Iterator<Item = Node<'tree>> + 's)> {
+        self.named_captures
+            .iter()
+            .find(|c| c.0 == name)
+            .map(|c| (c.1, self.mat.nodes_for_capture_index(c.2)))
+    }
+
+    /// Return an iterator over all capture names.
+    pub fn capture_names(&self) -> impl Iterator<Item = &String> {
+        self.named_captures.iter().map(|c| c.0)
+    }
+
+    /// Return the query location.
+    pub fn query_location(&self) -> &Location {
+        &self.query_location
+    }
 }
 
 /// Configuration for the execution of a File
···
 #[error("Cancelled at \"{0}\"")]
 pub struct CancellationError(pub &'static str);
 
-/// Get the value for the given capture, considering the suffix
-pub(self) fn query_capture_value<'tree>(
-    index: usize,
-    quantifier: CaptureQuantifier,
-    mat: &QueryMatch<'_, 'tree>,
-    graph: &mut Graph<'tree>,
-) -> Value {
-    let mut nodes = mat
-        .captures
-        .iter()
-        .filter(|c| c.index as usize == index)
-        .map(|c| c.node);
-    match quantifier {
-        CaptureQuantifier::Zero => unreachable!(),
-        CaptureQuantifier::One => {
-            let syntax_node = graph.add_syntax_node(nodes.next().expect("missing capture"));
-            syntax_node.into()
-        }
-        CaptureQuantifier::ZeroOrMore | CaptureQuantifier::OneOrMore => {
-            let syntax_nodes = nodes
-                .map(|n| graph.add_syntax_node(n.clone()).into())
-                .collect::<Vec<Value>>();
-            syntax_nodes.into()
-        }
-        CaptureQuantifier::ZeroOrOne => match nodes.next() {
-            None => Value::Null.into(),
-            Some(node) => {
-                let syntax_node = graph.add_syntax_node(node);
+impl Value {
+    pub fn from_nodes<'tree, NI: IntoIterator<Item = Node<'tree>>>(
+        graph: &mut Graph<'tree>,
+        nodes: NI,
+        quantifier: CaptureQuantifier,
+    ) -> Value {
+        let mut nodes = nodes.into_iter();
+        match quantifier {
+            CaptureQuantifier::Zero => unreachable!(),
+            CaptureQuantifier::One => {
+                let syntax_node = graph.add_syntax_node(nodes.next().expect("missing capture"));
                 syntax_node.into()
             }
-        },
+            CaptureQuantifier::ZeroOrMore | CaptureQuantifier::OneOrMore => {
+                let syntax_nodes = nodes
+                    .map(|n| graph.add_syntax_node(n.clone()).into())
+                    .collect::<Vec<Value>>();
+                syntax_nodes.into()
+            }
+            CaptureQuantifier::ZeroOrOne => match nodes.next() {
+                None => Value::Null.into(),
+                Some(node) => {
+                    let syntax_node = graph.add_syntax_node(node);
+                    syntax_node.into()
+                }
+            },
+        }
     }
 }

src/execution/error.rs

···
         Self {
             statement: format!("{}", stmt),
             statement_location: stmt.location(),
-            stanza_location: stanza.location,
+            stanza_location: stanza.range.start,
             source_location: Location::from(source_node.range().start_point),
             node_kind: source_node.kind().to_string(),
         }

src/execution/lazy.rs

···
 
 use std::collections::HashMap;
 
-use tree_sitter::CaptureQuantifier::One;
 use tree_sitter::QueryCursor;
 use tree_sitter::QueryMatch;
 use tree_sitter::Tree;
···
 use crate::execution::error::ExecutionError;
 use crate::execution::error::ResultWithExecutionError;
 use crate::execution::error::StatementContext;
-use crate::execution::query_capture_value;
 use crate::execution::ExecutionConfig;
 use crate::functions::Functions;
 use crate::graph;
 use crate::graph::Graph;
+use crate::graph::Value;
 use crate::variables::Globals;
 use crate::variables::MutVariables;
 use crate::variables::VariableMap;
···
         };
 
         let mut locals = VariableMap::new();
-        let mut cursor = QueryCursor::new();
         let mut store = LazyStore::new();
         let mut scoped_store = LazyScopedVariables::new();
         let mut lazy_graph = Vec::new();
         let mut function_parameters = Vec::new();
         let mut prev_element_debug_info = HashMap::new();
 
-        let query = &self.query.as_ref().unwrap();
-        let matches = cursor.matches(query, tree.root_node(), source.as_bytes());
-        for mat in matches {
+        self.try_visit_matches_lazy(tree, source, |stanza, mat| {
             cancellation_flag.check("processing matches")?;
-            let stanza = &self.stanzas[mat.pattern_index];
             stanza.execute_lazy(
                 source,
                 &mat,
···
                 &mut prev_element_debug_info,
                 &self.shorthands,
                 cancellation_flag,
-            )?;
-        }
+            )
+        })?;
 
         let mut exec = EvaluationContext {
             source,
···
 
         Ok(())
     }
+
+    pub(super) fn try_visit_matches_lazy<'tree, E, F>(
+        &self,
+        tree: &'tree Tree,
+        source: &'tree str,
+        mut visit: F,
+    ) -> Result<(), E>
+    where
+        F: FnMut(&ast::Stanza, QueryMatch<'_, 'tree>) -> Result<(), E>,
+    {
+        let mut cursor = QueryCursor::new();
+        let query = self.query.as_ref().unwrap();
+        let matches = cursor.matches(query, tree.root_node(), source.as_bytes());
+        for mat in matches {
+            let stanza = &self.stanzas[mat.pattern_index];
+            visit(stanza, mat)?;
+        }
+        Ok(())
+    }
 }
 
 /// Context for execution, which executes stanzas to build the lazy graph
···
     ) -> Result<(), ExecutionError> {
         let current_regex_captures = vec![];
         locals.clear();
-        let node = query_capture_value(self.full_match_file_capture_index, One, &mat, graph);
-        debug!("match {} at {}", node, self.location);
+        let node = mat
+            .nodes_for_capture_index(self.full_match_file_capture_index as u32)
+            .next()
+            .expect("missing capture for full match");
+        debug!("match {:?} at {}", node, self.range.start);
         trace!("{{");
         for statement in &self.statements {
-            let error_context = {
-                let node = mat
-                    .captures
-                    .iter()
-                    .find(|c| c.index as usize == self.full_match_file_capture_index)
-                    .expect("missing capture for full match")
-                    .node;
-                StatementContext::new(&statement, &self, &node)
-            };
+            let error_context = { StatementContext::new(&statement, &self, &node) };
             let mut exec = ExecutionContext {
                 source,
                 graph,
···
 
 impl ast::Capture {
     fn evaluate_lazy(&self, exec: &mut ExecutionContext) -> Result<LazyValue, ExecutionError> {
-        Ok(query_capture_value(
-            self.file_capture_index,
-            self.quantifier,
-            exec.mat,
+        Ok(Value::from_nodes(
             exec.graph,
+            exec.mat
+                .nodes_for_capture_index(self.file_capture_index as u32),
+            self.quantifier,
         )
         .into())
     }

src/execution/strict.rs

···
 use crate::ast::Variable;
 use crate::execution::error::ExecutionError;
 use crate::execution::error::ResultWithExecutionError;
-use crate::execution::query_capture_value;
+use crate::execution::error::StatementContext;
 use crate::execution::CancellationFlag;
 use crate::execution::ExecutionConfig;
 use crate::graph::Attributes;
···
 use crate::variables::Variables;
 use crate::Identifier;
 use crate::Location;
-
-use super::error::StatementContext;
 
 impl File {
     /// Executes this graph DSL file against a source file, saving the results into an existing
···
         let mut scoped = ScopedVariables::new();
         let current_regex_captures = Vec::new();
         let mut function_parameters = Vec::new();
-        let mut cursor = QueryCursor::new();
-        for stanza in &self.stanzas {
-            cancellation_flag.check("executing stanza")?;
+
+        self.try_visit_matches_strict(tree, source, |stanza, mat| {
             stanza.execute(
-                tree,
                 source,
+                &mat,
                 graph,
                 &mut config,
                 &mut locals,
                 &mut scoped,
                 &current_regex_captures,
                 &mut function_parameters,
-                &mut cursor,
                 &self.shorthands,
                 cancellation_flag,
-            )?;
+            )
+        })?;
+
+        Ok(())
+    }
+
+    pub(super) fn try_visit_matches_strict<'tree, E, F>(
+        &self,
+        tree: &'tree Tree,
+        source: &'tree str,
+        mut visit: F,
+    ) -> Result<(), E>
+    where
+        F: FnMut(&Stanza, QueryMatch<'_, 'tree>) -> Result<(), E>,
+    {
+        for stanza in &self.stanzas {
+            stanza.try_visit_matches_strict(tree, source, |mat| visit(stanza, mat))?;
         }
         Ok(())
     }
···
 impl Stanza {
     fn execute<'a, 'g, 'l, 's, 'tree>(
         &self,
-        tree: &'tree Tree,
         source: &'tree str,
+        mat: &QueryMatch<'_, 'tree>,
         graph: &mut Graph<'tree>,
         config: &ExecutionConfig<'_, 'g>,
         locals: &mut VariableMap<'l, Value>,
         scoped: &mut ScopedVariables<'s>,
         current_regex_captures: &Vec<String>,
         function_parameters: &mut Vec<Value>,
-        cursor: &mut QueryCursor,
         shorthands: &AttributeShorthands,
         cancellation_flag: &dyn CancellationFlag,
     ) -> Result<(), ExecutionError> {
+        locals.clear();
+        for statement in &self.statements {
+            let error_context = {
+                let node = mat
+                    .nodes_for_capture_index(self.full_match_stanza_capture_index as u32)
+                    .next()
+                    .expect("missing full capture");
+                StatementContext::new(&statement, &self, &node)
+            };
+            let mut exec = ExecutionContext {
+                source,
+                graph,
+                config,
+                locals,
+                scoped,
+                current_regex_captures,
+                function_parameters,
+                mat: &mat,
+                error_context,
+                shorthands,
+                cancellation_flag,
+            };
+            statement
+                .execute(&mut exec)
+                .with_context(|| exec.error_context.into())?;
+        }
+        Ok(())
+    }
+
+    pub(super) fn try_visit_matches_strict<'tree, E, F>(
+        &self,
+        tree: &'tree Tree,
+        source: &'tree str,
+        mut visit: F,
+    ) -> Result<(), E>
+    where
+        F: FnMut(QueryMatch<'_, 'tree>) -> Result<(), E>,
+    {
+        let mut cursor = QueryCursor::new();
         let matches = cursor.matches(&self.query, tree.root_node(), source.as_bytes());
         for mat in matches {
-            cancellation_flag.check("processing matches")?;
-            locals.clear();
-            for statement in &self.statements {
-                let error_context = {
-                    let node = mat
-                        .captures
-                        .iter()
-                        .find(|c| c.index as usize == self.full_match_stanza_capture_index)
-                        .expect("missing capture for full match")
-                        .node;
-                    StatementContext::new(&statement, &self, &node)
-                };
-                let mut exec = ExecutionContext {
-                    source,
-                    graph,
-                    config,
-                    locals,
-                    scoped,
-                    current_regex_captures,
-                    function_parameters,
-                    mat: &mat,
-                    error_context,
-                    shorthands,
-                    cancellation_flag,
-                };
-                statement
-                    .execute(&mut exec)
-                    .with_context(|| exec.error_context.into())?;
-            }
+            visit(mat)?;
         }
         Ok(())
     }
···
 
 impl Capture {
     fn evaluate(&self, exec: &mut ExecutionContext) -> Result<Value, ExecutionError> {
-        Ok(query_capture_value(
-            self.stanza_capture_index,
+        Ok(Value::from_nodes(
+            exec.graph,
+            exec.mat
+                .nodes_for_capture_index(self.stanza_capture_index as u32),
             self.quantifier,
-            exec.mat,
-            exec.graph,
-        ))
+        )
+        .into())
     }
 }
 

src/lib.rs

···
 pub use execution::CancellationError;
 pub use execution::CancellationFlag;
 pub use execution::ExecutionConfig;
+pub use execution::Match;
 pub use execution::NoCancellation;
 pub use parser::Location;
 pub use parser::ParseError;

src/parser.rs

···
 
 use std::fmt::Display;
 use std::iter::Peekable;
-use std::ops::Range;
 use std::path::Path;
 use std::str::Chars;
 
···
         }
     }
 
-    pub(crate) fn to_column_range(&self) -> Range<usize> {
+    pub(crate) fn to_column_range(&self) -> std::ops::Range<usize> {
         self.column..self.column + 1
     }
 }
···
 }
 
 // ----------------------------------------------------------------------------
+// Range
+
+/// The range of a graph DSL entity within its file
+#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
+pub struct Range {
+    pub start: Location,
+    pub end: Location,
+}
+
+impl Display for Range {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        write!(f, "{} - {}", self.start, self.end)
+    }
+}
+
+// ----------------------------------------------------------------------------
 // Parser
 
 struct Parser<'a> {
···
     }
 
     fn parse_stanza(&mut self, language: Language) -> Result<ast::Stanza, ParseError> {
-        let location = self.location;
+        let start = self.location;
         let (query, full_match_stanza_capture_index) = self.parse_query(language)?;
         self.consume_whitespace();
         let statements = self.parse_statements()?;
+        let end = self.location;
+        let range = Range { start, end };
         Ok(ast::Stanza {
             query,
             statements,
             full_match_stanza_capture_index,
             full_match_file_capture_index: usize::MAX, // set in checker
-            location,
+            range,
         })
     }