Fix if_then_some_else_none suggests wrongly when then ends with comment (rust-lang#16278)

Closes rust-lang#16269

changelog: [`if_then_some_else_none`] fix wrong suggestions when then
ends with comment

Author: dswij

clippy_lints/src/if_then_some_else_none.rs

@@ -109,13 +109,19 @@ impl<'tcx> LateLintPass<'tcx> for IfThenSomeElseNone {
                         .maybe_paren()
                         .to_string();
                     let arg_snip = snippet_with_context(cx, then_arg.span, ctxt, "[body]", &mut app).0;
-                    let method_body = if let Some(first_stmt) = then_block.stmts.first()
-                        && let Some(first_stmt_span) = walk_span_to_context(first_stmt.span, ctxt)
+                    let method_body = if let Some(_) = then_block.stmts.first()
+                        && let Some(then_span) = walk_span_to_context(then.span, ctxt)
                     {
-                        let block_snippet =
-                            snippet_with_applicability(cx, first_stmt_span.until(then_expr.span), "..", &mut app);
+                        let block_before_snippet =
+                            snippet_with_applicability(cx, then_span.until(then_expr.span), "..", &mut app);
+                        let block_after_snippet = snippet_with_applicability(
+                            cx,
+                            then_expr.span.shrink_to_hi().until(then_span.shrink_to_hi()),
+                            "..",
+                            &mut app,
+                        );
                         let closure = if method_name == sym::then { "|| " } else { "" };
-                        format!("{closure} {{ {} {arg_snip} }}", block_snippet.trim_end())
+                        format!("{closure}{block_before_snippet}{arg_snip}{block_after_snippet}")
                     } else if method_name == sym::then {
                         (std::borrow::Cow::Borrowed("|| ") + arg_snip).into_owned()
                     } else {

tests/ui/if_then_some_else_none.fixed

@@ -3,10 +3,20 @@
 
 fn main() {
     // Should issue an error.
-    let _ = foo().then(||  { println!("true!"); "foo" });
+    let _ = foo().then(|| {
+        //~^ if_then_some_else_none
+
+        println!("true!");
+        "foo"
+    });
 
     // Should issue an error when macros are used.
-    let _ = matches!(true, true).then(||  { println!("true!"); matches!(true, false) });
+    let _ = matches!(true, true).then(|| {
+        //~^ if_then_some_else_none
+
+        println!("true!");
+        matches!(true, false)
+    });
 
     // Should issue an error. Binary expression `o < 32` should be parenthesized.
     let x = Some(5);
@@ -71,7 +81,12 @@ fn _msrv_1_49() {
 
 #[clippy::msrv = "1.50"]
 fn _msrv_1_50() {
-    let _ = foo().then(||  { println!("true!"); 150 });
+    let _ = foo().then(|| {
+        //~^ if_then_some_else_none
+
+        println!("true!");
+        150
+    });
 }
 
 fn foo() -> bool {
@@ -182,7 +197,10 @@ fn issue15005() {
 
         fn next(&mut self) -> Option<Self::Item> {
             //~v if_then_some_else_none
-            (self.count < 5).then(||  { self.count += 1; self.count })
+            (self.count < 5).then(|| {
+                self.count += 1;
+                self.count
+            })
         }
     }
 }
@@ -195,7 +213,10 @@ fn statements_from_macro() {
         };
     }
     //~v if_then_some_else_none
-    let _ = true.then(||  { mac!(); 42 });
+    let _ = true.then(|| {
+        mac!();
+        42
+    });
 }
 
 fn dont_lint_inside_macros() {
@@ -228,3 +249,14 @@ mod issue16176 {
         if cond { Some(foo().await) } else { None } // OK
     }
 }
+
+fn issue16269() -> Option<i32> {
+    use std::cell::UnsafeCell;
+
+    //~v if_then_some_else_none
+    (1 <= 3).then(|| {
+        let a = UnsafeCell::new(1);
+        // SAFETY: `bytes` bytes starting at `new_end` were just reserved.
+        unsafe { *a.get() }
+    })
+}

tests/ui/if_then_some_else_none.rs

@@ -284,3 +284,16 @@ mod issue16176 {
         if cond { Some(foo().await) } else { None } // OK
     }
 }
+
+fn issue16269() -> Option<i32> {
+    use std::cell::UnsafeCell;
+
+    //~v if_then_some_else_none
+    if 1 <= 3 {
+        let a = UnsafeCell::new(1);
+        // SAFETY: `bytes` bytes starting at `new_end` were just reserved.
+        Some(unsafe { *a.get() })
+    } else {
+        None
+    }
+}

tests/ui/if_then_some_else_none.stderr

@@ -9,10 +9,19 @@ LL | |         println!("true!");
 ...  |
 LL | |         None
 LL | |     };
-   | |_____^ help: try: `foo().then(||  { println!("true!"); "foo" })`
+   | |_____^
    |
    = note: `-D clippy::if-then-some-else-none` implied by `-D warnings`
    = help: to override `-D warnings` add `#[allow(clippy::if_then_some_else_none)]`
+help: try
+   |
+LL ~     let _ = foo().then(|| {
+LL +
+LL + 
+LL +         println!("true!");
+LL +         "foo"
+LL ~     });
+   |
 
 error: this could be simplified with `bool::then`
   --> tests/ui/if_then_some_else_none.rs:16:13
@@ -25,7 +34,17 @@ LL | |         println!("true!");
 ...  |
 LL | |         None
 LL | |     };
-   | |_____^ help: try: `matches!(true, true).then(||  { println!("true!"); matches!(true, false) })`
+   | |_____^
+   |
+help: try
+   |
+LL ~     let _ = matches!(true, true).then(|| {
+LL +
+LL + 
+LL +         println!("true!");
+LL +         matches!(true, false)
+LL ~     });
+   |
 
 error: this could be simplified with `bool::then_some`
   --> tests/ui/if_then_some_else_none.rs:27:28
@@ -50,7 +69,17 @@ LL | |         println!("true!");
 ...  |
 LL | |         None
 LL | |     };
-   | |_____^ help: try: `foo().then(||  { println!("true!"); 150 })`
+   | |_____^
+   |
+help: try
+   |
+LL ~     let _ = foo().then(|| {
+LL +
+LL + 
+LL +         println!("true!");
+LL +         150
+LL ~     });
+   |
 
 error: this could be simplified with `bool::then`
   --> tests/ui/if_then_some_else_none.rs:138:5
@@ -125,7 +154,15 @@ LL | |                 Some(self.count)
 LL | |             } else {
 LL | |                 None
 LL | |             }
-   | |_____________^ help: try: `(self.count < 5).then(||  { self.count += 1; self.count })`
+   | |_____________^
+   |
+help: try
+   |
+LL ~             (self.count < 5).then(|| {
+LL +                 self.count += 1;
+LL +                 self.count
+LL +             })
+   |
 
 error: this could be simplified with `bool::then`
   --> tests/ui/if_then_some_else_none.rs:249:13
@@ -137,7 +174,36 @@ LL | |         Some(42)
 LL | |     } else {
 LL | |         None
 LL | |     };
-   | |_____^ help: try: `true.then(||  { mac!(); 42 })`
+   | |_____^
+   |
+help: try
+   |
+LL ~     let _ = true.then(|| {
+LL +         mac!();
+LL +         42
+LL ~     });
+   |
+
+error: this could be simplified with `bool::then`
+  --> tests/ui/if_then_some_else_none.rs:292:5
+   |
+LL | /     if 1 <= 3 {
+LL | |         let a = UnsafeCell::new(1);
+LL | |         // SAFETY: `bytes` bytes starting at `new_end` were just reserved.
+LL | |         Some(unsafe { *a.get() })
+LL | |     } else {
+LL | |         None
+LL | |     }
+   | |_____^
+   |
+help: try
+   |
+LL ~     (1 <= 3).then(|| {
+LL +         let a = UnsafeCell::new(1);
+LL +         // SAFETY: `bytes` bytes starting at `new_end` were just reserved.
+LL +         unsafe { *a.get() }
+LL +     })
+   |
 
-error: aborting due to 13 previous errors
+error: aborting due to 14 previous errors