update

src/main.rs

@@ -62,28 +62,42 @@ fn run(args: &Args) -> Result<u64> {
     } else {
         let searcher =
             try!(LineSearcherBuilder::new(&args.arg_pattern).create());
-        run_mmap(args, &searcher)
+        if args.flag_count {
+            run_mmap_count_only(args, &searcher)
+        } else {
+            run_mmap(args, &searcher)
+        }
     }
 }
 
+#[inline(never)]
 fn run_mmap(args: &Args, searcher: &LineSearcher) -> Result<u64> {
     use memmap::{Mmap, Protection};
 
     assert!(args.arg_file.len() == 1);
     let mut wtr = io::BufWriter::new(io::stdout());
-    let mut count = 0;
     let mmap = try!(Mmap::open_path(&args.arg_file[0], Protection::Read));
     let text = unsafe { mmap.as_slice() };
+
+    let mut count = 0;
     for m in searcher.search(text) {
-        if !args.flag_count {
-            try!(wtr.write(&text[m.start..m.end]));
-            try!(wtr.write(b"\n"));
-        }
+        try!(wtr.write(&text[m.start..m.end]));
+        try!(wtr.write(b"\n"));
         count += 1;
     }
-    if args.flag_count {
-        try!(writeln!(wtr, "{}", count));
-    }
+    Ok(count)
+}
+
+#[inline(never)]
+fn run_mmap_count_only(args: &Args, searcher: &LineSearcher) -> Result<u64> {
+    use memmap::{Mmap, Protection};
+
+    assert!(args.arg_file.len() == 1);
+    let mut wtr = io::BufWriter::new(io::stdout());
+    let mmap = try!(Mmap::open_path(&args.arg_file[0], Protection::Read));
+    let text = unsafe { mmap.as_slice() };
+    let count = searcher.search(text).last().map_or(0, |m| m.count + 1);
+    try!(writeln!(wtr, "{}", count));
     Ok(count)
 }
 

src/search.rs

@@ -1,3 +1,5 @@
+use std::cmp;
+
 use memchr::{memchr, memrchr};
 use regex::bytes::Regex;
 use syntax;
@@ -82,7 +84,7 @@ impl LineSearcher {
 pub struct Match {
     pub start: usize,
     pub end: usize,
-    pub count: usize,
+    pub count: u64,
     pub line: Option<usize>,
     pub locations: Vec<(usize, usize)>,
 }
@@ -91,14 +93,12 @@ pub struct Iter<'b, 's> {
     searcher: &'s LineSearcher,
     buf: &'b [u8],
     start: usize,
-    count: usize,
+    count: u64,
 }
 
 impl<'b, 's> Iter<'b, 's> {
+    #[inline(always)] // reduces constant overhead
     fn next_line_match(&mut self) -> Option<(usize, usize)> {
-        if self.start >= self.buf.len() {
-            return None;
-        }
         if let Some(ref req) = self.searcher.required {
             while self.start < self.buf.len() {
                 let e = match req.shortest_match(&self.buf[self.start..]) {
@@ -116,11 +116,9 @@ impl<'b, 's> Iter<'b, 's> {
             }
             None
         } else {
-            let e = match self.searcher.re.shortest_match(&self.buf[self.start..]) {
-                None => return None,
-                Some(e) => self.start + e,
-            };
-            Some(self.find_line(e, e))
+            self.searcher.re
+                .shortest_match(&self.buf[self.start..])
+                .map(|e| self.find_line(self.start + e, self.start + e))
         }
     }
 
@@ -140,21 +138,23 @@ impl<'b, 's> Iter<'b, 's> {
 impl<'b, 's> Iterator for Iter<'b, 's> {
     type Item = Match;
 
+    #[inline(always)] // reduces constant overhead
     fn next(&mut self) -> Option<Match> {
-        let (prevnl, nextnl) = match self.next_line_match() {
-            None => return None,
-            Some((s, e)) => (s, e),
-        };
-        let count = self.count;
-        self.start = nextnl + 1;
-        self.count += 1;
-        Some(Match {
-            start: prevnl,
-            end: nextnl,
-            count: count,
-            line: None,
-            locations: vec![],
-        })
+        match self.next_line_match() {
+            None => None,
+            Some((prevnl, nextnl)) => {
+                let count = self.count;
+                self.start = cmp::min(self.buf.len(), nextnl + 1);
+                self.count += 1;
+                Some(Match {
+                    start: prevnl,
+                    end: nextnl,
+                    count: count,
+                    line: None,
+                    locations: vec![],
+                })
+            }
+        }
     }
 }