1 files changed, 113 insertions, 13 deletions

diff --git a/crates/ra_assists/src/lib.rs b/crates/ra_assists/src/lib.rs
index 3337805a5..4029962f7 100644
--- a/crates/ra_assists/src/lib.rs
+++ b/crates/ra_assists/src/lib.rs
@@ -14,9 +14,9 @@ mod test_db;
 pub mod ast_transform;
 
 use either::Either;
-use hir::db::HirDatabase;
+use hir::{db::HirDatabase, InFile, ModPath, Module};
 use ra_db::FileRange;
-use ra_syntax::{TextRange, TextUnit};
+use ra_syntax::{ast::NameRef, TextRange, TextUnit};
 use ra_text_edit::TextEdit;
 
 pub(crate) use crate::assist_ctx::{Assist, AssistCtx};
@@ -77,6 +77,55 @@ where
     })
 }
 
+/// A functionality for locating imports for the given name.
+///
+/// Currently has to be a trait with the real implementation provided by the ra_ide_api crate,
+/// due to the search functionality located there.
+/// Later, this trait should be removed completely and the search functionality moved to a separate crate,
+/// accessible from the ra_assists crate.
+pub trait ImportsLocator<'a> {
+    /// Finds all imports for the given name and the module that contains this name.
+    fn find_imports(
+        &mut self,
+        name_to_import: InFile<&NameRef>,
+        module_with_name_to_import: Module,
+    ) -> Option<Vec<ModPath>>;
+}
+
+/// Return all the assists applicable at the given position
+/// and additional assists that need the imports locator functionality to work.
+///
+/// Assists are returned in the "resolved" state, that is with edit fully
+/// computed.
+pub fn assists_with_imports_locator<'a, H, F: 'a>(
+    db: &H,
+    range: FileRange,
+    mut imports_locator: F,
+) -> Vec<ResolvedAssist>
+where
+    H: HirDatabase + 'static,
+    F: ImportsLocator<'a>,
+{
+    AssistCtx::with_ctx(db, range, true, |ctx| {
+        let mut assists = assists::all()
+            .iter()
+            .map(|f| f(ctx.clone()))
+            .chain(
+                assists::all_with_imports_locator()
+                    .iter()
+                    .map(|f| f(ctx.clone(), &mut imports_locator)),
+            )
+            .filter_map(std::convert::identity)
+            .map(|a| match a {
+                Assist::Resolved { assist } => assist,
+                Assist::Unresolved { .. } => unreachable!(),
+            })
+            .collect();
+        sort_assists(&mut assists);
+        assists
+    })
+}
+
 /// Return all the assists applicable at the given position.
 ///
 /// Assists are returned in the "resolved" state, that is with edit fully
@@ -85,8 +134,6 @@ pub fn assists<H>(db: &H, range: FileRange) -> Vec<ResolvedAssist>
 where
     H: HirDatabase + 'static,
 {
-    use std::cmp::Ordering;
-
     AssistCtx::with_ctx(db, range, true, |ctx| {
         let mut a = assists::all()
             .iter()
@@ -95,19 +142,24 @@ where
                 Assist::Resolved { assist } => assist,
                 Assist::Unresolved { .. } => unreachable!(),
             })
-            .collect::<Vec<_>>();
-        a.sort_by(|a, b| match (a.get_first_action().target, b.get_first_action().target) {
-            (Some(a), Some(b)) => a.len().cmp(&b.len()),
-            (Some(_), None) => Ordering::Less,
-            (None, Some(_)) => Ordering::Greater,
-            (None, None) => Ordering::Equal,
-        });
+            .collect();
+        sort_assists(&mut a);
         a
     })
 }
 
+fn sort_assists(assists: &mut Vec<ResolvedAssist>) {
+    use std::cmp::Ordering;
+    assists.sort_by(|a, b| match (a.get_first_action().target, b.get_first_action().target) {
+        (Some(a), Some(b)) => a.len().cmp(&b.len()),
+        (Some(_), None) => Ordering::Less,
+        (None, Some(_)) => Ordering::Greater,
+        (None, None) => Ordering::Equal,
+    });
+}
+
 mod assists {
-    use crate::{Assist, AssistCtx};
+    use crate::{Assist, AssistCtx, ImportsLocator};
     use hir::db::HirDatabase;
 
     mod add_derive;
@@ -116,6 +168,7 @@ mod assists {
     mod add_custom_impl;
     mod add_new;
     mod apply_demorgan;
+    mod auto_import;
     mod invert_if;
     mod flip_comma;
     mod flip_binexpr;
@@ -168,6 +221,11 @@ mod assists {
             early_return::convert_to_guarded_return,
         ]
     }
+
+    pub(crate) fn all_with_imports_locator<'a, DB: HirDatabase, F: ImportsLocator<'a>>(
+    ) -> &'a [fn(AssistCtx<DB>, &mut F) -> Option<Assist>] {
+        &[auto_import::auto_import]
+    }
 }
 
 #[cfg(test)]
@@ -176,7 +234,7 @@ mod helpers {
     use ra_syntax::TextRange;
     use test_utils::{add_cursor, assert_eq_text, extract_offset, extract_range};
 
-    use crate::{test_db::TestDB, Assist, AssistCtx};
+    use crate::{test_db::TestDB, Assist, AssistCtx, ImportsLocator};
 
     pub(crate) fn check_assist(
         assist: fn(AssistCtx<TestDB>) -> Option<Assist>,
@@ -206,6 +264,35 @@ mod helpers {
         assert_eq_text!(after, &actual);
     }
 
+    pub(crate) fn check_assist_with_imports_locator<'a, F: ImportsLocator<'a>>(
+        assist: fn(AssistCtx<TestDB>, &mut F) -> Option<Assist>,
+        imports_locator: &mut F,
+        before: &str,
+        after: &str,
+    ) {
+        let (before_cursor_pos, before) = extract_offset(before);
+        let (db, file_id) = TestDB::with_single_file(&before);
+        let frange =
+            FileRange { file_id, range: TextRange::offset_len(before_cursor_pos, 0.into()) };
+        let assist = AssistCtx::with_ctx(&db, frange, true, |ctx| assist(ctx, imports_locator))
+            .expect("code action is not applicable");
+        let action = match assist {
+            Assist::Unresolved { .. } => unreachable!(),
+            Assist::Resolved { assist } => assist.get_first_action(),
+        };
+
+        let actual = action.edit.apply(&before);
+        let actual_cursor_pos = match action.cursor_position {
+            None => action
+                .edit
+                .apply_to_offset(before_cursor_pos)
+                .expect("cursor position is affected by the edit"),
+            Some(off) => off,
+        };
+        let actual = add_cursor(&actual, actual_cursor_pos);
+        assert_eq_text!(after, &actual);
+    }
+
     pub(crate) fn check_assist_range(
         assist: fn(AssistCtx<TestDB>) -> Option<Assist>,
         before: &str,
@@ -279,6 +366,19 @@ mod helpers {
         assert!(assist.is_none());
     }
 
+    pub(crate) fn check_assist_with_imports_locator_not_applicable<'a, F: ImportsLocator<'a>>(
+        assist: fn(AssistCtx<TestDB>, &mut F) -> Option<Assist>,
+        imports_locator: &mut F,
+        before: &str,
+    ) {
+        let (before_cursor_pos, before) = extract_offset(before);
+        let (db, file_id) = TestDB::with_single_file(&before);
+        let frange =
+            FileRange { file_id, range: TextRange::offset_len(before_cursor_pos, 0.into()) };
+        let assist = AssistCtx::with_ctx(&db, frange, true, |ctx| assist(ctx, imports_locator));
+        assert!(assist.is_none());
+    }
+
     pub(crate) fn check_assist_range_not_applicable(
         assist: fn(AssistCtx<TestDB>) -> Option<Assist>,
         before: &str,