Crate-level API changes for tree-sitter-lint

Cargo.toml

@@ -1,5 +1,5 @@
 [package]
-name = "tree-sitter-grep"
+name = "tree_sitter_lint_tree-sitter-grep"
 version = "0.1.0"
 edition = "2021"
 license = "Unlicense OR MIT"
@@ -8,6 +8,7 @@ authors = [
   "Peter Stuart <peter@peterstuart.org>"
 ]
 description = """
+(not-yet-landed version used by tree-sitter-lint)
 tree-sitter-grep is a grep-like search tool that
 recursively searches the current directory for a
 tree-sitter query pattern. Like ripgrep, it respects

examples/filter_before_line_number.rs

@@ -5,7 +5,7 @@ use std::{
 
 use libc::c_char;
 use tree_sitter::Node;
-use tree_sitter_grep::PluginInitializeReturn;
+use tree_sitter_lint_tree_sitter_grep::PluginInitializeReturn;
 
 static ROW_NUMBER: AtomicUsize = AtomicUsize::new(0);
 

examples/print_match_text.rs

@@ -1,12 +1,12 @@
 use clap::Parser;
-use tree_sitter_grep::{run_with_callback, Args};
+use tree_sitter_lint_tree_sitter_grep::{run_with_callback, Args};
 
 fn main() {
     let args = Args::parse_from(["tree_sitter_grep", "-q", "(function_item) @f"]);
-    run_with_callback(args, |node, file_contents, path| {
+    run_with_callback(args, |capture_info, file_contents, path| {
         println!(
             "Found match in {path:?}: {}",
-            std::str::from_utf8(&file_contents[node.byte_range()]).unwrap(),
+            std::str::from_utf8(&file_contents[capture_info.node.byte_range()]).unwrap(),
         );
     })
     .unwrap();

src/args.rs

@@ -23,7 +23,7 @@ use crate::{
 
 const ALL_NODES_QUERY: &str = "(_) @node";
 
-#[derive(Parser)]
+#[derive(Clone, Parser)]
 #[clap(group(
     ArgGroup::new("query_or_filter")
         .multiple(true)

src/bin/tree-sitter-grep.rs

@@ -1,7 +1,7 @@
 use std::process;
 
 use clap::Parser;
-use tree_sitter_grep::{run_print, Args, RunStatus};
+use tree_sitter_lint_tree_sitter_grep::{run_print, Args, RunStatus};
 
 pub fn main() {
     let args = Args::parse();

src/lib.rs

@@ -32,7 +32,8 @@ mod use_printer;
 mod use_searcher;
 
 pub use args::Args;
-use language::{BySupportedLanguage, SupportedLanguage};
+use language::BySupportedLanguage;
+pub use language::SupportedLanguage;
 pub use plugin::PluginInitializeReturn;
 use query_context::QueryContext;
 use treesitter::maybe_get_query;
@@ -74,6 +75,8 @@ pub enum Error {
     FilterPluginExpectedArgument,
     #[error("plugin couldn't parse argument {filter_arg:?}")]
     FilterPluginCouldntParseArgument { filter_arg: String },
+    #[error("language is required when passing a slice")]
+    LanguageMissingForSlice,
 }
 
 #[derive(Clone, Debug, Error)]
@@ -290,9 +293,14 @@ pub fn run_print(args: Args) -> Result<RunStatus, Error> {
     )
 }
 
+pub struct CaptureInfo<'node> {
+    pub node: Node<'node>,
+    pub pattern_index: usize,
+}
+
 pub fn run_with_callback(
     args: Args,
-    callback: impl Fn(Node, &[u8], &Path) + Sync,
+    callback: impl Fn(CaptureInfo, &[u8], &Path) + Sync,
 ) -> Result<RunStatus, Error> {
     run_for_context(
         args,
@@ -307,8 +315,8 @@ pub fn run_with_callback(
                 .search_path_callback::<_, io::Error>(
                     query_context,
                     path,
-                    |node: Node, file_contents: &[u8], path: &Path| {
-                        callback(node, file_contents, path);
+                    |capture_info: CaptureInfo, file_contents: &[u8], path: &Path| {
+                        callback(capture_info, file_contents, path);
                         matched.store(true, Ordering::SeqCst);
                     },
                 )
@@ -412,6 +420,50 @@ fn run_for_context<TContext: Sync>(
     })
 }
 
+pub fn run_for_slice_with_callback(
+    slice: &[u8],
+    args: Args,
+    mut callback: impl FnMut(CaptureInfo) + Sync,
+) -> Result<RunStatus, Error> {
+    let language = args.language.ok_or(Error::LanguageMissingForSlice)?;
+    let query_text = args.get_loaded_query_text()?;
+    let filter = args.get_loaded_filter()?;
+    let cached_queries: CachedQueries = Default::default();
+    let capture_index = CaptureIndex::default();
+    let matched = AtomicBool::new(false);
+    let non_fatal_errors: Arc<Mutex<Vec<NonFatalError>>> = Default::default();
+
+    let query = match cached_queries.get_and_cache_query_for_language(&query_text, language) {
+        Some(query) => query,
+        None => {
+            return Err(cached_queries
+                .error_if_no_successful_query_parsing()
+                .unwrap_err())
+        }
+    };
+    let capture_index = capture_index.get_or_init(&query, args.capture_name.as_deref())?;
+
+    let query_context = QueryContext::new(query, capture_index, language.language(), filter);
+
+    get_searcher(&args)
+        .borrow_mut()
+        .search_slice_callback_no_path(query_context, slice, |capture_info: CaptureInfo| {
+            callback(capture_info);
+            matched.store(true, Ordering::SeqCst);
+        })
+        .unwrap();
+
+    let non_fatal_errors = non_fatal_errors.lock().unwrap().clone();
+    if non_fatal_errors.is_empty() {
+        cached_queries.error_if_no_successful_query_parsing()?;
+    }
+
+    Ok(RunStatus {
+        matched: matched.load(Ordering::SeqCst),
+        non_fatal_errors,
+    })
+}
+
 fn for_each_project_file(
     args: &Args,
     non_fatal_errors: Arc<Mutex<Vec<NonFatalError>>>,

src/searcher/mod.rs

@@ -9,7 +9,7 @@ use std::{
 };
 
 use encoding_rs_io::DecodeReaderBytesBuilder;
-use tree_sitter::{Node, QueryCursor};
+use tree_sitter::QueryCursor;
 
 pub use self::mmap::MmapChoice;
 use crate::{
@@ -19,6 +19,7 @@ use crate::{
     searcher::glue::MultiLine,
     sink::{Sink, SinkError},
     treesitter::get_parser,
+    CaptureInfo,
 };
 
 mod core;
@@ -218,7 +219,7 @@ impl Searcher {
         &mut self,
         query_context: QueryContext,
         path: P,
-        callback: impl Fn(Node, &[u8], &Path),
+        callback: impl Fn(CaptureInfo, &[u8], &Path),
     ) -> Result<(), TError>
     where
         P: AsRef<Path>,
@@ -338,7 +339,7 @@ impl Searcher {
         &mut self,
         query_context: QueryContext,
         slice: &[u8],
-        callback: impl Fn(Node, &[u8], &Path),
+        callback: impl Fn(CaptureInfo, &[u8], &Path),
         path: &Path,
     ) -> Result<(), ConfigError> {
         self.check_config()?;
@@ -349,11 +350,53 @@ impl Searcher {
         Ok(())
     }
 
+    pub fn search_slice_callback_no_path(
+        &mut self,
+        query_context: QueryContext,
+        slice: &[u8],
+        mut callback: impl FnMut(CaptureInfo),
+    ) -> Result<(), ConfigError> {
+        self.check_config()?;
+
+        log::trace!("slice reader: searching via multiline strategy");
+        let mut query_cursor = QueryCursor::new();
+        let tree = get_parser(query_context.language)
+            .parse(slice, None)
+            .unwrap();
+        let query = &query_context.query;
+        let capture_index = query_context.capture_index;
+        let filter = &query_context.filter;
+        query_cursor
+            .captures(query, tree.root_node(), slice)
+            .filter_map(|(match_, index_into_query_match_captures)| {
+                let this_capture = &match_.captures[index_into_query_match_captures];
+                if this_capture.index != capture_index {
+                    return None;
+                }
+                let single_captured_node = this_capture.node;
+                match filter.as_ref() {
+                    None => Some(CaptureInfo {
+                        node: single_captured_node,
+                        pattern_index: match_.pattern_index,
+                    }),
+                    Some(filter) => filter.call(&single_captured_node).then_some(CaptureInfo {
+                        node: single_captured_node,
+                        pattern_index: match_.pattern_index,
+                    }),
+                }
+            })
+            .for_each(|capture_info| {
+                callback(capture_info);
+            });
+
+        Ok(())
+    }
+
     fn run_with_callback(
         &self,
         query_context: QueryContext,
         slice: &[u8],
-        callback: impl Fn(Node, &[u8], &Path),
+        callback: impl Fn(CaptureInfo, &[u8], &Path),
         path: &Path,
     ) {
         let mut query_cursor = QueryCursor::new();
@@ -365,26 +408,25 @@ impl Searcher {
         let filter = &query_context.filter;
         query_cursor
             .captures(query, tree.root_node(), slice)
-            .filter_map(|(match_, found_capture_index)| {
-                let found_capture_index = found_capture_index as u32;
-                if found_capture_index != capture_index {
+            .filter_map(|(match_, index_into_query_match_captures)| {
+                let this_capture = &match_.captures[index_into_query_match_captures];
+                if this_capture.index != capture_index {
                     return None;
                 }
-                let mut nodes_for_this_capture = match_.nodes_for_capture_index(capture_index);
-                let single_captured_node = nodes_for_this_capture.next().unwrap();
-                assert!(
-                    nodes_for_this_capture.next().is_none(),
-                    "I guess .captures() always wraps up the single capture like this?"
-                );
+                let single_captured_node = this_capture.node;
                 match filter.as_ref() {
-                    None => Some(single_captured_node),
-                    Some(filter) => filter
-                        .call(&single_captured_node)
-                        .then_some(single_captured_node),
+                    None => Some(CaptureInfo {
+                        node: single_captured_node,
+                        pattern_index: match_.pattern_index,
+                    }),
+                    Some(filter) => filter.call(&single_captured_node).then_some(CaptureInfo {
+                        node: single_captured_node,
+                        pattern_index: match_.pattern_index,
+                    }),
                 }
             })
-            .for_each(|node| {
-                callback(node, slice, path);
+            .for_each(|capture_info| {
+                callback(capture_info, slice, path);
             });
     }
 

src/use_searcher.rs

@@ -1,22 +1,16 @@
-use std::{
-    cell::{OnceCell, RefCell},
-    ptr,
-    rc::Rc,
-};
+use std::{cell::RefCell, collections::HashMap, rc::Rc};
 
 use crate::{searcher::Searcher, Args};
 
 thread_local! {
-    static SEARCHER: OnceCell<(Rc<RefCell<Searcher>>, *const Args)> = Default::default();
+    static SEARCHER_PER_ARGS_INSTANCE: RefCell<HashMap<*const Args, Rc<RefCell<Searcher>>>> = Default::default();
 }
 pub(crate) fn get_searcher(args: &Args) -> Rc<RefCell<Searcher>> {
-    SEARCHER.with(|searcher| {
-        let (searcher, args_when_initialized) =
-            searcher.get_or_init(|| (Rc::new(RefCell::new(args.get_searcher())), args));
-        assert!(
-            ptr::eq(*args_when_initialized, args),
-            "Using multiple instances of args not supported"
-        );
-        searcher.clone()
+    SEARCHER_PER_ARGS_INSTANCE.with(|searcher_per_args_instance| {
+        searcher_per_args_instance
+            .borrow_mut()
+            .entry(args)
+            .or_insert_with(|| Rc::new(RefCell::new(args.get_searcher())))
+            .clone()
     })
 }