Add 'remove this semicolon'

2 files changed, 40 insertions, 78 deletions

diff --git a/crates/ide_diagnostics/src/handlers/remove_this_semicolon.rs b/crates/ide_diagnostics/src/handlers/remove_this_semicolon.rs
deleted file mode 100644
index 141fbc42fa5..00000000000
--- a/crates/ide_diagnostics/src/handlers/remove_this_semicolon.rs
+++ /dev/null
@@ -1,76 +0,0 @@
-use ide_db::{
-    base_db::{FileLoader, FileRange},
-    source_change::SourceChange,
-};
-use syntax::{TextRange, TextSize};
-use text_edit::TextEdit;
-
-use crate::{fix, Assist, Diagnostic, DiagnosticsContext};
-
-// Diagnostic: remove-this-semicolon
-//
-// This diagnostic is triggered when there's an erroneous `;` at the end of the block.
-pub(crate) fn remove_this_semicolon(
-    ctx: &DiagnosticsContext<'_>,
-    d: &hir::RemoveThisSemicolon,
-) -> Diagnostic {
-    Diagnostic::new(
-        "remove-this-semicolon",
-        "remove this semicolon",
-        semicolon_range(ctx, d).unwrap_or_else(|it| it).range,
-    )
-    .with_fixes(fixes(ctx, d))
-}
-
-fn semicolon_range(
-    ctx: &DiagnosticsContext<'_>,
-    d: &hir::RemoveThisSemicolon,
-) -> Result<FileRange, FileRange> {
-    let expr_range = ctx.sema.diagnostics_display_range(d.expr.clone().map(|it| it.into()));
-    let file_text = ctx.sema.db.file_text(expr_range.file_id);
-    let range_end: usize = expr_range.range.end().into();
-    // FIXME: This doesn't handle whitespace and comments, but handling those in
-    // the presence of macros might prove tricky...
-    if file_text[range_end..].starts_with(';') {
-        Ok(FileRange {
-            file_id: expr_range.file_id,
-            range: TextRange::at(expr_range.range.end(), TextSize::of(';')),
-        })
-    } else {
-        Err(expr_range)
-    }
-}
-
-fn fixes(ctx: &DiagnosticsContext<'_>, d: &hir::RemoveThisSemicolon) -> Option<Vec<Assist>> {
-    let semicolon_range = semicolon_range(ctx, d).ok()?;
-
-    let edit = TextEdit::delete(semicolon_range.range);
-    let source_change = SourceChange::from_text_edit(semicolon_range.file_id, edit);
-
-    Some(vec![fix(
-        "remove_semicolon",
-        "Remove this semicolon",
-        source_change,
-        semicolon_range.range,
-    )])
-}
-
-#[cfg(test)]
-mod tests {
-    use crate::tests::{check_diagnostics, check_fix};
-
-    #[test]
-    fn missing_semicolon() {
-        check_diagnostics(
-            r#"
-fn test() -> i32 { 123; }
-                    //^ 💡 error: remove this semicolon
-"#,
-        );
-    }
-
-    #[test]
-    fn remove_semicolon() {
-        check_fix(r#"fn f() -> i32 { 92$0; }"#, r#"fn f() -> i32 { 92 }"#);
-    }
-}
diff --git a/crates/ide_diagnostics/src/handlers/type_mismatch.rs b/crates/ide_diagnostics/src/handlers/type_mismatch.rs
index 571605ef266..2c9f7839cee 100644
--- a/crates/ide_diagnostics/src/handlers/type_mismatch.rs
+++ b/crates/ide_diagnostics/src/handlers/type_mismatch.rs
@@ -1,9 +1,14 @@
 use hir::{db::AstDatabase, HirDisplay, Type, TypeInfo};
 use ide_db::{
-    famous_defs::FamousDefs, source_change::SourceChange,
+    base_db::{FileLoader, FileRange},
+    famous_defs::FamousDefs,
+    source_change::SourceChange,
     syntax_helpers::node_ext::for_each_tail_expr,
 };
-use syntax::{AstNode, TextRange};
+use syntax::{
+    ast::{BlockExpr, ExprStmt},
+    AstNode, TextRange, TextSize,
+};
 use text_edit::TextEdit;
 
 use crate::{fix, Assist, Diagnostic, DiagnosticsContext};
@@ -34,6 +39,7 @@ fn fixes(ctx: &DiagnosticsContext<'_>, d: &hir::TypeMismatch) -> Option<Vec<Assi
 
     add_reference(ctx, d, &mut fixes);
     add_missing_ok_or_some(ctx, d, &mut fixes);
+    remove_semicolon(ctx, d, &mut fixes);
 
     if fixes.is_empty() {
         None
@@ -110,6 +116,33 @@ fn add_missing_ok_or_some(
     Some(())
 }
 
+fn remove_semicolon(
+    ctx: &DiagnosticsContext<'_>,
+    d: &hir::TypeMismatch,
+    acc: &mut Vec<Assist>,
+) -> Option<()> {
+    let root = ctx.sema.db.parse_or_expand(d.expr.file_id)?;
+    let expr = d.expr.value.to_node(&root);
+    if !d.actual.is_unit() {
+        return None;
+    }
+    let block = BlockExpr::cast(expr.syntax().clone())?;
+    let expr_before_semi =
+        block.statements().last().and_then(|s| ExprStmt::cast(s.syntax().clone()))?;
+    let type_before_semi = ctx.sema.type_of_expr(&expr_before_semi.expr()?)?.original();
+    if !type_before_semi.could_coerce_to(ctx.sema.db, &d.expected) {
+        return None;
+    }
+    let semicolon_range = expr_before_semi.semicolon_token()?.text_range();
+
+    let edit = TextEdit::delete(semicolon_range);
+    let source_change =
+        SourceChange::from_text_edit(d.expr.file_id.original_file(ctx.sema.db), edit);
+
+    acc.push(fix("remove_semicolon", "Remove this semicolon", source_change, semicolon_range));
+    Some(())
+}
+
 #[cfg(test)]
 mod tests {
     use crate::tests::{check_diagnostics, check_fix, check_no_fix};
@@ -437,4 +470,9 @@ fn foo() -> SomeOtherEnum { 0$0 }
 "#,
         );
     }
+
+    #[test]
+    fn remove_semicolon() {
+        check_fix(r#"fn f() -> i32 { 92$0; }"#, r#"fn f() -> i32 { 92 }"#);
+    }
 }