From 2c3d71dd8fb71f379cab33a5ca8a5f307a31ef72 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cl=C3=A9ment=20Renault?= <renault.cle@gmail.com>
Date: Fri, 7 Dec 2018 11:53:17 +0100
Subject: [PATCH] fix: Improve the bucket sort algorithm

---
 src/index/mod.rs                                |  2 +-
 src/rank/mod.rs                                 |  4 ++--
 src/rank/{ranked_stream.rs => query_builder.rs} | 13 +++++++++----
 3 files changed, 12 insertions(+), 7 deletions(-)
 rename src/rank/{ranked_stream.rs => query_builder.rs} (94%)

diff --git a/src/index/mod.rs b/src/index/mod.rs
index 90a02e4d9..0329477b1 100644
--- a/src/index/mod.rs
+++ b/src/index/mod.rs
@@ -121,7 +121,7 @@ impl Index {
         let snapshot = self.database.snapshot();
 
         let builder = QueryBuilder::new(snapshot)?;
-        let documents = builder.query(query, 0..20);
+        let documents = builder.query(query, 20);
 
         Ok(documents)
     }
diff --git a/src/rank/mod.rs b/src/rank/mod.rs
index cee465cac..81b9463e2 100644
--- a/src/rank/mod.rs
+++ b/src/rank/mod.rs
@@ -1,10 +1,10 @@
 pub mod criterion;
-mod ranked_stream;
+mod query_builder;
 mod distinct_map;
 
 use crate::{Match, DocumentId};
 
-pub use self::ranked_stream::{QueryBuilder, DistinctQueryBuilder};
+pub use self::query_builder::{QueryBuilder, DistinctQueryBuilder};
 
 #[inline]
 fn match_query_index(a: &Match, b: &Match) -> bool {
diff --git a/src/rank/ranked_stream.rs b/src/rank/query_builder.rs
similarity index 94%
rename from src/rank/ranked_stream.rs
rename to src/rank/query_builder.rs
index 6dba99c9a..2f3642b37 100644
--- a/src/rank/ranked_stream.rs
+++ b/src/rank/query_builder.rs
@@ -109,23 +109,28 @@ impl<T, C> QueryBuilder<T, C>
 where T: Deref<Target=DB>,
       C: Criterion,
 {
-    pub fn query(&self, query: &str, range: Range<usize>) -> Vec<Document> {
+    pub fn query(&self, query: &str, limit: usize) -> Vec<Document> {
         let mut documents = self.query_all(query);
         let mut groups = vec![documents.as_mut_slice()];
 
-        for criterion in &self.criteria {
+        'group: for criterion in &self.criteria {
             let tmp_groups = mem::replace(&mut groups, Vec::new());
+            let mut computed = 0;
 
             for group in tmp_groups {
+
                 group.sort_unstable_by(|a, b| criterion.evaluate(a, b));
                 for group in GroupByMut::new(group, |a, b| criterion.eq(a, b)) {
+
+                    computed += group.len();
                     groups.push(group);
+                    if computed >= limit { break 'group }
                 }
             }
         }
 
-        let range = clamp_range(range, 0..documents.len());
-        documents[range].to_vec()
+        documents.truncate(limit);
+        documents
     }
 }