IDELightClassTest: provide tools to check laziness of light class construction

StubComputationTracker knows which context was used to construct light class

compiler/light-classes/src/org/jetbrains/kotlin/asJava/builder/LightClassBuilder.kt

@@ -80,7 +80,7 @@ fun buildLightClass(
             LOG.error("Unbalanced stack operations: " + pop)
         }
 
-        ServiceManager.getService(project, StubComputationTracker::class.java)?.onStubComputed(javaFileStub)
+        ServiceManager.getService(project, StubComputationTracker::class.java)?.onStubComputed(javaFileStub, context)
         return LightClassBuilderResult(javaFileStub, bindingContext, state.collectedExtraJvmDiagnostics)
     }
     catch (e: ProcessCanceledException) {

compiler/light-classes/src/org/jetbrains/kotlin/asJava/builder/LightClassConstructionContext.java

@@ -1,41 +0,0 @@
-/*
- * Copyright 2010-2016 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.kotlin.asJava.builder;
-
-import org.jetbrains.annotations.NotNull;
-import org.jetbrains.kotlin.descriptors.ModuleDescriptor;
-import org.jetbrains.kotlin.resolve.BindingContext;
-
-public class LightClassConstructionContext {
-    private final BindingContext bindingContext;
-    private final ModuleDescriptor module;
-
-    public LightClassConstructionContext(@NotNull BindingContext bindingContext, @NotNull ModuleDescriptor module) {
-        this.bindingContext = bindingContext;
-        this.module = module;
-    }
-
-    @NotNull
-    public BindingContext getBindingContext() {
-        return bindingContext;
-    }
-
-    @NotNull
-    public ModuleDescriptor getModule() {
-        return module;
-    }
-}

compiler/light-classes/src/org/jetbrains/kotlin/asJava/builder/LightClassConstructionContext.kt

@@ -0,0 +1,25 @@
+/*
+ * Copyright 2010-2016 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.kotlin.asJava.builder
+
+import org.jetbrains.kotlin.descriptors.ModuleDescriptor
+import org.jetbrains.kotlin.resolve.BindingContext
+
+open class LightClassConstructionContext(
+        val bindingContext: BindingContext,
+        val module: ModuleDescriptor
+)

compiler/light-classes/src/org/jetbrains/kotlin/asJava/builder/LightClassDataProvider.kt

@@ -212,5 +212,5 @@ sealed class LightClassDataProviderForFileFacade constructor(
 
 
 interface StubComputationTracker {
-    fun onStubComputed(javaFileStub: PsiJavaFileStub)
+    fun onStubComputed(javaFileStub: PsiJavaFileStub, context: LightClassConstructionContext)
 }

idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/resolve/lightClasses/IDELightClassContexts.kt

@@ -37,6 +37,8 @@ import org.jetbrains.kotlin.descriptors.impl.ModuleDescriptorImpl
 import org.jetbrains.kotlin.frontend.di.configureModule
 import org.jetbrains.kotlin.idea.caches.resolve.getModuleInfo
 import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade
+import org.jetbrains.kotlin.idea.caches.resolve.lightClasses.IDELightClassConstructionContext.Mode.EXACT
+import org.jetbrains.kotlin.idea.caches.resolve.lightClasses.IDELightClassConstructionContext.Mode.LIGHT
 import org.jetbrains.kotlin.idea.project.IdeaEnvironment
 import org.jetbrains.kotlin.idea.project.ResolveElementCache
 import org.jetbrains.kotlin.incremental.components.LookupTracker
@@ -68,6 +70,15 @@ import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.WrappedTypeFactory
 import org.jetbrains.kotlin.utils.sure
 
+
+class IDELightClassConstructionContext(bindingContext: BindingContext, module: ModuleDescriptor, val mode: Mode)
+    : LightClassConstructionContext(bindingContext, module) {
+    enum class Mode {
+        LIGHT,
+        EXACT
+    }
+}
+
 object IDELightClassContexts {
 
     private val LOG = Logger.getInstance(this::class.java)
@@ -87,7 +98,7 @@ object IDELightClassContexts {
             "Class descriptor was not found for ${classOrObject.getElementTextWithContext()}"
         }
         ForceResolveUtil.forceResolveAllContents(classDescriptor)
-        return LightClassConstructionContext(bindingContext, resolutionFacade.moduleDescriptor)
+        return IDELightClassConstructionContext(bindingContext, resolutionFacade.moduleDescriptor, EXACT)
     }
 
     fun contextForLocalClassOrObject(classOrObject: KtClassOrObject): LightClassConstructionContext {
@@ -98,12 +109,12 @@ object IDELightClassContexts {
 
         if (descriptor == null) {
             LOG.warn("No class descriptor in context for class: " + classOrObject.getElementTextWithContext())
-            return LightClassConstructionContext(bindingContext, resolutionFacade.moduleDescriptor)
+            return IDELightClassConstructionContext(bindingContext, resolutionFacade.moduleDescriptor, EXACT)
         }
 
         ForceResolveUtil.forceResolveAllContents(descriptor)
 
-        return LightClassConstructionContext(bindingContext, resolutionFacade.moduleDescriptor)
+        return IDELightClassConstructionContext(bindingContext, resolutionFacade.moduleDescriptor, EXACT)
     }
 
 
@@ -111,8 +122,7 @@ object IDELightClassContexts {
         val resolveSession = files.first().getResolutionFacade().getFrontendService(ResolveSession::class.java)
 
         forceResolvePackageDeclarations(files, resolveSession)
-
+        return IDELightClassConstructionContext(resolveSession.bindingContext, resolveSession.moduleDescriptor, EXACT)
-        return LightClassConstructionContext(resolveSession.bindingContext, resolveSession.moduleDescriptor)
     }
 
     fun lightContextForClassOrObject(classOrObject: KtClassOrObject): LightClassConstructionContext {
@@ -120,7 +130,7 @@ object IDELightClassContexts {
 
         ForceResolveUtil.forceResolveAllContents(resolveSession.resolveToDescriptor(classOrObject))
 
-        return LightClassConstructionContext(resolveSession.bindingContext, resolveSession.moduleDescriptor)
+        return IDELightClassConstructionContext(resolveSession.bindingContext, resolveSession.moduleDescriptor, LIGHT)
     }
 
     fun lightContextForFacade(files: List<KtFile>): LightClassConstructionContext {
@@ -129,7 +139,7 @@ object IDELightClassContexts {
 
         forceResolvePackageDeclarations(files, resolveSession)
 
-        return LightClassConstructionContext(resolveSession.bindingContext, resolveSession.moduleDescriptor)
+        return IDELightClassConstructionContext(resolveSession.bindingContext, resolveSession.moduleDescriptor, LIGHT)
     }
 
     fun forceResolvePackageDeclarations(files: Collection<KtFile>, session: ResolveSession) {

idea/idea-completion/tests/org/jetbrains/kotlin/idea/completion/test/LightClassComputationControl.kt

@@ -20,6 +20,7 @@ import com.intellij.openapi.components.ComponentManager
 import com.intellij.openapi.project.Project
 import com.intellij.psi.impl.java.stubs.PsiClassStub
 import com.intellij.psi.impl.java.stubs.PsiJavaFileStub
+import org.jetbrains.kotlin.asJava.builder.LightClassConstructionContext
 import org.jetbrains.kotlin.asJava.builder.StubComputationTracker
 import org.jetbrains.kotlin.test.InTextDirectivesUtils
 import org.junit.Assert
@@ -34,13 +35,13 @@ object LightClassComputationControl {
 
         val actualFqNames = ArrayList<String>()
         val stubComputationTracker = object : StubComputationTracker {
-            override fun onStubComputed(javaFileStub: PsiJavaFileStub) {
+            override fun onStubComputed(javaFileStub: PsiJavaFileStub, context: LightClassConstructionContext) {
                 val qualifiedName = (javaFileStub.childrenStubs.single() as PsiClassStub<*>).qualifiedName!!
                 actualFqNames.add(qualifiedName)
             }
         }
 
-        project.withServiceRegistered<StubComputationTracker>(stubComputationTracker) {
+        project.withServiceRegistered<StubComputationTracker, Unit>(stubComputationTracker) {
             testBody()
         }
 
@@ -58,13 +59,13 @@ object LightClassComputationControl {
     private fun List<String>.prettyToString() = if (isEmpty()) "<empty>" else joinToString()
 }
 
-private inline fun <reified T : Any> ComponentManager.withServiceRegistered(instance: T, body: () -> Unit) {
+inline fun <reified T : Any, R> ComponentManager.withServiceRegistered(instance: T, body: () -> R): R {
     val picoContainer = picoContainer as MutablePicoContainer
     val key = T::class.java.name
     try {
         picoContainer.unregisterComponent(key)
         picoContainer.registerComponentInstance(key, instance)
-        body()
+        return body()
     }
     finally {
         picoContainer.unregisterComponent(key)

idea/tests/org/jetbrains/kotlin/idea/caches/resolve/AbstractIdeLightClassTest.kt

@@ -21,11 +21,22 @@ import com.intellij.openapi.roots.ModuleRootModificationUtil
 import com.intellij.openapi.util.io.FileUtilRt
 import com.intellij.psi.JavaPsiFacade
 import com.intellij.psi.PsiClass
+import com.intellij.psi.PsiField
+import com.intellij.psi.PsiMethod
+import com.intellij.psi.impl.java.stubs.PsiJavaFileStub
 import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.psi.util.PsiTreeUtil
 import org.jetbrains.kotlin.asJava.LightClassTestCommon
+import org.jetbrains.kotlin.asJava.builder.LightClassConstructionContext
+import org.jetbrains.kotlin.asJava.builder.StubComputationTracker
+import org.jetbrains.kotlin.asJava.classes.KtLightClass
 import org.jetbrains.kotlin.asJava.classes.KtLightClassForSourceDeclaration
+import org.jetbrains.kotlin.asJava.elements.KtLightField
+import org.jetbrains.kotlin.asJava.elements.KtLightMethod
 import org.jetbrains.kotlin.idea.KotlinDaemonAnalyzerTestCase
+import org.jetbrains.kotlin.idea.caches.resolve.LightClassLazinessChecker.Tracker.Level.*
+import org.jetbrains.kotlin.idea.caches.resolve.lightClasses.IDELightClassConstructionContext
+import org.jetbrains.kotlin.idea.completion.test.withServiceRegistered
 import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCase
 import org.jetbrains.kotlin.idea.test.KotlinWithJdkAndRuntimeLightProjectDescriptor
 import org.jetbrains.kotlin.load.java.JvmAnnotationNames
@@ -33,13 +44,25 @@ import org.jetbrains.kotlin.psi.KtClassOrObject
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.test.KotlinTestUtils
 import org.jetbrains.kotlin.test.MockLibraryUtil
+import org.jetbrains.kotlin.utils.keysToMap
 import org.junit.Assert
 import java.io.File
+import kotlin.test.assertEquals
 
 abstract class AbstractIdeLightClassTest : KotlinLightCodeInsightFixtureTestCase() {
     fun doTest(testDataPath: String) {
         myFixture.configureByFile(testDataPath)
-        testLightClass(project, myFixture.file as KtFile, testDataPath, { LightClassTestCommon.removeEmptyDefaultImpls(it) })
+        val ktFile = myFixture.file as KtFile
+        testLightClass(testDataPath, { LightClassTestCommon.removeEmptyDefaultImpls(it) }) { fqName ->
+            val tracker = LightClassLazinessChecker.Tracker(fqName)
+            project.withServiceRegistered<StubComputationTracker, PsiClass?>(tracker) {
+                findClass(fqName, ktFile, project)?.apply {
+                    LightClassLazinessChecker.check(this as KtLightClass, tracker)
+                    tracker.allowLevel(EXACT)
+                    PsiElementChecker.checkPsiElementStructure(this)
+                }
+            }
+        }
     }
 
     override fun getProjectDescriptor() = KotlinWithJdkAndRuntimeLightProjectDescriptor.INSTANCE
@@ -64,26 +87,18 @@ abstract class AbstractIdeCompiledLightClassTest : KotlinDaemonAnalyzerTestCase(
     private fun libName() = "libFor" + getTestName(false)
 
     fun doTest(testDataPath: String) {
-        testLightClass(project, null, testDataPath, { it })
+        testLightClass(testDataPath, { it }) {
+            findClass(it, null, project)?.apply {
+                PsiElementChecker.checkPsiElementStructure(this)
+            }
+        }
     }
 }
 
-private fun testLightClass(project: Project, ktFile: KtFile?, testDataPath: String, normalize: (String) -> String) {
+private fun testLightClass(testDataPath: String, normalize: (String) -> String, findLightClass: (String) -> PsiClass?) {
     LightClassTestCommon.testLightClass(
             File(testDataPath),
-            findLightClass = { fqName ->
+            findLightClass,
-                var clazz: PsiClass? = JavaPsiFacade.getInstance(project).findClass(fqName, GlobalSearchScope.allScope(project))
-                if (clazz == null) {
-                    clazz = PsiTreeUtil.findChildrenOfType(ktFile, KtClassOrObject::class.java)
-                            .find { fqName.endsWith(it.nameAsName!!.asString()) }
-                            ?.let { KtLightClassForSourceDeclaration.create(it) }
-                }
-                if (clazz != null) {
-                    PsiElementChecker.checkPsiElementStructure(clazz)
-                }
-                clazz
-
-            },
             normalizeText = { text ->
                 //NOTE: ide and compiler differ in names generated for parameters with unspecified names
                 text
@@ -97,6 +112,98 @@ private fun testLightClass(project: Project, ktFile: KtFile?, testDataPath: Stri
     )
 }
 
+private fun findClass(fqName: String, ktFile: KtFile?, project: Project): PsiClass? {
+    return JavaPsiFacade.getInstance(project).findClass(fqName, GlobalSearchScope.allScope(project)) ?:
+           PsiTreeUtil.findChildrenOfType(ktFile, KtClassOrObject::class.java)
+                   .find { fqName.endsWith(it.nameAsName!!.asString()) }
+                   ?.let { KtLightClassForSourceDeclaration.create(it) }
+}
+
+object LightClassLazinessChecker {
+
+    class Tracker(private val fqName: String) : StubComputationTracker {
+
+        private var level = NONE
+            set(newLevel) {
+                if (newLevel.ordinal <= field.ordinal) {
+                    error("Level should not decrease at any point: $level -> $newLevel, allowed: $allowedLevel")
+                }
+                if (newLevel.ordinal > allowedLevel.ordinal) {
+                    error("Level increased before it was expected: $level -> $newLevel, allowed: $allowedLevel")
+                }
+                field = newLevel
+            }
+
+        private var allowedLevel = NONE
+
+        enum class Level {
+            NONE,
+            LIGHT,
+            EXACT
+        }
+
+        override fun onStubComputed(javaFileStub: PsiJavaFileStub, context: LightClassConstructionContext) {
+            if (fqName != javaFileStub.classes.single().qualifiedName!!) return
+            if (context !is IDELightClassConstructionContext) error("Unknown context ${context::class}")
+            level = when (context.mode) {
+                IDELightClassConstructionContext.Mode.LIGHT -> LIGHT
+                IDELightClassConstructionContext.Mode.EXACT -> EXACT
+            }
+        }
+
+        fun checkLevel(expectedLevel: Level) {
+            assert(level == expectedLevel)
+        }
+
+        fun allowLevel(newAllowed: Level) {
+            allowedLevel = newAllowed
+        }
+    }
+
+    fun check(lightClass: KtLightClass, tracker: Tracker) {
+        // lighter classes not implemented for locals
+        if (lightClass.kotlinOrigin?.isLocal ?: false) return
+
+        tracker.allowLevel(LIGHT)
+
+        // collect method class results on light members that should not trigger exact context evaluation
+        val fieldsToInfo = lightClass.fields.asList().keysToMap { fieldInfo(it) }
+        val methodsToInfo = lightClass.methods.asList().keysToMap { methodInfo(it) }
+
+        tracker.allowLevel(EXACT)
+
+        lightClass.clsDelegate // trigger exact context
+
+        tracker.checkLevel(EXACT)
+
+        // check collected data against delegates which should contain correct data
+        for ((field, lightFieldInfo) in fieldsToInfo) {
+            val delegate = (field as KtLightField).clsDelegate
+            assertEquals(fieldInfo(delegate), lightFieldInfo)
+        }
+        for ((method, lightMethodInfo) in methodsToInfo) {
+            val delegate = (method as KtLightMethod).clsDelegate
+            assertEquals(methodInfo(delegate), lightMethodInfo)
+        }
+    }
+
+    private data class FieldInfo(
+            val name: String
+    )
+
+    private fun fieldInfo(field: PsiField) = with(field) {
+        FieldInfo(name!!)
+    }
+
+    private data class MethodInfo(
+            val name: String
+    )
+
+    private fun methodInfo(method: PsiMethod) = with(method) {
+        MethodInfo(name)
+    }
+}
+
 private fun String.removeLinesStartingWith(prefix: String): String {
     return lines().filterNot { it.trimStart().startsWith(prefix) }.joinToString(separator = "\n")
 }