Separate tree structures for local and non-local root declarations

compiler/frontend/src/org/jetbrains/jet/lang/psi/JetPsiUtil.java

@@ -1030,7 +1030,7 @@ public class JetPsiUtil {
         return header != null ? header.getQualifiedName() : null;
     }
 
-    public static JetElement getLocalizingCodeBlock(@NotNull JetNamedDeclaration declaration) {
+    public static JetElement getEnclosingBlockForLocalDeclaration(@NotNull JetNamedDeclaration declaration) {
         //noinspection unchecked
         JetDeclaration container =
                 PsiTreeUtil.getParentOfType(declaration, JetNamedFunction.class, JetPropertyAccessor.class, JetClassInitializer.class);
@@ -1041,4 +1041,8 @@ public class JetPsiUtil {
                ? ((JetClassInitializer) container).getBody()
                : ((JetDeclarationWithBody) container).getBodyExpression();
     }
+
+    public static boolean isLocal(@NotNull JetNamedDeclaration declaration) {
+        return getEnclosingBlockForLocalDeclaration(declaration) != null;
+    }
 }

idea/src/org/jetbrains/jet/plugin/hierarchy/calls/KotlinCallHierarchyBrowser.java

@@ -31,7 +31,6 @@ import com.intellij.psi.PsiMethod;
 import com.intellij.ui.PopupHandler;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.jet.lang.psi.JetNamedDeclaration;
-import org.jetbrains.jet.plugin.hierarchy.HierarchyUtils;
 
 import javax.swing.*;
 import java.util.Comparator;
@@ -87,7 +86,7 @@ public class KotlinCallHierarchyBrowser extends CallHierarchyBrowserBase {
     @Override
     protected HierarchyTreeStructure createHierarchyTreeStructure(@NotNull String typeName, @NotNull PsiElement psiElement) {
         if (typeName.equals(CALLER_TYPE)) {
-            return new KotlinCallerMethodsTreeStructure(myProject, psiElement, getCurrentScopeType());
+            return KotlinCallerMethodsTreeStructure.newInstance(myProject, psiElement, getCurrentScopeType());
         }
 
         if (typeName.equals(CALLEE_TYPE)) {

idea/src/org/jetbrains/jet/plugin/hierarchy/calls/KotlinCallTreeStructure.java

@@ -1,5 +1,6 @@
 package org.jetbrains.jet.plugin.hierarchy.calls;
 
+import com.google.common.base.Predicate;
 import com.intellij.ide.hierarchy.HierarchyNodeDescriptor;
 import com.intellij.ide.hierarchy.HierarchyTreeStructure;
 import com.intellij.ide.hierarchy.call.CallHierarchyNodeDescriptor;
@@ -7,7 +8,6 @@ import com.intellij.openapi.project.Project;
 import com.intellij.psi.PsiClass;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiMethod;
-import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.util.containers.HashMap;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -19,30 +19,16 @@ import java.util.Map;
 
 public abstract class KotlinCallTreeStructure extends HierarchyTreeStructure {
     protected final String scopeType;
-    protected final JetElement localizingCodeBlock;
-    protected final PsiMethod basePsiMethod;
-    protected final PsiClass basePsiClass;
 
     public KotlinCallTreeStructure(@NotNull Project project, PsiElement element, String scopeType) {
         super(project, createNodeDescriptor(project, element, null));
-
         this.scopeType = scopeType;
+    }
 
-        localizingCodeBlock = element instanceof JetNamedDeclaration
+    protected static JetElement getEnclosingBlockForLocalDeclaration(PsiElement element) {
-                                   ? JetPsiUtil.getLocalizingCodeBlock((JetNamedDeclaration) element)
+        return element instanceof JetNamedDeclaration
+                                   ? JetPsiUtil.getEnclosingBlockForLocalDeclaration((JetNamedDeclaration) element)
                                    : null;
-
-        if (localizingCodeBlock == null) {
-            basePsiMethod = getPsiMethod(element);
-            assert basePsiMethod != null;
-
-            basePsiClass = basePsiMethod.getContainingClass();
-            assert basePsiClass != null;
-        }
-        else {
-            basePsiMethod = null;
-            basePsiClass = null;
-        }
     }
 
     protected static HierarchyNodeDescriptor createNodeDescriptor(Project project, PsiElement element, HierarchyNodeDescriptor parent) {
@@ -58,15 +44,46 @@ public abstract class KotlinCallTreeStructure extends HierarchyTreeStructure {
                                        : ((KotlinCallHierarchyNodeDescriptor)descriptor).getTargetElement();
     }
 
-    private static PsiMethod getClosestContainingClassConstructor(PsiElement element) {
+    private static final Predicate<PsiElement> IS_NON_LOCAL_DECLARATION = new Predicate<PsiElement>() {
-        while (element != null) {
+        @Override
-            JetClassOrObject classOrObject = PsiTreeUtil.getParentOfType(element, JetClassOrObject.class, false);
+        public boolean apply(@javax.annotation.Nullable PsiElement input) {
-            if (classOrObject == null) return null;
+            return input instanceof PsiMethod
+                   || ((input instanceof JetNamedFunction || input instanceof JetClassOrObject || input instanceof JetProperty)
+                       && !JetPsiUtil.isLocal((JetNamedDeclaration) input));
+        }
+    };
 
-            element = classOrObject.getParent();
+    @Nullable
+    protected static PsiMethod getRepresentativePsiMethod(PsiElement element) {
+        while (true) {
+            element = JetPsiUtil.getParentByTypeAndPredicate(element, PsiElement.class, IS_NON_LOCAL_DECLARATION, false);
+            if (element == null) return null;
 
-            PsiClass psiClass = LightClassUtil.getPsiClass(classOrObject);
+            PsiMethod method = getRepresentativePsiMethodForNonLocalDeclaration(element);
-            if (psiClass == null) continue;
+            if (method != null) return method;
+
+            element = element.getParent();
+        }
+    }
+
+    private static PsiMethod getRepresentativePsiMethodForNonLocalDeclaration(PsiElement element) {
+        if (element instanceof PsiMethod) {
+            return (PsiMethod) element;
+        }
+
+        if (element instanceof JetNamedFunction) {
+            return LightClassUtil.getLightClassMethod((JetNamedFunction) element);
+        }
+
+        if (element instanceof JetProperty) {
+            LightClassUtil.PropertyAccessorsPsiMethods propertyMethods =
+                    LightClassUtil.getLightClassPropertyMethods((JetProperty) element);
+            return (propertyMethods.getGetter() != null) ? propertyMethods.getGetter() : propertyMethods.getSetter();
+        }
+
+        if (element instanceof JetClassOrObject) {
+            PsiClass psiClass = LightClassUtil.getPsiClass((JetClassOrObject) element);
+            if (psiClass == null) return null;
 
             PsiMethod[] constructors = psiClass.getConstructors();
             if (constructors.length > 0) return constructors[0];
@@ -75,33 +92,6 @@ public abstract class KotlinCallTreeStructure extends HierarchyTreeStructure {
         return null;
     }
 
-    @Nullable
-    protected final PsiMethod getPsiMethod(PsiElement element) {
-        assert localizingCodeBlock == null : "Can't build light method for local declaration";
-
-        if (element instanceof PsiMethod) {
-            return (PsiMethod) element;
-        }
-
-        PsiMethod method = null;
-        if (element instanceof JetNamedFunction) {
-            method = LightClassUtil.getLightClassMethod((JetNamedFunction) element);
-        }
-        else if (element instanceof JetProperty) {
-            LightClassUtil.PropertyAccessorsPsiMethods propertyMethods =
-                    LightClassUtil.getLightClassPropertyMethods((JetProperty) element);
-
-            if (propertyMethods.getGetter() != null) {
-                method = propertyMethods.getGetter();
-            }
-            else {
-                method = propertyMethods.getSetter();
-            }
-        }
-
-        return method != null ? method : getClosestContainingClassConstructor(element);
-    }
-
     protected static CallHierarchyNodeDescriptor getJavaNodeDescriptor(HierarchyNodeDescriptor originalDescriptor) {
         if (originalDescriptor instanceof CallHierarchyNodeDescriptor) return (CallHierarchyNodeDescriptor) originalDescriptor;
 
@@ -109,7 +99,9 @@ public abstract class KotlinCallTreeStructure extends HierarchyTreeStructure {
         return ((KotlinCallHierarchyNodeDescriptor) originalDescriptor).getJavaDelegate();
     }
 
-    protected final Object[] collectNodeDescriptors(HierarchyNodeDescriptor descriptor, List<? extends PsiElement> calleeElements) {
+    protected Object[] collectNodeDescriptors(
+            HierarchyNodeDescriptor descriptor, List<? extends PsiElement> calleeElements, PsiClass basePsiClass
+    ) {
         HashMap<PsiElement, HierarchyNodeDescriptor> declarationToDescriptorMap = new HashMap<PsiElement, HierarchyNodeDescriptor>();
         for (PsiElement callee : calleeElements) {
             if (basePsiClass != null && !isInScope(basePsiClass, callee, scopeType)) continue;

idea/src/org/jetbrains/jet/plugin/hierarchy/calls/KotlinCalleeMethodsTreeStructure.java

@@ -35,10 +35,16 @@ import java.util.List;
 
 public class KotlinCalleeMethodsTreeStructure extends KotlinCallTreeStructure {
     private final CalleeMethodsTreeStructure javaTreeStructure;
+    private final PsiClass representativePsiClass;
 
     public KotlinCalleeMethodsTreeStructure(@NotNull Project project, @NotNull PsiElement element, String scopeType) {
         super(project, element, scopeType);
-        this.javaTreeStructure = basePsiMethod != null ? new CalleeMethodsTreeStructure(project, basePsiMethod, scopeType) : null;
+
+        PsiMethod representativePsiMethod = getRepresentativePsiMethod(element);
+        assert representativePsiMethod != null;
+
+        this.representativePsiClass = representativePsiMethod.getContainingClass();
+        this.javaTreeStructure = new CalleeMethodsTreeStructure(project, representativePsiMethod, scopeType);
     }
 
     private static List<? extends PsiElement> getCalleeElements(@NotNull JetElement rootElement, BindingContext bindingContext) {
@@ -126,6 +132,6 @@ public class KotlinCalleeMethodsTreeStructure extends KotlinCallTreeStructure {
         BindingContext bindingContext =
                 AnalyzerFacadeWithCache.analyzeFileWithCache((JetFile) targetElement.getContainingFile()).getBindingContext();
         List<? extends PsiElement> calleeDescriptors = getCalleeElements((JetElement) targetElement, bindingContext);
-        return collectNodeDescriptors(descriptor, calleeDescriptors);
+        return collectNodeDescriptors(descriptor, calleeDescriptors, representativePsiClass);
     }
 }

idea/src/org/jetbrains/jet/plugin/hierarchy/calls/KotlinCallerMethodsTreeStructure.java

@@ -43,31 +43,29 @@ import org.jetbrains.jet.plugin.references.JetPsiReference;
 import java.util.ArrayList;
 import java.util.Map;
 
-public class KotlinCallerMethodsTreeStructure extends KotlinCallTreeStructure {
+public abstract class KotlinCallerMethodsTreeStructure extends KotlinCallTreeStructure {
-    private final CallerMethodsTreeStructure javaTreeStructure;
+    private static class WithLocalRoot extends KotlinCallerMethodsTreeStructure {
+        private final JetElement codeBlockForLocalDeclaration;
 
-    public KotlinCallerMethodsTreeStructure(@NotNull Project project, @NotNull PsiElement element, String scopeType) {
+        private WithLocalRoot(
+                @NotNull Project project,
+                @NotNull PsiElement element,
+                String scopeType,
+                JetElement codeBlockForLocalDeclaration
+        ) {
             super(project, element, scopeType);
-        this.javaTreeStructure = basePsiMethod != null ? new CallerMethodsTreeStructure(project, basePsiMethod, scopeType) : null;
+            this.codeBlockForLocalDeclaration = codeBlockForLocalDeclaration;
         }
 
         @Override
         protected Object[] buildChildren(HierarchyNodeDescriptor descriptor) {
-        PsiElement targetElement = getTargetElement(descriptor);
+            final PsiElement element = getTargetElement(descriptor);
 
-        if (localizingCodeBlock != null) {
-            return buildLocalizedChildren(targetElement, (KotlinCallHierarchyNodeDescriptor) descriptor);
-        }
-
-        return processCallers(targetElement, descriptor);
-    }
-
-    private Object[] buildLocalizedChildren(final PsiElement element, KotlinCallHierarchyNodeDescriptor descriptor) {
             BindingContext bindingContext =
                     AnalyzerFacadeWithCache.analyzeFileWithCache((JetFile) element.getContainingFile()).getBindingContext();
 
             final ArrayList<PsiElement> result = new ArrayList<PsiElement>();
-        localizingCodeBlock.accept(new CalleeReferenceVisitorBase(bindingContext, true) {
+            codeBlockForLocalDeclaration.accept(new CalleeReferenceVisitorBase(bindingContext, true) {
                 @Override
                 protected void processDeclaration(JetReferenceExpression reference, PsiElement declaration) {
                     if (!declaration.equals(element)) return;
@@ -84,10 +82,25 @@ public class KotlinCallerMethodsTreeStructure extends KotlinCallTreeStructure {
                     }
                 }
             });
-        return collectNodeDescriptors(descriptor, result);
+            return collectNodeDescriptors(descriptor, result, null);
+        }
     }
 
-    private Object[] processCallers(PsiElement element, HierarchyNodeDescriptor descriptor) {
+    private static class WithNonLocalRoot extends KotlinCallerMethodsTreeStructure {
+        private final CallerMethodsTreeStructure javaTreeStructure;
+        private final PsiClass basePsiClass;
+
+        private WithNonLocalRoot(@NotNull Project project, @NotNull PsiElement element, String scopeType, PsiMethod basePsiMethod) {
+            super(project, element, scopeType);
+
+            this.basePsiClass = basePsiMethod.getContainingClass();
+            this.javaTreeStructure = new CallerMethodsTreeStructure(project, basePsiMethod, scopeType);
+        }
+
+        @Override
+        protected Object[] buildChildren(HierarchyNodeDescriptor descriptor) {
+            PsiElement element = getTargetElement(descriptor);
+
             SearchScope searchScope = getSearchScope(scopeType, basePsiClass);
             Map<PsiElement, HierarchyNodeDescriptor> methodToDescriptorMap =
                     new HashMap<PsiElement, HierarchyNodeDescriptor>();
@@ -196,3 +209,17 @@ public class KotlinCallerMethodsTreeStructure extends KotlinCallTreeStructure {
             };
         }
     }
+
+    private KotlinCallerMethodsTreeStructure(@NotNull Project project, @NotNull PsiElement element, String scopeType) {
+        super(project, element, scopeType);
+    }
+
+    public static KotlinCallerMethodsTreeStructure newInstance(@NotNull Project project, @NotNull PsiElement element, String scopeType) {
+        JetElement codeBlockForLocalDeclaration = getEnclosingBlockForLocalDeclaration(element);
+        if (codeBlockForLocalDeclaration != null) return new WithLocalRoot(project, element, scopeType, codeBlockForLocalDeclaration);
+
+        PsiMethod representativeMethod = getRepresentativePsiMethod(element);
+        assert representativeMethod != null;
+        return new WithNonLocalRoot(project, element, scopeType, representativeMethod);
+    }
+}

idea/tests/org/jetbrains/jet/plugin/hierarchy/AbstractHierarchyTest.java

@@ -120,7 +120,7 @@ public abstract class AbstractHierarchyTest extends HierarchyViewTestBase {
         return new Computable<HierarchyTreeStructure>() {
             @Override
             public HierarchyTreeStructure compute() {
-                return new KotlinCallerMethodsTreeStructure(
+                return KotlinCallerMethodsTreeStructure.newInstance(
                         getProject(),
                         getElementAtCaret(LanguageCallHierarchy.INSTANCE.forLanguage(getLanguage())),
                         HierarchyBrowserBaseEx.SCOPE_PROJECT