From d2d0873fdb072c39520c543e37b959fe9d94d46f Mon Sep 17 00:00:00 2001
From: Paul Masurel <paul@quickwit.io>
Date: Thu, 18 Mar 2021 17:28:09 +0900
Subject: [PATCH] Added support for Option<TCollector>.

---
 CHANGELOG.md                         |  1 +
 src/collector/histogram_collector.rs | 10 ++---
 src/collector/mod.rs                 | 55 ++++++++++++++++++++++++++++
 src/collector/tests.rs               | 33 +++++++++++++++--
 4 files changed, 91 insertions(+), 8 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 6c7b7d295..ed47485d7 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,6 +3,7 @@ Tantivy 0.15.0
 - API Changes. Using Range instead of (start, end) in the API and internals (`FileSlice`, `OwnedBytes`, `Snippets`, ...)
   This change is breaking but migration is trivial.
 - Added an Histogram collector. (@fulmicoton) #994
+- Added support for Option<TCollector>.  (@fulmicoton) 
 
 
 Tantivy 0.14.0
diff --git a/src/collector/histogram_collector.rs b/src/collector/histogram_collector.rs
index 4d28fe8c2..6db5c7167 100644
--- a/src/collector/histogram_collector.rs
+++ b/src/collector/histogram_collector.rs
@@ -1,8 +1,8 @@
-use fastdivide::DividerU64;
 use crate::collector::{Collector, SegmentCollector};
 use crate::fastfield::{FastFieldReader, FastValue};
 use crate::schema::{Field, Type};
 use crate::{DocId, Score};
+use fastdivide::DividerU64;
 
 /// Histogram builds an histogram of the values of a fastfield for the
 /// collected DocSet.
@@ -147,12 +147,12 @@ fn add_vecs(mut vals_list: Vec<Vec<u64>>, len: usize) -> Vec<u64> {
 
 #[cfg(test)]
 mod tests {
-    use super::{add_vecs, HistogramComputer, HistogramCollector};
-    use fastdivide::DividerU64;
-    use query::AllQuery;
+    use super::{add_vecs, HistogramCollector, HistogramComputer};
     use crate::chrono::{TimeZone, Utc};
     use crate::schema::{Schema, FAST};
-    use crate::{Index, doc, query};
+    use crate::{doc, query, Index};
+    use fastdivide::DividerU64;
+    use query::AllQuery;
 
     #[test]
     fn test_add_histograms_simple() {
diff --git a/src/collector/mod.rs b/src/collector/mod.rs
index 79f4efa12..f55b946ee 100644
--- a/src/collector/mod.rs
+++ b/src/collector/mod.rs
@@ -193,6 +193,61 @@ pub trait Collector: Sync + Send {
     }
 }
 
+impl<TSegmentCollector: SegmentCollector> SegmentCollector for Option<TSegmentCollector> {
+    type Fruit = Option<TSegmentCollector::Fruit>;
+
+    fn collect(&mut self, doc: DocId, score: Score) {
+        if let Some(segment_collector) = self {
+            segment_collector.collect(doc, score);
+        }
+    }
+
+    fn harvest(self) -> Self::Fruit {
+        self.map(|segment_collector| segment_collector.harvest())
+    }
+}
+
+impl<TCollector: Collector> Collector for Option<TCollector> {
+    type Fruit = Option<TCollector::Fruit>;
+
+    type Child = Option<<TCollector as Collector>::Child>;
+
+    fn for_segment(
+        &self,
+        segment_local_id: SegmentLocalId,
+        segment: &SegmentReader,
+    ) -> crate::Result<Self::Child> {
+        Ok(if let Some(inner) = self {
+            let inner_segment_collector = inner.for_segment(segment_local_id, segment)?;
+            Some(inner_segment_collector)
+        } else {
+            None
+        })
+    }
+
+    fn requires_scoring(&self) -> bool {
+        self.as_ref()
+            .map(|inner| inner.requires_scoring())
+            .unwrap_or(false)
+    }
+
+    fn merge_fruits(
+        &self,
+        segment_fruits: Vec<<Self::Child as SegmentCollector>::Fruit>,
+    ) -> crate::Result<Self::Fruit> {
+        if let Some(inner) = self.as_ref() {
+            let inner_segment_fruits: Vec<_> = segment_fruits
+                .into_iter()
+                .flat_map(|fruit_opt| fruit_opt.into_iter())
+                .collect();
+            let fruit = inner.merge_fruits(inner_segment_fruits)?;
+            Ok(Some(fruit))
+        } else {
+            Ok(None)
+        }
+    }
+}
+
 /// The `SegmentCollector` is the trait in charge of defining the
 /// collect operation at the scale of the segment.
 ///
diff --git a/src/collector/tests.rs b/src/collector/tests.rs
index ab03a723a..450d1cd90 100644
--- a/src/collector/tests.rs
+++ b/src/collector/tests.rs
@@ -3,13 +3,13 @@ use crate::core::SegmentReader;
 use crate::fastfield::BytesFastFieldReader;
 use crate::fastfield::FastFieldReader;
 use crate::schema::Field;
-use crate::DocAddress;
 use crate::DocId;
 use crate::Score;
 use crate::SegmentLocalId;
+use crate::{DocAddress, Document, Searcher};
 
-use crate::collector::{FilterCollector, TopDocs};
-use crate::query::QueryParser;
+use crate::collector::{Count, FilterCollector, TopDocs};
+use crate::query::{AllQuery, QueryParser};
 use crate::schema::{Schema, FAST, TEXT};
 use crate::DateTime;
 use crate::{doc, Index};
@@ -268,3 +268,30 @@ impl SegmentCollector for BytesFastFieldSegmentCollector {
         self.vals
     }
 }
+
+fn make_test_searcher() -> crate::Result<crate::LeasedItem<Searcher>> {
+    let schema = Schema::builder().build();
+    let index = Index::create_in_ram(schema);
+    let mut index_writer = index.writer_for_tests()?;
+    index_writer.add_document(Document::default());
+    index_writer.add_document(Document::default());
+    index_writer.commit()?;
+    Ok(index.reader()?.searcher())
+}
+
+#[test]
+fn test_option_collector_some() -> crate::Result<()> {
+    let searcher = make_test_searcher()?;
+    let counts = searcher.search(&AllQuery, &Some(Count))?;
+    assert_eq!(counts, Some(2));
+    Ok(())
+}
+
+#[test]
+fn test_option_collector_none() -> crate::Result<()> {
+    let searcher = make_test_searcher()?;
+    let none_collector: Option<Count> = None;
+    let counts = searcher.search(&AllQuery, &none_collector)?;
+    assert_eq!(counts, None);
+    Ok(())
+}