Repository: intellij-rust/intellij-rust
Issue: Unexpected difference in "Surround with unsafe block" behavior

## Environment

* **IntelliJ Rust plugin version:** 0.4.153.4027-212-nightly
* **Rust toolchain version:** 1.56.0-nightly (492723897 2021-07-29) x86_64-apple-darwin
* **IDE name and version:** CLion 2021.3 EAP (CL-213.2029)
* **Operating system:** macOS 11.4
* **Macro expansion engine:** new
* **Name resolution engine:** new

## Problem description

I have two seemingly identical mutable static usages though "Surround with unsafe block" yields a different result for them. Not sure why, I would expect the result to be the same to preserve the consistency.

## Steps to reproduce

```Rust
static mut ERROR: i32 = 0;

fn main() {

    if ERROR != 0 {
        panic!("an error has occurred while reading the file");
    }

    if ERROR != 0 {
        panic!("an error has occurred while closing the file");
    }
}
```

Call "Surround with unsafe block" on the first `ERROR` usage

```Rust
static mut ERROR: i32 = 0;

fn main() {
    unsafe {
        if ERROR != 0 {
            panic!("an error has occurred while reading the file");
        }
    }

    if ERROR != 0 {
        panic!("an error has occurred while closing the file");
    }
}
```

Call "Surround with unsafe block" on the second `ERROR` usage

```Rust
static mut ERROR: i32 = 0;

fn main() {
    unsafe {
        if ERROR != 0 {
            panic!("an error has occurred while reading the file");
        }
    }

    if unsafe { ERROR } != 0 {
        panic!("an error has occurred while closing the file");
    }
}
```

Changed files:
src/test/kotlin/org/rust/ide/annotator/fixes/AddUnsafeFixTest.kt, src/main/kotlin/org/rust/ide/annotator/fixes/SurroundWithUnsafeFix.kt

diff --git a/src/main/kotlin/org/rust/ide/annotator/fixes/SurroundWithUnsafeFix.kt b/src/main/kotlin/org/rust/ide/annotator/fixes/SurroundWithUnsafeFix.kt
index ca6b4a29b..ff77363c8 100644
--- a/src/main/kotlin/org/rust/ide/annotator/fixes/SurroundWithUnsafeFix.kt
+++ b/src/main/kotlin/org/rust/ide/annotator/fixes/SurroundWithUnsafeFix.kt
@@ -11,6 +11,8 @@ import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
 import com.intellij.psi.util.parentOfType
+import org.rust.ide.utils.getTopmostParentInside
+import org.rust.lang.core.psi.RsBlock
 import org.rust.lang.core.psi.RsExpr
 import org.rust.lang.core.psi.RsExprStmt
 import org.rust.lang.core.psi.RsPsiFactory
@@ -20,8 +22,14 @@ class SurroundWithUnsafeFix(expr: RsExpr) : LocalQuickFixAndIntentionActionOnPsi
     override fun getText() = "Surround with unsafe block"
 
     override fun invoke(project: Project, file: PsiFile, editor: Editor?, expr: PsiElement, endElement: PsiElement) {
-        val target = expr.parentOfType<RsExprStmt>() ?: expr
+        val target = getTopLevelStmtOrExpr(expr) ?: expr.parentOfType<RsExprStmt>() ?: expr
         val unsafe = RsPsiFactory(project).createUnsafeBlockExpr(target.text)
         target.replace(unsafe)
     }
 }
+
+private fun getTopLevelStmtOrExpr(expr: PsiElement): PsiElement? {
+    val block = expr.parentOfType<RsBlock>() ?: return null
+    return expr.getTopmostParentInside(block)
+        .takeIf { it is RsExprStmt || it == block.expr }
+}
diff --git a/src/test/kotlin/org/rust/ide/annotator/fixes/AddUnsafeFixTest.kt b/src/test/kotlin/org/rust/ide/annotator/fixes/AddUnsafeFixTest.kt
index 9312b177b..1efd7d943 100644
--- a/src/test/kotlin/org/rust/ide/annotator/fixes/AddUnsafeFixTest.kt
+++ b/src/test/kotlin/org/rust/ide/annotator/fixes/AddUnsafeFixTest.kt
@@ -257,6 +257,82 @@ class AddUnsafeFixTest : RsAnnotationTestBase() {
         }
     """)
 
+    fun `test wrap tail expression with an unsafe block 1`() = checkFixByText("Surround with unsafe block", """
+        static mut ERROR: i32 = 0;
+        fn main() {
+            if <error>ERROR/*caret*/</error> != 0 {
+                panic!();
+            }
+        }
+    """, """
+        static mut ERROR: i32 = 0;
+        fn main() {
+            unsafe {
+                if ERROR != 0 {
+                    panic!();
+                }
+            }
+        }
+    """)
+
+    fun `test wrap tail expression with an unsafe block 2`() = checkFixByText("Surround with unsafe block", """
+        static mut ERROR: i32 = 0;
+        fn main() {
+            if true {
+                if <error>ERROR/*caret*/</error> != 0 {
+                    panic!();
+                }
+            }
+        }
+    """, """
+        static mut ERROR: i32 = 0;
+        fn main() {
+            if true {
+                unsafe {
+                    if ERROR != 0 {
+                        panic!();
+                    }
+                }
+            }
+        }
+    """)
+
+    fun `test wrap tail expression inside closure with an unsafe block`() = checkFixByText("Surround with unsafe block", """
+        static mut ERROR: i32 = 0;
+        fn main() {
+            bar(|| {
+                if <error>ERROR/*caret*/</error> != 0 {
+                    panic!();
+                }
+            });
+        }
+        fn bar(f: fn()) {}
+    """, """
+        static mut ERROR: i32 = 0;
+        fn main() {
+            bar(|| {
+                unsafe {
+                    if ERROR != 0 {
+                        panic!();
+                    }
+                }
+            });
+        }
+        fn bar(f: fn()) {}
+    """)
+
+    fun `test wrap item inside block with an unsafe block`() = checkFixByText("Surround with unsafe block", """
+        static mut C1: i32 = 0;
+        fn main() {
+            static C2: i32 = <error>C1/*caret*/</error>;
+        }
+    """, """
+        static mut C1: i32 = 0;
+        fn main() {
+            static C2: i32 = unsafe { C1 };
+        }
+    """)
+
     fun `test add unsafe to impl`() = checkFixByText("Add unsafe to impl", """
         unsafe trait Trait {}