From c2fb0f5229b1bb72005da5177457fafb1560954a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tu=E1=BA=A5n-Anh=20Nguy=E1=BB=85n?= <ubolonton@gmail.com>
Date: Sun, 12 Jul 2020 20:45:17 +0700
Subject: [PATCH] cli: Add --byte-range flag to query command

---
 cli/src/main.rs  | 12 +++++++++++-
 cli/src/query.rs |  4 ++++
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/cli/src/main.rs b/cli/src/main.rs
index 713bf28f..a543202d 100644
--- a/cli/src/main.rs
+++ b/cli/src/main.rs
@@ -85,6 +85,12 @@ fn run() -> error::Result<()> {
                         .multiple(true)
                         .required(true),
                 )
+                .arg(
+                    Arg::with_name("beg>:<end")
+                        .help("The range of byte offsets in which the query will be executed")
+                        .long("byte-range")
+                        .takes_value(true)
+                )
                 .arg(Arg::with_name("scope").long("scope").takes_value(true))
                 .arg(Arg::with_name("captures").long("captures").short("c")),
         )
@@ -259,7 +265,11 @@ fn run() -> error::Result<()> {
             matches.value_of("scope"),
         )?;
         let query_path = Path::new(matches.value_of("query-path").unwrap());
-        query::query_files_at_paths(language, paths, query_path, ordered_captures)?;
+        let range = matches.value_of("beg>:<end").map(|br| {
+            let r: Vec<&str> = br.split(":").collect();
+            (r[0].parse().unwrap(), r[1].parse().unwrap())
+        });
+        query::query_files_at_paths(language, paths, query_path, ordered_captures, range)?;
     } else if let Some(matches) = matches.subcommand_matches("tags") {
         loader.find_all_languages(&config.parser_directories)?;
         let paths = collect_paths(matches.values_of("inputs").unwrap())?;
diff --git a/cli/src/query.rs b/cli/src/query.rs
index 47242273..8d097911 100644
--- a/cli/src/query.rs
+++ b/cli/src/query.rs
@@ -9,6 +9,7 @@ pub fn query_files_at_paths(
     paths: Vec<&Path>,
     query_path: &Path,
     ordered_captures: bool,
+    range: Option<(usize, usize)>,
 ) -> Result<()> {
     let stdout = io::stdout();
     let mut stdout = stdout.lock();
@@ -20,6 +21,9 @@ pub fn query_files_at_paths(
         .map_err(|e| Error::new(format!("Query compilation failed: {:?}", e)))?;
 
     let mut query_cursor = QueryCursor::new();
+    if let Some((beg, end)) = range {
+        query_cursor.set_byte_range(beg, end);
+    }
 
     let mut parser = Parser::new();
     parser.set_language(language).map_err(|e| e.to_string())?;