More advanced partial body resolve caching

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

@@ -47,7 +47,7 @@ public abstract class ElementResolver protected constructor(
         public val resolveSession: ResolveSession
 ) {
     public open fun getElementAdditionalResolve(resolveElement: JetElement, contextElement: JetElement, bodyResolveMode: BodyResolveMode): BindingContext {
-        return performElementAdditionalResolve(resolveElement, resolveElement, bodyResolveMode)
+        return performElementAdditionalResolve(resolveElement, resolveElement, bodyResolveMode).first
     }
 
     protected open fun probablyNothingCallableNames(): ProbablyNothingCallableNames
@@ -74,7 +74,7 @@ public abstract class ElementResolver protected constructor(
         return resolveSession.getBindingContext()
     }
 
-    protected fun findElementOfAdditionalResolve(element: JetElement): JetElement? {
+    private fun findElementOfAdditionalResolve(element: JetElement): JetElement? {
         val elementOfAdditionalResolve = JetPsiUtil.getTopmostParentOfTypes(
                 element,
                 javaClass<JetNamedFunction>(),
@@ -107,7 +107,7 @@ public abstract class ElementResolver protected constructor(
         }
     }
 
-    protected fun performElementAdditionalResolve(resolveElement: JetElement, contextElement: JetElement, bodyResolveMode: BodyResolveMode): BindingContext {
+    protected fun performElementAdditionalResolve(resolveElement: JetElement, contextElement: JetElement, bodyResolveMode: BodyResolveMode): Pair<BindingContext, StatementFilter> {
         val file = resolveElement.getContainingJetFile()
 
         val statementFilter = if (bodyResolveMode != BodyResolveMode.FULL && resolveElement is JetDeclaration)
@@ -158,7 +158,7 @@ public abstract class ElementResolver protected constructor(
         JetFlowInformationProvider(resolveElement, controlFlowTrace).checkDeclaration()
         controlFlowTrace.addOwnDataTo(trace, null, false)
 
-        return trace.getBindingContext()
+        return Pair(trace.getBindingContext(), statementFilter)
     }
 
     private fun packageRefAdditionalResolve(resolveSession: ResolveSession, jetElement: JetElement): BindingTrace {
@@ -447,8 +447,9 @@ public abstract class ElementResolver protected constructor(
         return null
     }
 
-    protected abstract fun createAdditionalCheckerProvider(jetFile: JetFile, module: ModuleDescriptor): AdditionalCheckerProvider
+    protected abstract fun createAdditionalCheckerProvider(file: JetFile, module: ModuleDescriptor): AdditionalCheckerProvider
-    protected abstract fun getDynamicTypesSettings(jetFile: JetFile): DynamicTypesSettings
+
+    protected abstract fun getDynamicTypesSettings(file: JetFile): DynamicTypesSettings
 
     private class BodyResolveContextForLazy(
             private val topDownAnalysisMode: TopDownAnalysisMode,

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

@@ -46,7 +46,10 @@ class PartialBodyResolveFilter(
     private val nothingFunctionNames = HashSet(probablyNothingCallableNames.functionNames())
     private val nothingVariableNames = HashSet(probablyNothingCallableNames.propertyNames())
 
-    override val filter: ((JetExpression) -> Boolean)? = { statementMarks.statementMark(it) != MarkLevel.SKIP }
+    override val filter: ((JetExpression) -> Boolean)? = { statementMarks.statementMark(it) != MarkLevel.NONE }
+
+    val allStatementsToResolve: Collection<JetExpression>
+        get() = statementMarks.allMarkedStatements()
 
     init {
         assert(declaration.isAncestor(elementToResolve))
@@ -532,14 +535,14 @@ class PartialBodyResolveFilter(
     }
 
     private enum class MarkLevel {
-        SKIP,
+        NONE,
         TAKE,
         NEED_REFERENCE_RESOLVE,
         NEED_COMPLETION
     }
 
     companion object {
-        public fun findResolveElement(element: JetElement, declaration: JetDeclaration): JetExpression? {
+        public fun findStatementToResolve(element: JetElement, declaration: JetDeclaration): JetExpression? {
             return element.parentsWithSelf.takeWhile { it != declaration }.firstOrNull { it.isStatement() } as JetExpression?
         }
 
@@ -628,7 +631,7 @@ class PartialBodyResolveFilter(
                 statementMarks[statement] = level
 
                 val block = statement.getParent() as JetBlockExpression
-                val currentBlockLevel = blockLevels[block] ?: MarkLevel.SKIP
+                val currentBlockLevel = blockLevels[block] ?: MarkLevel.NONE
                 if (currentBlockLevel < level) {
                     blockLevels[block] = level
                 }
@@ -636,10 +639,13 @@ class PartialBodyResolveFilter(
         }
 
         fun statementMark(statement: JetExpression): MarkLevel
-                = statementMarks[statement] ?: MarkLevel.SKIP
+                = statementMarks[statement] ?: MarkLevel.NONE
+
+        fun allMarkedStatements(): Collection<JetExpression>
+                = statementMarks.keySet()
 
         fun lastMarkedStatement(block: JetBlockExpression, minLevel: MarkLevel): JetExpression? {
-            val level = blockLevels[block] ?: MarkLevel.SKIP
+            val level = blockLevels[block] ?: MarkLevel.NONE
             if (level < minLevel) return null // optimization
             return block.getLastChild().siblings(forward = false)
                     .filterIsInstance<JetExpression>()

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

@@ -21,6 +21,7 @@ import com.intellij.psi.util.CachedValue
 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.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.idea.stubindex.JetProbablyNothingFunctionShortNameIndex
 import org.jetbrains.kotlin.idea.stubindex.JetProbablyNothingPropertyShortNameIndex
@@ -37,7 +38,7 @@ public class ResolveElementCache(resolveSession: ResolveSession, private val pro
                 override fun compute(): CachedValueProvider.Result<MemoizedFunctionToNotNull<JetElement, BindingContext>> {
                     val manager = resolveSession.getStorageManager()
                     val cacheFunction = manager.createSoftlyRetainedMemoizedFunction<JetElement, BindingContext> { element ->
-                        performElementAdditionalResolve(element, element, BodyResolveMode.FULL)
+                        performElementAdditionalResolve(element, element, BodyResolveMode.FULL).first
                     }
                     return CachedValueProvider.Result.create(cacheFunction,
                                                              PsiModificationTracker.MODIFICATION_COUNT,
@@ -46,15 +47,10 @@ public class ResolveElementCache(resolveSession: ResolveSession, private val pro
             },
             false)
 
-    private val partialBodyResolveCache: CachedValue<MemoizedFunctionToNotNull<JetExpression, BindingContext>> = CachedValuesManager.getManager(project).createCachedValue(
+    private val partialBodyResolveCache: CachedValue<MutableMap<JetExpression, BindingContext>> = CachedValuesManager.getManager(project).createCachedValue(
-            object : CachedValueProvider<MemoizedFunctionToNotNull<JetExpression, BindingContext>> {
+            object : CachedValueProvider<MutableMap<JetExpression, BindingContext>> {
-                override fun compute(): CachedValueProvider.Result<MemoizedFunctionToNotNull<JetExpression, BindingContext>> {
+                override fun compute(): CachedValueProvider.Result<MutableMap<JetExpression, BindingContext>> {
-                    val manager = resolveSession.getStorageManager()
+                    return CachedValueProvider.Result.create(ContainerUtil.createConcurrentSoftValueMap<JetExpression, BindingContext>(),
-                    val cacheFunction = manager.createSoftlyRetainedMemoizedFunction<JetExpression, BindingContext> { expression ->
-                        val resolveElement = findElementOfAdditionalResolve(expression)!!
-                        performElementAdditionalResolve(resolveElement, expression, BodyResolveMode.PARTIAL)
-                    }
-                    return CachedValueProvider.Result.create(cacheFunction,
                                                              PsiModificationTracker.MODIFICATION_COUNT,
                                                              resolveSession.getExceptionTracker())
                 }
@@ -64,27 +60,38 @@ public class ResolveElementCache(resolveSession: ResolveSession, private val pro
     override fun getElementAdditionalResolve(resolveElement: JetElement, contextElement: JetElement, bodyResolveMode: BodyResolveMode): BindingContext {
         if (bodyResolveMode != BodyResolveMode.FULL && !hasElementAdditionalResolveCached(resolveElement) && resolveElement is JetDeclaration) {
             if (bodyResolveMode == BodyResolveMode.PARTIAL) {
-                val partialResolveElement = PartialBodyResolveFilter.findResolveElement(contextElement, resolveElement)
+                val statementToResolve = PartialBodyResolveFilter.findStatementToResolve(contextElement, resolveElement)
-                if (partialResolveElement != null) {
+                if (statementToResolve != null) {
-                    return partialBodyResolveCache.getValue().invoke(partialResolveElement)
+                    val map = partialBodyResolveCache.getValue()
+                    map[statementToResolve]?.let { return it }
+
+                    val (bindingContext, statementFilter) = performElementAdditionalResolve(resolveElement, statementToResolve, BodyResolveMode.PARTIAL)
+
+                    for (statement in (statementFilter as PartialBodyResolveFilter).allStatementsToResolve) {
+                        if (!map.containsKey(statement)) {
+                            map[statement] = bindingContext
+                        }
+                    }
+
+                    return bindingContext
                 }
             }
             else {
-                return performElementAdditionalResolve(resolveElement, contextElement, bodyResolveMode)
+                return performElementAdditionalResolve(resolveElement, contextElement, bodyResolveMode).first
             }
         }
 
         return additionalResolveCache.getValue().invoke(resolveElement)
     }
 
-    private fun hasElementAdditionalResolveCached(jetElement: JetElement)
+    private fun hasElementAdditionalResolveCached(element: JetElement)
-            = additionalResolveCache.hasUpToDateValue() && additionalResolveCache.getValue().isComputed(jetElement)
+            = additionalResolveCache.hasUpToDateValue() && additionalResolveCache.getValue().isComputed(element)
 
-    override fun createAdditionalCheckerProvider(jetFile: JetFile, module: ModuleDescriptor)
+    override fun createAdditionalCheckerProvider(file: JetFile, module: ModuleDescriptor)
-            = TargetPlatformDetector.getPlatform(jetFile).createAdditionalCheckerProvider(module)
+            = TargetPlatformDetector.getPlatform(file).createAdditionalCheckerProvider(module)
 
-    override fun getDynamicTypesSettings(jetFile: JetFile)
+    override fun getDynamicTypesSettings(file: JetFile)
-            = TargetPlatformDetector.getPlatform(jetFile).getDynamicTypesSettings()
+            = TargetPlatformDetector.getPlatform(file).getDynamicTypesSettings()
 
     override fun probablyNothingCallableNames(): ProbablyNothingCallableNames {
         return object : ProbablyNothingCallableNames {

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

@@ -31,7 +31,8 @@ public class ResolveElementCacheTest : JetLightCodeInsightFixtureTestCase() {
 class C {
     fun a() {
         b(1, 2)
-        c()
+        val x = c()
+        d(x)
     }
 
     fun b() {
@@ -132,4 +133,15 @@ class C {
             assert(bindingContext1 === bindingContext2)
         }
     }
+
+    public fun testPartialResolveCachedForAllStatementsResolved() {
+        doTest {
+            val bindingContext1 = statements[2].analyze(BodyResolveMode.PARTIAL) // resolve 'd(x)'
+            val bindingContext2 = (statements[1] as JetVariableDeclaration).getInitializer()!!.analyze(BodyResolveMode.PARTIAL) // resolve initializer in 'val x = c()' - it required for resolved 'd(x)' and should be already resolved
+            assert(bindingContext1 === bindingContext2)
+
+            val bindingContext3 = statements[0].analyze(BodyResolveMode.PARTIAL)
+            assert(bindingContext3 !== bindingContext1)
+        }
+    }
 }