From b249b2a81b7c666727316ab680d649cdf404d18a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cl=C3=A9ment=20Renault?= <renault.cle@gmail.com>
Date: Mon, 24 Jun 2019 14:19:39 +0200
Subject: [PATCH] feat: Support removing specific synonym alternatives

---
 .../src/database/synonyms_addition.rs         | 11 +--
 .../src/database/synonyms_deletion.rs         | 78 ++++++++++++++++---
 2 files changed, 74 insertions(+), 15 deletions(-)

diff --git a/meilidb-data/src/database/synonyms_addition.rs b/meilidb-data/src/database/synonyms_addition.rs
index d0f6c2160..6e16ab97b 100644
--- a/meilidb-data/src/database/synonyms_addition.rs
+++ b/meilidb-data/src/database/synonyms_addition.rs
@@ -2,7 +2,6 @@ use std::collections::BTreeMap;
 use std::sync::Arc;
 
 use fst::{SetBuilder, set::OpBuilder};
-use meilidb_tokenizer::is_cjk;
 use meilidb_core::normalize_str;
 use sdset::SetBuf;
 
@@ -19,11 +18,13 @@ impl<'a> SynonymsAddition<'a> {
         SynonymsAddition { inner, synonyms: BTreeMap::new() }
     }
 
-    pub fn add_synonym<I>(&mut self, synonym: String, alternatives: I)
-    where I: Iterator<Item=String>,
+    pub fn add_synonym<S, T, I>(&mut self, synonym: S, alternatives: I)
+    where S: AsRef<str>,
+          T: AsRef<str>,
+          I: Iterator<Item=T>,
     {
-        let mut synonym = normalize_str(&synonym);
-        let alternatives = alternatives.map(|s| s.to_lowercase());
+        let synonym = normalize_str(synonym.as_ref());
+        let alternatives = alternatives.map(|s| s.as_ref().to_lowercase());
         self.synonyms.entry(synonym).or_insert_with(Vec::new).extend(alternatives);
     }
 
diff --git a/meilidb-data/src/database/synonyms_deletion.rs b/meilidb-data/src/database/synonyms_deletion.rs
index 8720d4b5c..c07e92859 100644
--- a/meilidb-data/src/database/synonyms_deletion.rs
+++ b/meilidb-data/src/database/synonyms_deletion.rs
@@ -1,23 +1,41 @@
-use std::collections::BTreeSet;
+use std::collections::BTreeMap;
+use std::iter::FromIterator;
 use std::sync::Arc;
 
 use fst::{SetBuilder, set::OpBuilder};
+use meilidb_core::normalize_str;
+use sdset::SetBuf;
 
 use crate::database::index::InnerIndex;
 use super::{Error, Index};
 
 pub struct SynonymsDeletion<'a> {
     inner: &'a Index,
-    synonyms: BTreeSet<String>,
+    synonyms: BTreeMap<String, Option<Vec<String>>>,
 }
 
 impl<'a> SynonymsDeletion<'a> {
     pub fn new(inner: &'a Index) -> SynonymsDeletion<'a> {
-        SynonymsDeletion { inner, synonyms: BTreeSet::new() }
+        SynonymsDeletion { inner, synonyms: BTreeMap::new() }
     }
 
-    pub fn delete_alternatives_of<I>(&mut self, synonym: String) {
-        self.synonyms.insert(synonym);
+    pub fn delete_all_alternatives_of<S: AsRef<str>>(&mut self, synonym: S) {
+        let synonym = normalize_str(synonym.as_ref());
+        self.synonyms.insert(synonym, None);
+    }
+
+    pub fn delete_specific_alternatives_of<S, T, I>(&mut self, synonym: S, alternatives: I)
+    where S: AsRef<str>,
+          T: AsRef<str>,
+          I: Iterator<Item=T>,
+    {
+        let synonym = normalize_str(synonym.as_ref());
+        let value = self.synonyms.entry(synonym).or_insert(None);
+        let alternatives = alternatives.map(|s| s.as_ref().to_lowercase());
+        match value {
+            Some(v) => v.extend(alternatives),
+            None => *value = Some(Vec::from_iter(alternatives)),
+        }
     }
 
     pub fn finalize(self) -> Result<(), Error> {
@@ -25,14 +43,54 @@ impl<'a> SynonymsDeletion<'a> {
         let synonyms = &lease_inner.raw.synonyms;
         let main = &lease_inner.raw.main;
 
-        let mut synonyms_builder = SetBuilder::memory();
+        let mut delete_whole_synonym_builder = SetBuilder::memory();
 
-        for synonym in self.synonyms {
-            synonyms_builder.insert(&synonym).unwrap();
-            synonyms.del_alternatives_of(synonym.as_bytes())?;
+        for (synonym, alternatives) in self.synonyms {
+            match alternatives {
+                Some(alternatives) => {
+                    let prev_alternatives = synonyms.alternatives_to(synonym.as_bytes())?;
+                    let prev_alternatives = match prev_alternatives {
+                        Some(alternatives) => alternatives,
+                        None => continue,
+                    };
+
+                    let delta_alternatives = {
+                        let alternatives = SetBuf::from_dirty(alternatives);
+                        let mut builder = SetBuilder::memory();
+                        builder.extend_iter(alternatives).unwrap();
+                        builder.into_inner()
+                            .and_then(fst::Set::from_bytes)
+                            .unwrap()
+                    };
+
+                    let op = OpBuilder::new()
+                        .add(prev_alternatives.stream())
+                        .add(delta_alternatives.stream())
+                        .difference();
+
+                    let (alternatives, empty_alternatives) = {
+                        let mut builder = SetBuilder::memory();
+                        let len = builder.get_ref().len();
+                        builder.extend_stream(op).unwrap();
+                        let is_empty = len == builder.get_ref().len();
+                        let alternatives = builder.into_inner().unwrap();
+                        (alternatives, is_empty)
+                    };
+
+                    if empty_alternatives {
+                        delete_whole_synonym_builder.insert(synonym.as_bytes())?;
+                    } else {
+                        synonyms.set_alternatives_to(synonym.as_bytes(), alternatives)?;
+                    }
+                },
+                None => {
+                    delete_whole_synonym_builder.insert(&synonym).unwrap();
+                    synonyms.del_alternatives_of(synonym.as_bytes())?;
+                }
+            }
         }
 
-        let delta_synonyms = synonyms_builder
+        let delta_synonyms = delete_whole_synonym_builder
             .into_inner()
             .and_then(fst::Set::from_bytes)
             .unwrap();