Add -m/--max-count flag.

This flag limits the number of matches printed *per file*.

Closes #159

8 files changed, 128 insertions, 3 deletions

diff --git a/doc/rg.1 b/doc/rg.1
index 3d7c1025..fa640e15 100644
--- a/doc/rg.1
+++ b/doc/rg.1
@@ -207,6 +207,11 @@ Follow symlinks.
 .RS
 .RE
 .TP
+.B \-m, \-\-max\-count NUM
+Limit the number of matching lines per file searched to NUM.
+.RS
+.RE
+.TP
 .B \-\-maxdepth \f[I]NUM\f[]
 Descend at most NUM directories below the command line arguments.
 A value of zero searches only the starting\-points themselves.
diff --git a/doc/rg.1.md b/doc/rg.1.md
index 67b3c33c..e2f42e73 100644
--- a/doc/rg.1.md
+++ b/doc/rg.1.md
@@ -135,6 +135,9 @@ Project home page: https://github.com/BurntSushi/ripgrep
 -L, --follow
 : Follow symlinks.
 
+-m, --max-count NUM
+: Limit the number of matching lines per file searched to NUM.
+
 --maxdepth *NUM*
 : Descend at most NUM directories below the command line arguments.
   A value of zero searches only the starting-points themselves.
diff --git a/src/args.rs b/src/args.rs
index f66667a9..66bb5cd9 100644
--- a/src/args.rs
+++ b/src/args.rs
@@ -141,6 +141,9 @@ Less common options:
     -L, --follow
         Follow symlinks.
 
+    -m, --max-count NUM
+        Limit the number of matching lines per file searched to NUM.
+
     --maxdepth NUM
         Descend at most NUM directories below the command line arguments.
         A value of zero only searches the starting-points themselves.
@@ -245,6 +248,7 @@ pub struct RawArgs {
     flag_invert_match: bool,
     flag_line_number: bool,
     flag_fixed_strings: bool,
+    flag_max_count: Option<usize>,
     flag_maxdepth: Option<usize>,
     flag_mmap: bool,
     flag_no_heading: bool,
@@ -296,6 +300,7 @@ pub struct Args {
     invert_match: bool,
     line_number: bool,
     line_per_match: bool,
+    max_count: Option<u64>,
     maxdepth: Option<usize>,
     mmap: bool,
     no_ignore: bool,
@@ -414,6 +419,7 @@ impl RawArgs {
             invert_match: self.flag_invert_match,
             line_number: !self.flag_no_line_number && self.flag_line_number,
             line_per_match: self.flag_vimgrep,
+            max_count: self.flag_max_count.map(|max| max as u64),
             maxdepth: self.flag_maxdepth,
             mmap: mmap,
             no_ignore: no_ignore,
@@ -629,6 +635,11 @@ impl Args {
         }
     }
 
+    /// Returns true if the given arguments are known to never produce a match.
+    pub fn never_match(&self) -> bool {
+        self.max_count == Some(0)
+    }
+
     /// Create a new buffer for use with searching.
     #[cfg(not(windows))]
     pub fn outbuf(&self) -> ColoredTerminal<term::TerminfoTerminal<Vec<u8>>> {
@@ -677,6 +688,7 @@ impl Args {
             .eol(self.eol)
             .line_number(self.line_number)
             .invert_match(self.invert_match)
+            .max_count(self.max_count)
             .mmap(self.mmap)
             .quiet(self.quiet)
             .text(self.text)
diff --git a/src/main.rs b/src/main.rs
index 276ee059..33f99ad9 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -73,6 +73,9 @@ fn main() {
 }
 
 fn run(args: Arc<Args>) -> Result<u64> {
+    if args.never_match() {
+        return Ok(0);
+    }
     {
         let args = args.clone();
         ctrlc::set_handler(move || {
diff --git a/src/search_buffer.rs b/src/search_buffer.rs
index 6a32a631..c7c3bca0 100644
--- a/src/search_buffer.rs
+++ b/src/search_buffer.rs
@@ -81,6 +81,14 @@ impl<'a, W: Send + Terminal> BufferSearcher<'a, W> {
         self
     }
 
+    /// Limit the number of matches to the given count.
+    ///
+    /// The default is None, which corresponds to no limit.
+    pub fn max_count(mut self, count: Option<u64>) -> Self {
+        self.opts.max_count = count;
+        self
+    }
+
     /// If enabled, don't show any output and quit searching after the first
     /// match is found.
     pub fn quiet(mut self, yes: bool) -> Self {
@@ -111,11 +119,11 @@ impl<'a, W: Send + Terminal> BufferSearcher<'a, W> {
                 self.print_match(m.start(), m.end());
             }
             last_end = m.end();
-            if self.opts.stop_after_first_match() {
+            if self.opts.terminate(self.match_count) {
                 break;
             }
         }
-        if self.opts.invert_match {
+        if self.opts.invert_match && !self.opts.terminate(self.match_count) {
             let upto = self.buf.len();
             self.print_inverted_matches(last_end, upto);
         }
@@ -146,6 +154,9 @@ impl<'a, W: Send + Terminal> 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) {
+                return;
+            }
             self.print_match(s, e);
         }
     }
@@ -267,6 +278,26 @@ and exhibited clearly, with a label attached.\
     }
 
     #[test]
+    fn max_count() {
+        let (count, out) = search(
+            "Sherlock", SHERLOCK, |s| s.max_count(Some(1)));
+        assert_eq!(1, count);
+        assert_eq!(out, "\
+/baz.rs:For the Doctor Watsons of this world, as opposed to the Sherlock
+");
+    }
+
+    #[test]
+    fn invert_match_max_count() {
+        let (count, out) = search(
+            "zzzz", SHERLOCK, |s| s.invert_match(true).max_count(Some(1)));
+        assert_eq!(1, count);
+        assert_eq!(out, "\
+/baz.rs:For the Doctor Watsons of this world, as opposed to the Sherlock
+");
+    }
+
+    #[test]
     fn invert_match() {
         let (count, out) = search(
             "Sherlock", SHERLOCK, |s| s.invert_match(true));
diff --git a/src/search_stream.rs b/src/search_stream.rs
index cbd7a63e..5b3880b6 100644
--- a/src/search_stream.rs
+++ b/src/search_stream.rs
@@ -86,6 +86,7 @@ pub struct Options {
     pub eol: u8,
     pub invert_match: bool,
     pub line_number: bool,
+    pub max_count: Option<u64>,
     pub quiet: bool,
     pub text: bool,
 }
@@ -100,6 +101,7 @@ impl Default for Options {
             eol: b'\n',
             invert_match: false,
             line_number: false,
+            max_count: None,
             quiet: false,
             text: false,
         }
@@ -119,6 +121,17 @@ impl Options {
     pub fn stop_after_first_match(&self) -> bool {
         self.files_with_matches || self.quiet
     }
+
+    /// Returns true if the search should terminate based on the match count.
+    pub fn terminate(&self, match_count: u64) -> bool {
+        if match_count > 0 && self.stop_after_first_match() {
+            return true;
+        }
+        if self.max_count.map_or(false, |max| match_count >= max) {
+            return true;
+        }
+        false
+    }
 }
 
 impl<'a, R: io::Read, W: Terminal + Send> Searcher<'a, R, W> {
@@ -207,6 +220,14 @@ impl<'a, R: io::Read, W: Terminal + Send> Searcher<'a, R, W> {
         self
     }
 
+    /// Limit the number of matches to the given count.
+    ///
+    /// The default is None, which corresponds to no limit.
+    pub fn max_count(mut self, count: Option<u64>) -> Self {
+        self.opts.max_count = count;
+        self
+    }
+
     /// If enabled, don't show any output and quit searching after the first
     /// match is found.
     pub fn quiet(mut self, yes: bool) -> Self {
@@ -282,7 +303,7 @@ impl<'a, R: io::Read, W: Terminal + Send> Searcher<'a, R, W> {
 
     #[inline(always)]
     fn terminate(&self) -> bool {
-        self.match_count > 0 && self.opts.stop_after_first_match()
+        self.opts.terminate(self.match_count)
     }
 
     #[inline(always)]
@@ -319,6 +340,9 @@ impl<'a, R: io::Read, W: Terminal + Send> Searcher<'a, R, W> {
         debug_assert!(self.opts.invert_match);
         let mut it = IterLines::new(self.opts.eol, self.inp.pos);
         while let Some((start, end)) = it.next(&self.inp.buf[..upto]) {
+            if self.terminate() {
+                return;
+            }
             self.print_match(start, end);
             self.inp.pos = end;
         }
@@ -963,6 +987,26 @@ fn main() {
     }
 
     #[test]
+    fn max_count() {
+        let (count, out) = search_smallcap(
+            "Sherlock", SHERLOCK, |s| s.max_count(Some(1)));
+        assert_eq!(1, count);
+        assert_eq!(out, "\
+/baz.rs:For the Doctor Watsons of this world, as opposed to the Sherlock
+");
+    }
+
+    #[test]
+    fn invert_match_max_count() {
+        let (count, out) = search(
+            "zzzz", SHERLOCK, |s| s.invert_match(true).max_count(Some(1)));
+        assert_eq!(1, count);
+        assert_eq!(out, "\
+/baz.rs:For the Doctor Watsons of this world, as opposed to the Sherlock
+");
+    }
+
+    #[test]
     fn invert_match() {
         let (count, out) = search_smallcap(
             "Sherlock", SHERLOCK, |s| s.invert_match(true));
diff --git a/src/worker.rs b/src/worker.rs
index 797fe9d7..bc8a62b3 100644
--- a/src/worker.rs
+++ b/src/worker.rs
@@ -34,6 +34,7 @@ struct Options {
     eol: u8,
     invert_match: bool,
     line_number: bool,
+    max_count: Option<u64>,
     quiet: bool,
     text: bool,
 }
@@ -49,6 +50,7 @@ impl Default for Options {
             eol: b'\n',
             invert_match: false,
             line_number: false,
+            max_count: None,
             quiet: false,
             text: false,
         }
@@ -128,6 +130,14 @@ impl WorkerBuilder {
         self
     }
 
+    /// Limit the number of matches to the given count.
+    ///
+    /// The default is None, which corresponds to no limit.
+    pub fn max_count(mut self, count: Option<u64>) -> Self {
+        self.opts.max_count = count;
+        self
+    }
+
     /// If enabled, try to use memory maps for searching if possible.
     pub fn mmap(mut self, yes: bool) -> Self {
         self.opts.mmap = yes;
@@ -217,6 +227,7 @@ impl Worker {
             .eol(self.opts.eol)
             .line_number(self.opts.line_number)
             .invert_match(self.opts.invert_match)
+            .max_count(self.opts.max_count)
             .quiet(self.opts.quiet)
             .text(self.opts.text)
             .run()
@@ -246,6 +257,7 @@ impl Worker {
             .eol(self.opts.eol)
             .line_number(self.opts.line_number)
             .invert_match(self.opts.invert_match)
+            .max_count(self.opts.max_count)
             .quiet(self.opts.quiet)
             .text(self.opts.text)
             .run())
diff --git a/tests/tests.rs b/tests/tests.rs
index bf6f471d..59cefb59 100644
--- a/tests/tests.rs
+++ b/tests/tests.rs
@@ -1071,6 +1071,21 @@ clean!(feature_109_case_sensitive_part2, "test", ".",
     wd.assert_err(&mut cmd);
 });
 
+// See: https://github.com/BurntSushi/ripgrep/issues/159
+clean!(feature_159_works, "test", ".", |wd: WorkDir, mut cmd: Command| {
+    wd.create("foo", "test\ntest");
+    cmd.arg("-m1");
+    let lines: String = wd.stdout(&mut cmd);
+    assert_eq!(lines, "foo:test\n");
+});
+
+// See: https://github.com/BurntSushi/ripgrep/issues/159
+clean!(feature_159_zero_max, "test", ".", |wd: WorkDir, mut cmd: Command| {
+    wd.create("foo", "test\ntest");
+    cmd.arg("-m0");
+    wd.assert_err(&mut cmd);
+});
+
 #[test]
 fn binary_nosearch() {
     let wd = WorkDir::new("binary_nosearch");