ensure assert inserted in correct place (IDEA-123296)

java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/DataFlowInspectionBase.java

@@ -215,7 +215,7 @@ public class DataFlowInspectionBase extends BaseJavaBatchLocalInspectionTool {
         final PsiElementFactory elementFactory = JavaPsiFacade.getInstance(project).getElementFactory();
         final PsiBinaryExpression binary = (PsiBinaryExpression)elementFactory.createExpressionFromText("a != null", null);
         binary.getLOperand().replace(qualifier);
-        fixes.add(new AddAssertStatementFix(binary));
+        ContainerUtil.addIfNotNull(fixes, createAssertFix(binary));
       }
 
       addSurroundWithIfFix(qualifier, fixes, onTheFly);
@@ -231,6 +231,10 @@ public class DataFlowInspectionBase extends BaseJavaBatchLocalInspectionTool {
     }
   }
 
+  protected LocalQuickFix createAssertFix(PsiBinaryExpression binary) {
+    return null;
+  }
+
   protected void addSurroundWithIfFix(PsiExpression qualifier, List<LocalQuickFix> fixes, boolean onTheFly) {
   }
 

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

@@ -20,6 +20,7 @@ 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.intellij.util.IncorrectOperationException;
 import org.jetbrains.annotations.NonNls;
 import org.jetbrains.annotations.NotNull;
@@ -51,7 +52,7 @@ public class AddAssertStatementFix implements LocalQuickFix {
     if (expressionToAssert == null) return;
     if (!FileModificationService.getInstance().preparePsiElementForWrite(descriptor.getPsiElement())) return;
     PsiElement element = descriptor.getPsiElement();
-    PsiElement anchorElement = PsiTreeUtil.getParentOfType(element, PsiStatement.class);
+    PsiElement anchorElement = RefactoringUtil.getParentStatement(element, false);
     LOG.assertTrue(anchorElement != null);
     final PsiElement tempParent = anchorElement.getParent();
     if (tempParent instanceof PsiForStatement && !PsiTreeUtil.isAncestor(((PsiForStatement)tempParent).getBody(), anchorElement, false)) {
@@ -75,11 +76,7 @@ public class AddAssertStatementFix implements LocalQuickFix {
         parent.addBefore(assertStatement, anchorElement);
       }
       else {
-        PsiBlockStatement blockStatement = (PsiBlockStatement)factory.createStatementFromText("{}", null);
-        final PsiCodeBlock block = blockStatement.getCodeBlock();
-        block.add(assertStatement);
-        block.add(anchorElement);
-        anchorElement.replace(blockStatement);
+        RefactoringUtil.putStatementInLoopBody(assertStatement, parent, anchorElement);
       }
     }
     catch (IncorrectOperationException e) {

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

@@ -16,6 +16,7 @@
 package com.intellij.codeInspection.dataFlow;
 
 import com.intellij.codeInsight.NullableNotNullDialog;
+import com.intellij.codeInspection.AddAssertStatementFix;
 import com.intellij.codeInspection.InspectionsBundle;
 import com.intellij.codeInspection.LocalQuickFix;
 import com.intellij.codeInspection.SurroundWithIfFix;
@@ -26,6 +27,7 @@ import com.intellij.openapi.project.Project;
 import com.intellij.openapi.project.ProjectManager;
 import com.intellij.psi.JavaTokenType;
 import com.intellij.psi.PsiAssignmentExpression;
+import com.intellij.psi.PsiBinaryExpression;
 import com.intellij.psi.PsiExpression;
 import com.intellij.psi.tree.IElementType;
 
@@ -60,6 +62,11 @@ public class DataFlowInspection extends DataFlowInspectionBase {
     return new OptionsPanel();
   }
 
+  @Override
+  protected AddAssertStatementFix createAssertFix(PsiBinaryExpression binary) {
+    return new AddAssertStatementFix(binary);
+  }
+
   private class OptionsPanel extends JPanel {
     private final JCheckBox myIgnoreAssertions;
     private final JCheckBox myReportConstantReferences;

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

@@ -0,0 +1,10 @@
+// "Assert 'container != null'" "true"
+class A{
+  void test(){
+    Object container = null;
+    Runnable r = () -> {
+        assert container != null;
+        container == null ? container.toString() : "";
+    };
+  }
+}

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

@@ -0,0 +1,7 @@
+// "Assert 'container != null'" "true"
+class A{
+  void test(){
+    Object container = null;
+    Runnable r = () -> container == null ? conta<caret>iner.toString() : "";
+  }
+}