[java][switch completion] IDEA-270439 Code completion for pattern matching in switch

Reuse the SealedUtils#findSameFileInheritorsClasses method in order to eliminate the duplicating code and reduce the searching scope. Enable "return" when there is a lambda in a switch label

GitOrigin-RevId: 6e6dbcb8d9b1b16828fcdd29c1e59b2f47a71bb5

java/java-impl/src/com/intellij/codeInsight/completion/JavaKeywordCompletion.java

@@ -20,7 +20,6 @@ import com.intellij.psi.*;
 import com.intellij.psi.codeStyle.CodeStyleManager;
 import com.intellij.psi.filters.FilterPositionUtil;
 import com.intellij.psi.javadoc.PsiDocComment;
-import com.intellij.psi.search.searches.DirectClassInheritorsSearch;
 import com.intellij.psi.templateLanguages.OuterLanguageElement;
 import com.intellij.psi.tree.IElementType;
 import com.intellij.psi.util.InheritanceUtil;
@@ -29,11 +28,15 @@ import com.intellij.psi.util.PsiUtil;
 import com.intellij.util.Consumer;
 import com.intellij.util.ObjectUtils;
 import com.intellij.util.containers.ContainerUtil;
-import one.util.streamex.StreamEx;
+import com.siyeh.ig.psiutils.SealedUtils;
+import org.jetbrains.annotations.Contract;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
-import java.util.*;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Locale;
+import java.util.Set;
 
 import static com.intellij.openapi.util.Conditions.notInstanceOf;
 import static com.intellij.patterns.PsiJavaPatterns.*;
@@ -216,7 +219,7 @@ public class JavaKeywordCompletion {
       addKeyword(new OverridableSpace(createKeyword(PsiKeyword.ASSERT), TailType.INSERT_SPACE));
     }
 
-    if (!psiElement().inside(PsiSwitchExpression.class).accepts(myPosition)) {
+    if (!psiElement().inside(PsiSwitchExpression.class).accepts(myPosition) || psiElement().inside(PsiLambdaExpression.class).accepts(myPosition)) {
       TailType returnTail = getReturnTail(myPosition);
       LookupElement ret = createKeyword(PsiKeyword.RETURN);
       if (returnTail != TailType.NONE) {
@@ -294,6 +297,12 @@ public class JavaKeywordCompletion {
   private void addCaseNullToSwitch() {
     if (!isInsideCaseLabel()) return;
 
+    final PsiSwitchBlock switchBlock = PsiTreeUtil.getParentOfType(myPosition, PsiSwitchBlock.class, false, PsiMember.class);
+    if (switchBlock == null) return;
+
+    final PsiType selectorType = getSelectorType(switchBlock);
+    if (selectorType instanceof PsiPrimitiveType) return;
+
     addKeyword(createKeyword(PsiKeyword.NULL));
   }
 
@@ -375,47 +384,39 @@ public class JavaKeywordCompletion {
     PsiSwitchBlock switchBlock = getSwitchFromLabelPosition(myPosition);
     if (switchBlock == null) return;
 
-    final PsiExpression selector = switchBlock.getExpression();
+    final PsiType selectorType = getSelectorType(switchBlock);
-    if (selector == null) return;
-
-    final PsiType selectorType = selector.getType();
     if (selectorType == null || selectorType instanceof PsiPrimitiveType) return;
 
     final TailType caseRuleTail = TailTypes.forSwitchLabel(switchBlock);
     addKeyword(createCaseRule(PsiKeyword.NULL, caseRuleTail));
     addKeyword(createCaseRule(PsiKeyword.DEFAULT, caseRuleTail));
 
-    addSealedHierarchyCases(selector, caseRuleTail);
+    addSealedHierarchyCases(selectorType);
   }
 
-  private void addSealedHierarchyCases(@Nullable PsiExpression expression, @NotNull TailType caseRuleTail) {
+  @Contract(pure = true)
-    if (expression == null) return;
+  private static @Nullable PsiType getSelectorType(@NotNull PsiSwitchBlock switchBlock) {
 
-    final PsiClass aClass = PsiUtil.resolveClassInClassTypeOnly(expression.getType());
+    final PsiExpression selector = switchBlock.getExpression();
-    if (aClass == null || aClass.isEnum()) return;
+    if (selector == null) return null;
 
-    if (!aClass.hasModifierProperty(PsiModifier.SEALED)) return;
+    return selector.getType();
+  }
 
-    final PsiReferenceList list = aClass.getPermitsList();
+  private void addSealedHierarchyCases(@NotNull PsiType type) {
-    final StreamEx<String> inheritors = list != null ? getFromPermittedList(list) : getFromHierarchy(aClass);
+    final PsiResolveHelper resolver = JavaPsiFacade.getInstance(myPosition.getProject()).getResolveHelper();
 
-    for (String inheritor : inheritors) {
+    final PsiClass aClass = resolver.resolveReferencedClass(type.getCanonicalText(), null);
-      final LookupElement rule = createCaseRule(inheritor, caseRuleTail);
+    if (aClass == null || aClass.isEnum() || !aClass.hasModifierProperty(PsiModifier.SEALED)) return;
-      addKeyword(rule);
+
+    for (PsiClass inheritor : SealedUtils.findSameFileInheritorsClasses(aClass)) {
+      final JavaPsiClassReferenceElement item = AllClassesGetter.createLookupItem(inheritor, AllClassesGetter.TRY_SHORTENING);
+      item.setForcedPresentableName("case " + inheritor.getName());
+      addKeyword(item);
     }
   }
 
-  private static @NotNull StreamEx<String> getFromHierarchy(@NotNull PsiClass aClass) {
+  private static @NotNull LookupElement createCaseRule(@NotNull String caseRuleName, TailType tailType) {
-    final Collection<PsiClass> inheritors = DirectClassInheritorsSearch.search(aClass).findAll();
-    return StreamEx.of(inheritors).map(PsiNamedElement::getName);
-  }
-
-  private static @NotNull StreamEx<String> getFromPermittedList(@NotNull PsiReferenceList list) {
-    final PsiJavaCodeReferenceElement[] elements = list.getReferenceElements();
-    return StreamEx.of(elements).map(PsiQualifiedReference::getReferenceName);
-  }
-
-  private static @NotNull LookupElement createCaseRule(String caseRuleName, TailType tailType) {
     final String prefix = "case ";
 
     final LookupElement lookupElement = LookupElementBuilder

java/java-tests/testData/codeInsight/completion/normal/CompleteReturnInSwitch.java

@@ -0,0 +1,26 @@
+
+class Main {
+
+  private static interface Empty {
+    void f();
+  }
+
+  private Empty f(Object o) {
+    return switch (o) {
+      default -> () -> {
+        retu<caret>
+      };
+    };
+  }
+
+  private Empty g(Object o1, Object o2) {
+    return switch (o1) {
+      default -> () -> {
+        switch (o2) {
+          case null -> { retu<caret> }
+          case default -> { retu<caret> }
+        }
+      };
+    };
+  }
+}

java/java-tests/testData/codeInsight/completion/normal/CompleteReturnInSwitch_after.java

@@ -0,0 +1,26 @@
+
+class Main {
+
+  private static interface Empty {
+    void f();
+  }
+
+  private Empty f(Object o) {
+    return switch (o) {
+      default -> () -> {
+        return;<caret>
+      };
+    };
+  }
+
+  private Empty g(Object o1, Object o2) {
+    return switch (o1) {
+      default -> () -> {
+        switch (o2) {
+          case null -> { return;<caret> }
+          case default -> { return;<caret> }
+        }
+      };
+    };
+  }
+}

java/java-tests/testData/codeInsight/completion/normal/variants/CompletionSealedHierarchyExpr.java

@@ -5,8 +5,8 @@ class Main {
       <caret>
     }
   }
-}
 
-sealed interface Sealed permits Variant1, Variant2 { }
+  private static sealed interface Sealed permits Variant1, Variant2 { }
-final class Variant1 implements Sealed { }
+  private static final class Variant1 implements Sealed { }
-final class Variant2 implements Sealed { }
+  private static final class Variant2 implements Sealed { }
+}

java/java-tests/testData/codeInsight/completion/normal/variants/CompletionSealedHierarchyStmt.java

@@ -5,8 +5,8 @@ class Main {
       <caret>
     }
   }
-}
 
-sealed interface Sealed permits Variant1, Variant2 { }
+  private static sealed interface Sealed permits Variant1, Variant2 { }
-final class Variant1 implements Sealed { }
+  private static final class Variant1 implements Sealed { }
-final class Variant2 implements Sealed { }
+  private static final class Variant2 implements Sealed { }
+}

java/java-tests/testSrc/com/intellij/java/codeInsight/completion/NormalSwitchCompletionTest.groovy

@@ -26,6 +26,8 @@ class NormalSwitchCompletionTest extends NormalCompletionTestCase {
   void testCompletePatternVariableInSwitchExpr() { doTest() }
   void testCompletePatternVariableInSwitchStmt() { doTest() }
 
+  void testCompleteReturnInSwitch() { doTest() }
+
   void testCompleteConstantInSwitchExpr() { doTest() }
   void testCompleteConstantInSwitchStmt() { doTest() }
 

java/java-tests/testSrc/com/intellij/java/codeInsight/completion/NormalSwitchCompletionVariantsTest.groovy

@@ -5,6 +5,7 @@ import com.intellij.JavaTestUtil
 import com.intellij.codeInsight.completion.CompletionType
 import com.intellij.codeInsight.completion.LightFixtureCompletionTestCase
 import com.intellij.testFramework.LightProjectDescriptor
+import com.intellij.testFramework.NeedsIndex
 import groovy.transform.CompileStatic
 import org.jetbrains.annotations.NotNull
 
@@ -28,7 +29,11 @@ class NormalSwitchCompletionVariantsTest extends LightFixtureCompletionTestCase
   void testCompletionPrimitiveTypeStmt() { doTest(COMMON_VARIANTS) }
   void testCompletionVariantsInStmt() { doTest(COMMON_OBJECT_VARIANTS) }
   void testCompletionVariantsInExpr() { doTest(COMMON_OBJECT_VARIANTS) }
+
+  @NeedsIndex.Full
   void testCompletionSealedHierarchyStmt() { doTest(COMMON_OBJECT_VARIANTS + ["case Variant1", "case Variant2"]) }
+
+  @NeedsIndex.Full
   void testCompletionSealedHierarchyExpr() { doTest(COMMON_OBJECT_VARIANTS + ["case Variant1", "case Variant2"]) }
 
   private void doTest(String[] variants) {

plugins/InspectionGadgets/InspectionGadgetsAnalysis/src/com/siyeh/ig/psiutils/SealedUtils.java

@@ -1,4 +1,4 @@
-// Copyright 2000-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+// Copyright 2000-2021 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
 package com.siyeh.ig.psiutils;
 
 import com.intellij.codeInsight.daemon.impl.analysis.JavaModuleGraphUtil;
@@ -18,6 +18,7 @@ import org.jetbrains.annotations.PropertyKey;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Objects;
+import java.util.function.Function;
 import java.util.stream.Stream;
 
 import static com.intellij.util.ObjectUtils.tryCast;
@@ -55,11 +56,21 @@ public final class SealedUtils {
       .anyMatch(parent -> parent != null && parent.hasModifierProperty(PsiModifier.SEALED));
   }
 
+  public static Collection<PsiClass> findSameFileInheritorsClasses(@NotNull PsiClass psiClass, PsiClass @NotNull ... classesToExclude) {
+    return getClasses(psiClass, Function.identity(), classesToExclude);
+  }
+
   public static Collection<String> findSameFileInheritors(@NotNull PsiClass psiClass, PsiClass @NotNull ... classesToExclude) {
+    return getClasses(psiClass, PsiClass::getQualifiedName, classesToExclude);
+  }
+
+  private static @NotNull <T> Collection<T> getClasses(@NotNull PsiClass psiClass,
+                                                       Function<PsiClass, T> mapper,
+                                                       PsiClass @NotNull ... classesToExclude) {
     GlobalSearchScope fileScope = GlobalSearchScope.fileScope(psiClass.getContainingFile());
     return DirectClassInheritorsSearch.search(psiClass, fileScope)
       .filtering(inheritor -> !ArrayUtil.contains(inheritor, classesToExclude))
-      .mapping(inheritor -> inheritor.getQualifiedName())
+      .mapping(mapper)
       .findAll();
   }