diff --git a/idea/src/org/jetbrains/jet/plugin/JetIconProvider.java b/idea/src/org/jetbrains/jet/plugin/JetIconProvider.java index afdbac7b6e5..b64f5e5b531 100644 --- a/idea/src/org/jetbrains/jet/plugin/JetIconProvider.java +++ b/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 classes = ContainerUtil.filter(file.getDeclarations(), new Condition() { - @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 getIcon(classes.get(0), flags); - } - } - return JetIcons.FILE; + JetClassOrObject mainClass = getMainClass(file); + return mainClass != null ? getIcon(mainClass, flags) : 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 classes = ContainerUtil.filter(file.getDeclarations(), new Condition() { + @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; + } } diff --git a/idea/src/org/jetbrains/jet/plugin/projectView/JetProjectViewProvider.java b/idea/src/org/jetbrains/jet/plugin/projectView/JetProjectViewProvider.java index 7f23c4ab221..31f8e8c54a1 100644 --- a/idea/src/org/jetbrains/jet/plugin/projectView/JetProjectViewProvider.java +++ b/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 declarations = file.getDeclarations(); - ArrayList classDeclarations = new ArrayList(); - for (JetDeclaration declaration : declarations) { - if (declaration instanceof JetClassOrObject) { - 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)); + JetClassOrObject mainClass = JetIconProvider.getMainClass(file); + if (mainClass != null && (!settings.isShowMembers() || declarations.size() == 1)) { + result.add(new JetClassOrObjectTreeNode(file.getProject(), mainClass, settings)); } else { result.add(new JetFileTreeNode(file.getProject(), file, settings));