Java: Infer nullability of extracted method's parameter in presence of a lambda or an anonymous class (IDEA-175727)

java/java-impl/src/com/intellij/refactoring/extractMethod/ExtractMethodProcessor.java

@@ -1487,37 +1487,78 @@ public class ExtractMethodProcessor implements MatchProvider {
 
   @Nullable
   private static Boolean isNotNullAt(@NotNull PsiVariable variable, PsiElement startElement) {
+    String variableName = variable.getName();
+    if (variableName == null) return null;
+
     final PsiElement methodOrLambdaBody = getSurroundingMethodOrLambdaBody(variable);
     if (methodOrLambdaBody instanceof PsiCodeBlock) {
-      final Set<PsiReferenceExpression> firstReadUsages = findFirstReadUsagesAt(variable, startElement);
-      return firstReadUsages != null &&
-             firstReadUsages.stream()
-               .map(firstReadUsage -> DfaUtil.checkNullness(variable, firstReadUsage, methodOrLambdaBody))
-               .allMatch(nullness -> nullness == Nullness.NOT_NULL);
+      PsiElement topmostLambdaOrAnonymousClass = null;
+      for (PsiElement element = startElement; element != null && element != methodOrLambdaBody; element = element.getParent()) {
+        if (element instanceof PsiLambdaExpression || element instanceof PsiAnonymousClass) {
+          topmostLambdaOrAnonymousClass = element;
+        }
+      }
+      if (topmostLambdaOrAnonymousClass != null) {
+        startElement = topmostLambdaOrAnonymousClass;
+      }
+
+      Project project = methodOrLambdaBody.getProject();
+      PsiFile file = methodOrLambdaBody.getContainingFile();
+      final PsiFile copy = PsiFileFactory.getInstance(project)
+        .createFileFromText(file.getName(), file.getFileType(), file.getText(), file.getModificationStamp(), false);
+
+      PsiCodeBlock bodyCopy = findCopy(copy, methodOrLambdaBody, PsiCodeBlock.class);
+      PsiVariable variableCopy = findCopy(copy, variable, PsiVariable.class);
+      if (startElement instanceof PsiExpression) {
+        startElement = PsiTreeUtil.getParentOfType(startElement, PsiStatement.class);
+      }
+      if (startElement instanceof PsiStatement) {
+        PsiStatement startStatementCopy = findCopy(copy, startElement, PsiStatement.class);
+
+        try {
+          PsiElementFactory factory = JavaPsiFacade.getElementFactory(project);
+          startStatementCopy = wrapWithBlockStatementIfNeeded(startStatementCopy, factory);
+          PsiDeclarationStatement declarationStatement = (PsiDeclarationStatement)factory.createStatementFromText(
+            CommonClassNames.JAVA_LANG_OBJECT + " _Dummy_ = " + variableName + ";", startStatementCopy);
+
+          PsiElement parent = startStatementCopy.getParent();
+          declarationStatement = (PsiDeclarationStatement)parent.addBefore(declarationStatement, startStatementCopy);
+          PsiElement[] declaredElements = declarationStatement.getDeclaredElements();
+          PsiExpression initializer = ((PsiVariable)declaredElements[0]).getInitializer();
+
+          Nullness nullness = DfaUtil.checkNullness(variableCopy, initializer, bodyCopy);
+          return nullness == Nullness.NOT_NULL;
+        }
+        catch (IncorrectOperationException e) {
+          LOG.debug(e);
+          return null;
+        }
+      }
     }
     return null;
   }
 
-  @Nullable
-  private static Set<PsiReferenceExpression> findFirstReadUsagesAt(@NotNull PsiVariable variable, PsiElement startElement) {
-    final PsiCodeBlock closestCodeBlock = PsiTreeUtil.getParentOfType(startElement, PsiCodeBlock.class);
-    if (closestCodeBlock != null) {
-      try {
-        final ControlFlow controlFlow = ControlFlowFactory.getInstance(closestCodeBlock.getProject())
-          .getControlFlow(closestCodeBlock, AllVariablesControlFlowPolicy.getInstance(), false, false);
+  private static <T extends PsiElement> T findCopy(@NotNull PsiFile copy, @NotNull PsiElement element, @NotNull Class<T> clazz) {
+    TextRange range = element.getTextRange();
+    return CodeInsightUtil.findElementInRange(copy, range.getStartOffset(), range.getEndOffset(), clazz);
+  }
 
-        final int startOffset = controlFlow.getStartOffset(startElement);
-        final List<PsiReferenceExpression> readBeforeWrite = ControlFlowUtil.getReadBeforeWrite(controlFlow, startOffset);
-        final Set<PsiReferenceExpression> result = StreamEx.of(readBeforeWrite)
-          .filter(referenceExpression -> referenceExpression.isReferenceTo(variable))
-          .toSet();
-        return !result.isEmpty() ? result : null;
-      }
-      catch (AnalysisCanceledException e) {
-        return null;
+  private static PsiStatement wrapWithBlockStatementIfNeeded(@NotNull PsiStatement statement, @NotNull PsiElementFactory factory) {
+    PsiElement parent = statement.getParent();
+    if (parent instanceof PsiLoopStatement && ((PsiLoopStatement)parent).getBody() == statement ||
+        parent instanceof PsiIfStatement &&
+        (((PsiIfStatement)parent).getThenBranch() == statement || ((PsiIfStatement)parent).getElseBranch() == statement)) {
+      PsiBlockStatement blockStatement = (PsiBlockStatement)factory.createStatementFromText("{}", statement);
+      blockStatement.getCodeBlock().add(statement);
+      blockStatement = (PsiBlockStatement)statement.replace(blockStatement);
+      return blockStatement.getCodeBlock().getStatements()[0];
+    }
+    if (parent instanceof PsiForStatement) {
+      if (((PsiForStatement)parent).getInitialization() == statement || ((PsiForStatement)parent).getUpdate() == statement) {
+        return wrapWithBlockStatementIfNeeded((PsiForStatement)parent, factory);
       }
     }
-    return null;
+    return statement;
   }
 
   @NotNull

java/java-tests/testData/refactoring/extractMethod/NotNullArgumentAnonymousClassBare.java

@@ -0,0 +1,13 @@
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+class X {
+    void foo(@NotNull Object o) {
+        <selection>Runnable r = new Runnable() {
+            @Override
+            public void run() {
+                System.out.println(o);
+            }
+        };</selection>
+    }
+}

java/java-tests/testData/refactoring/extractMethod/NotNullArgumentAnonymousClassBare_after.java

@@ -0,0 +1,17 @@
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+class X {
+    void foo(@NotNull Object o) {
+        newMethod(o);
+    }
+
+    private void newMethod(@NotNull Object o) {
+        Runnable r = new Runnable() {
+            @Override
+            public void run() {
+                System.out.println(o);
+            }
+        };
+    }
+}

java/java-tests/testData/refactoring/extractMethod/NotNullArgumentAnonymousClassInIf.java

@@ -0,0 +1,18 @@
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+class C {
+    void foo(@Nullable Object o) {
+        if (o != null) {
+            <selection>Runnable r = new Runnable() {
+                @Override
+                public void run() {
+                    bar(o);
+                }
+            };</selection>
+        }
+    }
+
+    void bar(@NotNull Object o) {
+    }
+}

java/java-tests/testData/refactoring/extractMethod/NotNullArgumentAnonymousClassInIfNoBlock.java

@@ -0,0 +1,17 @@
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+class C {
+    void foo(@Nullable Object o) {
+        if (o != null)
+            <selection>new Runnable() {
+                @Override
+                public void run() {
+                    bar(o);
+                }
+            }</selection>.run();
+    }
+
+    void bar(@NotNull Object o) {
+    }
+}

java/java-tests/testData/refactoring/extractMethod/NotNullArgumentAnonymousClassInIfNoBlock_after.java

@@ -0,0 +1,22 @@
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+class C {
+    void foo(@Nullable Object o) {
+        if (o != null)
+            newMethod(o).run();
+    }
+
+    @NotNull
+    private Runnable newMethod(@NotNull Object o) {
+        return new Runnable() {
+            @Override
+            public void run() {
+                bar(o);
+            }
+        };
+    }
+
+    void bar(@NotNull Object o) {
+    }
+}

java/java-tests/testData/refactoring/extractMethod/NotNullArgumentAnonymousClassInIf_after.java

@@ -0,0 +1,22 @@
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+class C {
+    void foo(@Nullable Object o) {
+        if (o != null) {
+            newMethod(o);
+        }
+    }
+
+    private void newMethod(@NotNull Object o) {
+        Runnable r = new Runnable() {
+            @Override
+            public void run() {
+                bar(o);
+            }
+        };
+    }
+
+    void bar(@NotNull Object o) {
+    }
+}

java/java-tests/testData/refactoring/extractMethod/NotNullArgumentLambdaBare.java

@@ -0,0 +1,11 @@
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+class C {
+    void foo(@NotNull Object o) {
+        <selection>Runnable r = () -> bar(o);</selection>
+    }
+
+    void bar(@NotNull Object o) {
+    }
+}

java/java-tests/testData/refactoring/extractMethod/NotNullArgumentLambdaBare_after.java

@@ -0,0 +1,15 @@
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+class C {
+    void foo(@NotNull Object o) {
+        newMethod(o);
+    }
+
+    private void newMethod(@NotNull Object o) {
+        Runnable r = () -> bar(o);
+    }
+
+    void bar(@NotNull Object o) {
+    }
+}

java/java-tests/testData/refactoring/extractMethod/NotNullArgumentLambdaInForInitializer.java

@@ -0,0 +1,12 @@
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+class C {
+    void foo(@Nullable Object o) {
+        if (o != null)
+            for(Runnable r = () -> <selection>bar(o)</selection>; ; r.run()) {}
+    }
+
+    void bar(@NotNull Object o) {
+    }
+}

java/java-tests/testData/refactoring/extractMethod/NotNullArgumentLambdaInForInitializer_after.java

@@ -0,0 +1,16 @@
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+class C {
+    void foo(@Nullable Object o) {
+        if (o != null)
+            for(Runnable r = () -> newMethod(o); ; r.run()) {}
+    }
+
+    private void newMethod(@NotNull Object o) {
+        bar(o);
+    }
+
+    void bar(@NotNull Object o) {
+    }
+}

java/java-tests/testData/refactoring/extractMethod/NotNullArgumentLambdaInIf.java

@@ -0,0 +1,13 @@
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+class C {
+    void foo(@Nullable Object o) {
+        if (o != null) {
+            <selection>Runnable r = () -> bar(o);</selection>
+        }
+    }
+
+    void bar(@NotNull Object o) {
+    }
+}

java/java-tests/testData/refactoring/extractMethod/NotNullArgumentLambdaInIfInsideBody.java

@@ -0,0 +1,12 @@
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+class C {
+    void foo(@Nullable Object o) {
+        if (o != null)
+            ((Runnable) (() -> <selection>bar(o)</selection>)).run();
+    }
+
+    void bar(@NotNull Object o) {
+    }
+}

java/java-tests/testData/refactoring/extractMethod/NotNullArgumentLambdaInIfInsideBody_after.java

@@ -0,0 +1,16 @@
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+class C {
+    void foo(@Nullable Object o) {
+        if (o != null)
+            ((Runnable) (() -> newMethod(o))).run();
+    }
+
+    private void newMethod(@NotNull Object o) {
+        bar(o);
+    }
+
+    void bar(@NotNull Object o) {
+    }
+}

java/java-tests/testData/refactoring/extractMethod/NotNullArgumentLambdaInIfNoBlock.java

@@ -0,0 +1,12 @@
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+class C {
+    void foo(@Nullable Object o) {
+        if (o != null)
+            <selection>((Runnable)(() -> bar(o)))</selection>.run();
+    }
+
+    void bar(@NotNull Object o) {
+    }
+}

java/java-tests/testData/refactoring/extractMethod/NotNullArgumentLambdaInIfNoBlock_after.java

@@ -0,0 +1,17 @@
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+class C {
+    void foo(@Nullable Object o) {
+        if (o != null)
+            newMethod(o).run();
+    }
+
+    @NotNull
+    private Runnable newMethod(@NotNull Object o) {
+        return (Runnable)(() -> bar(o));
+    }
+
+    void bar(@NotNull Object o) {
+    }
+}

java/java-tests/testData/refactoring/extractMethod/NotNullArgumentLambdaInIf_after.java

@@ -0,0 +1,17 @@
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+class C {
+    void foo(@Nullable Object o) {
+        if (o != null) {
+            newMethod(o);
+        }
+    }
+
+    private void newMethod(@NotNull Object o) {
+        Runnable r = () -> bar(o);
+    }
+
+    void bar(@NotNull Object o) {
+    }
+}

java/java-tests/testData/refactoring/extractMethod/NotNullArgumentLambdaInWhileNoBlock.java

@@ -0,0 +1,12 @@
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+class C {
+    void foo(@Nullable Object o) {
+        while (o != null)
+            <selection>((Runnable)(() -> bar(o)))</selection>.run();
+    }
+
+    void bar(@NotNull Object o) {
+    }
+}

java/java-tests/testData/refactoring/extractMethod/NotNullArgumentLambdaInWhileNoBlock_after.java

@@ -0,0 +1,17 @@
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+class C {
+    void foo(@Nullable Object o) {
+        while (o != null)
+            newMethod(o).run();
+    }
+
+    @NotNull
+    private Runnable newMethod(@NotNull Object o) {
+        return (Runnable)(() -> bar(o));
+    }
+
+    void bar(@NotNull Object o) {
+    }
+}

java/java-tests/testData/refactoring/extractMethod/NotNullArgumentLambdaInsideBody.java

@@ -0,0 +1,11 @@
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+class C {
+    void foo(@NotNull Object o) {
+        Runnable r = () -> <selection>bar(o)</selection>;
+    }
+
+    void bar(@NotNull Object o) {
+    }
+}

java/java-tests/testData/refactoring/extractMethod/NotNullArgumentLambdaInsideBody_after.java

@@ -0,0 +1,15 @@
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+class C {
+    void foo(@NotNull Object o) {
+        Runnable r = () -> newMethod(o);
+    }
+
+    private void newMethod(@NotNull Object o) {
+        bar(o);
+    }
+
+    void bar(@NotNull Object o) {
+    }
+}

java/java-tests/testSrc/com/intellij/java/refactoring/ExtractMethodTest.java

@@ -963,6 +963,46 @@ public class ExtractMethodTest extends LightCodeInsightTestCase {
     doTest();
   }
 
+  public void testNotNullArgumentLambdaBare() throws Exception {
+    doTest();
+  }
+
+  public void testNotNullArgumentLambdaInIf() throws Exception {
+    doTest();
+  }
+
+  public void testNotNullArgumentLambdaInIfNoBlock() throws Exception {
+    doTest();
+  }
+
+  public void testNotNullArgumentLambdaInWhileNoBlock() throws Exception {
+    doTest();
+  }
+
+  public void testNotNullArgumentLambdaInsideBody() throws Exception {
+    doTest();
+  }
+
+  public void testNotNullArgumentLambdaInIfInsideBody() throws Exception {
+    doTest();
+  }
+
+  public void testNotNullArgumentAnonymousClassBare() throws Exception {
+    doTest();
+  }
+
+  public void testNotNullArgumentAnonymousClassInIf() throws Exception {
+    doTest();
+  }
+
+  public void testNotNullArgumentAnonymousClassInIfNoBlock() throws Exception {
+    doTest();
+  }
+
+  public void testNotNullArgumentLambdaInForInitializer() throws Exception {
+    doTest();
+  }
+
   public void testQualifyWhenConflictingNamePresent() throws Exception {
     final CodeStyleSettings settings = CodeStyleSettingsManager.getSettings(getProject());
     settings.ELSE_ON_NEW_LINE = true;