Do not invalidate resolve of everything after change inside a code block

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

@@ -0,0 +1,81 @@
+/*
+ * Copyright 2010-2015 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.idea.caches.resolve
+
+import com.intellij.openapi.util.Key
+import com.intellij.psi.PsiElement
+import com.intellij.psi.impl.PsiTreeChangeEventImpl
+import com.intellij.psi.impl.PsiTreeChangePreprocessor
+import org.jetbrains.kotlin.psi.JetFile
+import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
+
+public class SubtreeModificationCountUpdater : PsiTreeChangePreprocessor {
+    override fun treeChanged(event: PsiTreeChangeEventImpl) {
+        if (event.file !is JetFile) return
+
+        when (event.code!!) {
+            PsiTreeChangeEventImpl.PsiEventType.BEFORE_CHILDREN_CHANGE,
+            PsiTreeChangeEventImpl.PsiEventType.BEFORE_PROPERTY_CHANGE,
+            PsiTreeChangeEventImpl.PsiEventType.BEFORE_CHILD_MOVEMENT,
+            PsiTreeChangeEventImpl.PsiEventType.BEFORE_CHILD_REPLACEMENT,
+            PsiTreeChangeEventImpl.PsiEventType.BEFORE_CHILD_ADDITION,
+            PsiTreeChangeEventImpl.PsiEventType.BEFORE_CHILD_REMOVAL,
+            PsiTreeChangeEventImpl.PsiEventType.PROPERTY_CHANGED -> {
+                // skip
+            }
+
+            PsiTreeChangeEventImpl.PsiEventType.CHILD_ADDED,
+            PsiTreeChangeEventImpl.PsiEventType.CHILD_REMOVED,
+            PsiTreeChangeEventImpl.PsiEventType.CHILD_REPLACED -> {
+                incCounters(event.parent)
+            }
+
+            PsiTreeChangeEventImpl.PsiEventType.CHILDREN_CHANGED -> {
+                if (!event.isGenericChange) {
+                    incCounters(event.parent)
+                }
+            }
+
+            PsiTreeChangeEventImpl.PsiEventType.CHILD_MOVED -> {
+                incCounters(event.oldParent)
+                incCounters(event.newParent)
+            }
+
+            else -> error("Unknown code:${event.code}")
+        }
+    }
+
+    private fun incCounters(element: PsiElement?) {
+        element?.parentsWithSelf?.forEach {
+            val count = it.getUserData(MODIFICATION_COUNT_KEY)
+            if (count != null) {
+                it.putUserData(MODIFICATION_COUNT_KEY, count + 1)
+            }
+        }
+    }
+
+    companion object {
+        private val MODIFICATION_COUNT_KEY = Key<Int>("SubtreeModificationCountUpdater.MODIFICATION_COUNT_KEY")
+
+        public fun getModificationCount(element: PsiElement): Int {
+            element.getUserData(MODIFICATION_COUNT_KEY)?.let { return it }
+
+            element.putUserData(MODIFICATION_COUNT_KEY, 0)
+            return 0
+        }
+    }
+}

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

@@ -22,6 +22,9 @@ import com.intellij.psi.util.CachedValueProvider
 import com.intellij.psi.util.CachedValuesManager
 import com.intellij.psi.util.PsiModificationTracker
 import com.intellij.util.containers.ContainerUtil
+import org.jetbrains.kotlin.asJava.KotlinCodeBlockModificationListener
+import org.jetbrains.kotlin.descriptors.ModuleDescriptor
+import org.jetbrains.kotlin.idea.caches.resolve.SubtreeModificationCountUpdater
 import org.jetbrains.kotlin.idea.stubindex.JetProbablyNothingFunctionShortNameIndex
 import org.jetbrains.kotlin.idea.stubindex.JetProbablyNothingPropertyShortNameIndex
 import org.jetbrains.kotlin.psi.*
@@ -32,13 +35,24 @@ import org.jetbrains.kotlin.resolve.TargetPlatform
 import org.jetbrains.kotlin.resolve.lazy.*
 
 public class ResolveElementCache(resolveSession: ResolveSession, private val project: Project) : ElementResolver(resolveSession), BodyResolveCache {
-    // Recreate internal cache after change of modification count
-    private val fullResolveCache: CachedValue<MutableMap<JetElement, BindingContext>> = CachedValuesManager.getManager(project).createCachedValue(
-            object : CachedValueProvider<MutableMap<JetElement, BindingContext>> {
-                override fun compute(): CachedValueProvider.Result<MutableMap<JetElement, BindingContext>> {
-                    return CachedValueProvider.Result.create(ContainerUtil.createConcurrentSoftValueMap<JetElement, BindingContext>(),
-                                                             PsiModificationTracker.MODIFICATION_COUNT,
-                                                             resolveSession.getExceptionTracker())
+    private class CachedResolve(val bindingContext: BindingContext, resolveElement: JetElement) {
+        private val elementModificationCount: Int? =
+                if (resolveElement is JetDeclaration && KotlinCodeBlockModificationListener.isBlockDeclaration(resolveElement))
+                    SubtreeModificationCountUpdater.getModificationCount(resolveElement)
+                else
+                    null
+
+        fun isUpToDate(resolveElement: JetElement)
+                = elementModificationCount == null || elementModificationCount == SubtreeModificationCountUpdater.getModificationCount(resolveElement)
+    }
+
+    // drop whole cache after change "out of code block"
+    private val fullResolveCache: CachedValue<MutableMap<JetElement, CachedResolve>> = CachedValuesManager.getManager(project).createCachedValue(
+            object : CachedValueProvider<MutableMap<JetElement, CachedResolve>> {
+                override fun compute(): CachedValueProvider.Result<MutableMap<JetElement, CachedResolve>> {
+                    return CachedValueProvider.Result.create(ContainerUtil.createConcurrentSoftValueMap<JetElement, CachedResolve>(),
+                                                             PsiModificationTracker.OUT_OF_CODE_BLOCK_MODIFICATION_COUNT,
+                                                             resolveSession.exceptionTracker)
                 }
             },
             false)
@@ -48,19 +62,28 @@ public class ResolveElementCache(resolveSession: ResolveSession, private val pro
                 override fun compute(): CachedValueProvider.Result<MutableMap<JetExpression, BindingContext>> {
                     return CachedValueProvider.Result.create(ContainerUtil.createConcurrentSoftValueMap<JetExpression, BindingContext>(),
                                                              PsiModificationTracker.MODIFICATION_COUNT,
-                                                             resolveSession.getExceptionTracker())
+                                                             resolveSession.exceptionTracker)
                 }
             },
             false)
 
     override fun getElementAdditionalResolve(resolveElement: JetElement, contextElement: JetElement, bodyResolveMode: BodyResolveMode): BindingContext {
-        val fullResolveMap = fullResolveCache.getValue()
-        fullResolveMap[resolveElement]?.let { return it } // check if full additional resolve already performed
+        // check if full additional resolve already performed and is up-to-date
+        val fullResolveMap = fullResolveCache.value
+        val cachedFullResolve = fullResolveMap[resolveElement]
+        if (cachedFullResolve != null) {
+            if (cachedFullResolve.isUpToDate(resolveElement)) {
+                return cachedFullResolve.bindingContext
+            }
+            else {
+                fullResolveMap.remove(resolveElement) // remove outdated cache entry
+            }
+        }
 
         when (bodyResolveMode) {
             BodyResolveMode.FULL -> {
                 val bindingContext = performElementAdditionalResolve(resolveElement, resolveElement, BodyResolveMode.FULL).first
-                fullResolveMap[resolveElement] = bindingContext
+                fullResolveMap[resolveElement] = CachedResolve(bindingContext, resolveElement)
                 return bindingContext
             }
 
@@ -70,14 +93,13 @@ public class ResolveElementCache(resolveSession: ResolveSession, private val pro
                 }
 
                 val statementToResolve = PartialBodyResolveFilter.findStatementToResolve(contextElement, resolveElement)
-                val partialResolveMap = partialBodyResolveCache.getValue()
+                val partialResolveMap = partialBodyResolveCache.value
                 partialResolveMap[statementToResolve ?: resolveElement]?.let { return it } // partial resolve is already cached for this statement
 
                 val (bindingContext, statementFilter) = performElementAdditionalResolve(resolveElement, contextElement, BodyResolveMode.PARTIAL)
 
-                if (statementFilter == StatementFilter.NONE) {
-                    // partial resolve is not supported for the given declaration - full resolve performed instead
-                    fullResolveMap[resolveElement] = bindingContext
+                if (statementFilter == StatementFilter.NONE) { // partial resolve is not supported for the given declaration - full resolve performed instead
+                    fullResolveMap[resolveElement] = CachedResolve(bindingContext, resolveElement)
                     return bindingContext
                 }