From 90a58ee75c3377b123c86318a5f0256a675aad85 Mon Sep 17 00:00:00 2001
From: Tagir Valeev <Tagir.Valeev@jetbrains.com>
Date: Tue, 4 Feb 2020 15:47:49 +0700
Subject: [PATCH] DfaAssist: fix analysis of switch expressions (EA-221783)

GitOrigin-RevId: a721cc9c54ec3487ed87049b4f6fddb7021848a0
---
 .../intellij/debugger/engine/dfaassist/DfaAssist.java  |  5 ++++-
 .../codeInspection/dataFlow/ControlFlowAnalyzer.java   | 10 +++++++++-
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/java/debugger/impl/src/com/intellij/debugger/engine/dfaassist/DfaAssist.java b/java/debugger/impl/src/com/intellij/debugger/engine/dfaassist/DfaAssist.java
index 7ec4360af30e..247371d442ef 100644
--- a/java/debugger/impl/src/com/intellij/debugger/engine/dfaassist/DfaAssist.java
+++ b/java/debugger/impl/src/com/intellij/debugger/engine/dfaassist/DfaAssist.java
@@ -256,7 +256,10 @@ public class DfaAssist implements DebuggerContextListener {
         if (parent instanceof PsiMethod || parent instanceof PsiLambdaExpression || parent instanceof PsiClassInitializer ||
             // We cannot properly restore context if we started from finally, so let's analyze just finally block
             parent instanceof PsiTryStatement && ((PsiTryStatement)parent).getFinallyBlock() == e ||
-            parent instanceof PsiBlockStatement && parent.getParent() instanceof PsiLoopStatement) {
+            parent instanceof PsiBlockStatement && 
+            (parent.getParent() instanceof PsiLoopStatement || 
+             parent.getParent() instanceof PsiSwitchLabeledRuleStatement && 
+             ((PsiSwitchLabeledRuleStatement)parent.getParent()).getEnclosingSwitchBlock() instanceof PsiSwitchExpression)) {
           if (parent.getParent() instanceof PsiDoWhileStatement) {
             return parent.getParent();
           }
diff --git a/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/ControlFlowAnalyzer.java b/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/ControlFlowAnalyzer.java
index 126d21e76194..a7273f4a03db 100644
--- a/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/ControlFlowAnalyzer.java
+++ b/java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/ControlFlowAnalyzer.java
@@ -391,8 +391,16 @@ public class ControlFlowAnalyzer extends JavaElementVisitor {
   public void visitYieldStatement(PsiYieldStatement statement) {
     startElement(statement);
     PsiSwitchExpression enclosing = statement.findEnclosingExpression();
+    PsiExpression expression = statement.getExpression();
     if (enclosing != null && myExpressionBlockContext != null && myExpressionBlockContext.myCodeBlock == enclosing.getBody()) {
-      myExpressionBlockContext.generateReturn(statement.getExpression(), this);
+      myExpressionBlockContext.generateReturn(expression, this);
+    } else {
+      // yield in incorrect location or only part of switch is analyzed
+      if (expression != null) {
+        expression.accept(this);
+        addInstruction(new PopInstruction());
+      }
+      jumpOut(enclosing);
     }
     finishElement(statement);
   }