Caching for PARTIAL_FOR_COMPLETION resolve too (because it's used not only in completion)

idea/ide-common/src/org/jetbrains/kotlin/resolve/lazy/BodyResolveMode.kt

@@ -18,19 +18,15 @@ package org.jetbrains.kotlin.resolve.lazy
 
 import org.jetbrains.kotlin.resolve.BindingTraceFilter
 
-open class BodyResolveMode {
+enum class BodyResolveMode(val bindingTraceFilter: BindingTraceFilter) {
-    open val bindingTraceFilter: BindingTraceFilter = BindingTraceFilter.ACCEPT_ALL
+    FULL(BindingTraceFilter.ACCEPT_ALL),
+    PARTIAL_FOR_COMPLETION(BindingTraceFilter.NO_DIAGNOSTICS),
+    PARTIAL_WITH_DIAGNOSTICS(BindingTraceFilter.ACCEPT_ALL),
+    PARTIAL(BindingTraceFilter.NO_DIAGNOSTICS)
 
-    companion object {
+    ;
-        @JvmField val FULL = BodyResolveMode()
 
-        @JvmField val PARTIAL = object : BodyResolveMode() {
+    fun doesNotLessThan(other: BodyResolveMode): Boolean {
-            override val bindingTraceFilter: BindingTraceFilter
+        return this <= other && this.bindingTraceFilter.includesEverythingIn(other.bindingTraceFilter)
-                get() = BindingTraceFilter.NO_DIAGNOSTICS
-        }
-
-        @JvmField val PARTIAL_WITH_DIAGNOSTICS = BodyResolveMode()
-
-        @JvmField val PARTIAL_FOR_COMPLETION = BodyResolveMode()
     }
 }

idea/idea-analysis/src/org/jetbrains/kotlin/idea/project/ResolveElementCache.kt

@@ -83,11 +83,11 @@ class ResolveElementCache(
             },
             false)
 
-    private class CachedPartialResolve(val bindingContext: BindingContext, file: KtFile, val filter: BindingTraceFilter) {
+    private class CachedPartialResolve(val bindingContext: BindingContext, file: KtFile, val mode: BodyResolveMode) {
         private val modificationStamp: Long? = modificationStamp(file)
 
-        fun isUpToDate(file: KtFile, newFilter: BindingTraceFilter) =
+        fun isUpToDate(file: KtFile, newMode: BodyResolveMode) =
-                modificationStamp == modificationStamp(file) && filter.includesEverythingIn(newFilter)
+                modificationStamp == modificationStamp(file) && mode.doesNotLessThan(newMode)
 
         private fun modificationStamp(file: KtFile): Long? {
             return if (!file.isPhysical) // for non-physical file we don't get MODIFICATION_COUNT increased and must reset data on any modification of the file
@@ -150,15 +150,6 @@ class ResolveElementCache(
                 return bindingContext
             }
 
-            BodyResolveMode.PARTIAL_FOR_COMPLETION -> {
-                if (resolveElement !is KtDeclaration) {
-                    return getElementsAdditionalResolve(resolveElement, null, BodyResolveMode.FULL)
-                }
-
-                // not cached
-                return performElementAdditionalResolve(resolveElement, contextElements, bodyResolveMode).first
-            }
-
             else -> {
                 if (resolveElement !is KtDeclaration) {
                     return getElementsAdditionalResolve(resolveElement, null, BodyResolveMode.FULL)
@@ -168,7 +159,7 @@ class ResolveElementCache(
                 val statementsToResolve = contextElements!!.map { PartialBodyResolveFilter.findStatementToResolve(it, resolveElement) }.distinct()
                 val partialResolveMap = partialBodyResolveCache.value
                 val cachedResults = statementsToResolve.map { partialResolveMap[it ?: resolveElement] }
-                if (cachedResults.all { it != null && it.isUpToDate(file, bodyResolveMode.bindingTraceFilter) }) { // partial resolve is already cached for these statements
+                if (cachedResults.all { it != null && it.isUpToDate(file, bodyResolveMode) }) { // partial resolve is already cached for these statements
                     return CompositeBindingContext.create(cachedResults.map { it!!.bindingContext }.distinct())
                 }
 
@@ -179,7 +170,7 @@ class ResolveElementCache(
                     return bindingContext
                 }
 
-                val resolveToCache = CachedPartialResolve(bindingContext, file, bodyResolveMode.bindingTraceFilter)
+                val resolveToCache = CachedPartialResolve(bindingContext, file, bodyResolveMode)
 
                 for (statement in (statementFilter as PartialBodyResolveFilter).allStatementsToResolve) {
                     if (!partialResolveMap.containsKey(statement) && bindingContext[BindingContext.PROCESSED, statement] == true) {
@@ -408,7 +399,7 @@ class ResolveElementCache(
         val modifierList = ktAnnotationEntry.getParentOfType<KtModifierList>(true)
         val declaration = modifierList?.getParentOfType<KtDeclaration>(true)
         if (declaration != null) {
-            doResolveAnnotations(getAnnotationsByDeclaration(resolveSession, modifierList!!, declaration))
+            doResolveAnnotations(getAnnotationsByDeclaration(resolveSession, modifierList, declaration))
         }
         else {
             val fileAnnotationList = ktAnnotationEntry.getParentOfType<KtFileAnnotationList>(true)

idea/tests/org/jetbrains/kotlin/idea/ResolveElementCacheTest.kt

@@ -193,33 +193,40 @@ class C(param1: String = "", param2: Int = 0) {
 
             val bindingContext3 = statement2.analyze(BodyResolveMode.PARTIAL_FOR_COMPLETION)
             assert(bindingContext3 === bindingContext1)
+
+            val bindingContext4 = statement2.analyze(BodyResolveMode.PARTIAL_WITH_DIAGNOSTICS)
+            assert(bindingContext4 === bindingContext1)
         }
     }
 
     fun testPartialResolveCaching() {
-        doTest { this.testPartialResolveCaching() }
+        doTest { this.testPartialResolveCaching(BodyResolveMode.PARTIAL) }
     }
 
-    private fun Data.testPartialResolveCaching() {
+    fun testPartialForCompletionResolveCaching() {
+        doTest { this.testPartialResolveCaching(BodyResolveMode.PARTIAL_FOR_COMPLETION) }
+    }
+
+    private fun Data.testPartialResolveCaching(mode: BodyResolveMode) {
         val statement1 = statements[0]
         val statement2 = statements[1]
-        val bindingContext1 = statement1.analyze(BodyResolveMode.PARTIAL)
+        val bindingContext1 = statement1.analyze(mode)
-        val bindingContext2 = statement2.analyze(BodyResolveMode.PARTIAL)
+        val bindingContext2 = statement2.analyze(mode)
         assert(bindingContext1 !== bindingContext2)
 
-        val bindingContext3 = statement1.analyze(BodyResolveMode.PARTIAL)
+        val bindingContext3 = statement1.analyze(mode)
-        val bindingContext4 = statement2.analyze(BodyResolveMode.PARTIAL)
+        val bindingContext4 = statement2.analyze(mode)
         assert(bindingContext3 === bindingContext1)
         assert(bindingContext4 === bindingContext2)
 
         file.add(factory.createFunction("fun foo(){}"))
 
-        val bindingContext5 = statement1.analyze(BodyResolveMode.PARTIAL)
+        val bindingContext5 = statement1.analyze(mode)
         assert(bindingContext5 !== bindingContext1)
 
         statement1.parent.addAfter(factory.createExpression("x()"), statement1)
 
-        val bindingContext6 = statement1.analyze(BodyResolveMode.PARTIAL)
+        val bindingContext6 = statement1.analyze(mode)
         assert(bindingContext6 !== bindingContext5)
     }
 
@@ -227,7 +234,7 @@ class C(param1: String = "", param2: Int = 0) {
         doTest {
             val nonPhysicalFile = KtPsiFactory(project).createAnalyzableFile("NonPhysical.kt", FILE_TEXT, file)
             val nonPhysicalData = extractData(nonPhysicalFile)
-            nonPhysicalData.testPartialResolveCaching()
+            nonPhysicalData.testPartialResolveCaching(BodyResolveMode.PARTIAL)
         }
     }
 
@@ -265,6 +272,30 @@ class C(param1: String = "", param2: Int = 0) {
         }
     }
 
+    fun testPartialForCompletionAndPartialAfter() {
+        doTest {
+            val statement = statements[0]
+            val bindingContext1 = statement.analyze(BodyResolveMode.PARTIAL_FOR_COMPLETION)
+            val bindingContext2 = statement.analyze(BodyResolveMode.PARTIAL)
+            assert(bindingContext2 === bindingContext1)
+
+            val bindingContext3 = statement.analyze(BodyResolveMode.PARTIAL_WITH_DIAGNOSTICS)
+            assert(bindingContext3 !== bindingContext1)
+        }
+    }
+
+    fun testPartialWithDiagnosticsAndPartialAfter() {
+        doTest {
+            val statement = statements[0]
+            val bindingContext1 = statement.analyze(BodyResolveMode.PARTIAL_WITH_DIAGNOSTICS)
+            val bindingContext2 = statement.analyze(BodyResolveMode.PARTIAL)
+            assert(bindingContext2 === bindingContext1)
+
+            val bindingContext3 = statement.analyze(BodyResolveMode.PARTIAL_FOR_COMPLETION)
+            assert(bindingContext3 !== bindingContext1)
+        }
+    }
+
     fun testFullResolvedCachedWhenPartialForConstructorInvoked() {
         doTest {
             val defaultValue1 = klass.getPrimaryConstructorParameters()[0].defaultValue!!