Java: Support more variants with switch expressions in '.switch' template (IDEA-204010)

java/java-impl/src/com/intellij/codeInsight/template/postfix/templates/SwitchStatementPostfixTemplate.java

@@ -140,7 +140,10 @@ public class SwitchStatementPostfixTemplate extends SurroundPostfixTemplateBase
           if (parent instanceof PsiExpressionStatement) {
             result.add(element);
           }
-          else if (isEnhancedSwitchAvailable && (isVariableInitializer(element, parent) || isRightSideOfAssignment(element, parent))) {
+          else if (isEnhancedSwitchAvailable && (isVariableInitializer(element, parent) ||
+                                                 isRightSideOfAssignment(element, parent) ||
+                                                 isReturnValue(element, parent) ||
+                                                 isArgumentList(parent))) {
             result.add(element);
           }
         }
@@ -165,6 +168,14 @@ public class SwitchStatementPostfixTemplate extends SurroundPostfixTemplateBase
       private boolean isRightSideOfAssignment(PsiElement element, PsiElement parent) {
         return parent instanceof PsiAssignmentExpression && ((PsiAssignmentExpression)parent).getRExpression() == element;
       }
+
+      private boolean isReturnValue(PsiElement element, PsiElement parent) {
+        return parent instanceof PsiReturnStatement && ((PsiReturnStatement)parent).getReturnValue() == element;
+      }
+
+      private boolean isArgumentList(PsiElement parent) {
+        return parent instanceof PsiExpressionList && parent.getParent() instanceof PsiCall;
+      }
     };
   }
 }

java/java-tests/testData/codeInsight/template/postfix/templates/switch/charExprArg.java

@@ -0,0 +1,5 @@
+public class Foo {
+    void f(char x) {
+        System.out.println(x.switch<caret>);
+    }
+}

java/java-tests/testData/codeInsight/template/postfix/templates/switch/charExprArg_after.java

@@ -0,0 +1,7 @@
+public class Foo {
+    void f(char x) {
+        System.out.println(switch (x) {
+            <caret>
+        });
+    }
+}

java/java-tests/testData/codeInsight/template/postfix/templates/switch/charExprReturn.java

@@ -0,0 +1,5 @@
+public class Foo {
+    String f(char x) {
+        return x.switch<caret>
+    }
+}

java/java-tests/testData/codeInsight/template/postfix/templates/switch/charExprReturn_after.java

@@ -0,0 +1,7 @@
+public class Foo {
+    String f(char x) {
+        return switch (x) {
+            <caret>
+        }
+    }
+}

java/java-tests/testData/codeInsight/template/postfix/templates/switch/shortExprArg.java

@@ -0,0 +1,5 @@
+public class Foo {
+    void f(short x) {
+        System.out.println(1 + x.switch<caret>);
+    }
+}

java/java-tests/testData/codeInsight/template/postfix/templates/switch/shortExprArg_after.java

@@ -0,0 +1,7 @@
+public class Foo {
+    void f(short x) {
+        System.out.println(switch (1 + x) {
+            <caret>
+        });
+    }
+}

java/java-tests/testData/codeInsight/template/postfix/templates/switch/shortExprReturn.java

@@ -0,0 +1,5 @@
+public class Foo {
+    String f(short x) {
+        return 1 + x.switch<caret>
+    }
+}

java/java-tests/testData/codeInsight/template/postfix/templates/switch/shortExprReturn_after.java

@@ -0,0 +1,7 @@
+public class Foo {
+    String f(short x) {
+        return switch (1 + x) {
+            <caret>
+        }
+    }
+}

java/java-tests/testSrc/com/intellij/java/codeInsight/template/postfix/templates/SwitchExpressionPostfixTemplateTest.java

@@ -33,6 +33,22 @@ public class SwitchExpressionPostfixTemplateTest extends PostfixTemplateTestCase
     doTest();
   }
 
+  public void testCharExprReturn() {
+    doTest();
+  }
+
+  public void testCharExprArg() {
+    doTest();
+  }
+
+  public void testShortExprReturn() {
+    doTest();
+  }
+
+  public void testShortExprArg() {
+    doTest();
+  }
+
   public void testEnumExprInit() {
     doTest();
   }