Initial auto import action implementation

5 files changed, 317 insertions, 15 deletions

diff --git a/crates/ra_assists/src/assist_ctx.rs b/crates/ra_assists/src/assist_ctx.rs
index 43f0d664b..2ab65ab99 100644
--- a/crates/ra_assists/src/assist_ctx.rs
+++ b/crates/ra_assists/src/assist_ctx.rs
@@ -101,7 +101,6 @@ impl<'a, DB: HirDatabase> AssistCtx<'a, DB> {
         Some(assist)
     }
 
-    #[allow(dead_code)] // will be used for auto import assist with multiple actions
     pub(crate) fn add_assist_group(
         self,
         id: AssistId,
@@ -168,7 +167,6 @@ pub(crate) struct ActionBuilder {
 }
 
 impl ActionBuilder {
-    #[allow(dead_code)]
     /// Adds a custom label to the action, if it needs to be different from the assist label
     pub(crate) fn label(&mut self, label: impl Into<String>) {
         self.label = Some(label.into())
diff --git a/crates/ra_assists/src/assists/auto_import.rs b/crates/ra_assists/src/assists/auto_import.rs
new file mode 100644
index 000000000..fe226521e
--- /dev/null
+++ b/crates/ra_assists/src/assists/auto_import.rs
@@ -0,0 +1,181 @@
+use hir::db::HirDatabase;
+use ra_syntax::{
+    ast::{self, AstNode},
+    SmolStr, SyntaxElement,
+    SyntaxKind::{NAME_REF, USE_ITEM},
+    SyntaxNode,
+};
+
+use crate::{
+    assist_ctx::{ActionBuilder, Assist, AssistCtx},
+    auto_import_text_edit, AssistId, ImportsLocator,
+};
+
+// Assist: auto_import
+//
+// If the name is unresolved, provides all possible imports for it.
+//
+// ```
+// fn main() {
+//     let map = HashMap<|>::new();
+// }
+// ```
+// ->
+// ```
+// use std::collections::HashMap;
+//
+// fn main() {
+//     let map = HashMap<|>::new();
+// }
+// ```
+pub(crate) fn auto_import<'a, F: ImportsLocator<'a>>(
+    ctx: AssistCtx<impl HirDatabase>,
+    imports_locator: &mut F,
+) -> Option<Assist> {
+    let path: ast::Path = ctx.find_node_at_offset()?;
+    let module = path.syntax().ancestors().find_map(ast::Module::cast);
+    let position = match module.and_then(|it| it.item_list()) {
+        Some(item_list) => item_list.syntax().clone(),
+        None => {
+            let current_file = path.syntax().ancestors().find_map(ast::SourceFile::cast)?;
+            current_file.syntax().clone()
+        }
+    };
+
+    let module_with_name_to_import = ctx.source_analyzer(&position, None).module()?;
+    let name_to_import = hir::InFile {
+        file_id: ctx.frange.file_id.into(),
+        value: &find_applicable_name_ref(ctx.covering_element())?,
+    };
+
+    let proposed_imports =
+        imports_locator.find_imports(name_to_import, module_with_name_to_import)?;
+    if proposed_imports.is_empty() {
+        return None;
+    }
+
+    ctx.add_assist_group(AssistId("auto_import"), "auto import", || {
+        proposed_imports
+            .into_iter()
+            .map(|import| import_to_action(import.to_string(), &position, &path))
+            .collect()
+    })
+}
+
+fn find_applicable_name_ref(element: SyntaxElement) -> Option<ast::NameRef> {
+    if element.ancestors().find(|ancestor| ancestor.kind() == USE_ITEM).is_some() {
+        None
+    } else if element.kind() == NAME_REF {
+        Some(element.as_node().cloned().and_then(ast::NameRef::cast)?)
+    } else {
+        let parent = element.parent()?;
+        if parent.kind() == NAME_REF {
+            Some(ast::NameRef::cast(parent)?)
+        } else {
+            None
+        }
+    }
+}
+
+fn import_to_action(import: String, position: &SyntaxNode, path: &ast::Path) -> ActionBuilder {
+    let mut action_builder = ActionBuilder::default();
+    action_builder.label(format!("Import `{}`", &import));
+    auto_import_text_edit(
+        position,
+        &path.syntax().clone(),
+        &[SmolStr::new(import)],
+        action_builder.text_edit_builder(),
+    );
+    action_builder
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::helpers::{
+        check_assist_with_imports_locator, check_assist_with_imports_locator_not_applicable,
+    };
+    use hir::Name;
+
+    #[derive(Clone)]
+    struct TestImportsLocator<'a> {
+        import_path: &'a [Name],
+    }
+
+    impl<'a> TestImportsLocator<'a> {
+        fn new(import_path: &'a [Name]) -> Self {
+            TestImportsLocator { import_path }
+        }
+    }
+
+    impl<'a> ImportsLocator<'_> for TestImportsLocator<'_> {
+        fn find_imports(
+            &mut self,
+            _: hir::InFile<&ast::NameRef>,
+            _: hir::Module,
+        ) -> Option<Vec<hir::ModPath>> {
+            if self.import_path.is_empty() {
+                None
+            } else {
+                Some(vec![hir::ModPath {
+                    kind: hir::PathKind::Plain,
+                    segments: self.import_path.to_owned(),
+                }])
+            }
+        }
+    }
+
+    #[test]
+    fn applicable_when_found_an_import() {
+        let import_path = &[hir::name::known::std, hir::name::known::ops, hir::name::known::Debug];
+        let mut imports_locator = TestImportsLocator::new(import_path);
+        check_assist_with_imports_locator(
+            auto_import,
+            &mut imports_locator,
+            "
+            fn main() {
+            }
+
+            Debug<|>",
+            &format!(
+                "
+            use {};
+
+            fn main() {{
+            }}
+
+            Debug<|>",
+                import_path
+                    .into_iter()
+                    .map(|name| name.to_string())
+                    .collect::<Vec<String>>()
+                    .join("::")
+            ),
+        );
+    }
+
+    #[test]
+    fn not_applicable_when_no_imports_found() {
+        let mut imports_locator = TestImportsLocator::new(&[]);
+        check_assist_with_imports_locator_not_applicable(
+            auto_import,
+            &mut imports_locator,
+            "
+            fn main() {
+            }
+
+            Debug<|>",
+        );
+    }
+
+    #[test]
+    fn not_applicable_in_import_statements() {
+        let import_path = &[hir::name::known::std, hir::name::known::ops, hir::name::known::Debug];
+        let mut imports_locator = TestImportsLocator::new(import_path);
+        check_assist_with_imports_locator_not_applicable(
+            auto_import,
+            &mut imports_locator,
+            "use Debug<|>;",
+        );
+    }
+}
diff --git a/crates/ra_assists/src/doc_tests.rs b/crates/ra_assists/src/doc_tests.rs
index 5dc1ee233..65d51428b 100644
--- a/crates/ra_assists/src/doc_tests.rs
+++ b/crates/ra_assists/src/doc_tests.rs
@@ -11,6 +11,10 @@ use test_utils::{assert_eq_text, extract_range_or_offset};
 use crate::test_db::TestDB;
 
 fn check(assist_id: &str, before: &str, after: &str) {
+    // FIXME we cannot get the imports search functionality here yet, but still need to generate a test and a doc for an assist
+    if assist_id == "auto_import" {
+        return;
+    }
     let (selection, before) = extract_range_or_offset(before);
     let (db, file_id) = TestDB::with_single_file(&before);
     let frange = FileRange { file_id, range: selection.into() };
diff --git a/crates/ra_assists/src/doc_tests/generated.rs b/crates/ra_assists/src/doc_tests/generated.rs
index 7d84dc8fb..ec4587ce7 100644
--- a/crates/ra_assists/src/doc_tests/generated.rs
+++ b/crates/ra_assists/src/doc_tests/generated.rs
@@ -215,6 +215,25 @@ fn main() {
 }
 
 #[test]
+fn doctest_auto_import() {
+    check(
+        "auto_import",
+        r#####"
+fn main() {
+    let map = HashMap<|>::new();
+}
+"#####,
+        r#####"
+use std::collections::HashMap;
+
+fn main() {
+    let map = HashMap<|>::new();
+}
+"#####,
+    )
+}
+
+#[test]
 fn doctest_change_visibility() {
     check(
         "change_visibility",
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,