Limit import map queries

2 files changed, 51 insertions, 7 deletions

diff --git a/crates/ra_hir/src/code_model.rs b/crates/ra_hir/src/code_model.rs
index c8329d971..a55fe03a6 100644
--- a/crates/ra_hir/src/code_model.rs
+++ b/crates/ra_hir/src/code_model.rs
@@ -104,12 +104,16 @@ impl Crate {
         db: &dyn DefDatabase,
         query: &str,
     ) -> impl Iterator<Item = Either<ModuleDef, MacroDef>> {
-        import_map::search_dependencies(db, self.into(), import_map::Query::new(query).anchor_end())
-            .into_iter()
-            .map(|item| match item {
-                ItemInNs::Types(mod_id) | ItemInNs::Values(mod_id) => Either::Left(mod_id.into()),
-                ItemInNs::Macros(mac_id) => Either::Right(mac_id.into()),
-            })
+        import_map::search_dependencies(
+            db,
+            self.into(),
+            import_map::Query::new(query).anchor_end().limit(40),
+        )
+        .into_iter()
+        .map(|item| match item {
+            ItemInNs::Types(mod_id) | ItemInNs::Values(mod_id) => Either::Left(mod_id.into()),
+            ItemInNs::Macros(mac_id) => Either::Right(mac_id.into()),
+        })
     }
 
     pub fn all(db: &dyn HirDatabase) -> Vec<Crate> {
diff --git a/crates/ra_hir_def/src/import_map.rs b/crates/ra_hir_def/src/import_map.rs
index f2e4ca2db..70368d8df 100644
--- a/crates/ra_hir_def/src/import_map.rs
+++ b/crates/ra_hir_def/src/import_map.rs
@@ -178,11 +178,12 @@ fn cmp((_, lhs): &(&ItemInNs, &ModPath), (_, rhs): &(&ItemInNs, &ModPath)) -> Or
 pub struct Query {
     query: String,
     anchor_end: bool,
+    limit: usize,
 }
 
 impl Query {
     pub fn new(query: &str) -> Self {
-        Self { query: query.to_lowercase(), anchor_end: false }
+        Self { query: query.to_lowercase(), anchor_end: false, limit: usize::max_value() }
     }
 
     /// Only returns items whose paths end with the (case-insensitive) query string as their last
@@ -190,6 +191,11 @@ impl Query {
     pub fn anchor_end(self) -> Self {
         Self { anchor_end: true, ..self }
     }
+
+    /// Limits the returned number of items to `limit`.
+    pub fn limit(self, limit: usize) -> Self {
+        Self { limit, ..self }
+    }
 }
 
 /// Searches dependencies of `krate` for an importable path matching `query`.
@@ -237,6 +243,11 @@ pub fn search_dependencies<'a>(
                 let item_path = &import_map.map[item];
                 fst_path(item_path) == fst_path(path)
             }));
+
+            if res.len() >= query.limit {
+                res.truncate(query.limit);
+                return res;
+            }
         }
     }
 
@@ -570,4 +581,33 @@ mod tests {
         dep::Fmt (m)
         "###);
     }
+
+    #[test]
+    fn search_limit() {
+        let res = search_dependencies_of(
+            r#"
+        //- /main.rs crate:main deps:dep
+        //- /dep.rs crate:dep
+        pub mod fmt {
+            pub trait Display {
+                fn fmt();
+            }
+        }
+        #[macro_export]
+        macro_rules! Fmt {
+            () => {};
+        }
+        pub struct Fmt;
+
+        pub fn format() {}
+        pub fn no() {}
+    "#,
+            "main",
+            Query::new("").limit(2),
+        );
+        assert_snapshot!(res, @r###"
+        dep::fmt (t)
+        dep::Fmt (t)
+        "###);
+    }
 }