From 1a6b4ca813d6450a2e5bac3b9c22e6ec62406665 Mon Sep 17 00:00:00 2001
From: Anna Kozlova <anna.kozlova@jetbrains.com>
Date: Fri, 3 Jul 2015 14:54:30 +0300
Subject: [PATCH] restore coverage information for groovy files and for
 multiple classes in one java file (IDEA-142057)

---
 .../view/JavaCoverageViewExtension.java       | 30 +++++++++----------
 1 file changed, 15 insertions(+), 15 deletions(-)

diff --git a/plugins/coverage/src/com/intellij/coverage/view/JavaCoverageViewExtension.java b/plugins/coverage/src/com/intellij/coverage/view/JavaCoverageViewExtension.java
index c6a8c108b734..99799ff40d8f 100644
--- a/plugins/coverage/src/com/intellij/coverage/view/JavaCoverageViewExtension.java
+++ b/plugins/coverage/src/com/intellij/coverage/view/JavaCoverageViewExtension.java
@@ -251,21 +251,7 @@ public class JavaCoverageViewExtension extends CoverageViewExtension {
             }
           });
           for (final PsiFile file : childFiles) {
-            if (file instanceof PsiJavaFile) {
-              PsiClass[] classes = ApplicationManager.getApplication().runReadAction(new Computable<PsiClass[]>() {
-                public PsiClass[] compute() {
-                  return file.isValid() ? ((PsiJavaFile) file).getClasses() : PsiClass.EMPTY_ARRAY;
-                }
-              });
-              if (classes.length > 0) {
-                PsiClass aClass = classes[0];
-                if (!(node instanceof CoverageListRootNode) && getClassCoverageInfo(aClass) == null) continue;
-                children.add(new CoverageListNode(myProject, aClass, mySuitesBundle, myStateBean));
-              }
-            }
-            else if (file instanceof PsiClassOwner) {
-              children.add(new CoverageListNode(myProject, file, mySuitesBundle, myStateBean));
-            }
+            collectFileChildren(file, node, children);
           }
         }
         else if (!myStateBean.myFlattenPackages) {
@@ -287,6 +273,20 @@ public class JavaCoverageViewExtension extends CoverageViewExtension {
     return children;
   }
 
+  protected void collectFileChildren(final PsiFile file, AbstractTreeNode node, List<AbstractTreeNode> children) {
+    if (file instanceof PsiClassOwner) {
+      PsiClass[] classes = ApplicationManager.getApplication().runReadAction(new Computable<PsiClass[]>() {
+        public PsiClass[] compute() {
+          return file.isValid() ? ((PsiClassOwner) file).getClasses() : PsiClass.EMPTY_ARRAY;
+        }
+      });
+      for (PsiClass aClass : classes) {
+        if (!(node instanceof CoverageListRootNode) && getClassCoverageInfo(aClass) == null) continue;
+        children.add(new CoverageListNode(myProject, aClass, mySuitesBundle, myStateBean));
+      }
+    }
+  }
+
   @Nullable
   private PackageAnnotator.ClassCoverageInfo getClassCoverageInfo(final PsiClass aClass) {
     return myAnnotator.getClassCoverageInfo(ApplicationManager.getApplication().runReadAction(new NullableComputable<String>() {