1 files changed, 223 insertions, 10 deletions

diff --git a/crates/ra_ide_api/src/diagnostics.rs b/crates/ra_ide_api/src/diagnostics.rs
index 98b840b26..1a4882824 100644
--- a/crates/ra_ide_api/src/diagnostics.rs
+++ b/crates/ra_ide_api/src/diagnostics.rs
@@ -9,7 +9,7 @@ use ra_assists::ast_editor::{AstBuilder, AstEditor};
 use ra_db::SourceDatabase;
 use ra_prof::profile;
 use ra_syntax::{
-    ast::{self, AstNode, NamedField},
+    ast::{self, AstNode, RecordField},
     Location, SyntaxNode, TextRange, T,
 };
 use ra_text_edit::{TextEdit, TextEditBuilder};
@@ -62,7 +62,7 @@ pub(crate) fn diagnostics(db: &RootDatabase, file_id: FileId) -> Vec<Diagnostic>
         let node = d.ast(db);
         let mut ast_editor = AstEditor::new(node);
         for f in d.missed_fields.iter() {
-            ast_editor.append_field(&AstBuilder::<NamedField>::from_name(f));
+            ast_editor.append_field(&AstBuilder::<RecordField>::from_name(f));
         }
 
         let mut builder = TextEditBuilder::default();
@@ -75,6 +75,19 @@ pub(crate) fn diagnostics(db: &RootDatabase, file_id: FileId) -> Vec<Diagnostic>
             severity: Severity::Error,
             fix: Some(fix),
         })
+    })
+    .on::<hir::diagnostics::MissingOkInTailExpr, _>(|d| {
+        let node = d.ast(db);
+        let mut builder = TextEditBuilder::default();
+        let replacement = format!("Ok({})", node.syntax());
+        builder.replace(node.syntax().text_range(), replacement);
+        let fix = SourceChange::source_file_edit_from("wrap with ok", file_id, builder.finish());
+        res.borrow_mut().push(Diagnostic {
+            range: d.highlight_range(),
+            message: d.message(),
+            severity: Severity::Error,
+            fix: Some(fix),
+        })
     });
     if let Some(m) = source_binder::module_from_file_id(db, file_id) {
         m.diagnostics(db, &mut sink);
@@ -141,20 +154,20 @@ fn check_struct_shorthand_initialization(
     file_id: FileId,
     node: &SyntaxNode,
 ) -> Option<()> {
-    let struct_lit = ast::StructLit::cast(node.clone())?;
-    let named_field_list = struct_lit.named_field_list()?;
-    for named_field in named_field_list.fields() {
-        if let (Some(name_ref), Some(expr)) = (named_field.name_ref(), named_field.expr()) {
+    let record_lit = ast::RecordLit::cast(node.clone())?;
+    let record_field_list = record_lit.record_field_list()?;
+    for record_field in record_field_list.fields() {
+        if let (Some(name_ref), Some(expr)) = (record_field.name_ref(), record_field.expr()) {
             let field_name = name_ref.syntax().text().to_string();
             let field_expr = expr.syntax().text().to_string();
             if field_name == field_expr {
                 let mut edit_builder = TextEditBuilder::default();
-                edit_builder.delete(named_field.syntax().text_range());
-                edit_builder.insert(named_field.syntax().text_range().start(), field_name);
+                edit_builder.delete(record_field.syntax().text_range());
+                edit_builder.insert(record_field.syntax().text_range().start(), field_name);
                 let edit = edit_builder.finish();
 
                 acc.push(Diagnostic {
-                    range: named_field.syntax().text_range(),
+                    range: record_field.syntax().text_range(),
                     message: "Shorthand struct initialization".to_string(),
                     severity: Severity::WeakWarning,
                     fix: Some(SourceChange::source_file_edit(
@@ -171,10 +184,11 @@ fn check_struct_shorthand_initialization(
 #[cfg(test)]
 mod tests {
     use insta::assert_debug_snapshot_matches;
+    use join_to_string::join;
     use ra_syntax::SourceFile;
     use test_utils::assert_eq_text;
 
-    use crate::mock_analysis::single_file;
+    use crate::mock_analysis::{analysis_and_position, single_file};
 
     use super::*;
 
@@ -203,6 +217,48 @@ mod tests {
         assert_eq_text!(after, &actual);
     }
 
+    /// Takes a multi-file input fixture with annotated cursor positions,
+    /// and checks that:
+    ///  * a diagnostic is produced
+    ///  * this diagnostic touches the input cursor position
+    ///  * that the contents of the file containing the cursor match `after` after the diagnostic fix is applied
+    fn check_apply_diagnostic_fix_from_position(fixture: &str, after: &str) {
+        let (analysis, file_position) = analysis_and_position(fixture);
+        let diagnostic = analysis.diagnostics(file_position.file_id).unwrap().pop().unwrap();
+        let mut fix = diagnostic.fix.unwrap();
+        let edit = fix.source_file_edits.pop().unwrap().edit;
+        let target_file_contents = analysis.file_text(file_position.file_id).unwrap();
+        let actual = edit.apply(&target_file_contents);
+
+        // Strip indent and empty lines from `after`, to match the behaviour of
+        // `parse_fixture` called from `analysis_and_position`.
+        let margin = fixture
+            .lines()
+            .filter(|it| it.trim_start().starts_with("//-"))
+            .map(|it| it.len() - it.trim_start().len())
+            .next()
+            .expect("empty fixture");
+        let after = join(after.lines().filter_map(|line| {
+            if line.len() > margin {
+                Some(&line[margin..])
+            } else {
+                None
+            }
+        }))
+        .separator("\n")
+        .suffix("\n")
+        .to_string();
+
+        assert_eq_text!(&after, &actual);
+        assert!(
+            diagnostic.range.start() <= file_position.offset
+                && diagnostic.range.end() >= file_position.offset,
+            "diagnostic range {} does not touch cursor position {}",
+            diagnostic.range,
+            file_position.offset
+        );
+    }
+
     fn check_apply_diagnostic_fix(before: &str, after: &str) {
         let (analysis, file_id) = single_file(before);
         let diagnostic = analysis.diagnostics(file_id).unwrap().pop().unwrap();
@@ -212,6 +268,14 @@ mod tests {
         assert_eq_text!(after, &actual);
     }
 
+    /// Takes a multi-file input fixture with annotated cursor position and checks that no diagnostics
+    /// apply to the file containing the cursor.
+    fn check_no_diagnostic_for_target_file(fixture: &str) {
+        let (analysis, file_position) = analysis_and_position(fixture);
+        let diagnostics = analysis.diagnostics(file_position.file_id).unwrap();
+        assert_eq!(diagnostics.len(), 0);
+    }
+
     fn check_no_diagnostic(content: &str) {
         let (analysis, file_id) = single_file(content);
         let diagnostics = analysis.diagnostics(file_id).unwrap();
@@ -219,6 +283,155 @@ mod tests {
     }
 
     #[test]
+    fn test_wrap_return_type() {
+        let before = r#"
+            //- /main.rs
+            use std::{string::String, result::Result::{self, Ok, Err}};
+
+            fn div(x: i32, y: i32) -> Result<i32, String> {
+                if y == 0 {
+                    return Err("div by zero".into());
+                }
+                x / y<|>
+            }
+
+            //- /std/lib.rs
+            pub mod string {
+                pub struct String { }
+            }
+            pub mod result {
+                pub enum Result<T, E> { Ok(T), Err(E) }
+            }
+        "#;
+        let after = r#"
+            use std::{string::String, result::Result::{self, Ok, Err}};
+
+            fn div(x: i32, y: i32) -> Result<i32, String> {
+                if y == 0 {
+                    return Err("div by zero".into());
+                }
+                Ok(x / y)
+            }
+        "#;
+        check_apply_diagnostic_fix_from_position(before, after);
+    }
+
+    #[test]
+    fn test_wrap_return_type_handles_generic_functions() {
+        let before = r#"
+            //- /main.rs
+            use std::result::Result::{self, Ok, Err};
+
+            fn div<T>(x: T) -> Result<T, i32> {
+                if x == 0 {
+                    return Err(7);
+                }
+                <|>x
+            }
+
+            //- /std/lib.rs
+            pub mod result {
+                pub enum Result<T, E> { Ok(T), Err(E) }
+            }
+        "#;
+        let after = r#"
+            use std::result::Result::{self, Ok, Err};
+
+            fn div<T>(x: T) -> Result<T, i32> {
+                if x == 0 {
+                    return Err(7);
+                }
+                Ok(x)
+            }
+        "#;
+        check_apply_diagnostic_fix_from_position(before, after);
+    }
+
+    #[test]
+    fn test_wrap_return_type_handles_type_aliases() {
+        let before = r#"
+            //- /main.rs
+            use std::{string::String, result::Result::{self, Ok, Err}};
+
+            type MyResult<T> = Result<T, String>;
+
+            fn div(x: i32, y: i32) -> MyResult<i32> {
+                if y == 0 {
+                    return Err("div by zero".into());
+                }
+                x <|>/ y
+            }
+
+            //- /std/lib.rs
+            pub mod string {
+                pub struct String { }
+            }
+            pub mod result {
+                pub enum Result<T, E> { Ok(T), Err(E) }
+            }
+        "#;
+        let after = r#"
+            use std::{string::String, result::Result::{self, Ok, Err}};
+
+            type MyResult<T> = Result<T, String>;
+            fn div(x: i32, y: i32) -> MyResult<i32> {
+                if y == 0 {
+                    return Err("div by zero".into());
+                }
+                Ok(x / y)
+            }
+        "#;
+        check_apply_diagnostic_fix_from_position(before, after);
+    }
+
+    #[test]
+    fn test_wrap_return_type_not_applicable_when_expr_type_does_not_match_ok_type() {
+        let content = r#"
+            //- /main.rs
+            use std::{string::String, result::Result::{self, Ok, Err}};
+
+            fn foo() -> Result<String, i32> {
+                0<|>
+            }
+
+            //- /std/lib.rs
+            pub mod string {
+                pub struct String { }
+            }
+            pub mod result {
+                pub enum Result<T, E> { Ok(T), Err(E) }
+            }
+        "#;
+        check_no_diagnostic_for_target_file(content);
+    }
+
+    #[test]
+    fn test_wrap_return_type_not_applicable_when_return_type_is_not_result() {
+        let content = r#"
+            //- /main.rs
+            use std::{string::String, result::Result::{self, Ok, Err}};
+
+            enum SomeOtherEnum {
+                Ok(i32),
+                Err(String),
+            }
+
+            fn foo() -> SomeOtherEnum {
+                0<|>
+            }
+
+            //- /std/lib.rs
+            pub mod string {
+                pub struct String { }
+            }
+            pub mod result {
+                pub enum Result<T, E> { Ok(T), Err(E) }
+            }
+        "#;
+        check_no_diagnostic_for_target_file(content);
+    }
+
+    #[test]
     fn test_fill_struct_fields_empty() {
         let before = r"
             struct TestStruct {