Add explicit type assist.

2 files changed, 97 insertions, 0 deletions

diff --git a/crates/ra_assists/src/add_explicit_type.rs b/crates/ra_assists/src/add_explicit_type.rs
new file mode 100644
index 000000000..dec4f68ee
--- /dev/null
+++ b/crates/ra_assists/src/add_explicit_type.rs
@@ -0,0 +1,95 @@
+use hir::{
+    HirDisplay, Ty,
+    db::HirDatabase,
+    source_binder::function_from_child_node,
+};
+use ra_syntax::{
+    SyntaxKind,
+    ast::{LetStmt, PatKind, NameOwner, AstNode}
+};
+
+use crate::{AssistCtx, Assist, AssistId};
+
+/// Add explicit type assist.
+pub(crate) fn add_explicit_type(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
+    let stmt = ctx.node_at_offset::<LetStmt>()?;
+    let expr = stmt.initializer()?;
+    let pat = stmt.pat()?;
+    // Must be a binding
+    let pat = match pat.kind() {
+        PatKind::BindPat(bind_pat) => bind_pat,
+        _ => {
+            return None;
+        }
+    };
+    let pat_range = pat.syntax().range();
+    // The binding must have a name
+    let name = pat.name()?;
+    let name_range = name.syntax().range();
+    // Assist not applicable if the type has already been specified
+    if stmt.syntax().children_with_tokens().any(|child| child.kind() == SyntaxKind::COLON) {
+        return None;
+    }
+    // Infer type
+    let func = function_from_child_node(ctx.db, ctx.frange.file_id, pat.syntax())?;
+    let inference_res = func.infer(ctx.db);
+    let source_map = func.body_source_map(ctx.db);
+    let expr_id = source_map.node_expr(expr.into())?;
+    let ty = inference_res[expr_id].clone();
+    // Assist not applicable if the type is unknown
+    if is_unknown(&ty) {
+        return None;
+    }
+    let ty_str = ty.display(ctx.db).to_string();
+
+    ctx.add_action(AssistId("add_explicit_type"), "add explicit type", |edit| {
+        edit.target(pat_range);
+        edit.insert(name_range.end(), format!(": {}", ty_str));
+    });
+    ctx.build()
+}
+
+/// Returns true if any type parameter is unknown
+fn is_unknown(ty: &Ty) -> bool {
+    match ty {
+        Ty::Unknown => true,
+        Ty::Apply(a_ty) => a_ty.parameters.iter().any(is_unknown),
+        _ => false,
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    use crate::helpers::{ check_assist, check_assist_target, check_assist_not_applicable };
+
+    #[test]
+    fn add_explicit_type_target() {
+        check_assist_target(add_explicit_type, "fn f() { let a<|> = 1; }", "a");
+    }
+
+    #[test]
+    fn add_explicit_type_works_for_simple_expr() {
+        check_assist(
+            add_explicit_type,
+            "fn f() { let a<|> = 1; }",
+            "fn f() { let a<|>: i32 = 1; }",
+        );
+    }
+
+    #[test]
+    fn add_explicit_type_not_applicable_if_ty_not_inferred() {
+        check_assist_not_applicable(add_explicit_type, "fn f() { let a<|> = None; }");
+    }
+
+    #[test]
+    fn add_explicit_type_not_applicable_if_ty_already_specified() {
+        check_assist_not_applicable(add_explicit_type, "fn f() { let a<|>: i32 = 1; }");
+    }
+
+    #[test]
+    fn add_explicit_type_not_applicable_if_specified_ty_is_tuple() {
+        check_assist_not_applicable(add_explicit_type, "fn f() { let a<|>: (i32, i32) = (3, 4); }");
+    }
+}
diff --git a/crates/ra_assists/src/lib.rs b/crates/ra_assists/src/lib.rs
index c1514f8e5..ded401b63 100644
--- a/crates/ra_assists/src/lib.rs
+++ b/crates/ra_assists/src/lib.rs
@@ -86,6 +86,7 @@ where
 }
 
 mod add_derive;
+mod add_explicit_type;
 mod add_impl;
 mod flip_comma;
 mod flip_binexpr;
@@ -103,6 +104,7 @@ mod add_missing_impl_members;
 fn all_assists<DB: HirDatabase>() -> &'static [fn(AssistCtx<DB>) -> Option<Assist>] {
     &[
         add_derive::add_derive,
+        add_explicit_type::add_explicit_type,
         add_impl::add_impl,
         change_visibility::change_visibility,
         fill_match_arms::fill_match_arms,