Using indexes to find subpackages

compiler/cli/src/org/jetbrains/jet/cli/jvm/compiler/CliLightClassGenerationSupport.java

@@ -24,19 +24,19 @@ import com.intellij.openapi.project.Project;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.search.GlobalSearchScope;
 import com.intellij.psi.search.PsiSearchScopeUtil;
+import com.intellij.util.Function;
 import com.intellij.util.SmartList;
+import com.intellij.util.containers.ContainerUtil;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.jet.asJava.LightClassConstructionContext;
 import org.jetbrains.jet.asJava.LightClassGenerationSupport;
 import org.jetbrains.jet.lang.descriptors.ClassDescriptor;
+import org.jetbrains.jet.lang.descriptors.DeclarationDescriptor;
 import org.jetbrains.jet.lang.descriptors.NamespaceDescriptor;
 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.lang.resolve.BindingContext;
+import org.jetbrains.jet.lang.resolve.*;
-import org.jetbrains.jet.lang.resolve.BindingContextUtils;
-import org.jetbrains.jet.lang.resolve.BindingTrace;
-import org.jetbrains.jet.lang.resolve.BindingTraceContext;
 import org.jetbrains.jet.lang.resolve.name.FqName;
 
 import java.util.Collection;
@@ -139,4 +139,22 @@ public class CliLightClassGenerationSupport extends LightClassGenerationSupport
     ) {
         return trace.get(BindingContext.FQNAME_TO_NAMESPACE_DESCRIPTOR, fqName) != null;
     }
+
+    @NotNull
+    @Override
+    public Collection<FqName> getSubPackages(@NotNull FqName fqn, @NotNull GlobalSearchScope scope) {
+        NamespaceDescriptor namespaceDescriptor = trace.get(BindingContext.FQNAME_TO_NAMESPACE_DESCRIPTOR, fqn);
+        if (namespaceDescriptor == null) return Collections.emptyList();
+
+        Collection<DeclarationDescriptor> allDescriptors = namespaceDescriptor.getMemberScope().getAllDescriptors();
+        return ContainerUtil.mapNotNull(allDescriptors, new Function<DeclarationDescriptor, FqName>() {
+            @Override
+            public FqName fun(DeclarationDescriptor input) {
+                if (input instanceof NamespaceDescriptor) {
+                    return DescriptorUtils.getFQName(input).toSafe();
+                }
+                return null;
+            }
+        });
+    }
 }

compiler/jet.as.java.psi/src/org/jetbrains/jet/asJava/JavaElementFinder.java

@@ -16,6 +16,8 @@
 
 package org.jetbrains.jet.asJava;
 
+import com.google.common.base.Function;
+import com.google.common.collect.Collections2;
 import com.google.common.collect.Sets;
 import com.intellij.openapi.extensions.Extensions;
 import com.intellij.openapi.project.Project;
@@ -37,7 +39,10 @@ import org.jetbrains.jet.lang.resolve.name.FqName;
 import org.jetbrains.jet.lang.resolve.name.Name;
 import org.jetbrains.jet.util.QualifiedNamesUtil;
 
-import java.util.*;
+import java.util.Collection;
+import java.util.List;
+import java.util.Set;
+import java.util.WeakHashMap;
 
 import static org.jetbrains.jet.codegen.CodegenUtil.getLocalNameForObject;
 
@@ -216,19 +221,17 @@ public class JavaElementFinder extends PsiElementFinder implements JavaPsiFacade
 
     @NotNull
     @Override
-    public PsiPackage[] getSubPackages(@NotNull PsiPackage psiPackage, @NotNull GlobalSearchScope scope) {
+    public PsiPackage[] getSubPackages(@NotNull PsiPackage psiPackage, @NotNull final GlobalSearchScope scope) {
-        final Collection<JetFile> psiFiles = collectProjectJetFiles(project, GlobalSearchScope.allScope(project));
+        FqName packageFQN = new FqName(psiPackage.getQualifiedName());
 
-        Set<PsiPackage> answer = new HashSet<PsiPackage>();
+        Collection<FqName> subpackages = lightClassGenerationSupport.getSubPackages(packageFQN, scope);
 
-        for (JetFile psiFile : psiFiles) {
+        Collection<PsiPackage> answer = Collections2.transform(subpackages, new Function<FqName, PsiPackage>() {
-            FqName jetRootNamespace = JetPsiUtil.getFQName(psiFile);
+            @Override
-
+            public PsiPackage apply(@Nullable FqName input) {
-            final FqName subPackageFQN = QualifiedNamesUtil.plusOneSegment(new FqName(psiPackage.getQualifiedName()), jetRootNamespace);
+                return new JetLightPackage(psiManager, input, scope);
-            if (subPackageFQN != null) {
-                answer.add(new JetLightPackage(psiManager, subPackageFQN, scope));
             }
-        }
+        });
 
         return answer.toArray(new PsiPackage[answer.size()]);
     }
@@ -263,7 +266,7 @@ public class JavaElementFinder extends PsiElementFinder implements JavaPsiFacade
     @Deprecated
     private synchronized Collection<JetFile> collectProjectJetFiles(final Project project, @NotNull final GlobalSearchScope scope) {
         Collection<JetFile> cachedFiles = jetFiles.get(scope);
-        
+
         if (cachedFiles == null) {
             cachedFiles = JetFilesProvider.getInstance(project).allInScope(scope);
             jetFiles.put(scope, cachedFiles);

compiler/jet.as.java.psi/src/org/jetbrains/jet/asJava/LightClassGenerationSupport.java

@@ -40,7 +40,11 @@ public abstract class LightClassGenerationSupport {
     public abstract Collection<JetClassOrObject> findClassOrObjectDeclarations(@NotNull FqName fqName, @NotNull GlobalSearchScope searchScope);
 
     /*
-     * Returns empty collection for absent package
+     * Finds files whose package declaration is exactly {@code fqName}. For example, if a file declares
+     *     package a.b.c
+     * it will not be returned for fqName "a.b"
+     *
+     * If the resulting collection is empty, it means that this package has not other declarations than sub-packages
      */
     @NotNull
     public abstract Collection<JetFile> findFilesForPackage(@NotNull FqName fqName, @NotNull GlobalSearchScope searchScope);
@@ -53,4 +57,7 @@ public abstract class LightClassGenerationSupport {
     );
 
     public abstract boolean packageExists(@NotNull FqName fqName, @NotNull GlobalSearchScope scope);
+
+    @NotNull
+    public abstract Collection<FqName> getSubPackages(@NotNull FqName fqn, @NotNull GlobalSearchScope scope);
 }

idea/src/org/jetbrains/jet/plugin/caches/resolve/IDELightClassGenerationSupport.java

@@ -16,6 +16,7 @@
 
 package org.jetbrains.jet.plugin.caches.resolve;
 
+import com.google.common.collect.Sets;
 import com.intellij.openapi.project.Project;
 import com.intellij.psi.search.GlobalSearchScope;
 import org.jetbrains.annotations.NotNull;
@@ -23,13 +24,16 @@ import org.jetbrains.jet.asJava.LightClassConstructionContext;
 import org.jetbrains.jet.asJava.LightClassGenerationSupport;
 import org.jetbrains.jet.lang.psi.JetClassOrObject;
 import org.jetbrains.jet.lang.psi.JetFile;
+import org.jetbrains.jet.lang.psi.JetPsiUtil;
 import org.jetbrains.jet.lang.resolve.name.FqName;
 import org.jetbrains.jet.plugin.stubindex.JetAllPackagesIndex;
 import org.jetbrains.jet.plugin.stubindex.JetClassByPackageIndex;
 import org.jetbrains.jet.plugin.stubindex.JetFullClassNameIndex;
 import org.jetbrains.jet.plugin.stubindex.JetPackageDeclarationIndex;
+import org.jetbrains.jet.util.QualifiedNamesUtil;
 
 import java.util.Collection;
+import java.util.Set;
 
 public class IDELightClassGenerationSupport extends LightClassGenerationSupport {
 
@@ -73,4 +77,25 @@ public class IDELightClassGenerationSupport extends LightClassGenerationSupport
     ) {
         return !JetAllPackagesIndex.getInstance().get(fqName.getFqName(), project, scope).isEmpty();
     }
+
+    @NotNull
+    @Override
+    public Collection<FqName> getSubPackages(@NotNull FqName fqn, @NotNull GlobalSearchScope scope) {
+        Collection<JetFile> files = JetAllPackagesIndex.getInstance().get(fqn.getFqName(), project, scope);
+
+        Set<FqName> result = Sets.newHashSet();
+        for (JetFile file : files) {
+            FqName fqName = JetPsiUtil.getFQName(file);
+
+            assert QualifiedNamesUtil.isSubpackageOf(fqn, fqName) : "Registered package is not a subpackage of actually declared package:\n" +
+                                                                    "in index: " + fqn + "\n" +
+                                                                    "declared: " + fqName;
+            FqName subpackage = QualifiedNamesUtil.plusOneSegment(fqn, fqName);
+            if (subpackage != null) {
+                result.add(subpackage);
+            }
+        }
+
+        return result;
+    }
 }