Merge pull request tree-sitter#2085 from tree-sitter/faster-query-in-…

…range

Improve the performance of running a query in a small range of a large file

cli/Cargo.toml

@@ -3,7 +3,7 @@ name = "tree-sitter-cli"
 description = "CLI tool for developing, testing, and using Tree-sitter parsers"
 version = "0.20.7"
 authors = ["Max Brunsfeld <[email protected]>"]
-edition = "2018"
+edition = "2021"
 license = "MIT"
 readme = "README.md"
 keywords = ["incremental", "parsing"]
@@ -73,6 +73,7 @@ rand = "0.8"
 tempfile = "3"
 pretty_assertions = "0.7.2"
 ctor = "0.1"
+unindent = "0.2"
 
 [build-dependencies]
 toml = "0.5"

cli/src/main.rs

@@ -3,6 +3,7 @@ use clap::{App, AppSettings, Arg, SubCommand};
 use glob::glob;
 use std::path::{Path, PathBuf};
 use std::{env, fs, u64};
+use tree_sitter::Point;
 use tree_sitter_cli::parse::ParseOutput;
 use tree_sitter_cli::{
     generate, highlight, logger, parse, playground, query, tags, test, test_highlight, test_tags,
@@ -173,6 +174,8 @@ fn run() -> Result<()> {
                         .index(1)
                         .required(true),
                 )
+                .arg(&time_arg)
+                .arg(&quiet_arg)
                 .arg(&paths_file_arg)
                 .arg(&paths_arg.clone().index(2))
                 .arg(
@@ -181,6 +184,12 @@ fn run() -> Result<()> {
                         .long("byte-range")
                         .takes_value(true),
                 )
+                .arg(
+                    Arg::with_name("row-range")
+                        .help("The range of rows in which the query will be executed")
+                        .long("row-range")
+                        .takes_value(true),
+                )
                 .arg(&scope_arg)
                 .arg(Arg::with_name("captures").long("captures").short("c"))
                 .arg(Arg::with_name("test").long("test")),
@@ -456,6 +465,8 @@ fn run() -> Result<()> {
 
         ("query", Some(matches)) => {
             let ordered_captures = matches.values_of("captures").is_some();
+            let quiet = matches.values_of("quiet").is_some();
+            let time = matches.values_of("time").is_some();
             let paths = collect_paths(matches.value_of("paths-file"), matches.values_of("paths"))?;
             let loader_config = config.get()?;
             loader.find_all_languages(&loader_config)?;
@@ -465,18 +476,29 @@ fn run() -> Result<()> {
                 matches.value_of("scope"),
             )?;
             let query_path = Path::new(matches.value_of("query-path").unwrap());
-            let range = matches.value_of("byte-range").map(|br| {
-                let r: Vec<&str> = br.split(":").collect();
-                r[0].parse().unwrap()..r[1].parse().unwrap()
+            let byte_range = matches.value_of("byte-range").and_then(|arg| {
+                let mut parts = arg.split(":");
+                let start = parts.next()?.parse().ok()?;
+                let end = parts.next().unwrap().parse().ok()?;
+                Some(start..end)
+            });
+            let point_range = matches.value_of("row-range").and_then(|arg| {
+                let mut parts = arg.split(":");
+                let start = parts.next()?.parse().ok()?;
+                let end = parts.next().unwrap().parse().ok()?;
+                Some(Point::new(start, 0)..Point::new(end, 0))
             });
             let should_test = matches.is_present("test");
             query::query_files_at_paths(
                 language,
                 paths,
                 query_path,
                 ordered_captures,
-                range,
+                byte_range,
+                point_range,
                 should_test,
+                quiet,
+                time,
             )?;
         }
 

cli/src/query.rs

@@ -5,16 +5,20 @@ use std::{
     io::{self, Write},
     ops::Range,
     path::Path,
+    time::Instant,
 };
-use tree_sitter::{Language, Parser, Query, QueryCursor};
+use tree_sitter::{Language, Parser, Point, Query, QueryCursor};
 
 pub fn query_files_at_paths(
     language: Language,
     paths: Vec<String>,
     query_path: &Path,
     ordered_captures: bool,
-    range: Option<Range<usize>>,
+    byte_range: Option<Range<usize>>,
+    point_range: Option<Range<Point>>,
     should_test: bool,
+    quiet: bool,
+    print_time: bool,
 ) -> Result<()> {
     let stdout = io::stdout();
     let mut stdout = stdout.lock();
@@ -24,9 +28,12 @@ pub fn query_files_at_paths(
     let query = Query::new(language, &query_source).with_context(|| "Query compilation failed")?;
 
     let mut query_cursor = QueryCursor::new();
-    if let Some(range) = range {
+    if let Some(range) = byte_range {
         query_cursor.set_byte_range(range);
     }
+    if let Some(range) = point_range {
+        query_cursor.set_point_range(range);
+    }
 
     let mut parser = Parser::new();
     parser.set_language(language)?;
@@ -40,22 +47,25 @@ pub fn query_files_at_paths(
             fs::read(&path).with_context(|| format!("Error reading source file {:?}", path))?;
         let tree = parser.parse(&source_code, None).unwrap();
 
+        let start = Instant::now();
         if ordered_captures {
             for (mat, capture_index) in
                 query_cursor.captures(&query, tree.root_node(), source_code.as_slice())
             {
                 let capture = mat.captures[capture_index];
                 let capture_name = &query.capture_names()[capture.index as usize];
-                writeln!(
-                    &mut stdout,
-                    "    pattern: {:>2}, capture: {} - {}, start: {}, end: {}, text: `{}`",
-                    mat.pattern_index,
-                    capture.index,
-                    capture_name,
-                    capture.node.start_position(),
-                    capture.node.end_position(),
-                    capture.node.utf8_text(&source_code).unwrap_or("")
-                )?;
+                if !quiet {
+                    writeln!(
+                        &mut stdout,
+                        "    pattern: {:>2}, capture: {} - {}, start: {}, end: {}, text: `{}`",
+                        mat.pattern_index,
+                        capture.index,
+                        capture_name,
+                        capture.node.start_position(),
+                        capture.node.end_position(),
+                        capture.node.utf8_text(&source_code).unwrap_or("")
+                    )?;
+                }
                 results.push(query_testing::CaptureInfo {
                     name: capture_name.to_string(),
                     start: capture.node.start_position(),
@@ -64,27 +74,31 @@ pub fn query_files_at_paths(
             }
         } else {
             for m in query_cursor.matches(&query, tree.root_node(), source_code.as_slice()) {
-                writeln!(&mut stdout, "  pattern: {}", m.pattern_index)?;
+                if !quiet {
+                    writeln!(&mut stdout, "  pattern: {}", m.pattern_index)?;
+                }
                 for capture in m.captures {
                     let start = capture.node.start_position();
                     let end = capture.node.end_position();
                     let capture_name = &query.capture_names()[capture.index as usize];
-                    if end.row == start.row {
-                        writeln!(
-                            &mut stdout,
-                            "    capture: {} - {}, start: {}, end: {}, text: `{}`",
-                            capture.index,
-                            capture_name,
-                            start,
-                            end,
-                            capture.node.utf8_text(&source_code).unwrap_or("")
-                        )?;
-                    } else {
-                        writeln!(
-                            &mut stdout,
-                            "    capture: {}, start: {}, end: {}",
-                            capture_name, start, end,
-                        )?;
+                    if !quiet {
+                        if end.row == start.row {
+                            writeln!(
+                                &mut stdout,
+                                "    capture: {} - {}, start: {}, end: {}, text: `{}`",
+                                capture.index,
+                                capture_name,
+                                start,
+                                end,
+                                capture.node.utf8_text(&source_code).unwrap_or("")
+                            )?;
+                        } else {
+                            writeln!(
+                                &mut stdout,
+                                "    capture: {}, start: {}, end: {}",
+                                capture_name, start, end,
+                            )?;
+                        }
                     }
                     results.push(query_testing::CaptureInfo {
                         name: capture_name.to_string(),
@@ -103,6 +117,9 @@ pub fn query_files_at_paths(
         if should_test {
             query_testing::assert_expected_captures(results, path, &mut parser, language)?
         }
+        if print_time {
+            writeln!(&mut stdout, "{:?}", start.elapsed())?;
+        }
     }
 
     Ok(())