Blinking class object tests fixed

Now lazy descriptors for all class objects are retained by the owning class

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

@@ -18,7 +18,6 @@ package org.jetbrains.jet.lang.psi;
 
 import com.intellij.lang.ASTNode;
 import com.intellij.psi.PsiElement;
-import com.intellij.psi.stubs.StubElement;
 import com.intellij.psi.tree.TokenSet;
 import com.intellij.psi.util.PsiTreeUtil;
 import org.jetbrains.annotations.NotNull;
@@ -65,6 +64,12 @@ public class JetClassBody extends JetElementImplStub<PsiJetClassBodyStub> implem
         return (JetClassObject) findChildByType(JetNodeTypes.CLASS_OBJECT);
     }
 
+    @NotNull
+    public List<JetClassObject> getAllClassObjects() {
+        //noinspection unchecked
+        return (List) findChildrenByType(JetNodeTypes.CLASS_OBJECT);
+    }
+
     @Nullable
     public PsiElement getRBrace() {
         ASTNode[] children = getNode().getChildren(TokenSet.create(JetTokens.RBRACE));

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

@@ -22,7 +22,6 @@ import com.intellij.psi.PsiElement;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.util.Function;
 import com.intellij.util.containers.ContainerUtil;
-import kotlin.Function0;
 import kotlin.Function1;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -40,7 +39,6 @@ import org.jetbrains.jet.lang.resolve.lazy.descriptors.LazyPackageDescriptor;
 import org.jetbrains.jet.lang.resolve.lazy.descriptors.LazyScriptDescriptor;
 import org.jetbrains.jet.lang.resolve.name.FqName;
 import org.jetbrains.jet.lang.resolve.name.Name;
-import org.jetbrains.jet.lang.resolve.name.SpecialNames;
 import org.jetbrains.jet.lang.resolve.scopes.JetScope;
 import org.jetbrains.jet.renderer.DescriptorRenderer;
 import org.jetbrains.jet.storage.*;
@@ -299,7 +297,7 @@ public class ResolveSession implements KotlinCodeAnalyzer {
     /*package*/ LazyClassDescriptor getClassObjectDescriptor(@NotNull JetClassObject classObject) {
         JetClass aClass = PsiTreeUtil.getParentOfType(classObject, JetClass.class);
 
-        final LazyClassDescriptor parentClassDescriptor;
+        LazyClassDescriptor parentClassDescriptor;
 
         if (aClass != null) {
             parentClassDescriptor = (LazyClassDescriptor) getClassDescriptor(aClass);
@@ -313,27 +311,11 @@ public class ResolveSession implements KotlinCodeAnalyzer {
 
         // Activate resolution and writing to trace
         parentClassDescriptor.getClassObjectDescriptor();
-        DeclarationDescriptor declaration = getBindingContext().get(BindingContext.DECLARATION_TO_DESCRIPTOR, classObject.getObjectDeclaration());
+        parentClassDescriptor.getDescriptorsForExtraClassObjects();
+        DeclarationDescriptor classObjectDescriptor = getBindingContext().get(BindingContext.DECLARATION_TO_DESCRIPTOR, classObject.getObjectDeclaration());
+        assert classObjectDescriptor != null : "No descriptor found for " + JetPsiUtil.getElementTextWithContext(classObject);
 
-        if (declaration == null) {
+        return (LazyClassDescriptor) classObjectDescriptor;
-            // It's possible that there are several class objects and another class object is taking part in lazy resolve. We still want to
-            // build descriptors for such class objects.
-            final JetClassLikeInfo classObjectInfo = parentClassDescriptor.getClassObjectInfo(classObject);
-            assert classObjectInfo != null :
-                    String.format("Failed to find class object info for existent class object declaration: %s",
-                                  JetPsiUtil.getElementTextWithContext(classObject));
-
-            final Name name = SpecialNames.getClassObjectName(parentClassDescriptor.getName());
-            return storageManager.compute(new Function0<LazyClassDescriptor>() {
-                @Override
-                public LazyClassDescriptor invoke() {
-                    // Create under lock to avoid premature access to published 'this'
-                    return new LazyClassDescriptor(ResolveSession.this, parentClassDescriptor, name, classObjectInfo);
-                }
-            });
-        }
-
-        return (LazyClassDescriptor) declaration;
     }
 
     @Override

compiler/frontend/src/org/jetbrains/jet/lang/resolve/lazy/data/JetClassLikeInfo.java

@@ -37,6 +37,10 @@ public interface JetClassLikeInfo extends JetDeclarationContainer {
     @Nullable
     JetClassObject getClassObject();
 
+    @NotNull
+    @ReadOnly
+    List<JetClassObject> getClassObjects();
+
     // This element is used to identify resolution scope for the class
     @NotNull
     PsiElement getScopeAnchor();

compiler/frontend/src/org/jetbrains/jet/lang/resolve/lazy/data/JetClassOrObjectInfo.java

@@ -23,6 +23,7 @@ import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.psi.*;
 import org.jetbrains.jet.lang.resolve.name.FqName;
 
+import java.util.Collections;
 import java.util.List;
 
 public abstract class JetClassOrObjectInfo<E extends JetClassOrObject> implements JetClassLikeInfo {
@@ -49,6 +50,16 @@ public abstract class JetClassOrObjectInfo<E extends JetClassOrObject> implement
         return element.getDeclarations();
     }
 
+    @NotNull
+    @Override
+    public List<JetClassObject> getClassObjects() {
+        JetClassBody body = element.getBody();
+        if (body == null) {
+            return Collections.emptyList();
+        }
+        return body.getAllClassObjects();
+    }
+
     @NotNull
     @Override
     public PsiElement getScopeAnchor() {

compiler/frontend/src/org/jetbrains/jet/lang/resolve/lazy/data/JetScriptInfo.kt

@@ -21,11 +21,11 @@ import org.jetbrains.jet.lang.psi.JetParameter
 import org.jetbrains.jet.lang.descriptors.ClassKind
 import org.jetbrains.jet.lang.psi.JetScript
 import org.jetbrains.jet.lang.resolve.ScriptNameUtil
-import org.jetbrains.jet.lang.psi.JetFile
 import org.jetbrains.jet.lang.psi.JetNamedFunction
 import org.jetbrains.jet.lang.psi.JetDeclaration
 import org.jetbrains.jet.lang.psi.JetCallableDeclaration
 import org.jetbrains.jet.lang.resolve.name.FqName
+import org.jetbrains.jet.lang.psi.JetClassObject
 
 public class JetScriptInfo(
         val script: JetScript
@@ -34,6 +34,7 @@ public class JetScriptInfo(
     override fun getContainingPackageFqName() = fqName.parent()
     override fun getModifierList() = null
     override fun getClassObject() = null
+    override fun getClassObjects() = listOf<JetClassObject>()
     override fun getScopeAnchor() = script
     override fun getCorrespondingClassOrObject() = null
     override fun getTypeParameters() = listOf<JetTypeParameter>()

compiler/frontend/src/org/jetbrains/jet/lang/resolve/lazy/data/SyntheticClassObjectInfo.java

@@ -60,6 +60,12 @@ public class SyntheticClassObjectInfo implements JetClassLikeInfo {
         return null;
     }
 
+    @NotNull
+    @Override
+    public List<JetClassObject> getClassObjects() {
+        return Collections.emptyList();
+    }
+
     @NotNull
     @Override
     public PsiElement getScopeAnchor() {

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

@@ -22,16 +22,19 @@ import com.google.common.collect.Lists;
 import com.intellij.psi.PsiElement;
 import kotlin.Function0;
 import kotlin.Function1;
-import kotlin.Unit;
 import kotlin.KotlinPackage;
+import kotlin.Unit;
 import org.jetbrains.annotations.Mutable;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import org.jetbrains.annotations.ReadOnly;
 import org.jetbrains.jet.lang.descriptors.*;
 import org.jetbrains.jet.lang.descriptors.annotations.Annotations;
 import org.jetbrains.jet.lang.descriptors.impl.ClassDescriptorBase;
 import org.jetbrains.jet.lang.psi.*;
-import org.jetbrains.jet.lang.resolve.*;
+import org.jetbrains.jet.lang.resolve.BindingContext;
+import org.jetbrains.jet.lang.resolve.DescriptorUtils;
+import org.jetbrains.jet.lang.resolve.TypeHierarchyResolver;
 import org.jetbrains.jet.lang.resolve.lazy.ForceResolveUtil;
 import org.jetbrains.jet.lang.resolve.lazy.LazyEntity;
 import org.jetbrains.jet.lang.resolve.lazy.ResolveSession;
@@ -46,6 +49,7 @@ import org.jetbrains.jet.lang.types.JetType;
 import org.jetbrains.jet.lang.types.TypeConstructor;
 import org.jetbrains.jet.lang.types.TypeUtils;
 import org.jetbrains.jet.lang.types.lang.KotlinBuiltIns;
+import org.jetbrains.jet.storage.MemoizedFunctionToNotNull;
 import org.jetbrains.jet.storage.NotNullLazyValue;
 import org.jetbrains.jet.storage.NullableLazyValue;
 import org.jetbrains.jet.storage.StorageManager;
@@ -78,6 +82,7 @@ public class LazyClassDescriptor extends ClassDescriptorBase implements ClassDes
 
     private final Annotations annotations;
     private final NullableLazyValue<ClassDescriptorWithResolutionScopes> classObjectDescriptor;
+    private final MemoizedFunctionToNotNull<JetClassObject, ClassDescriptorWithResolutionScopes> extraClassObjectDescriptors;
 
     private final LazyClassMemberScope unsubstitutedMemberScope;
 
@@ -151,7 +156,13 @@ public class LazyClassDescriptor extends ClassDescriptorBase implements ClassDes
         this.classObjectDescriptor = storageManager.createNullableLazyValue(new Function0<ClassDescriptorWithResolutionScopes>() {
             @Override
             public ClassDescriptorWithResolutionScopes invoke() {
-                return computeClassObjectDescriptor();
+                return computeClassObjectDescriptor(declarationProvider.getOwnerInfo().getClassObject());
+            }
+        });
+        this.extraClassObjectDescriptors = storageManager.createMemoizedFunction(new Function1<JetClassObject, ClassDescriptorWithResolutionScopes>() {
+            @Override
+            public ClassDescriptorWithResolutionScopes invoke(JetClassObject classObject) {
+                return computeClassObjectDescriptor(classObject);
             }
         });
         this.scopeForClassHeaderResolution = storageManager.createLazyValue(new Function0<JetScope>() {
@@ -305,10 +316,29 @@ public class LazyClassDescriptor extends ClassDescriptorBase implements ClassDes
         return classObjectDescriptor.invoke();
     }
 
-    @Nullable
+    @NotNull
-    private ClassDescriptorWithResolutionScopes computeClassObjectDescriptor() {
+    @ReadOnly
-        JetClassObject classObject = declarationProvider.getOwnerInfo().getClassObject();
+    public List<ClassDescriptor> getDescriptorsForExtraClassObjects() {
+        return KotlinPackage.map(
+                KotlinPackage.filter(
+                        declarationProvider.getOwnerInfo().getClassObjects(),
+                        new Function1<JetClassObject, Boolean>() {
+                            @Override
+                            public Boolean invoke(JetClassObject classObject) {
+                                return classObject != declarationProvider.getOwnerInfo().getClassObject();
+                            }
+                        }),
+                new Function1<JetClassObject, ClassDescriptor>() {
+                    @Override
+                    public ClassDescriptor invoke(JetClassObject classObject) {
+                        return extraClassObjectDescriptors.invoke(classObject);
+                    }
+                }
+        );
+    }
 
+    @Nullable
+    private ClassDescriptorWithResolutionScopes computeClassObjectDescriptor(@Nullable JetClassObject classObject) {
         JetClassLikeInfo classObjectInfo = getClassObjectInfo(classObject);
         if (classObjectInfo != null) {
             return new LazyClassDescriptor(resolveSession, this, getClassObjectName(getName()), classObjectInfo);
@@ -317,7 +347,7 @@ public class LazyClassDescriptor extends ClassDescriptorBase implements ClassDes
     }
 
     @Nullable
-    public JetClassLikeInfo getClassObjectInfo(JetClassObject classObject) {
+    private JetClassLikeInfo getClassObjectInfo(@Nullable JetClassObject classObject) {
         if (classObject != null) {
             if (getKind() != ClassKind.CLASS && getKind() != ClassKind.TRAIT && getKind() != ClassKind.ANNOTATION_CLASS || isInner()) {
                 resolveSession.getTrace().report(CLASS_OBJECT_NOT_ALLOWED.on(classObject));
@@ -381,6 +411,8 @@ public class LazyClassDescriptor extends ClassDescriptorBase implements ClassDes
             ForceResolveUtil.forceResolveAllContents(classObjectDescriptor);
         }
 
+        ForceResolveUtil.forceResolveAllContents(getDescriptorsForExtraClassObjects());
+
         getClassObjectType();
         ForceResolveUtil.forceResolveAllContents(getConstructors());
         getContainingDeclaration();