crates/ide/src/diagnostics/fixes/wrap_tail_expr.rs

   1 use hir::{db::AstDatabase, diagnostics::MissingOkOrSomeInTailExpr, Semantics};
   2 use ide_assists::{Assist, AssistResolveStrategy};
   3 use ide_db::{source_change::SourceChange, RootDatabase};
   4 use syntax::AstNode;
   5 use text_edit::TextEdit;
   6
   7 use crate::diagnostics::{fix, DiagnosticWithFix};
   8
   9 impl DiagnosticWithFix for MissingOkOrSomeInTailExpr {
  10     fn fix(
  11         &self,
  12         sema: &Semantics<RootDatabase>,
  13         _resolve: &AssistResolveStrategy,
  14     ) -> Option<Assist> {
  15         let root = sema.db.parse_or_expand(self.file)?;
  16         let tail_expr = self.expr.to_node(&root);
  17         let tail_expr_range = tail_expr.syntax().text_range();
  18         let replacement = format!("{}({})", self.required, tail_expr.syntax());
  19         let edit = TextEdit::replace(tail_expr_range, replacement);
  20         let source_change = SourceChange::from_text_edit(self.file.original_file(sema.db), edit);
  21         let name = if self.required == "Ok" { "Wrap with Ok" } else { "Wrap with Some" };
  22         Some(fix("wrap_tail_expr", name, source_change, tail_expr_range))
  23     }
  24 }
  25
  26 #[cfg(test)]
  27 mod tests {
  28     use crate::diagnostics::tests::{check_fix, check_no_diagnostics};
  29
  30     #[test]
  31     fn test_wrap_return_type_option() {
  32         check_fix(
  33             r#"
  34 //- /main.rs crate:main deps:core
  35 use core::option::Option::{self, Some, None};
  36
  37 fn div(x: i32, y: i32) -> Option<i32> {
  38     if y == 0 {
  39         return None;
  40     }
  41     x / y$0
  42 }
  43 //- /core/lib.rs crate:core
  44 pub mod result {
  45     pub enum Result<T, E> { Ok(T), Err(E) }
  46 }
  47 pub mod option {
  48     pub enum Option<T> { Some(T), None }
  49 }
  50 "#,
  51             r#"
  52 use core::option::Option::{self, Some, None};
  53
  54 fn div(x: i32, y: i32) -> Option<i32> {
  55     if y == 0 {
  56         return None;
  57     }
  58     Some(x / y)
  59 }
  60 "#,
  61         );
  62     }
  63
  64     #[test]
  65     fn test_wrap_return_type() {
  66         check_fix(
  67             r#"
  68 //- /main.rs crate:main deps:core
  69 use core::result::Result::{self, Ok, Err};
  70
  71 fn div(x: i32, y: i32) -> Result<i32, ()> {
  72     if y == 0 {
  73         return Err(());
  74     }
  75     x / y$0
  76 }
  77 //- /core/lib.rs crate:core
  78 pub mod result {
  79     pub enum Result<T, E> { Ok(T), Err(E) }
  80 }
  81 pub mod option {
  82     pub enum Option<T> { Some(T), None }
  83 }
  84 "#,
  85             r#"
  86 use core::result::Result::{self, Ok, Err};
  87
  88 fn div(x: i32, y: i32) -> Result<i32, ()> {
  89     if y == 0 {
  90         return Err(());
  91     }
  92     Ok(x / y)
  93 }
  94 "#,
  95         );
  96     }
  97
  98     #[test]
  99     fn test_wrap_return_type_handles_generic_functions() {
 100         check_fix(
 101             r#"
 102 //- /main.rs crate:main deps:core
 103 use core::result::Result::{self, Ok, Err};
 104
 105 fn div<T>(x: T) -> Result<T, i32> {
 106     if x == 0 {
 107         return Err(7);
 108     }
 109     $0x
 110 }
 111 //- /core/lib.rs crate:core
 112 pub mod result {
 113     pub enum Result<T, E> { Ok(T), Err(E) }
 114 }
 115 pub mod option {
 116     pub enum Option<T> { Some(T), None }
 117 }
 118 "#,
 119             r#"
 120 use core::result::Result::{self, Ok, Err};
 121
 122 fn div<T>(x: T) -> Result<T, i32> {
 123     if x == 0 {
 124         return Err(7);
 125     }
 126     Ok(x)
 127 }
 128 "#,
 129         );
 130     }
 131
 132     #[test]
 133     fn test_wrap_return_type_handles_type_aliases() {
 134         check_fix(
 135             r#"
 136 //- /main.rs crate:main deps:core
 137 use core::result::Result::{self, Ok, Err};
 138
 139 type MyResult<T> = Result<T, ()>;
 140
 141 fn div(x: i32, y: i32) -> MyResult<i32> {
 142     if y == 0 {
 143         return Err(());
 144     }
 145     x $0/ y
 146 }
 147 //- /core/lib.rs crate:core
 148 pub mod result {
 149     pub enum Result<T, E> { Ok(T), Err(E) }
 150 }
 151 pub mod option {
 152     pub enum Option<T> { Some(T), None }
 153 }
 154 "#,
 155             r#"
 156 use core::result::Result::{self, Ok, Err};
 157
 158 type MyResult<T> = Result<T, ()>;
 159
 160 fn div(x: i32, y: i32) -> MyResult<i32> {
 161     if y == 0 {
 162         return Err(());
 163     }
 164     Ok(x / y)
 165 }
 166 "#,
 167         );
 168     }
 169
 170     #[test]
 171     fn test_wrap_return_type_not_applicable_when_expr_type_does_not_match_ok_type() {
 172         check_no_diagnostics(
 173             r#"
 174 //- /main.rs crate:main deps:core
 175 use core::result::Result::{self, Ok, Err};
 176
 177 fn foo() -> Result<(), i32> { 0 }
 178
 179 //- /core/lib.rs crate:core
 180 pub mod result {
 181     pub enum Result<T, E> { Ok(T), Err(E) }
 182 }
 183 pub mod option {
 184     pub enum Option<T> { Some(T), None }
 185 }
 186 "#,
 187         );
 188     }
 189
 190     #[test]
 191     fn test_wrap_return_type_not_applicable_when_return_type_is_not_result_or_option() {
 192         check_no_diagnostics(
 193             r#"
 194 //- /main.rs crate:main deps:core
 195 use core::result::Result::{self, Ok, Err};
 196
 197 enum SomeOtherEnum { Ok(i32), Err(String) }
 198
 199 fn foo() -> SomeOtherEnum { 0 }
 200
 201 //- /core/lib.rs crate:core
 202 pub mod result {
 203     pub enum Result<T, E> { Ok(T), Err(E) }
 204 }
 205 pub mod option {
 206     pub enum Option<T> { Some(T), None }
 207 }
 208 "#,
 209         );
 210     }
 211 }