From 1e47f9b3ffa7f55071274d95d542bb183c7862cb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cl=C3=A9ment=20Renault?= <clement@meilisearch.com>
Date: Fri, 19 Feb 2021 15:32:14 +0100
Subject: [PATCH] Introduce the Words criterion

---
 milli/src/search/criteria/mod.rs   |   1 +
 milli/src/search/criteria/typo.rs  |   4 +-
 milli/src/search/criteria/words.rs | 109 +++++++++++++++++++++++++++++
 3 files changed, 112 insertions(+), 2 deletions(-)
 create mode 100644 milli/src/search/criteria/words.rs

diff --git a/milli/src/search/criteria/mod.rs b/milli/src/search/criteria/mod.rs
index e7549cae1..5cc803dee 100644
--- a/milli/src/search/criteria/mod.rs
+++ b/milli/src/search/criteria/mod.rs
@@ -8,6 +8,7 @@ use roaring::RoaringBitmap;
 use super::query_tree::{Operation, Query, QueryKind};
 
 pub mod typo;
+pub mod words;
 
 pub trait Criterion {
     fn next(&mut self) -> anyhow::Result<Option<CriterionResult>>;
diff --git a/milli/src/search/criteria/typo.rs b/milli/src/search/criteria/typo.rs
index 659292619..a62616f08 100644
--- a/milli/src/search/criteria/typo.rs
+++ b/milli/src/search/criteria/typo.rs
@@ -13,7 +13,7 @@ pub struct Typo<'t> {
     number_typos: u8,
     candidates: Candidates,
     bucket_candidates: RoaringBitmap,
-    parent: Option<Box<dyn Criterion>>,
+    parent: Option<Box<dyn Criterion + 't>>,
     candidates_cache: HashMap<(Operation, u8), RoaringBitmap>,
     typo_cache: HashMap<(String, bool, u8), Vec<(String, u8)>>,
 }
@@ -39,7 +39,7 @@ impl<'t> Typo<'t> {
 
     pub fn new(
         ctx: &'t dyn Context,
-        parent: Box<dyn Criterion>,
+        parent: Box<dyn Criterion + 't>,
     ) -> anyhow::Result<Self> where Self: Sized
     {
         Ok(Typo {
diff --git a/milli/src/search/criteria/words.rs b/milli/src/search/criteria/words.rs
new file mode 100644
index 000000000..ade370fda
--- /dev/null
+++ b/milli/src/search/criteria/words.rs
@@ -0,0 +1,109 @@
+use std::collections::HashMap;
+use std::mem::take;
+
+use roaring::RoaringBitmap;
+
+use crate::search::query_tree::Operation;
+use super::{Candidates, Criterion, CriterionResult, Context};
+
+pub struct Words<'t> {
+    ctx: &'t dyn Context,
+    query_trees: Vec<Operation>,
+    candidates: Candidates,
+    bucket_candidates: RoaringBitmap,
+    parent: Option<Box<dyn Criterion + 't>>,
+    candidates_cache: HashMap<(Operation, u8), RoaringBitmap>,
+}
+
+impl<'t> Words<'t> {
+    pub fn initial(
+        ctx: &'t dyn Context,
+        query_tree: Option<Operation>,
+        candidates: Option<RoaringBitmap>,
+    ) -> anyhow::Result<Self> where Self: Sized
+    {
+        Ok(Words {
+            ctx,
+            query_trees: query_tree.map(explode_query_tree).unwrap_or_default(),
+            candidates: candidates.map_or_else(Candidates::default, Candidates::Allowed),
+            bucket_candidates: RoaringBitmap::new(),
+            parent: None,
+            candidates_cache: HashMap::default(),
+        })
+    }
+
+    pub fn new(
+        ctx: &'t dyn Context,
+        parent: Box<dyn Criterion + 't>,
+    ) -> anyhow::Result<Self> where Self: Sized
+    {
+        Ok(Words {
+            ctx,
+            query_trees: Vec::default(),
+            candidates: Candidates::default(),
+            bucket_candidates: RoaringBitmap::new(),
+            parent: Some(parent),
+            candidates_cache: HashMap::default(),
+        })
+    }
+}
+
+impl<'t> Criterion for Words<'t> {
+    fn next(&mut self) -> anyhow::Result<Option<CriterionResult>> {
+        use Candidates::{Allowed, Forbidden};
+
+        loop {
+            match (self.query_trees.pop(), &mut self.candidates) {
+                (_, Allowed(candidates)) if candidates.is_empty() => {
+                    self.query_trees = Vec::new();
+                    self.candidates = Candidates::default();
+                },
+                (Some(qt), Allowed(candidates)) => {
+                    let bucket_candidates = match self.parent {
+                        Some(_) => take(&mut self.bucket_candidates),
+                        None => candidates.clone(),
+                    };
+
+                    return Ok(Some(CriterionResult {
+                        query_tree: Some(qt),
+                        candidates: candidates.clone(),
+                        bucket_candidates,
+                    }));
+                },
+                (Some(_qt), Forbidden(candidates)) => {
+                    todo!()
+                },
+                (None, Allowed(_)) => {
+                    let candidates = take(&mut self.candidates).into_inner();
+                    return Ok(Some(CriterionResult {
+                        query_tree: None,
+                        candidates: candidates.clone(),
+                        bucket_candidates: candidates,
+                    }));
+                },
+                (None, Forbidden(_)) => {
+                    match self.parent.as_mut() {
+                        Some(parent) => {
+                            match parent.next()? {
+                                Some(CriterionResult { query_tree, candidates, bucket_candidates }) => {
+                                    self.query_trees = query_tree.map(explode_query_tree).unwrap_or_default();
+                                    self.candidates = Candidates::Allowed(candidates);
+                                    self.bucket_candidates = bucket_candidates;
+                                },
+                                None => return Ok(None),
+                            }
+                        },
+                        None => return Ok(None),
+                    }
+                },
+            }
+        }
+    }
+}
+
+fn explode_query_tree(query_tree: Operation) -> Vec<Operation> {
+    match query_tree {
+        Operation::Or(true, ops) => ops,
+        otherwise => vec![otherwise],
+    }
+}