Extracted getting all library source dirs and building binding context for library from JetSourceNavigationHelper.getSourceClass().

idea/src/org/jetbrains/jet/plugin/libraries/JetSourceNavigationHelper.java

@@ -17,6 +17,7 @@
 package org.jetbrains.jet.plugin.libraries;
 
 import com.google.common.base.Predicates;
+import com.intellij.openapi.project.Project;
 import com.intellij.openapi.roots.OrderEntry;
 import com.intellij.openapi.roots.OrderRootType;
 import com.intellij.openapi.roots.ProjectFileIndex;
@@ -24,6 +25,7 @@ import com.intellij.openapi.vfs.VfsUtil;
 import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
+import com.intellij.psi.PsiManager;
 import com.intellij.util.Processor;
 import jet.Tuple0;
 import jet.Tuple1;
@@ -38,6 +40,7 @@ import org.jetbrains.jet.lang.resolve.BindingContext;
 import org.jetbrains.jet.plugin.JetFileType;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 
 /**
@@ -50,42 +53,52 @@ public class JetSourceNavigationHelper {
 
     @Nullable
     public static JetClass getSourceClass(final @NotNull JetClass decompiledClass) {
-        VirtualFile decompiledFile = decompiledClass.getContainingFile().getVirtualFile();
+        for (VirtualFile sourceDir : getAllSourceDirs(decompiledClass)) {
-        if (decompiledFile == null) {
+            BindingContext bindingContext = analyzeLibrary(sourceDir, decompiledClass.getProject());
-            return null;
+            ClassDescriptor cd = bindingContext.get(BindingContext.FQNAME_TO_CLASS_DESCRIPTOR, JetPsiUtil.getFQName(decompiledClass));
-        }
+            if (cd != null) {
-        ProjectFileIndex projectFileIndex = ProjectFileIndex.SERVICE.getInstance(decompiledClass.getProject());
+                PsiElement declaration = bindingContext.get(BindingContext.DESCRIPTOR_TO_DECLARATION, cd);
-        List<OrderEntry> orderEntries = projectFileIndex.getOrderEntriesForFile(decompiledFile);
+                assert declaration instanceof JetClass;
-        for (OrderEntry orderEntry : orderEntries) {
+                return (JetClass) declaration;
-            // TODO this is too slow, needs optimizing
-            VirtualFile[] sourceDirs = orderEntry.getFiles(OrderRootType.SOURCES);
-            for (VirtualFile sourceDir : sourceDirs) {
-                final List<JetFile> libraryFiles = new ArrayList<JetFile>();
-                VfsUtil.processFilesRecursively(sourceDir, new Processor<VirtualFile>() {
-                    @Override
-                    public boolean process(VirtualFile virtualFile) {
-                        if (virtualFile.getFileType() == JetFileType.INSTANCE) {
-                            libraryFiles.add((JetFile) decompiledClass.getManager().findFile(virtualFile));
-                        }
-                        return true;
-                    }
-                });
-                BindingContext bindingContext = AnalyzingUtils.analyzeFiles(decompiledClass.getProject(),
-                                                                            ModuleConfiguration.EMPTY,
-                                                                            libraryFiles,
-                                                                            Predicates.<PsiFile>alwaysTrue(),
-                                                                            JetControlFlowDataTraceFactory.EMPTY);
-                ClassDescriptor cd = bindingContext.get(BindingContext.FQNAME_TO_CLASS_DESCRIPTOR, JetPsiUtil.getFQName(decompiledClass));
-                if (cd != null) {
-                    PsiElement declaration = bindingContext.get(BindingContext.DESCRIPTOR_TO_DECLARATION, cd);
-                    assert declaration instanceof JetClass;
-                    return (JetClass) declaration;
-                }
             }
         }
         return null;
     }
 
+    private static BindingContext analyzeLibrary(VirtualFile sourceDir, final Project project) {
+        final List<JetFile> libraryFiles = new ArrayList<JetFile>();
+        VfsUtil.processFilesRecursively(sourceDir, new Processor<VirtualFile>() {
+            @Override
+            public boolean process(VirtualFile virtualFile) {
+                if (virtualFile.getFileType() == JetFileType.INSTANCE) {
+                    libraryFiles.add((JetFile) PsiManager.getInstance(project).findFile(virtualFile));
+                }
+                return true;
+            }
+        });
+        return AnalyzingUtils.analyzeFiles(project,
+                                           ModuleConfiguration.EMPTY,
+                                           libraryFiles,
+                                           Predicates.<PsiFile>alwaysTrue(),
+                                           JetControlFlowDataTraceFactory.EMPTY);
+    }
+
+    @NotNull
+    private static List<VirtualFile> getAllSourceDirs(@NotNull JetClass decompiledClass) {
+        List<VirtualFile> allSourceDirs = new ArrayList<VirtualFile>();
+
+        VirtualFile decompiledFile = decompiledClass.getContainingFile().getVirtualFile();
+        if (decompiledFile != null) {
+            ProjectFileIndex projectFileIndex = ProjectFileIndex.SERVICE.getInstance(decompiledClass.getProject());
+            List<OrderEntry> orderEntries = projectFileIndex.getOrderEntriesForFile(decompiledFile);
+            for (OrderEntry orderEntry : orderEntries) {
+                Collections.addAll(allSourceDirs, orderEntry.getFiles(OrderRootType.SOURCES));
+            }
+        }
+
+        return allSourceDirs;
+    }
+
     @Nullable
     public static JetProperty getSourceProperty(final @NotNull JetProperty decompiledProperty) {
         String propertyName = decompiledProperty.getName();