SimplifyOptionalCallChainInspection: support another case for unwrapping folded ifPresent

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

@@ -7,6 +7,7 @@ import com.intellij.codeInspection.dataFlow.value.DfaFactMapValue;
 import com.intellij.codeInspection.dataFlow.value.DfaValue;
 import com.intellij.codeInspection.util.LambdaGenerationUtil;
 import com.intellij.codeInspection.util.OptionalRefactoringUtil;
+import com.intellij.codeInspection.util.OptionalUtil;
 import com.intellij.openapi.project.Project;
 import com.intellij.pom.java.LanguageLevel;
 import com.intellij.psi.*;
@@ -831,60 +832,96 @@ public class SimplifyOptionalCallChainsInspection extends AbstractBaseJavaLocalI
     @NotNull
     @Override
     public String getName(@NotNull Context context) {
-      return CommonQuickFixBundle.message("fix.replace.map.with.flat.map.name");
+      return CommonQuickFixBundle.message("fix.eliminate.folded.if.present.name");
     }
 
     @NotNull
     @Override
     public String getDescription(@NotNull Context context) {
-      return CommonQuickFixBundle.message("fix.replace.map.with.flat.map.description");
+      return CommonQuickFixBundle.message("fix.eliminate.folded.if.present.description");
     }
 
     @Nullable
     @Override
     public Context extractContext(@NotNull Project project, @NotNull PsiMethodCallExpression call) {
-      PsiLambdaExpression mapLambda = getLambda(call.getArgumentList().getExpressions()[0]);
-      if (mapLambda == null) return null;
-      PsiMethodCallExpression maybeIfPresentCall = ExpressionUtils.getCallForQualifier(call);
-      if (!OPTIONAL_IF_PRESENT.test(maybeIfPresentCall)) return null;
-      PsiLambdaExpression ifPresentOuterLambda = getLambda(maybeIfPresentCall.getArgumentList().getExpressions()[0]);
-      if (ifPresentOuterLambda == null) return null;
-      PsiExpression argument = LambdaUtil.extractSingleExpressionFromBody(ifPresentOuterLambda.getBody());
-      PsiMethodCallExpression insideLambdaCall = tryCast(argument, PsiMethodCallExpression.class);
-      if (insideLambdaCall == null) return null;
-      if (!OPTIONAL_IF_PRESENT.test(insideLambdaCall)) return null;
-      return new Context(mapLambda, call, insideLambdaCall, maybeIfPresentCall);
+      PsiExpression outerIfPresentQualifier = call.getMethodExpression().getQualifierExpression();
+      PsiMethodCallExpression qualifierCall = tryCast(outerIfPresentQualifier, PsiMethodCallExpression.class);
+
+      PsiLambdaExpression outerIfPresentArgument = tryCast(call.getArgumentList().getExpressions()[0], PsiLambdaExpression.class);
+      if (outerIfPresentArgument == null) return null;
+      if (outerIfPresentArgument.getParameterList().getParametersCount() != 1) return null;
+      PsiParameter parameter = outerIfPresentArgument.getParameterList().getParameters()[0];
+      if (parameter == null) return null;
+      String outerIfPresentParameterName = parameter.getName();
+      if (outerIfPresentParameterName == null) return null;
+      PsiExpression outerIfPresentBodyExpr = LambdaUtil.extractSingleExpressionFromBody(outerIfPresentArgument.getBody());
+      PsiMethodCallExpression outerIfPresentBody = tryCast(outerIfPresentBodyExpr, PsiMethodCallExpression.class);
+      if (!OPTIONAL_IF_PRESENT.test(outerIfPresentBody)) return null;
+      PsiExpression innerIfPresentQualifier = outerIfPresentBody.getMethodExpression().getQualifierExpression();
+      PsiExpression nonTrivialQualifier = ExpressionUtils.isReferenceTo(innerIfPresentQualifier, parameter) ? null : innerIfPresentQualifier;
+      PsiExpression innerIfPresentArgument = outerIfPresentBody.getArgumentList().getExpressions()[0];
+
+      PsiMethodCallExpression mapBefore = null;
+      if (OPTIONAL_MAP.test(qualifierCall)) {
+        // case when map(Value::getOptional).ifPresent(p -> p.ifPresent(...))
+        if (isOptionalTypeParameter(qualifierCall.getType())) {
+          mapBefore = qualifierCall;
+        }
+      }
+      return new Context(mapBefore, nonTrivialQualifier, outerIfPresentParameterName, innerIfPresentArgument);
+    }
+
+    private static boolean isOptionalTypeParameter(@Nullable PsiType type) {
+      PsiClassType classType = tryCast(type, PsiClassType.class);
+      if (classType == null) return false;
+      if (classType.getParameterCount() != 1) return false;
+      PsiType typeParameter = classType.getParameters()[0];
+      PsiClass parameterClass = PsiUtil.resolveClassInClassTypeOnly(typeParameter);
+      if (parameterClass == null) return false;
+      return JAVA_UTIL_OPTIONAL.equals(parameterClass.getQualifiedName());
     }
 
     @Override
     public void apply(@NotNull Project project, @NotNull PsiMethodCallExpression call, @NotNull Context context) {
+      PsiMethodCallExpression mapBefore = context.myMapBefore;
       CommentTracker ct = new CommentTracker();
-      PsiExpression qualifier = Objects.requireNonNull(context.myMapCall.getMethodExpression().getQualifierExpression());
-      String text = ct.text(qualifier) + ".flatMap(" + ct.text(context.myMapLambda) + ")" + ".ifPresent" + ct.text(context.myInnerIfPresentCall.getArgumentList());
-      PsiElement result = ct.replaceAndRestoreComments(context.myOuterIfPresentCall, text);
+      StringBuilder sb = new StringBuilder();
+      PsiExpression qualifer = call.getMethodExpression().getQualifierExpression();
+      assert qualifer != null;
+      sb.append(ct.text(qualifer)).append(".");
+      if (mapBefore != null) {
+        PsiExpression mapArgument = mapBefore.getArgumentList().getExpressions()[0];
+        sb.append("flatMap(").append(ct.text(mapArgument)).append(").");
+      }
+      PsiExpression lambdaBodyAfter = context.myMapLambdaBodyAfter;
+      if (lambdaBodyAfter != null) {
+        sb.append("flatMap(").append(context.myOuterIfPresentVarName).append("->").append(ct.text(lambdaBodyAfter)).append(").");
+      }
+      sb.append("ifPresent(").append(ct.text(context.myInnerIfPresentArgument)).append(")");
+      PsiElement result = ct.replaceAndRestoreComments(call, sb.toString());
       LambdaCanBeMethodReferenceInspection.replaceAllLambdasWithMethodReferences(result);
     }
 
     @NotNull
     @Override
     public CallMatcher getMatcher() {
-      return OPTIONAL_MAP;
+      return OPTIONAL_IF_PRESENT;
     }
 
-    private static class Context {
-      private final @NotNull PsiExpression myMapLambda;
-      private final @NotNull PsiMethodCallExpression myMapCall;
-      private final @NotNull PsiMethodCallExpression myInnerIfPresentCall;
-      private final @NotNull PsiMethodCallExpression myOuterIfPresentCall;
+    class Context {
+      @Nullable PsiMethodCallExpression myMapBefore;
+      @Nullable PsiExpression myMapLambdaBodyAfter;
+      @NotNull String myOuterIfPresentVarName;
+      @NotNull PsiExpression myInnerIfPresentArgument;
 
-      private Context(@NotNull PsiExpression lambda,
-                      @NotNull PsiMethodCallExpression mapCall,
-                      @NotNull PsiMethodCallExpression innerIfPresentCall,
-                      @NotNull PsiMethodCallExpression outerIfPresentCall) {
-        myMapLambda = lambda;
-        myMapCall = mapCall;
-        myInnerIfPresentCall = innerIfPresentCall;
-        myOuterIfPresentCall = outerIfPresentCall;
+      Context(@Nullable PsiMethodCallExpression mapBefore,
+              @Nullable PsiExpression mapLambdaBodyAfter,
+              @NotNull String outerIfPresentVarName,
+              @NotNull PsiExpression innerIfPresentArgument) {
+        myMapBefore = mapBefore;
+        myMapLambdaBodyAfter = mapLambdaBodyAfter;
+        myOuterIfPresentVarName = outerIfPresentVarName;
+        myInnerIfPresentArgument = innerIfPresentArgument;
       }
     }
   }

java/java-tests/testData/inspection/optionalChains/afterIfPresentUnwrapping.java

@@ -0,0 +1,31 @@
+// "Remove folded 'ifPresent' call" "true"
+
+import java.util.Optional;
+
+public class After {
+  static void setDefaultReviewerForTaskOldStyle(Task task) {
+    task.getResponsible().flatMap(Person::getManager).ifPresent(task::setReviewer);
+  }
+
+  class Person {
+    private final Department department; // can be null
+
+    Person(Department department) {
+      this.department = department;
+    }
+
+    private Optional<Person> getManager() {
+      return Optional.ofNullable(department).map(Department::getManager);
+    }
+  }
+
+  interface Department {
+    Person getManager(); // never null
+  }
+
+  interface Task {
+    Optional<Person> getResponsible();
+
+    void setReviewer(Person person);
+  }
+}

java/java-tests/testData/inspection/optionalChains/afterMapOptionalIfPresentUnwrapping.java

@@ -1,10 +1,10 @@
-// "Replace 'map()' with 'flatMap()'" "true"
+// "Remove folded 'ifPresent' call" "true"
 import java.util.*;
 
 class Test {
   native Optional<String> getOptional();
 
   void test(Optional<Test> opt) {
-    opt.flatMap(Test::getOptional).ifPresent(System.out::println);
+    opt.map(Test::getOptional).flatMap(Test::getOptional).ifPresent(System.out::println);
   }
 }

java/java-tests/testData/inspection/optionalChains/afterMapOptionalIfPresentUnwrapping2.java

@@ -0,0 +1,12 @@
+// "Remove folded 'ifPresent' call" "true"
+import java.util.*;
+
+class Test {
+  native String getNonOptional();
+
+  static native Optional<String> getOptional(Test v);
+
+  void test(Optional<Test> opt) {
+    opt.map(Test::getNonOptional).flatMap(x -> getOptional(x)).ifPresent(System.out::println);
+  }
+}

java/java-tests/testData/inspection/optionalChains/beforeIfPresentUnwrapping.java

@@ -0,0 +1,31 @@
+// "Remove folded 'ifPresent' call" "true"
+
+import java.util.Optional;
+
+public class After {
+  static void setDefaultReviewerForTaskOldStyle(Task task) {
+    task.getResponsible().<caret>ifPresent(person -> person.getManager().ifPresent(task::setReviewer));
+  }
+
+  class Person {
+    private final Department department; // can be null
+
+    Person(Department department) {
+      this.department = department;
+    }
+
+    private Optional<Person> getManager() {
+      return Optional.ofNullable(department).map(Department::getManager);
+    }
+  }
+
+  interface Department {
+    Person getManager(); // never null
+  }
+
+  interface Task {
+    Optional<Person> getResponsible();
+
+    void setReviewer(Person person);
+  }
+}

java/java-tests/testData/inspection/optionalChains/beforeMapOptionalIfPresentUnwrapping.java

@@ -1,10 +1,10 @@
-// "Replace 'map()' with 'flatMap()'" "true"
+// "Remove folded 'ifPresent' call" "true"
 import java.util.*;
 
 class Test {
   native Optional<String> getOptional();
 
   void test(Optional<Test> opt) {
-    opt.map<caret>(Test::getOptional).ifPresent(x -> x.ifPresent(System.out::println));
+    opt.map(Test::getOptional).<caret>ifPresent(x -> x.ifPresent(System.out::println));
   }
 }

java/java-tests/testData/inspection/optionalChains/beforeMapOptionalIfPresentUnwrapping2.java

@@ -0,0 +1,12 @@
+// "Remove folded 'ifPresent' call" "true"
+import java.util.*;
+
+class Test {
+  native String getNonOptional();
+
+  static native Optional<String> getOptional(Test v);
+
+  void test(Optional<Test> opt) {
+    opt.map(Test::getNonOptional).<caret>ifPresent(x -> getOptional(x).ifPresent(System.out::println));
+  }
+}

platform/platform-resources-en/src/messages/CommonQuickFixBundle.properties

@@ -17,6 +17,9 @@ fix.replace.x.with.y=Replace ''{0}'' with ''{1}''
 fix.replace.map.with.flat.map.name=Replace 'map()' with 'flatMap()'
 fix.replace.map.with.flat.map.description='map()' can be replaced with 'flatMap()'
 
+fix.eliminate.folded.if.present.name=Remove folded 'ifPresent' call
+fix.eliminate.folded.if.present.description=Folded 'ifPresent' call can be eliminated
+
 fix.call=Call ''{0}''
 
 fix.simplify=Simplify