From 976a5b5417506a7a9b36f255204fabefdd534ba0 Mon Sep 17 00:00:00 2001
From: Andrey Cherkasov <andrey.cherkasov@jetbrains.com>
Date: Wed, 23 Nov 2022 17:59:00 +0400
Subject: [PATCH] com.siyeh.ig.psiutils.ExpressionUtils#nonStructuralChildren:
 switch expressions support

GitOrigin-RevId: ea7039750712eab413f72d4b88240a9fee5d0dea
---
 .../siyeh/ig/psiutils/ExpressionUtils.java    | 36 ++++++++++++++++++-
 1 file changed, 35 insertions(+), 1 deletion(-)

diff --git a/plugins/InspectionGadgets/InspectionGadgetsAnalysis/src/com/siyeh/ig/psiutils/ExpressionUtils.java b/plugins/InspectionGadgets/InspectionGadgetsAnalysis/src/com/siyeh/ig/psiutils/ExpressionUtils.java
index 49da0c7998a3..945288ac2f72 100644
--- a/plugins/InspectionGadgets/InspectionGadgetsAnalysis/src/com/siyeh/ig/psiutils/ExpressionUtils.java
+++ b/plugins/InspectionGadgets/InspectionGadgetsAnalysis/src/com/siyeh/ig/psiutils/ExpressionUtils.java
@@ -221,9 +221,31 @@ public final class ExpressionUtils {
         if (e instanceof PsiParenthesizedExpression parens) {
           return StreamEx.ofNullable(parens.getExpression());
         }
+        if (e instanceof PsiSwitchExpression switchExpression) {
+          PsiCodeBlock switchBody = switchExpression.getBody();
+          if (switchBody == null) return StreamEx.empty();
+          List<PsiExpression> result = new ArrayList<>();
+          PsiStatement[] statements = switchBody.getStatements();
+          for (PsiStatement statement : statements) {
+            if (statement instanceof PsiSwitchLabeledRuleStatement rule) {
+              PsiStatement ruleBody = rule.getBody();
+              if (ruleBody instanceof PsiBlockStatement blockStatement) {
+                collectYieldExpressions(blockStatement, switchExpression, result);
+              }
+              else if (ruleBody instanceof PsiExpressionStatement expr) {
+                result.add(expr.getExpression());
+              }
+            }
+            else {
+              collectYieldExpressions(statement, switchExpression, result);
+            }
+          }
+          return StreamEx.of(result);
+        }
         return null;
       }).remove(e -> e instanceof PsiConditionalExpression ||
-                   e instanceof PsiParenthesizedExpression)
+                     e instanceof PsiParenthesizedExpression ||
+                     e instanceof PsiSwitchExpression)
       .map(e -> {
         if (e instanceof PsiTypeCastExpression cast) {
           PsiExpression operand = cast.getOperand();
@@ -244,6 +266,18 @@ public final class ExpressionUtils {
       });
   }
 
+  private static void collectYieldExpressions(@NotNull PsiStatement statement,
+                                              @NotNull PsiSwitchExpression switchExpression,
+                                              @NotNull Collection<PsiExpression> result) {
+    Collection<PsiYieldStatement> blockStatement = statement instanceof PsiYieldStatement yieldStatement
+                                                   ? Collections.singleton(yieldStatement)
+                                                   : PsiTreeUtil.findChildrenOfType(statement, PsiYieldStatement.class);
+    List<PsiYieldStatement> myYields = ContainerUtil.filter(blockStatement, st -> st.findEnclosingExpression() == switchExpression);
+    for (PsiYieldStatement yield : myYields) {
+      ContainerUtil.addIfNotNull(result, yield.getExpression());
+    }
+  }
+
   public static boolean isZero(@Nullable PsiExpression expression) {
     if (expression == null) {
       return false;