From 2b52caf1feefeb8e2af433da91688b437c0cbedb Mon Sep 17 00:00:00 2001
From: Evgeny Gerashchenko <Evgeny.Gerashchenko@jetbrains.com>
Date: Wed, 14 Mar 2012 16:38:21 +0400
Subject: [PATCH] Added navigation to source classes.

---
 .../libraries/JetClsNavigationPolicy.java     | 10 ++-
 .../libraries/JetSourceNavigationHelper.java  | 87 +++++++++++++++++++
 2 files changed, 96 insertions(+), 1 deletion(-)
 create mode 100644 idea/src/org/jetbrains/jet/plugin/libraries/JetSourceNavigationHelper.java

diff --git a/idea/src/org/jetbrains/jet/plugin/libraries/JetClsNavigationPolicy.java b/idea/src/org/jetbrains/jet/plugin/libraries/JetClsNavigationPolicy.java
index 0a824090c0a..5ecf6b4cb93 100644
--- a/idea/src/org/jetbrains/jet/plugin/libraries/JetClsNavigationPolicy.java
+++ b/idea/src/org/jetbrains/jet/plugin/libraries/JetClsNavigationPolicy.java
@@ -21,6 +21,7 @@ import com.intellij.psi.PsiElement;
 import com.intellij.psi.impl.compiled.*;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import org.jetbrains.jet.lang.psi.JetClass;
 import org.jetbrains.jet.lang.psi.JetDeclaration;
 
 /**
@@ -31,7 +32,14 @@ public class JetClsNavigationPolicy implements ClsCustomNavigationPolicy {
     @Override
     @Nullable
     public PsiElement getNavigationElement(@NotNull ClsClassImpl clsClass) {
-        return getJetDeclarationByClsElement(clsClass);
+        JetClass jetClass = (JetClass) getJetDeclarationByClsElement(clsClass);
+        if (jetClass != null) {
+            JetClass sourceClass = JetSourceNavigationHelper.getSourceClass(jetClass);
+            if (sourceClass != null) {
+                return sourceClass;
+            }
+        }
+        return jetClass;
     }
 
     @Override
diff --git a/idea/src/org/jetbrains/jet/plugin/libraries/JetSourceNavigationHelper.java b/idea/src/org/jetbrains/jet/plugin/libraries/JetSourceNavigationHelper.java
new file mode 100644
index 00000000000..fa7050629ef
--- /dev/null
+++ b/idea/src/org/jetbrains/jet/plugin/libraries/JetSourceNavigationHelper.java
@@ -0,0 +1,87 @@
+/*
+ * Copyright 2010-2012 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.jet.plugin.libraries;
+
+import com.google.common.base.Predicates;
+import com.intellij.openapi.roots.OrderEntry;
+import com.intellij.openapi.roots.OrderRootType;
+import com.intellij.openapi.roots.ProjectFileIndex;
+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.util.Processor;
+import org.jetbrains.annotations.Nullable;
+import org.jetbrains.jet.lang.ModuleConfiguration;
+import org.jetbrains.jet.lang.cfg.pseudocode.JetControlFlowDataTraceFactory;
+import org.jetbrains.jet.lang.descriptors.ClassDescriptor;
+import org.jetbrains.jet.lang.psi.JetClass;
+import org.jetbrains.jet.lang.psi.JetFile;
+import org.jetbrains.jet.lang.psi.JetPsiUtil;
+import org.jetbrains.jet.lang.resolve.AnalyzingUtils;
+import org.jetbrains.jet.lang.resolve.BindingContext;
+import org.jetbrains.jet.plugin.JetFileType;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * @author Evgeny Gerashchenko
+ * @since 3/13/12
+ */
+public class JetSourceNavigationHelper {
+    private JetSourceNavigationHelper() {
+    }
+
+    @Nullable
+    public static JetClass getSourceClass(final JetClass decompiledClass) {
+        VirtualFile decompiledFile = decompiledClass.getContainingFile().getVirtualFile();
+        if (decompiledFile == null) {
+            return null;
+        }
+        ProjectFileIndex projectFileIndex = ProjectFileIndex.SERVICE.getInstance(decompiledClass.getProject());
+        List<OrderEntry> orderEntries = projectFileIndex.getOrderEntriesForFile(decompiledFile);
+        for (OrderEntry orderEntry : orderEntries) {
+            // 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;
+    }
+}