From 59879d73ad270285cd70d2b894eea7619187779b Mon Sep 17 00:00:00 2001
From: Alexey Utkin <alexey.utkin@jetbrains.com>
Date: Sun, 21 Jan 2018 17:43:16 +0300
Subject: [PATCH] CPP-2906  Code Fold of control flow statements (if/else,
 do/while, for, switch)

Fix take into account the injected languages (Doxygen comment in C++ as an example)
---
 .../impl/actions/CollapseBlockAction.java     | 36 +++++++++++++++----
 .../impl/CodeInsightTestFixtureImpl.java      | 16 +++++++--
 2 files changed, 43 insertions(+), 9 deletions(-)

diff --git a/platform/lang-impl/src/com/intellij/codeInsight/folding/impl/actions/CollapseBlockAction.java b/platform/lang-impl/src/com/intellij/codeInsight/folding/impl/actions/CollapseBlockAction.java
index cc99902c34fd..dee0f8672b5f 100644
--- a/platform/lang-impl/src/com/intellij/codeInsight/folding/impl/actions/CollapseBlockAction.java
+++ b/platform/lang-impl/src/com/intellij/codeInsight/folding/impl/actions/CollapseBlockAction.java
@@ -6,6 +6,8 @@ package com.intellij.codeInsight.folding.impl.actions;
 import com.intellij.codeInsight.CodeInsightActionHandler;
 import com.intellij.codeInsight.actions.BaseCodeInsightAction;
 import com.intellij.codeInsight.folding.CollapseBlockHandler;
+import com.intellij.injected.editor.EditorWindow;
+import com.intellij.lang.injection.InjectedLanguageManager;
 import com.intellij.openapi.editor.Editor;
 import com.intellij.openapi.project.Project;
 import com.intellij.psi.PsiFile;
@@ -13,23 +15,45 @@ import org.jetbrains.annotations.NotNull;
 
 import java.util.List;
 
-public class CollapseBlockAction  extends BaseCodeInsightAction {
+public class CollapseBlockAction extends BaseCodeInsightAction {
   @NotNull
   @Override
   protected CodeInsightActionHandler getHandler() {
     return new CodeInsightActionHandler() {
       public void invoke(@NotNull final Project project, @NotNull final Editor editor, @NotNull final PsiFile file) {
-        getHandlersForFile(file).forEach(handler->handler.invoke(project, editor, file));
+        executor(project, editor, file, true);
       }
     };
   }
 
   protected boolean isValidForFile(@NotNull Project project, @NotNull Editor editor, @NotNull PsiFile file) {
-    return !getHandlersForFile(file).isEmpty();
+    return executor(project, editor, file, false);
   }
 
-  @NotNull
-  private static List<CollapseBlockHandler> getHandlersForFile(@NotNull PsiFile file) {
-    return CollapseBlockHandler.EP_NAME.allForLanguage(file.getLanguage());
+  private static boolean executor(@NotNull final Project project,
+                                  @NotNull Editor editor,
+                                  @NotNull PsiFile file,
+                                  boolean executeAction) {
+    final InjectedLanguageManager instance = InjectedLanguageManager.getInstance(project);
+    while (true) {
+      final List<CollapseBlockHandler> handlers = CollapseBlockHandler.EP_NAME.allForLanguage(file.getLanguage());
+      if (handlers.isEmpty()) {
+        if (!instance.isInjectedFragment(file) || !(editor instanceof EditorWindow)) {
+          return false;
+        }
+        file = instance.getTopLevelFile(file);
+        if (file == null) {
+          return false;
+        }
+        editor = ((EditorWindow)editor).getDelegate();
+        continue;
+      }
+      if (executeAction) {
+        final Editor finalEditor = editor;
+        final PsiFile finalFile = file;
+        handlers.forEach(handler -> handler.invoke(project, finalEditor, finalFile));
+      }
+      return true;
+    }
   }
 }
diff --git a/platform/testFramework/src/com/intellij/testFramework/fixtures/impl/CodeInsightTestFixtureImpl.java b/platform/testFramework/src/com/intellij/testFramework/fixtures/impl/CodeInsightTestFixtureImpl.java
index b3818c7e47de..2cb50d8bec18 100644
--- a/platform/testFramework/src/com/intellij/testFramework/fixtures/impl/CodeInsightTestFixtureImpl.java
+++ b/platform/testFramework/src/com/intellij/testFramework/fixtures/impl/CodeInsightTestFixtureImpl.java
@@ -137,6 +137,7 @@ import com.intellij.util.indexing.FileBasedIndex;
 import com.intellij.util.io.ReadOnlyAttributeUtil;
 import com.intellij.util.ui.UIUtil;
 import junit.framework.ComparisonFailure;
+import org.jetbrains.annotations.Contract;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.annotations.TestOnly;
@@ -1642,14 +1643,23 @@ public class CodeInsightTestFixtureImpl extends BaseFixture implements CodeInsig
 
   @NotNull
   public String getFoldingDescription(boolean withCollapseStatus) {
-    CodeFoldingManager.getInstance(getProject()).buildInitialFoldings(myEditor);
-    return getTagsFromSegments(myEditor.getDocument().getText(),
-                               Arrays.asList(myEditor.getFoldingModel().getAllFoldRegions()),
+    final Editor topEditor = getTopEditor(myEditor);
+    CodeFoldingManager.getInstance(getProject()).buildInitialFoldings(topEditor);
+    return getTagsFromSegments(topEditor.getDocument().getText(),
+                               Arrays.asList(topEditor.getFoldingModel().getAllFoldRegions()),
                                FOLD,
                                foldRegion -> "text=\'" + foldRegion.getPlaceholderText() + "\'"
                                              + (withCollapseStatus ? " expand=\'" + foldRegion.isExpanded() + "\'" : ""));
   }
 
+  @Contract("null -> null")
+  private static Editor getTopEditor(Editor editor) {
+    while(editor instanceof EditorWindow) {
+      editor = ((EditorWindow)editor).getDelegate();
+    }
+    return editor;
+  }
+
   @NotNull
   public static <T extends Segment> String getTagsFromSegments(@NotNull String text,
                                                                @NotNull Collection<T> segments,