search: add support for counting individual matches (Fixes #566)

complete/_rg

@@ -26,7 +26,8 @@ _rg() {
     '--column[show column numbers]'
     '(-A -B -C --after-context --before-context --context)'{-C+,--context=}'[specify lines to show before and after each match]:number of lines'
     '--context-separator=[specify string used to separate non-continuous context lines in output]:separator'
-    '(-c --count --passthrough --passthru)'{-c,--count}'[only show count of matches for each file]'
+    '(-c --count --count-matches --passthrough --passthru)'{-c,--count}'[only show count of matching lines for each file]'
+    '(--count-matches -c --count --passthrough --passthru)--count-matches[only show count of individual matches for each file]'
     '--debug[show debug messages]'
     '--dfa-size-limit=[specify upper size limit of generated DFA]:DFA size'
     '(-E --encoding)'{-E+,--encoding=}'[specify text encoding of files to search]: :_rg_encodings'

src/app.rs

@@ -516,6 +516,7 @@ pub fn all_args_and_flags() -> Vec<RGArg> {
     flag_context(&mut args);
     flag_context_separator(&mut args);
     flag_count(&mut args);
+    flag_count_matches(&mut args);
     flag_debug(&mut args);
     flag_dfa_size_limit(&mut args);
     flag_encoding(&mut args);
@@ -758,7 +759,7 @@ sequences like \\x7F or \\t may be used. The default value is --.
 }
 
 fn flag_count(args: &mut Vec<RGArg>) {
-    const SHORT: &str = "Only show the count of matches for each file.";
+    const SHORT: &str = "Only show the count of matching lines for each file.";
     const LONG: &str = long!("\
 This flag suppresses normal output and shows the number of lines that match
 the given patterns for each file searched. Each file containing a match has its
@@ -768,9 +769,31 @@ that match and not the total number of matches.
 If only one file is given to ripgrep, then only the count is printed if there
 is a match. The --with-filename flag can be used to force printing the file
 path in this case.
+
+This overrides the --count-matches flag.
 ");
     let arg = RGArg::switch("count").short("c")
-        .help(SHORT).long_help(LONG);
+        .help(SHORT).long_help(LONG).overrides("count-matches");
+    args.push(arg);
+}
+
+fn flag_count_matches(args: &mut Vec<RGArg>) {
+    const SHORT: &str = "Only show the count of individual matches for each file.";
+    const LONG: &str = long!("\
+This flag suppresses normal output and shows the number of individual
+matches of the given patterns for each file searched. Each file
+containing matches has its path and match count printed on each line.
+Note that this reports the total number of individual matches and not
+the number of lines that match.
+
+If only one file is given to ripgrep, then only the count is printed if there
+is a match. The --with-filename flag can be used to force printing the file
+path in this case.
+
+This overrides the --count flag.
+");
+    let arg = RGArg::switch("count-matches")
+        .help(SHORT).long_help(LONG).overrides("count");
     args.push(arg);
 }
 

src/args.rs

@@ -40,6 +40,7 @@ pub struct Args {
     column: bool,
     context_separator: Vec<u8>,
     count: bool,
+    count_matches: bool,
     encoding: Option<&'static Encoding>,
     files_with_matches: bool,
     files_without_matches: bool,
@@ -199,6 +200,7 @@ impl Args {
     pub fn file_separator(&self) -> Option<Vec<u8>> {
         let contextless =
             self.count
+            || self.count_matches
             || self.files_with_matches
             || self.files_without_matches;
         let use_heading_sep = self.heading && !contextless;
@@ -260,6 +262,7 @@ impl Args {
             .after_context(self.after_context)
             .before_context(self.before_context)
             .count(self.count)
+            .count_matches(self.count_matches)
             .encoding(self.encoding)
             .files_with_matches(self.files_with_matches)
             .files_without_matches(self.files_without_matches)
@@ -356,6 +359,7 @@ impl<'a> ArgMatches<'a> {
         let mmap = self.mmap(&paths)?;
         let with_filename = self.with_filename(&paths);
         let (before_context, after_context) = self.contexts()?;
+        let (count, count_matches) = self.counts();
         let quiet = self.is_present("quiet");
         let args = Args {
             paths: paths,
@@ -365,7 +369,8 @@ impl<'a> ArgMatches<'a> {
             colors: self.color_specs()?,
             column: self.column(),
             context_separator: self.context_separator(),
-            count: self.is_present("count"),
+            count: count,
+            count_matches: count_matches,
             encoding: self.encoding()?,
             files_with_matches: self.is_present("files-with-matches"),
             files_without_matches: self.is_present("files-without-match"),
@@ -729,6 +734,22 @@ impl<'a> ArgMatches<'a> {
         })
     }
 
+    /// Returns whether the -c/--count or the --count-matches flags were
+    /// passed from the command line.
+    ///
+    /// If --count-matches and --invert-match were passed in, behave
+    /// as if --count and --invert-match were passed in (i.e. rg will
+    /// count inverted matches as per existing behavior).
+    fn counts(&self) -> (bool, bool) {
+        let count = self.is_present("count");
+        let count_matches = self.is_present("count-matches");
+        let invert_matches = self.is_present("invert-match");
+        if count_matches && invert_matches {
+            return (true, false);
+        }
+        (count, count_matches)
+    }
+
     /// Returns the user's color choice based on command line parameters and
     /// environment.
     fn color_choice(&self) -> termcolor::ColorChoice {

src/main.rs

@@ -88,13 +88,13 @@ fn run_parallel(args: &Arc<Args>) -> Result<u64> {
     let bufwtr = Arc::new(args.buffer_writer());
     let quiet_matched = args.quiet_matched();
     let paths_searched = Arc::new(AtomicUsize::new(0));
-    let match_count = Arc::new(AtomicUsize::new(0));
+    let match_line_count = Arc::new(AtomicUsize::new(0));
 
     args.walker_parallel().run(|| {
         let args = Arc::clone(args);
         let quiet_matched = quiet_matched.clone();
         let paths_searched = paths_searched.clone();
-        let match_count = match_count.clone();
+        let match_line_count = match_line_count.clone();
         let bufwtr = Arc::clone(&bufwtr);
         let mut buf = bufwtr.buffer();
         let mut worker = args.worker();
@@ -125,7 +125,7 @@ fn run_parallel(args: &Arc<Args>) -> Result<u64> {
                     } else {
                         worker.run(&mut printer, Work::DirEntry(dent))
                     };
-                match_count.fetch_add(count as usize, Ordering::SeqCst);
+                match_line_count.fetch_add(count as usize, Ordering::SeqCst);
                 if quiet_matched.set_match(count > 0) {
                     return Quit;
                 }
@@ -141,15 +141,15 @@ fn run_parallel(args: &Arc<Args>) -> Result<u64> {
             eprint_nothing_searched();
         }
     }
-    Ok(match_count.load(Ordering::SeqCst) as u64)
+    Ok(match_line_count.load(Ordering::SeqCst) as u64)
 }
 
 fn run_one_thread(args: &Arc<Args>) -> Result<u64> {
     let stdout = args.stdout();
     let mut stdout = stdout.lock();
     let mut worker = args.worker();
     let mut paths_searched: u64 = 0;
-    let mut match_count = 0;
+    let mut match_line_count = 0;
     for result in args.walker() {
         let dent = match get_or_log_dir_entry(
             result,
@@ -161,7 +161,7 @@ fn run_one_thread(args: &Arc<Args>) -> Result<u64> {
             Some(dent) => dent,
         };
         let mut printer = args.printer(&mut stdout);
-        if match_count > 0 {
+        if match_line_count > 0 {
             if args.quiet() {
                 break;
             }
@@ -170,7 +170,7 @@ fn run_one_thread(args: &Arc<Args>) -> Result<u64> {
             }
         }
         paths_searched += 1;
-        match_count +=
+        match_line_count +=
             if dent.is_stdin() {
                 worker.run(&mut printer, Work::Stdin)
             } else {
@@ -182,7 +182,7 @@ fn run_one_thread(args: &Arc<Args>) -> Result<u64> {
             eprint_nothing_searched();
         }
     }
-    Ok(match_count)
+    Ok(match_line_count)
 }
 
 fn run_files_parallel(args: Arc<Args>) -> Result<u64> {

src/search_buffer.rs

@@ -21,7 +21,8 @@ pub struct BufferSearcher<'a, W: 'a> {
     grep: &'a Grep,
     path: &'a Path,
     buf: &'a [u8],
-    match_count: u64,
+    match_line_count: u64,
+    match_count: Option<u64>,
     line_count: Option<u64>,
     last_line: usize,
 }
@@ -39,7 +40,8 @@ impl<'a, W: WriteColor> BufferSearcher<'a, W> {
             grep: grep,
             path: path,
             buf: buf,
-            match_count: 0,
+            match_line_count: 0,
+            match_count: None,
             line_count: None,
             last_line: 0,
         }
@@ -53,6 +55,15 @@ impl<'a, W: WriteColor> BufferSearcher<'a, W> {
         self
     }
 
+    /// If enabled, searching will print the count of individual matches
+    /// instead of each match.
+    ///
+    /// Disabled by default.
+    pub fn count_matches(mut self, yes: bool) -> Self {
+        self.opts.count_matches = yes;
+        self
+    }
+
     /// If enabled, searching will print the path instead of each match.
     ///
     /// Disabled by default.
@@ -118,8 +129,9 @@ impl<'a, W: WriteColor> BufferSearcher<'a, W> {
             return 0;
         }
 
-        self.match_count = 0;
+        self.match_line_count = 0;
         self.line_count = if self.opts.line_number { Some(0) } else { None };
+        self.match_count = if self.opts.count_matches { Some(0) } else { None };
         let mut last_end = 0;
         for m in self.grep.iter(self.buf) {
             if self.opts.invert_match {
@@ -128,29 +140,43 @@ impl<'a, W: WriteColor> BufferSearcher<'a, W> {
                 self.print_match(m.start(), m.end());
             }
             last_end = m.end();
-            if self.opts.terminate(self.match_count) {
+            if self.opts.terminate(self.match_line_count) {
                 break;
             }
         }
-        if self.opts.invert_match && !self.opts.terminate(self.match_count) {
+        if self.opts.invert_match && !self.opts.terminate(self.match_line_count) {
             let upto = self.buf.len();
             self.print_inverted_matches(last_end, upto);
         }
-        if self.opts.count && self.match_count > 0 {
-            self.printer.path_count(self.path, self.match_count);
+        if self.opts.count && self.match_line_count > 0 {
+            self.printer.path_count(self.path, self.match_line_count);
+        } else if self.opts.count_matches
+            && self.match_count.map_or(false, |c| c > 0)
+        {
+            self.printer.path_count(self.path, self.match_count.unwrap());
         }
-        if self.opts.files_with_matches && self.match_count > 0 {
+        if self.opts.files_with_matches && self.match_line_count > 0 {
             self.printer.path(self.path);
         }
-        if self.opts.files_without_matches && self.match_count == 0 {
+        if self.opts.files_without_matches && self.match_line_count == 0 {
             self.printer.path(self.path);
         }
-        self.match_count
+        self.match_line_count
+    }
+
+    #[inline(always)]
+    fn count_individual_matches(&mut self, start: usize, end: usize) {
+        if let Some(ref mut count) = self.match_count {
+            for _ in self.grep.regex().find_iter(&self.buf[start..end]) {
+                *count += 1;
+            }
+        }
     }
 
     #[inline(always)]
     pub fn print_match(&mut self, start: usize, end: usize) {
-        self.match_count += 1;
+        self.match_line_count += 1;
+        self.count_individual_matches(start, end);
         if self.opts.skip_matches() {
             return;
         }
@@ -166,7 +192,7 @@ impl<'a, W: WriteColor> BufferSearcher<'a, W> {
         debug_assert!(self.opts.invert_match);
         let mut it = IterLines::new(self.opts.eol, start);
         while let Some((s, e)) = it.next(&self.buf[..end]) {
-            if self.opts.terminate(self.match_count) {
+            if self.opts.terminate(self.match_line_count) {
                 return;
             }
             self.print_match(s, e);
@@ -279,6 +305,13 @@ and exhibited clearly, with a label attached.\
         assert_eq!(out, "/baz.rs:2\n");
     }
 
+    #[test]
+    fn count_matches() {
+        let (_, out) = search(
+            "the", SHERLOCK, |s| s.count_matches(true));
+        assert_eq!(out, "/baz.rs:4\n");
+    }
+
     #[test]
     fn files_with_matches() {
         let (count, out) = search(