From 5a56ac4b72e7f57fb20d49bcf531611f144cc0e0 Mon Sep 17 00:00:00 2001
From: Aleksey Kladov <aleksey.kladov@gmail.com>
Date: Mon, 13 Aug 2018 17:19:27 +0300
Subject: [PATCH] Better limiting

---
 crates/libanalysis/src/lib.rs           | 20 +++++++++++++-------
 crates/libanalysis/src/symbol_index.rs  | 18 ++++++++++++++----
 crates/server/src/main_loop/handlers.rs |  3 ++-
 3 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/crates/libanalysis/src/lib.rs b/crates/libanalysis/src/lib.rs
index 97b6dfca689..562fab29084 100644
--- a/crates/libanalysis/src/lib.rs
+++ b/crates/libanalysis/src/lib.rs
@@ -94,12 +94,12 @@ pub fn file_line_index(&self, path: &Path) -> Result<LineIndex> {
         Ok(index.clone())
     }
 
-    pub fn world_symbols<'a>(&'a self, query: Query) -> impl Iterator<Item=(&'a Path, &'a FileSymbol)> + 'a {
+    pub fn world_symbols<'a>(&'a self, mut query: Query) -> impl Iterator<Item=(&'a Path, &'a FileSymbol)> + 'a {
         self.data.file_map.iter()
             .flat_map(move |(path, data)| {
                 let path: &'a Path = path.as_path();
-                let symbols = data.symbols(path);
-                query.process(symbols).map(move |s| (path, s))
+                let symbols = data.symbols();
+                query.process(symbols).into_iter().map(move |s| (path, s))
             })
     }
 
@@ -125,7 +125,8 @@ pub fn approximately_resolve_symbol<'a>(
 
         let mut query = Query::new(name.to_string());
         query.exact();
-        Ok(self.world_symbols(query).take(4).collect())
+        query.limit(4);
+        Ok(self.world_symbols(query).collect())
     }
 
     fn file_data(&self, path: &Path) -> Result<Arc<FileData>> {
@@ -178,9 +179,14 @@ fn syntax(&self, path: &Path) -> &ast::File {
             })
     }
 
-    fn symbols(&self, path: &Path) -> &FileSymbols {
-        let syntax = self.syntax(path);
+    fn syntax_transient(&self) -> ast::File {
+        self.syntax.get().map(|s| s.clone())
+            .unwrap_or_else(|| ast::File::parse(&self.text))
+    }
+
+    fn symbols(&self) -> &FileSymbols {
+        let syntax = self.syntax_transient();
         self.symbols
-            .get_or_init(|| FileSymbols::new(syntax))
+            .get_or_init(|| FileSymbols::new(&syntax))
     }
 }
diff --git a/crates/libanalysis/src/symbol_index.rs b/crates/libanalysis/src/symbol_index.rs
index 88d5c499520..a7ae197e01b 100644
--- a/crates/libanalysis/src/symbol_index.rs
+++ b/crates/libanalysis/src/symbol_index.rs
@@ -35,6 +35,7 @@ pub struct Query {
     lowercased: String,
     only_types: bool,
     exact: bool,
+    limit: usize,
 }
 
 impl Query {
@@ -45,6 +46,7 @@ pub fn new(query: String) -> Query {
             lowercased,
             only_types: false,
             exact: false,
+            limit: usize::max_value()
         }
     }
 
@@ -56,10 +58,14 @@ pub fn exact(&mut self) {
         self.exact = true;
     }
 
+    pub fn limit(&mut self, limit: usize) {
+        self.limit = limit
+    }
+
     pub(crate) fn process<'a>(
-        &self,
+        &mut self,
         file: &'a FileSymbols,
-    ) -> impl Iterator<Item=&'a FileSymbol> + 'a {
+    ) -> Vec<&'a FileSymbol> {
         fn is_type(kind: SyntaxKind) -> bool {
             match kind {
                 STRUCT | ENUM | TRAIT | TYPE_ITEM => true,
@@ -70,6 +76,9 @@ fn is_type(kind: SyntaxKind) -> bool {
         let mut stream = file.map.search(automaton).into_stream();
         let mut res = Vec::new();
         while let Some((_, idx)) = stream.next() {
+            if self.limit == 0 {
+                break;
+            }
             let idx = idx as usize;
             let symbol = &file.symbols[idx];
             if self.only_types && !is_type(symbol.kind) {
@@ -78,9 +87,10 @@ fn is_type(kind: SyntaxKind) -> bool {
             if self.exact && symbol.name != self.query {
                 continue;
             }
-            res.push(symbol)
+            res.push(symbol);
+            self.limit -= 1;
         }
-        res.into_iter()
+        res
     }
 }
 
diff --git a/crates/server/src/main_loop/handlers.rs b/crates/server/src/main_loop/handlers.rs
index e9dc7842052..bfdfcb51e64 100644
--- a/crates/server/src/main_loop/handlers.rs
+++ b/crates/server/src/main_loop/handlers.rs
@@ -110,10 +110,11 @@ pub fn handle_workspace_symbol(
         if !all_symbols {
             q.only_types();
         }
+        q.limit(128);
         q
     };
 
-    for (path, symbol) in world.world_symbols(query).take(128) {
+    for (path, symbol) in world.world_symbols(query) {
         let line_index = world.file_line_index(path)?;
         let info = SymbolInformation {
             name: symbol.name.to_string(),