From d1e59be46b73bb371020a0dee2a18b8a02a27a5a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cl=C3=A9ment=20Renault?= <clement@meilisearch.com>
Date: Thu, 10 Oct 2019 10:35:17 +0200
Subject: [PATCH] Introduce the DocumentsIds iterator

---
 meilidb-core/src/store/documents_fields.rs | 37 ++++++++++++++++++++++
 1 file changed, 37 insertions(+)

diff --git a/meilidb-core/src/store/documents_fields.rs b/meilidb-core/src/store/documents_fields.rs
index 804508f05..a8e489f24 100644
--- a/meilidb-core/src/store/documents_fields.rs
+++ b/meilidb-core/src/store/documents_fields.rs
@@ -100,6 +100,15 @@ impl DocumentsFields {
         let iter = self.documents_fields.iter_from(reader, document_id_bytes)?;
         Ok(DocumentFieldsIter { document_id, iter })
     }
+
+    pub fn documents_ids<'r, T: rkv::Readable>(
+        &self,
+        reader: &'r T,
+    ) -> Result<DocumentsIdsIter<'r>, rkv::StoreError>
+    {
+        let iter = self.documents_fields.iter_start(reader)?;
+        Ok(DocumentsIdsIter { last_seen_id: None, iter })
+    }
 }
 
 pub struct DocumentFieldsIter<'r> {
@@ -125,3 +134,31 @@ impl<'r> Iterator for DocumentFieldsIter<'r> {
         }
     }
 }
+
+pub struct DocumentsIdsIter<'r> {
+    last_seen_id: Option<DocumentId>,
+    iter: rkv::store::single::Iter<'r>,
+}
+
+impl<'r> Iterator for DocumentsIdsIter<'r> {
+    type Item = Result<DocumentId, rkv::StoreError>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        for result in self.iter.next() {
+            match result {
+                Ok((key, Some(rkv::Value::Blob(bytes)))) => {
+                    let array = TryFrom::try_from(key).unwrap();
+                    let (document_id, attr) = document_attribute_from_key(array);
+                    if Some(document_id) != self.last_seen_id {
+                        self.last_seen_id = Some(document_id);
+                        return Some(Ok(document_id))
+                    }
+                },
+                Ok((key, data)) => panic!("{:?}, {:?}", key, data),
+                Err(e) => return Some(Err(e)),
+            }
+        }
+
+        None
+    }
+}