Got rid of some code duplication in JetProjectViewProvider.

idea/src/org/jetbrains/jet/plugin/JetIconProvider.java

@@ -24,8 +24,8 @@ import com.intellij.psi.PsiElement;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.util.PlatformIcons;
 import com.intellij.util.containers.ContainerUtil;
-import com.intellij.util.containers.FilteringIterator;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.psi.*;
 import org.jetbrains.jet.lexer.JetTokens;
 
@@ -43,18 +43,8 @@ public class JetIconProvider extends IconProvider {
     public Icon getIcon(@NotNull PsiElement psiElement, int flags) {
         if (psiElement instanceof JetFile) {
             JetFile file = (JetFile) psiElement;
-            List<JetDeclaration> classes = ContainerUtil.filter(file.getDeclarations(), new Condition<JetDeclaration>() {
+            JetClassOrObject mainClass = getMainClass(file);
-                @Override
+            return mainClass != null ? getIcon(mainClass, flags) : JetIcons.FILE;
-                public boolean value(JetDeclaration jetDeclaration) {
-                    return jetDeclaration instanceof JetClassOrObject;
-                }
-            });
-            if (classes.size() == 1) {
-                if (StringUtil.getPackageName(file.getName()).equals(classes.get(0).getName())) {
-                    return getIcon(classes.get(0), flags);
-                }
-            }
-            return JetIcons.FILE;
         }
         if (psiElement instanceof JetNamespaceHeader) {
             return (flags & Iconable.ICON_FLAG_OPEN) != 0 ? PlatformIcons.PACKAGE_OPEN_ICON : PlatformIcons.PACKAGE_ICON;
@@ -98,4 +88,20 @@ public class JetIconProvider extends IconProvider {
         }
         return null;
     }
+
+    @Nullable
+    public static JetClassOrObject getMainClass(@NotNull JetFile file) {
+        List<JetDeclaration> classes = ContainerUtil.filter(file.getDeclarations(), new Condition<JetDeclaration>() {
+            @Override
+            public boolean value(JetDeclaration jetDeclaration) {
+                return jetDeclaration instanceof JetClassOrObject;
+            }
+        });
+        if (classes.size() == 1) {
+            if (StringUtil.getPackageName(file.getName()).equals(classes.get(0).getName())) {
+                return (JetClassOrObject) classes.get(0);
+            }
+        }
+        return null;
+    }
 }

idea/src/org/jetbrains/jet/plugin/projectView/JetProjectViewProvider.java

@@ -30,6 +30,7 @@ import org.jetbrains.jet.lang.psi.JetClassBody;
 import org.jetbrains.jet.lang.psi.JetClassOrObject;
 import org.jetbrains.jet.lang.psi.JetDeclaration;
 import org.jetbrains.jet.lang.psi.JetFile;
+import org.jetbrains.jet.plugin.JetIconProvider;
 
 import java.util.ArrayList;
 import java.util.Collection;
@@ -57,17 +58,9 @@ public class JetProjectViewProvider implements SelectableTreeStructureProvider,
                 JetFile file = (JetFile) childValue;
                 List<JetDeclaration> declarations = file.getDeclarations();
 
-                ArrayList<JetClassOrObject> classDeclarations = new ArrayList<JetClassOrObject>();
+                JetClassOrObject mainClass = JetIconProvider.getMainClass(file);
-                for (JetDeclaration declaration : declarations) {
+                if (mainClass != null && (!settings.isShowMembers() || declarations.size() == 1)) {
-                    if (declaration instanceof JetClassOrObject) {
+                    result.add(new JetClassOrObjectTreeNode(file.getProject(), mainClass, settings));
-                        classDeclarations.add((JetClassOrObject) declaration);
-                    }
-                }
-                VirtualFile virtualFile = file.getVirtualFile();
-                String nameWithoutExtension = virtualFile != null ? virtualFile.getNameWithoutExtension() : file.getName();
-                if (classDeclarations.size() == 1 && (!settings.isShowMembers() || declarations.size() == 1) &&
-                    nameWithoutExtension.equals(classDeclarations.get(0).getName())) {
-                    result.add(new JetClassOrObjectTreeNode(file.getProject(), classDeclarations.get(0), settings));
                 }
                 else {
                     result.add(new JetFileTreeNode(file.getProject(), file, settings));