Add limit option (#40)

* Add test for limit option

* Add limit option

* Adjust test to use number instad of NonZero limit

* Make limit usize instead of NonZero

* Make limit option work

* Fix which loop gets break

* Ignore send error inside thread

* Add crossbeam-channel

* Formatting

* Add ThreadPool.stop and stop when limit reached

* Improve limit tests to have multiple cases

* Fix limit condition (it was backwards)

Author: sirbrillig

Cargo.lock

@@ -11,6 +11,7 @@ documentation = "https://docs.rs/grepdef/latest/grepdef"
 [dependencies]
 clap = { version = "4.5.9", features = ["derive"] }
 colored = "2.1.0"
+crossbeam-channel = "0.5.14"
 exitcode = "1.1.2"
 ignore = "0.4.22"
 memchr = "2.7.4"

Cargo.toml

@@ -112,6 +112,10 @@ pub struct Args {
     #[arg(long = "no-color")]
     pub no_color: bool,
 
+    /// Limit the number of results
+    #[arg(short = 'l', long = "limit")]
+    pub limit: Option<usize>,
+
     /// (Advanced) Print debugging information
     #[arg(long = "debug")]
     pub debug: bool,
@@ -192,6 +196,9 @@ struct Config {
     /// Output debugging info during search if true
     debug: bool,
 
+    /// Limit the number of results
+    limit: Option<usize>,
+
     /// Explicitly disable color output if true
     no_color: bool,
 
@@ -242,6 +249,7 @@ impl Config {
             no_color: args.no_color,
             color,
             search_method: args.search_method.unwrap_or_default(),
+            limit: args.limit,
             num_threads,
             format: args.format.unwrap_or_default(),
         };
@@ -438,10 +446,13 @@ impl Searcher {
     /// Perform the search and return formatted strings
     pub fn search_and_format(&self) -> Result<Vec<String>, Box<dyn Error>> {
         let results = self.search()?;
-        Ok(results.iter().map(|result| match self.config.format {
-            SearchResultFormat::Grep => result.to_grep(),
-            SearchResultFormat::JsonPerMatch => result.to_json_per_match(),
-        }).collect())
+        Ok(results
+            .iter()
+            .map(|result| match self.config.format {
+                SearchResultFormat::Grep => result.to_grep(),
+                SearchResultFormat::JsonPerMatch => result.to_json_per_match(),
+            })
+            .collect())
     }
 
     /// Perform the search and run a callback for each formatted string
@@ -468,7 +479,7 @@ impl Searcher {
         };
         let re = query_regex::get_regex_for_query(&self.config.query, &self.config.file_type);
         let file_type_re = file_type::get_regexp_for_file_type(&self.config.file_type);
-        let mut pool = threads::ThreadPool::new(self.config.num_threads);
+        let mut pool = threads::ThreadPool::new(self.config.num_threads, self.config.debug);
 
         match self.config.color {
             ColorOption::ALWAYS => colored::control::set_override(true),
@@ -517,10 +528,12 @@ impl Searcher {
                             &re1,
                             &path1,
                             &config1,
-                            // NOTE: it would be nice to have better error handling for if this
-                            // message send fails, but since error handling would happen through
-                            // message sending, I don't know what else to do other than panic.
-                            move |file_results: Vec<SearchResult>| tx1.send(file_results).unwrap(),
+                            move |file_results: Vec<SearchResult>| {
+                                // NOTE: it would be nice to have better error handling for if this
+                                // message send fails, but since normal error handling would happen through
+                                // message sending, I don't know what else to do.
+                                let _ = tx1.send(file_results);
+                            },
                         );
                     })
                 }
@@ -529,15 +542,25 @@ impl Searcher {
         };
 
         self.debug("Listening to searcher results");
-        for received_results in rx {
+        let mut result_counter: usize = 0;
+        'all_results: for received_results in rx {
             for received_result in received_results {
+                result_counter += 1;
                 callback(received_result);
                 // Don't try to even calculate elapsed time if we are not going to print it
                 if let (true, Some(start)) = (self.config.debug, start) {
                     self.debug(
                         format!("Found a result in {} ms", start.elapsed().as_millis()).as_str(),
                     );
                 }
+                if let Some(i) = self.config.limit {
+                    self.debug(format!("This is result {}; limit {}", result_counter, i).as_str());
+                    if result_counter >= i {
+                        self.debug("Limit reached");
+                        pool.stop();
+                        break 'all_results;
+                    }
+                }
             }
         }
 

src/lib.rs

@@ -1,10 +1,12 @@
+use colored::Colorize;
 use std::num::NonZero;
 use std::sync::mpsc;
 use std::sync::Arc;
 use std::sync::Mutex;
 use std::thread;
 
 type WorkerReceiver = Arc<Mutex<mpsc::Receiver<Job>>>;
+type ThreadControllerReceiver = crossbeam_channel::Receiver<usize>;
 type Job = Box<dyn FnOnce() + Send + 'static>;
 
 struct Worker {
@@ -13,21 +15,44 @@ struct Worker {
 }
 
 impl Worker {
-    pub fn new(id: usize, receiver: WorkerReceiver) -> Worker {
+    pub fn new(
+        id: usize,
+        debug: bool,
+        receiver: WorkerReceiver,
+        controller: ThreadControllerReceiver,
+    ) -> Worker {
         let thread = thread::spawn(move || loop {
+            // If the controller says to stop, ignore remaining jobs on the receiver and just stop.
+            if ! controller.is_empty() {
+                if debug {
+                    println!("thread {}: {}", id, "Controller said to stop".red());
+                }
+                break;
+            }
+
             // recv will block until the next job is sent.
+            if debug {
+                println!("thread {}: {}", id, "Waiting for job".green());
+            }
             let message = receiver
                 .lock()
                 .expect("Worker thread could not get message from main thread")
                 .recv();
 
             match message {
                 Ok(job) => {
+                    if debug {
+                        println!("thread {}: {}", id, "Job received; running".green());
+                    }
+                    // Run the job inside this thread.
                     job();
                 }
                 // The thread will stop when the job channel is sent an Err, which will happen when
                 // the channel is closed.
                 Err(_) => {
+                    if debug {
+                        println!("thread {}: {}", id, "No more jobs; stopping".red());
+                    }
                     break;
                 }
             }
@@ -42,24 +67,33 @@ impl Worker {
 pub struct ThreadPool {
     workers: Vec<Worker>,
     sender: Option<mpsc::Sender<Job>>,
+    thread_control_tx: crossbeam_channel::Sender<usize>,
 }
 
 impl ThreadPool {
-    pub fn new(count: NonZero<usize>) -> ThreadPool {
+    pub fn new(count: NonZero<usize>, debug: bool) -> ThreadPool {
         // This channel is used to send Jobs to each thread.
         let (sender, receiver) = mpsc::channel();
+        // This channel is used to send a stop signal to each thread.
+        let (thread_control_tx, thread_control_rx) = crossbeam_channel::unbounded();
 
         let receiver = Arc::new(Mutex::new(receiver));
 
         let mut workers = Vec::with_capacity(count.into());
 
         for id in 0..count.into() {
-            workers.push(Worker::new(id, Arc::clone(&receiver)));
+            workers.push(Worker::new(
+                id,
+                debug,
+                Arc::clone(&receiver),
+                thread_control_rx.clone(),
+            ));
         }
 
         ThreadPool {
             workers,
             sender: Some(sender),
+            thread_control_tx,
         }
     }
 
@@ -87,6 +121,12 @@ impl ThreadPool {
             }
         }
     }
+
+    pub fn stop(&mut self) {
+        self.thread_control_tx
+            .send(1)
+            .expect("Unable to send stop message to threads");
+    }
 }
 
 impl Drop for ThreadPool {

src/threads.rs

@@ -13,6 +13,7 @@ pub fn make_args(
         },
         file_type: file_type_string,
         line_number: true,
+        limit: None,
         search_method: None,
         debug: false,
         no_color: false,

tests/common/mod.rs

@@ -50,3 +50,6 @@ public function similarMethod() {}
 class ClassWithNonUniqueMethod2 {
 	public function similarMethod() {}
 }
+class ClassWithNonUniqueMethod3 {
+	public function similarMethod() {}
+}

tests/fixtures/by-language/php-fixture.php

@@ -406,11 +406,11 @@ fn search_returns_expected_line_number_for_file_type(
 }
 
 #[rstest]
-#[case::php_non_unique_method(String::from("similarMethod"), String::from("php"), [48, 51])]
+#[case::php_non_unique_method(String::from("similarMethod"), String::from("php"), [48, 51, 54])]
 fn search_returns_expected_line_number_group_for_file_type(
     #[case] query: String,
     #[case] file_type_string: String,
-    #[case] line_numbers: [usize;2],
+    #[case] line_numbers: [usize;3],
 ) {
     let file_path =
         common::get_default_fixture_for_file_type_string(file_type_string.as_str()).unwrap();
@@ -426,6 +426,31 @@ fn search_returns_expected_line_number_group_for_file_type(
     }
 }
 
+#[rstest]
+#[case::php_non_unique_method_limit_1(String::from("similarMethod"), String::from("php"), [48, 51, 54], 1)]
+#[case::php_non_unique_method_limit_2(String::from("similarMethod"), String::from("php"), [48, 51, 54], 2)]
+#[case::php_non_unique_method_limit_3(String::from("similarMethod"), String::from("php"), [48, 51, 54], 3)]
+fn search_returns_limited_expected_line_number_group_for_file_type(
+    #[case] query: String,
+    #[case] file_type_string: String,
+    #[case] line_numbers: [usize;3],
+    #[case] limit: usize,
+) {
+    let file_path =
+        common::get_default_fixture_for_file_type_string(file_type_string.as_str()).unwrap();
+    let mut args = common::make_args(query, Some(file_path), Some(file_type_string));
+    args.limit = Some(limit);
+    let actual = common::do_search(args);
+    assert_eq!(limit, actual.len(), "Did not find expected number of matches");
+    for (i, result) in actual.iter().enumerate() {
+        assert_eq!(
+            line_numbers[i],
+            result.line_number.unwrap(),
+            "Match found but it has the wrong line number"
+        );
+    }
+}
+
 #[rstest]
 fn search_returns_matching_js_function_line_for_recursive() {
     let file_path = String::from("./tests");