From 09e1f9c330ec99de06c834b2de95c378bc6d00d3 Mon Sep 17 00:00:00 2001
From: Daniel <kingdread@gmx.de>
Date: Fri, 28 Aug 2020 15:49:12 +0200
Subject: add --count/-n flag

---
 CHANGELOG.md              |  2 ++
 src/main.rs               |  8 ++++++++
 src/output/aggregators.rs | 30 +++++++++++++++++++++++++++++-
 src/output/mod.rs         |  4 +++-
 4 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 5755644..6dca5e2 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,6 +3,8 @@
 All notable changes to this project will be documented in this file.
 
 ## Unreleased
+### Added
+- The `--count`/`-n` flag to only output the number of matching logs.
 
 ## 1.3.0 - 2020-07-24
 ### Added
diff --git a/src/main.rs b/src/main.rs
index 0915f42..0ff9e44 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -92,6 +92,14 @@ pub struct Opt {
     #[structopt(short = "l", long = "files-with-matches")]
     file_name_only: bool,
 
+    /// Only output the number of matching logs.
+    #[structopt(
+        short = "n",
+        long = "count",
+        conflicts_with_all = &["sorting", "file-name-only"],
+    )]
+    count: bool,
+
     /// Disable colored output.
     #[structopt(long = "no-color")]
     no_color: bool,
diff --git a/src/output/aggregators.rs b/src/output/aggregators.rs
index 83171eb..34ecea4 100644
--- a/src/output/aggregators.rs
+++ b/src/output/aggregators.rs
@@ -8,7 +8,11 @@
 //! an Aggregator must make sure that the data is protected by a mutex or similar.
 use super::{super::LogResult, formats::Format, sorting::Sorting};
 
-use std::{io::Write, sync::Mutex};
+use std::{
+    io::Write,
+    sync::atomic::{AtomicU32, Ordering},
+    sync::Mutex,
+};
 
 pub trait Aggregator: Sync {
     fn push_item(&self, item: LogResult, format: &dyn Format, stream: &mut dyn Write);
@@ -62,3 +66,27 @@ impl Aggregator for SortedOutput {
         stream.flush().unwrap();
     }
 }
+
+/// An aggregator that just counts how many logs have been found.
+#[derive(Debug)]
+pub struct CountingOutput {
+    counter: AtomicU32,
+}
+
+impl CountingOutput {
+    pub fn new() -> Self {
+        CountingOutput {
+            counter: AtomicU32::new(0),
+        }
+    }
+}
+
+impl Aggregator for CountingOutput {
+    fn push_item(&self, _: LogResult, _: &dyn Format, _: &mut dyn Write) {
+        self.counter.fetch_add(1, Ordering::SeqCst);
+    }
+
+    fn finish(self: Box<Self>, _: &dyn Format, stream: &mut dyn Write) {
+        writeln!(stream, "{}", self.counter.into_inner()).unwrap();
+    }
+}
diff --git a/src/output/mod.rs b/src/output/mod.rs
index c1cd787..e0b28cc 100644
--- a/src/output/mod.rs
+++ b/src/output/mod.rs
@@ -14,7 +14,9 @@ use self::{aggregators::Aggregator, formats::Format};
 /// Build an pipeline for the given command line options.
 pub fn build_pipeline(opt: &Opt) -> Pipeline {
     let stream = io::stdout();
-    let aggregator: Box<dyn Aggregator> = if let Some(sorting) = &opt.sorting {
+    let aggregator: Box<dyn Aggregator> = if opt.count {
+        Box::new(aggregators::CountingOutput::new())
+    } else if let Some(sorting) = &opt.sorting {
         Box::new(aggregators::SortedOutput::new(sorting.clone()))
     } else {
         Box::new(aggregators::WriteThrough)
-- 
cgit v1.2.3