From fe82516f9f9314d5615bdcc3e2981a97802a2a1a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cl=C3=A9ment=20Renault?= <clement@meilisearch.com>
Date: Sun, 22 Nov 2020 17:28:41 +0100
Subject: [PATCH] Use the ExternalDocumentsIds in the Index struct

---
 src/external_documents_ids.rs |  9 ++++++++-
 src/index.rs                  | 37 +++++++++++++++++++++++++----------
 2 files changed, 35 insertions(+), 11 deletions(-)

diff --git a/src/external_documents_ids.rs b/src/external_documents_ids.rs
index f8765b57e..7c81cdde8 100644
--- a/src/external_documents_ids.rs
+++ b/src/external_documents_ids.rs
@@ -12,7 +12,14 @@ impl<'a> ExternalDocumentsIds<'a> {
         ExternalDocumentsIds { hard, soft }
     }
 
-    pub fn get<A: AsRef<str>>(&self, external_id: A) -> Option<u32> {
+    pub fn into_static(self) -> ExternalDocumentsIds<'static> {
+        ExternalDocumentsIds {
+            hard: self.hard.map_data(|c| Cow::Owned(c.into_owned())).unwrap(),
+            soft: self.soft.map_data(|c| Cow::Owned(c.into_owned())).unwrap(),
+        }
+    }
+
+    pub fn get<A: AsRef<[u8]>>(&self, external_id: A) -> Option<u32> {
         let external_id = external_id.as_ref();
         match self.soft.get(external_id).or_else(|| self.hard.get(external_id)) {
             // u64 MAX means deleted in the soft fst map
diff --git a/src/index.rs b/src/index.rs
index 8da7940a7..76ddd4fef 100644
--- a/src/index.rs
+++ b/src/index.rs
@@ -7,6 +7,7 @@ use heed::types::*;
 use heed::{PolyDatabase, Database, RwTxn, RoTxn};
 use roaring::RoaringBitmap;
 
+use crate::external_documents_ids::ExternalDocumentsIds;
 use crate::facet::FacetType;
 use crate::fields_ids_map::FieldsIdsMap;
 use crate::Search;
@@ -22,7 +23,8 @@ pub const FACETED_FIELDS_KEY: &str = "faceted-fields";
 pub const FIELDS_IDS_MAP_KEY: &str = "fields-ids-map";
 pub const PRIMARY_KEY_KEY: &str = "primary-key";
 pub const SEARCHABLE_FIELDS_KEY: &str = "searchable-fields";
-pub const EXTERNAL_DOCUMENTS_IDS_KEY: &str = "external-documents-ids";
+pub const HARD_EXTERNAL_DOCUMENTS_IDS_KEY: &str = "hard-external-documents-ids";
+pub const SOFT_EXTERNAL_DOCUMENTS_IDS_KEY: &str = "soft-external-documents-ids";
 pub const WORDS_FST_KEY: &str = "words-fst";
 
 #[derive(Clone)]
@@ -121,18 +123,33 @@ impl Index {
 
     /* external documents ids */
 
-    /// Writes the external documents ids, it is a byte slice (i.e. `[u8]`)
-    /// and refers to an internal id (i.e. `u32`).
-    pub fn put_external_documents_ids<A: AsRef<[u8]>>(&self, wtxn: &mut RwTxn, fst: &fst::Map<A>) -> heed::Result<()> {
-        self.main.put::<_, Str, ByteSlice>(wtxn, EXTERNAL_DOCUMENTS_IDS_KEY, fst.as_fst().as_bytes())
+    /// Writes the external documents ids and internal ids (i.e. `u32`).
+    pub fn put_external_documents_ids<'a>(
+        &self,
+        wtxn: &mut RwTxn,
+        external_documents_ids: &ExternalDocumentsIds<'a>,
+    ) -> heed::Result<()>
+    {
+        let ExternalDocumentsIds { hard, soft } = external_documents_ids;
+        let hard = hard.as_fst().as_bytes();
+        let soft = soft.as_fst().as_bytes();
+        self.main.put::<_, Str, ByteSlice>(wtxn, HARD_EXTERNAL_DOCUMENTS_IDS_KEY, hard)?;
+        self.main.put::<_, Str, ByteSlice>(wtxn, SOFT_EXTERNAL_DOCUMENTS_IDS_KEY, soft)?;
+        Ok(())
     }
 
-    /// Returns the external documents ids map which associate the external ids (i.e. `[u8]`)
+    /// Returns the external documents ids map which associate the external ids
     /// with the internal ids (i.e. `u32`).
-    pub fn external_documents_ids<'t>(&self, rtxn: &'t RoTxn) -> anyhow::Result<fst::Map<Cow<'t, [u8]>>> {
-        match self.main.get::<_, Str, ByteSlice>(rtxn, EXTERNAL_DOCUMENTS_IDS_KEY)? {
-            Some(bytes) => Ok(fst::Map::new(bytes)?.map_data(Cow::Borrowed)?),
-            None => Ok(fst::Map::default().map_data(Cow::Owned)?),
+    pub fn external_documents_ids<'t>(&self, rtxn: &'t RoTxn) -> anyhow::Result<ExternalDocumentsIds<'t>> {
+        let hard = self.main.get::<_, Str, ByteSlice>(rtxn, HARD_EXTERNAL_DOCUMENTS_IDS_KEY)?;
+        let soft = self.main.get::<_, Str, ByteSlice>(rtxn, SOFT_EXTERNAL_DOCUMENTS_IDS_KEY)?;
+        match hard.zip(soft) {
+            Some((hard, soft)) => {
+                let hard = fst::Map::new(hard)?.map_data(Cow::Borrowed)?;
+                let soft = fst::Map::new(soft)?.map_data(Cow::Borrowed)?;
+                Ok(ExternalDocumentsIds::new(hard, soft))
+            },
+            None => Ok(ExternalDocumentsIds::default()),
         }
     }