Store information about package declaration in context

This information is used for resolving references in plugin

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

@@ -896,4 +896,18 @@ public class JetPsiUtil {
         }
         return null;
     }
+
+    public static NavigatablePsiElement getPackageReference(@NotNull JetFile file, int partIndex) {
+        JetNamespaceHeader header = file.getNamespaceHeader();
+        if (header == null) {
+            throw new IllegalArgumentException("Should be called only for files with namespace: " + file);
+        }
+
+        List<JetSimpleNameExpression> names = header.getParentNamespaceNames();
+        if (!(0 <= partIndex && partIndex < names.size() + 1)) {
+            throw new IndexOutOfBoundsException(String.format("%s index for file with header %s is out of range", partIndex, header.getText()));
+        }
+
+        return (names.size() > partIndex) ? names.get(partIndex) : header.getLastPartExpression();
+    }
 }

compiler/frontend/src/org/jetbrains/jet/lang/resolve/lazy/declarations/AbstractPsiBasedDeclarationProvider.java

@@ -65,7 +65,6 @@ public abstract class AbstractPsiBasedDeclarationProvider implements Declaration
                 throw new IllegalArgumentException("Unknown declaration: " + declaration);
             }
         }
-
     }
 
     private final NotNullLazyValue<Index> index;

compiler/frontend/src/org/jetbrains/jet/lang/resolve/lazy/declarations/EmptyPackageMemberDeclarationProvider.java

@@ -16,6 +16,7 @@
 
 package org.jetbrains.jet.lang.resolve.lazy.declarations;
 
+import com.intellij.psi.NavigatablePsiElement;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.jet.lang.psi.JetClassOrObject;
 import org.jetbrains.jet.lang.psi.JetDeclaration;
@@ -44,6 +45,12 @@ public class EmptyPackageMemberDeclarationProvider implements PackageMemberDecla
         return Collections.emptyList();
     }
 
+    @NotNull
+    @Override
+    public Collection<NavigatablePsiElement> getPackageDeclarations(FqName fqName) {
+        return Collections.emptyList();
+    }
+
     @Override
     public List<JetDeclaration> getAllDeclarations() {
         return Collections.emptyList();

compiler/frontend/src/org/jetbrains/jet/lang/resolve/lazy/declarations/FileBasedDeclarationProviderFactory.java

@@ -18,23 +18,25 @@ package org.jetbrains.jet.lang.resolve.lazy.declarations;
 
 import com.google.common.base.Predicate;
 import com.google.common.base.Predicates;
-import com.google.common.collect.Collections2;
+import com.google.common.collect.*;
-import com.google.common.collect.HashMultimap;
-import com.google.common.collect.Multimap;
-import com.google.common.collect.Sets;
 import com.intellij.openapi.util.Computable;
+import com.intellij.psi.NavigatablePsiElement;
 import com.intellij.util.Function;
+import com.intellij.util.containers.ContainerUtil;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.psi.JetFile;
 import org.jetbrains.jet.lang.psi.JetNamespaceHeader;
+import org.jetbrains.jet.lang.psi.JetPsiUtil;
 import org.jetbrains.jet.lang.resolve.lazy.data.JetClassLikeInfo;
 import org.jetbrains.jet.lang.resolve.lazy.storage.MemoizedFunctionToNullable;
 import org.jetbrains.jet.lang.resolve.lazy.storage.NotNullLazyValue;
 import org.jetbrains.jet.lang.resolve.lazy.storage.StorageManager;
 import org.jetbrains.jet.lang.resolve.name.FqName;
+import org.jetbrains.jet.util.QualifiedNamesUtil;
 
 import java.util.Collection;
+import java.util.Collections;
 import java.util.Set;
 
 public class FileBasedDeclarationProviderFactory implements DeclarationProviderFactory {
@@ -116,6 +118,27 @@ public class FileBasedDeclarationProviderFactory implements DeclarationProviderF
         });
     }
 
+    /*package*/ Collection<NavigatablePsiElement> getPackageDeclarations(@NotNull final FqName fqName) {
+        if (fqName.isRoot()) {
+            return Collections.emptyList();
+        }
+
+        Collection<NavigatablePsiElement> resultElements = Lists.newArrayList();
+        for (FqName declaredPackage : index.compute().filesByPackage.keys()) {
+            if (QualifiedNamesUtil.isSubpackageOf(declaredPackage, fqName)) {
+                Collection<JetFile> files = index.compute().filesByPackage.get(declaredPackage);
+                resultElements.addAll(ContainerUtil.map(files, new Function<JetFile, NavigatablePsiElement>() {
+                    @Override
+                    public NavigatablePsiElement fun(JetFile file) {
+                        return JetPsiUtil.getPackageReference(file, QualifiedNamesUtil.numberOfSegments(fqName) - 1);
+                    }
+                }));
+            }
+        }
+
+        return resultElements;
+    }
+
     @Override
     public PackageMemberDeclarationProvider getPackageMemberDeclarationProvider(@NotNull FqName packageFqName) {
         return packageDeclarationProviders.fun(packageFqName);

compiler/frontend/src/org/jetbrains/jet/lang/resolve/lazy/declarations/FileBasedPackageMemberDeclarationProvider.java

@@ -17,8 +17,10 @@
 package org.jetbrains.jet.lang.resolve.lazy.declarations;
 
 import com.intellij.openapi.util.Computable;
+import com.intellij.psi.NavigatablePsiElement;
 import org.jetbrains.annotations.NotNull;
-import org.jetbrains.jet.lang.psi.*;
+import org.jetbrains.jet.lang.psi.JetDeclaration;
+import org.jetbrains.jet.lang.psi.JetFile;
 import org.jetbrains.jet.lang.resolve.lazy.storage.NotNullLazyValue;
 import org.jetbrains.jet.lang.resolve.lazy.storage.StorageManager;
 import org.jetbrains.jet.lang.resolve.name.FqName;
@@ -30,7 +32,7 @@ public class FileBasedPackageMemberDeclarationProvider extends AbstractPsiBasedD
 
     private final FqName fqName;
     private final FileBasedDeclarationProviderFactory factory;
-    private final Collection<JetFile> allFiles;
+    private final Collection<JetFile> packageFiles;
     private final NotNullLazyValue<Collection<FqName>> allDeclaredPackages;
 
 
@@ -38,12 +40,12 @@ public class FileBasedPackageMemberDeclarationProvider extends AbstractPsiBasedD
             @NotNull StorageManager storageManager,
             @NotNull FqName _fqName,
             @NotNull FileBasedDeclarationProviderFactory _factory,
-            @NotNull Collection<JetFile> allFiles
+            @NotNull Collection<JetFile> packageFiles
     ) {
         super(storageManager);
         this.fqName = _fqName;
         this.factory = _factory;
-        this.allFiles = allFiles;
+        this.packageFiles = packageFiles;
         this.allDeclaredPackages = storageManager.createLazyValue(new Computable<Collection<FqName>>() {
             @Override
             public Collection<FqName> compute() {
@@ -54,8 +56,9 @@ public class FileBasedPackageMemberDeclarationProvider extends AbstractPsiBasedD
 
     @Override
     protected void doCreateIndex(@NotNull Index index) {
-        for (JetFile file : allFiles) {
+        for (JetFile file : packageFiles) {
             for (JetDeclaration declaration : file.getDeclarations()) {
+                assert fqName.asString().equals(file.getPackageName()) : "Files declaration utils contains file with invalid namespace";
                 index.putToIndex(declaration);
             }
         }
@@ -71,6 +74,12 @@ public class FileBasedPackageMemberDeclarationProvider extends AbstractPsiBasedD
         return allDeclaredPackages.compute();
     }
 
+    @NotNull
+    @Override
+    public Collection<NavigatablePsiElement> getPackageDeclarations(FqName fqName) {
+        return factory.getPackageDeclarations(fqName);
+    }
+
     @Override
     public String toString() {
         return "Declarations for package " + fqName;

compiler/frontend/src/org/jetbrains/jet/lang/resolve/lazy/declarations/PackageMemberDeclarationProvider.java

@@ -16,6 +16,7 @@
 
 package org.jetbrains.jet.lang.resolve.lazy.declarations;
 
+import com.intellij.psi.NavigatablePsiElement;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.jet.lang.resolve.name.FqName;
 import org.jetbrains.jet.lang.resolve.name.Name;
@@ -27,4 +28,7 @@ public interface PackageMemberDeclarationProvider extends DeclarationProvider {
     boolean isPackageDeclared(@NotNull Name name);
 
     Collection<FqName> getAllDeclaredPackages();
+
+    @NotNull
+    Collection<NavigatablePsiElement> getPackageDeclarations(FqName fqName);
 }

compiler/frontend/src/org/jetbrains/jet/lang/resolve/lazy/descriptors/LazyPackageDescriptor.java

@@ -16,11 +16,14 @@
 
 package org.jetbrains.jet.lang.resolve.lazy.descriptors;
 
+import com.google.common.collect.Iterables;
+import com.intellij.psi.NavigatablePsiElement;
 import org.jetbrains.annotations.NotNull;
-import org.jetbrains.jet.lang.descriptors.impl.AbstractNamespaceDescriptorImpl;
 import org.jetbrains.jet.lang.descriptors.NamespaceDescriptor;
-import org.jetbrains.jet.lang.descriptors.impl.NamespaceDescriptorParent;
 import org.jetbrains.jet.lang.descriptors.annotations.AnnotationDescriptor;
+import org.jetbrains.jet.lang.descriptors.impl.AbstractNamespaceDescriptorImpl;
+import org.jetbrains.jet.lang.descriptors.impl.NamespaceDescriptorParent;
+import org.jetbrains.jet.lang.resolve.BindingContext;
 import org.jetbrains.jet.lang.resolve.DescriptorUtils;
 import org.jetbrains.jet.lang.resolve.lazy.ForceResolveUtil;
 import org.jetbrains.jet.lang.resolve.lazy.LazyDescriptor;
@@ -51,6 +54,11 @@ public class LazyPackageDescriptor extends AbstractNamespaceDescriptorImpl imple
         this.lazyScope = new LazyPackageMemberScope(resolveSession, declarationProvider, this);
 
         this.memberScope = new ChainedScope(this, "Lazy package members scope: " + name, lazyScope, scope);
+
+        NavigatablePsiElement declaration = Iterables.getFirst(declarationProvider.getPackageDeclarations(getFqName()), null);
+        if (declaration != null) {
+            resolveSession.getTrace().record(BindingContext.NAMESPACE, declaration, this);
+        }
     }
 
     @NotNull

compiler/frontend/src/org/jetbrains/jet/util/QualifiedNamesUtil.java

@@ -42,7 +42,7 @@ public final class QualifiedNamesUtil {
         String subpackageNameStr = subpackageName.asString();
         String packageNameStr = packageName.asString();
 
-        return (subpackageNameStr.startsWith(packageNameStr) && subpackageNameStr.charAt(packageNameStr.length()) == '.');
+        return isSubpackageOf(subpackageNameStr, packageNameStr);
     }
 
     public static boolean isOneSegmentFQN(@NotNull String fqn) {
@@ -78,6 +78,14 @@ public final class QualifiedNamesUtil {
         return new FqName(fqNameStr.substring(fqNameStr.indexOf('.'), fqNameStr.length()));
     }
 
+    public static int numberOfSegments(@NotNull FqName fqName) {
+        if (fqName.isRoot()) {
+            return 0;
+        }
+
+        return 1 + numberOfSegments(fqName.parent());
+    }
+
     @NotNull
     public static FqName combine(@NotNull FqName first, @NotNull Name second) {
         return first.child(second);
@@ -186,4 +194,9 @@ public final class QualifiedNamesUtil {
 
         return state != AFTER_DOT;
     }
+
+    public static boolean isSubpackageOf(String subpackageNameStr, String packageNameStr) {
+        return subpackageNameStr.equals(packageNameStr) ||
+                (subpackageNameStr.startsWith(packageNameStr) && subpackageNameStr.charAt(packageNameStr.length()) == '.');
+    }
 }

idea/src/org/jetbrains/jet/plugin/stubindex/resolve/StubPackageMemberDeclarationProvider.java

@@ -17,6 +17,7 @@
 package org.jetbrains.jet.plugin.stubindex.resolve;
 
 import com.intellij.openapi.project.Project;
+import com.intellij.psi.NavigatablePsiElement;
 import com.intellij.psi.search.GlobalSearchScope;
 import com.intellij.util.Function;
 import com.intellij.util.containers.ContainerUtil;
@@ -25,14 +26,17 @@ import org.jetbrains.jet.asJava.LightClassGenerationSupport;
 import org.jetbrains.jet.lang.psi.JetFile;
 import org.jetbrains.jet.lang.psi.JetNamedFunction;
 import org.jetbrains.jet.lang.psi.JetProperty;
+import org.jetbrains.jet.lang.psi.JetPsiUtil;
 import org.jetbrains.jet.lang.resolve.lazy.declarations.PackageMemberDeclarationProvider;
 import org.jetbrains.jet.lang.resolve.name.FqName;
 import org.jetbrains.jet.lang.resolve.name.Name;
 import org.jetbrains.jet.plugin.stubindex.JetAllPackagesIndex;
 import org.jetbrains.jet.plugin.stubindex.JetTopLevelFunctionsFqnNameIndex;
 import org.jetbrains.jet.plugin.stubindex.JetTopLevelPropertiesFqnNameIndex;
+import org.jetbrains.jet.util.QualifiedNamesUtil;
 
 import java.util.Collection;
+import java.util.Collections;
 
 public class StubPackageMemberDeclarationProvider extends AbstractStubDeclarationProvider implements PackageMemberDeclarationProvider {
     private final FqName fqName;
@@ -76,4 +80,20 @@ public class StubPackageMemberDeclarationProvider extends AbstractStubDeclaratio
             }
         });
     }
+
+    @NotNull
+    @Override
+    public Collection<NavigatablePsiElement> getPackageDeclarations(final FqName fqName) {
+        if (fqName.isRoot()) {
+            return Collections.emptyList();
+        }
+
+        Collection<JetFile> files = JetAllPackagesIndex.getInstance().get(fqName.asString(), project, searchScope);
+        return ContainerUtil.map(files, new Function<JetFile, NavigatablePsiElement>() {
+            @Override
+            public NavigatablePsiElement fun(JetFile file) {
+                return JetPsiUtil.getPackageReference(file, QualifiedNamesUtil.numberOfSegments(fqName) - 1);
+            }
+        });
+    }
 }

idea/tests/org/jetbrains/jet/plugin/stubs/resolve/JetStubResolveTest.java

@@ -19,6 +19,7 @@ package org.jetbrains.jet.plugin.stubs.resolve;
 import com.google.common.base.Function;
 import com.google.common.collect.Collections2;
 import com.google.common.collect.Lists;
+import com.intellij.psi.NavigatablePsiElement;
 import com.intellij.psi.search.GlobalSearchScope;
 import com.intellij.testFramework.fixtures.LightCodeInsightFixtureTestCase;
 import org.jetbrains.jet.lang.psi.JetNamedFunction;
@@ -80,4 +81,16 @@ public class JetStubResolveTest extends LightCodeInsightFixtureTestCase {
         assertSize(1, testProperties);
         assertEquals("val test = 12", testProperties.get(0).getText());
     }
+
+    public void testGetPackageByIntermediatePartition() {
+        myFixture.configureByText(JetFileType.INSTANCE, "package first.second.third");
+
+        StubPackageMemberDeclarationProvider provider = new StubPackageMemberDeclarationProvider(
+                new FqName("first.second"), getProject(), GlobalSearchScope.projectScope(getProject()));
+
+        Collection<NavigatablePsiElement> packageDeclarations = provider.getPackageDeclarations(new FqName("first.second"));
+
+        assertSize(1, packageDeclarations);
+        assertEquals("second", packageDeclarations.iterator().next().getText());
+    }
 }