IDEA-235296 Refactor ensureCodeBlock: AddAssertStatementFix updated

GitOrigin-RevId: e1e2f6021664be235b663876fe6629729fd0a4ef

java/java-impl/src/com/intellij/codeInspection/AddAssertStatementFix.java

@@ -16,11 +16,10 @@
 package com.intellij.codeInspection;
 
 import com.intellij.java.JavaBundle;
-import com.intellij.openapi.diagnostic.Logger;
 import com.intellij.openapi.project.Project;
 import com.intellij.psi.*;
 import com.intellij.psi.util.PsiTreeUtil;
-import com.intellij.refactoring.util.RefactoringUtil;
+import com.siyeh.ig.psiutils.CodeBlockSurrounder;
 import org.jetbrains.annotations.NonNls;
 import org.jetbrains.annotations.NotNull;
 
@@ -29,7 +28,6 @@ import org.jetbrains.annotations.NotNull;
  * @author ven
  */
 public class AddAssertStatementFix implements LocalQuickFix {
-  private static final Logger LOG = Logger.getInstance(AddAssertStatementFix.class);
   private final String myText;
 
   public AddAssertStatementFix(@NotNull String text) {
@@ -46,10 +44,11 @@ public class AddAssertStatementFix implements LocalQuickFix {
   public void applyFix(@NotNull Project project, @NotNull ProblemDescriptor descriptor) {
     PsiExpression element = PsiTreeUtil.getParentOfType(descriptor.getPsiElement(), PsiExpression.class);
     if (element == null) return;
-    element = RefactoringUtil.ensureCodeBlock(element);
-    if (element == null) return;
-    PsiElement anchorElement = RefactoringUtil.getParentStatement(element, false);
-    LOG.assertTrue(anchorElement != null);
+    CodeBlockSurrounder surrounder = CodeBlockSurrounder.forExpression(element);
+    if (surrounder == null) return;
+    CodeBlockSurrounder.SurroundResult result = surrounder.surround();
+    element = result.getExpression();
+    PsiElement anchorElement = result.getAnchor();
     PsiElement prev = PsiTreeUtil.skipWhitespacesBackward(anchorElement);
     if (prev instanceof PsiComment && JavaSuppressionUtil.getSuppressedInspectionIdsIn(prev) != null) {
       anchorElement = prev;

java/java-impl/src/com/intellij/codeInspection/dataFlow/DataFlowInspection.java

@@ -34,10 +34,7 @@ import com.intellij.util.SmartList;
 import com.intellij.util.containers.ContainerUtil;
 import com.intellij.util.ui.JBUI;
 import com.siyeh.ig.fixes.IntroduceVariableFix;
-import com.siyeh.ig.psiutils.ControlFlowUtils;
-import com.siyeh.ig.psiutils.ExpressionUtils;
-import com.siyeh.ig.psiutils.ParenthesesUtils;
-import com.siyeh.ig.psiutils.SideEffectChecker;
+import com.siyeh.ig.psiutils.*;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -144,7 +141,7 @@ public class DataFlowInspection extends DataFlowInspectionBase {
     PsiExpression operand = castExpression.getOperand();
     PsiTypeElement typeElement = castExpression.getCastType();
     if (typeElement != null && operand != null) {
-      if (!alwaysFails && !SideEffectChecker.mayHaveSideEffects(operand) && ControlFlowUtils.canExtractStatement(castExpression)) {
+      if (!alwaysFails && !SideEffectChecker.mayHaveSideEffects(operand) && CodeBlockSurrounder.canSurround(castExpression)) {
         String suffix = " instanceof " + typeElement.getText();
         fixes.add(new AddAssertStatementFix(ParenthesesUtils.getText(operand, PsiPrecedenceUtil.RELATIONAL_PRECEDENCE) + suffix));
         if (onTheFly && SurroundWithIfFix.isAvailable(operand)) {
@@ -187,8 +184,7 @@ public class DataFlowInspection extends DataFlowInspectionBase {
       }
       else if (!ExpressionUtils.isNullLiteral(qualifier) && !SideEffectChecker.mayHaveSideEffects(qualifier))  {
         String suffix = " != null";
-        if (PsiUtil.getLanguageLevel(qualifier).isAtLeast(LanguageLevel.JDK_1_4) &&
-            ControlFlowUtils.canExtractStatement(expression)) {
+        if (PsiUtil.getLanguageLevel(qualifier).isAtLeast(LanguageLevel.JDK_1_4) && CodeBlockSurrounder.canSurround(expression)) {
           String replacement = ParenthesesUtils.getText(qualifier, ParenthesesUtils.EQUALITY_PRECEDENCE) + suffix;
           fixes.add(new AddAssertStatementFix(replacement));
         }

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addAssert/afterLambda.java

@@ -1,13 +1,15 @@
 // "Assert 'container != null'" "true"
+import java.util.function.Supplier;
+
 class A{
   void test(){
     Object container = null;
-    Runnable r = () -> {
+    Supplier<String> r = () -> {
         if (container == null) {
             assert container != null;
-            container.toString();
+            return container.toString();
         } else {
-            "";
+            return "";
         }
     };
   }

java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/addAssert/beforeLambda.java

@@ -1,7 +1,9 @@
 // "Assert 'container != null'" "true"
+import java.util.function.Supplier;
+
 class A{
   void test(){
     Object container = null;
-    Runnable r = () -> container == null ? container.toS<caret>tring() : "";
+    Supplier<String> r = () -> container == null ? container.toS<caret>tring() : "";
   }
 }