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

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

@@ -17,14 +17,17 @@
 package org.jetbrains.kotlin.asJava
 
 import com.intellij.openapi.diagnostic.Logger
-import com.intellij.psi.*
+import com.intellij.psi.PsiClass
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiFileSystemItem
+import com.intellij.psi.PsiInvalidElementAccessException
 import com.intellij.psi.impl.PsiModificationTrackerImpl
-import com.intellij.psi.impl.PsiTreeChangeEventImpl.PsiEventType.*
 import com.intellij.psi.impl.PsiTreeChangeEventImpl
+import com.intellij.psi.impl.PsiTreeChangeEventImpl.PsiEventType.*
 import com.intellij.psi.impl.PsiTreeChangePreprocessor
 import com.intellij.psi.util.PsiModificationTracker
 import org.jetbrains.kotlin.psi.*
-import com.intellij.psi.util.PsiTreeUtil
+import org.jetbrains.kotlin.psi.psiUtil.isAncestor
 
 /**
  * Tested in OutOfBlockModificationTestGenerated
@@ -33,9 +36,9 @@ public class KotlinCodeBlockModificationListener(modificationTracker: PsiModific
     private val myModificationTracker = modificationTracker as PsiModificationTrackerImpl
 
     override fun treeChanged(event: PsiTreeChangeEventImpl) {
-        if (event.getFile() !is JetFile) return
+        if (event.file !is JetFile) return
 
-        when (event.getCode()) {
+        when (event.code) {
             BEFORE_CHILDREN_CHANGE,
             BEFORE_PROPERTY_CHANGE,
             BEFORE_CHILD_MOVEMENT,
@@ -48,12 +51,12 @@ public class KotlinCodeBlockModificationListener(modificationTracker: PsiModific
             CHILD_ADDED,
             CHILD_REMOVED,
             CHILD_REPLACED -> {
-                processChange(event.getParent(), event.getOldChild(), event.getChild())
+                processChange(event.parent, event.oldChild, event.child)
             }
 
             CHILDREN_CHANGED -> {
-                if (!event.isGenericChange()) {
+                if (!event.isGenericChange) {
-                    processChange(event.getParent(), event.getParent(), null)
+                    processChange(event.parent, event.parent, null)
                 }
             }
 
@@ -62,14 +65,14 @@ public class KotlinCodeBlockModificationListener(modificationTracker: PsiModific
                 myModificationTracker.incCounter()
             }
 
-            else -> LOG.error("Unknown code:" + event.getCode())
+            else -> LOG.error("Unknown code:" + event.code)
         }
     }
 
     private fun processChange(parent: PsiElement?, child1: PsiElement?, child2: PsiElement?) {
         try {
             if (!isInsideCodeBlock(parent)) {
-                if (parent != null && parent.getContainingFile() is JetFile) {
+                if (parent != null && parent.containingFile is JetFile) {
                     myModificationTracker.incCounter()
                 }
                 else {
@@ -88,16 +91,16 @@ public class KotlinCodeBlockModificationListener(modificationTracker: PsiModific
     }
 
     companion object {
-        private val LOG = Logger.getInstance("#org.jetbrains.kotlin.asJava.JetCodeBlockModificationListener")
+        private val LOG = Logger.getInstance("#org.jetbrains.kotlin.asJava.KotlinCodeBlockModificationListener")
 
         private fun containsClassesInside(element: PsiElement?): Boolean {
             if (element == null) return false
             if (element is PsiClass) return true
 
-            var child = element.getFirstChild()
+            var child = element.firstChild
             while (child != null) {
                 if (containsClassesInside(child)) return true
-                child = child.getNextSibling()
+                child = child.nextSibling
             }
 
             return false
@@ -105,38 +108,42 @@ public class KotlinCodeBlockModificationListener(modificationTracker: PsiModific
 
         public fun isInsideCodeBlock(element: PsiElement?): Boolean {
             if (element is PsiFileSystemItem) return false
-            if (element == null || element.getParent() == null) return true
+            if (element == null || element.parent == null) return true
 
-            val blockDeclarationCandidate = JetPsiUtil.getTopmostParentOfTypes(
+            //TODO: other types
-                    element,
+            val blockDeclaration = JetPsiUtil.getTopmostParentOfTypes(element, *BLOCK_DECLARATION_TYPES) ?: return false
-                    javaClass<JetProperty>(),
-                    javaClass<JetNamedFunction>())
 
-            when (blockDeclarationCandidate) {
+            when (blockDeclaration) {
                 is JetNamedFunction -> {
-                    val function = blockDeclarationCandidate : JetNamedFunction
+                    if (blockDeclaration.hasBlockBody()) {
-                    if (function.hasBlockBody() && function.getBodyExpression().isAncestorOf(element)) {
+                        return blockDeclaration.bodyExpression.isAncestor(element)
-                        return true
                     }
-
+                    else if (blockDeclaration.hasDeclaredReturnType()) {
-                    if (function.hasDeclaredReturnType() && function.getInitializer().isAncestorOf(element)) {
+                        return blockDeclaration.initializer.isAncestor(element)
-                        return true
                     }
                 }
+
                 is JetProperty -> {
-                    val property = blockDeclarationCandidate : JetProperty
+                    for (accessor in blockDeclaration.accessors) {
-                    for (accessor in property.getAccessors()) {
+                        if (accessor.initializer.isAncestor(element) || accessor.bodyExpression.isAncestor(element)) {
-                        when {
+                            return true
-                            accessor.getInitializer().isAncestorOf(element) -> return true
-                            accessor.getBodyExpression().isAncestorOf(element) -> return true
                         }
                     }
                 }
+
+                else -> throw IllegalStateException()
             }
 
             return false
         }
 
-        private fun PsiElement?.isAncestorOf(element: PsiElement) = PsiTreeUtil.isAncestor(this, element, false)
+        public fun isBlockDeclaration(declaration: JetDeclaration): Boolean {
+            return BLOCK_DECLARATION_TYPES.any { it.isInstance(declaration) }
+        }
+
+        private val BLOCK_DECLARATION_TYPES = arrayOf<Class<out JetDeclaration>>(
+                javaClass<JetProperty>(),
+                javaClass<JetNamedFunction>()
+        )
     }
 }

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 class CachedResolve(val bindingContext: BindingContext, resolveElement: JetElement) {
-    private val fullResolveCache: CachedValue<MutableMap<JetElement, BindingContext>> = CachedValuesManager.getManager(project).createCachedValue(
+        private val elementModificationCount: Int? =
-            object : CachedValueProvider<MutableMap<JetElement, BindingContext>> {
+                if (resolveElement is JetDeclaration && KotlinCodeBlockModificationListener.isBlockDeclaration(resolveElement))
-                override fun compute(): CachedValueProvider.Result<MutableMap<JetElement, BindingContext>> {
+                    SubtreeModificationCountUpdater.getModificationCount(resolveElement)
-                    return CachedValueProvider.Result.create(ContainerUtil.createConcurrentSoftValueMap<JetElement, BindingContext>(),
+                else
-                                                             PsiModificationTracker.MODIFICATION_COUNT,
+                    null
-                                                             resolveSession.getExceptionTracker())
+
+        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()
+        // check if full additional resolve already performed and is up-to-date
-        fullResolveMap[resolveElement]?.let { return it } // check if full additional resolve already performed
+        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) {
+                if (statementFilter == StatementFilter.NONE) { // partial resolve is not supported for the given declaration - full resolve performed instead
-                    // partial resolve is not supported for the given declaration - full resolve performed instead
+                    fullResolveMap[resolveElement] = CachedResolve(bindingContext, resolveElement)
-                    fullResolveMap[resolveElement] = bindingContext
                     return bindingContext
                 }
 

idea/src/META-INF/plugin.xml

@@ -498,6 +498,7 @@
     <editorNotificationProvider implementation="org.jetbrains.kotlin.idea.highlighter.ErrorDuringFileAnalyzeNotificationProvider"/>
 
     <psi.treeChangePreprocessor implementation="org.jetbrains.kotlin.asJava.KotlinCodeBlockModificationListener"/>
+    <psi.treeChangePreprocessor implementation="org.jetbrains.kotlin.idea.caches.resolve.SubtreeModificationCountUpdater"/>
 
     <referencesSearch implementation="org.jetbrains.kotlin.idea.search.ideaExtensions.KotlinReferencesSearcher"/>
     <directClassInheritorsSearch implementation="org.jetbrains.kotlin.idea.search.ideaExtensions.KotlinDirectInheritorsSearcher"/>

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

@@ -16,11 +16,13 @@
 
 package org.jetbrains.kotlin.idea
 
+import com.intellij.psi.PsiDocumentManager
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.test.JetLightCodeInsightFixtureTestCase
 import org.jetbrains.kotlin.idea.test.JetLightProjectDescriptor
 import org.jetbrains.kotlin.idea.util.application.executeWriteCommand
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.startOffset
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 
@@ -65,6 +67,7 @@ class C(param1: String = "", param2: Int = 0) {
 
     public fun testResolveCaching() {
         doTest {
+            // resolve statements in "a()"
             val statement1 = statements[0]
             val statement2 = statements[1]
             val bindingContext1 = statement1.analyze(BodyResolveMode.FULL)
@@ -74,15 +77,60 @@ class C(param1: String = "", param2: Int = 0) {
             val bindingContext3 = statement1.analyze(BodyResolveMode.FULL)
             assert(bindingContext3 === bindingContext1)
 
-            file.add(factory.createFunction("fun foo(){}"))
+            // modify body of "b()"
+            val bFun = members[1] as JetNamedFunction
+            val bBody = bFun.bodyExpression as JetBlockExpression
+            bBody.addAfter(factory.createExpression("x()"), bBody.lBrace)
 
             val bindingContext4 = statement1.analyze(BodyResolveMode.FULL)
-            assert(bindingContext4 !== bindingContext1)
+            assert(bindingContext4 === bindingContext1) // change in other function's body should not affect resolve of other one
 
-            statement1.getParent().addAfter(factory.createExpression("x()"), statement1)
+            // add parameter to "b()" this should invalidate all resolve
+            bFun.valueParameterList!!.addParameter(factory.createParameter("p: Int"))
 
             val bindingContext5 = statement1.analyze(BodyResolveMode.FULL)
-            assert(bindingContext5 !== bindingContext4)
+            assert(bindingContext5 !== bindingContext1)
+
+            statement1.parent.addAfter(factory.createExpression("x()"), statement1)
+
+            val bindingContext6 = statement1.analyze(BodyResolveMode.FULL)
+            assert(bindingContext6 !== bindingContext5)
+        }
+    }
+
+    public fun testResolveSurvivesTypingInCodeBlock() {
+        doTest {
+            val statement = statements[0]
+            val bindingContext1 = statement.analyze(BodyResolveMode.FULL)
+
+            val classConstructorParamTypeRef = klass.getPrimaryConstructor()!!.valueParameters.first().typeReference!!
+            val bindingContext2 = classConstructorParamTypeRef.analyze(BodyResolveMode.FULL)
+
+            val documentManager = PsiDocumentManager.getInstance(getProject())
+            val document = documentManager.getDocument(file)!!
+            documentManager.doPostponedOperationsAndUnblockDocument(document)
+
+            // modify body of "b()" via document
+            val bFun = members[1] as JetNamedFunction
+            val bBody = bFun.bodyExpression as JetBlockExpression
+            document.insertString(bBody.lBrace!!.startOffset + 1, "x()")
+            documentManager.commitAllDocuments()
+
+            val bindingContext3 = statement.analyze(BodyResolveMode.FULL)
+            assert(bindingContext3 === bindingContext1)
+
+            val bindingContext4 = classConstructorParamTypeRef.analyze(BodyResolveMode.FULL)
+            assert(bindingContext4 === bindingContext2)
+
+            // insert statement in "a()"
+            document.insertString(statement.startOffset, "x()\n")
+            documentManager.commitAllDocuments()
+
+            val bindingContext5 = (members[0] as JetNamedFunction).bodyExpression!!.analyze(BodyResolveMode.FULL)
+            assert(bindingContext5 !== bindingContext1)
+
+            val bindingContext6 = classConstructorParamTypeRef.analyze(BodyResolveMode.FULL)
+            assert(bindingContext6 === bindingContext2)
         }
     }