From 9b12e25cfab2461c4fc3eed78c9534d1fe647f3d Mon Sep 17 00:00:00 2001
From: Roman Shevchenko <roman.shevchenko@jetbrains.com>
Date: Mon, 24 Jun 2019 15:03:15 +0200
Subject: [PATCH] [java] exit point highlighting for yields (IDEA-CR-48919)

GitOrigin-RevId: d8b4687e331744d40c254ee261ffd12163d0f640
---
 .../HighlightBreakOutsHandler.java            | 15 +++++++++++---
 .../HighlightExitPointsHandlerFactory.java    | 20 +++----------------
 2 files changed, 15 insertions(+), 20 deletions(-)

diff --git a/java/java-impl/src/com/intellij/codeInsight/highlighting/HighlightBreakOutsHandler.java b/java/java-impl/src/com/intellij/codeInsight/highlighting/HighlightBreakOutsHandler.java
index a81066a2e899..7017cf879149 100644
--- a/java/java-impl/src/com/intellij/codeInsight/highlighting/HighlightBreakOutsHandler.java
+++ b/java/java-impl/src/com/intellij/codeInsight/highlighting/HighlightBreakOutsHandler.java
@@ -1,4 +1,4 @@
-// Copyright 2000-2018 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-2019 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.
 package com.intellij.codeInsight.highlighting;
 
 import com.intellij.openapi.editor.Editor;
@@ -26,6 +26,7 @@ public class HighlightBreakOutsHandler extends HighlightUsagesHandlerBase<PsiEle
   }
 
   @Override
+  @SuppressWarnings("BoundedWildcard")
   protected void selectTargets(List<PsiElement> targets, Consumer<List<PsiElement>> selectionConsumer) {
     selectionConsumer.consume(targets);
   }
@@ -49,6 +50,13 @@ public class HighlightBreakOutsHandler extends HighlightUsagesHandlerBase<PsiEle
         collectSiblings((PsiStatement)parent, exitedElement, exitedElement);
       }
     }
+    else if (parent instanceof PsiYieldStatement) {
+      PsiSwitchExpression enclosingExpression = ((PsiYieldStatement)parent).findEnclosingExpression();
+      if (enclosingExpression != null) {
+        addOccurrence(enclosingExpression.getFirstChild());
+        collectSiblings((PsiStatement)parent, enclosingExpression, enclosingExpression);
+      }
+    }
     addOccurrence(myTarget);
   }
 
@@ -69,7 +77,8 @@ public class HighlightBreakOutsHandler extends HighlightUsagesHandlerBase<PsiEle
       for (PsiStatement psiStatement: statements) {
         if (currentStatement == psiStatement) continue;
         if (psiStatement instanceof PsiContinueStatement && ((PsiContinueStatement)psiStatement).findContinuedStatement() == container ||
-            psiStatement instanceof PsiBreakStatement && ((PsiBreakStatement)psiStatement).findExitedElement() == container) {
+            psiStatement instanceof PsiBreakStatement && ((PsiBreakStatement)psiStatement).findExitedElement() == container ||
+            psiStatement instanceof PsiYieldStatement && ((PsiYieldStatement)psiStatement).findEnclosingExpression() == container) {
           addOccurrence(psiStatement.getFirstChild());
         }
       }
@@ -78,7 +87,7 @@ public class HighlightBreakOutsHandler extends HighlightUsagesHandlerBase<PsiEle
   }
 
   private void highlightLoopDeclaration(PsiLoopStatement statement) {
-    
+
     if (statement instanceof PsiDoWhileStatement) {
       PsiKeyword whileKeyword = ((PsiDoWhileStatement)statement).getWhileKeyword();
       if (whileKeyword != null) {
diff --git a/java/java-impl/src/com/intellij/codeInsight/highlighting/HighlightExitPointsHandlerFactory.java b/java/java-impl/src/com/intellij/codeInsight/highlighting/HighlightExitPointsHandlerFactory.java
index 9a9d455486f2..117fb6e226fb 100644
--- a/java/java-impl/src/com/intellij/codeInsight/highlighting/HighlightExitPointsHandlerFactory.java
+++ b/java/java-impl/src/com/intellij/codeInsight/highlighting/HighlightExitPointsHandlerFactory.java
@@ -1,18 +1,4 @@
-/*
- * Copyright 2000-2014 JetBrains s.r.o.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
+// Copyright 2000-2019 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.
 package com.intellij.codeInsight.highlighting;
 
 import com.intellij.openapi.editor.Editor;
@@ -31,10 +17,10 @@ public class HighlightExitPointsHandlerFactory extends HighlightUsagesHandlerFac
       if (PsiKeyword.RETURN.equals(target.getText()) || PsiKeyword.THROW.equals(target.getText())) {
         return new HighlightExitPointsHandler(editor, file, target);
       }
-      if (PsiKeyword.CONTINUE.equals(target.getText()) || PsiKeyword.BREAK.equals(target.getText())) {
+      if (PsiKeyword.CONTINUE.equals(target.getText()) || PsiKeyword.BREAK.equals(target.getText()) || PsiKeyword.YIELD.equals(target.getText())) {
         return new HighlightBreakOutsHandler(editor, file, target);
       }
     }
     return null;
   }
-}
+}
\ No newline at end of file