Perform and cache full resolve when partial is not supported

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

@@ -110,19 +110,28 @@ public abstract class ElementResolver protected constructor(
     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)
-            PartialBodyResolveFilter(contextElement, resolveElement, probablyNothingCallableNames(), bodyResolveMode == BodyResolveMode.PARTIAL_FOR_COMPLETION)
-        else
-            StatementFilter.NONE
+        var statementFilterUsed = StatementFilter.NONE
+
+        fun createStatementFilter(): StatementFilter {
+            assert(resolveElement is JetDeclaration)
+            if (bodyResolveMode != BodyResolveMode.FULL) {
+                statementFilterUsed = PartialBodyResolveFilter(
+                        contextElement,
+                        resolveElement as JetDeclaration,
+                        probablyNothingCallableNames(),
+                        bodyResolveMode == BodyResolveMode.PARTIAL_FOR_COMPLETION)
+            }
+            return statementFilterUsed
+        }
 
         val trace : BindingTrace = when (resolveElement) {
-            is JetNamedFunction -> functionAdditionalResolve(resolveSession, resolveElement, file, statementFilter)
+            is JetNamedFunction -> functionAdditionalResolve(resolveSession, resolveElement, file, createStatementFilter())
 
-            is JetClassInitializer -> initializerAdditionalResolve(resolveSession, resolveElement, file, statementFilter)
+            is JetClassInitializer -> initializerAdditionalResolve(resolveSession, resolveElement, file, createStatementFilter())
 
-            is JetSecondaryConstructor -> secondaryConstructorAdditionalResolve(resolveSession, resolveElement, file, statementFilter)
+            is JetSecondaryConstructor -> secondaryConstructorAdditionalResolve(resolveSession, resolveElement, file, createStatementFilter())
 
-            is JetProperty -> propertyAdditionalResolve(resolveSession, resolveElement, file, statementFilter)
+            is JetProperty -> propertyAdditionalResolve(resolveSession, resolveElement, file, createStatementFilter())
 
             is JetDelegationSpecifierList -> delegationSpecifierAdditionalResolve(resolveSession, resolveElement, resolveElement.getParent() as JetClassOrObject, file)
 
@@ -136,7 +145,7 @@ public abstract class ElementResolver protected constructor(
 
             is JetAnnotationEntry -> annotationAdditionalResolve(resolveSession, resolveElement)
 
-            is JetClass -> constructorAdditionalResolve(resolveSession, resolveElement, file, statementFilter)
+            is JetClass -> constructorAdditionalResolve(resolveSession, resolveElement, file)
 
             is JetTypeParameter -> typeParameterAdditionalResolve(resolveSession, resolveElement)
 
@@ -158,7 +167,7 @@ public abstract class ElementResolver protected constructor(
         JetFlowInformationProvider(resolveElement, controlFlowTrace).checkDeclaration()
         controlFlowTrace.addOwnDataTo(trace, null, false)
 
-        return Pair(trace.getBindingContext(), statementFilter)
+        return Pair(trace.getBindingContext(), statementFilterUsed)
     }
 
     private fun packageRefAdditionalResolve(resolveSession: ResolveSession, jetElement: JetElement): BindingTrace {
@@ -339,7 +348,7 @@ public abstract class ElementResolver protected constructor(
         return trace
     }
 
-    private fun constructorAdditionalResolve(resolveSession: ResolveSession, klass: JetClass, file: JetFile, statementFilter: StatementFilter): BindingTrace {
+    private fun constructorAdditionalResolve(resolveSession: ResolveSession, klass: JetClass, file: JetFile): BindingTrace {
         val trace = createDelegatingTrace(klass)
         val scope = resolveSession.getDeclarationScopeProvider().getResolutionScopeForDeclaration(klass)
 
@@ -347,7 +356,7 @@ public abstract class ElementResolver protected constructor(
         val constructorDescriptor = classDescriptor.getUnsubstitutedPrimaryConstructor()
                                     ?: error("Can't get primary constructor for descriptor '$classDescriptor' in from class '${klass.getElementTextWithContext()}'")
 
-        val bodyResolver = createBodyResolver(resolveSession, trace, file, statementFilter)
+        val bodyResolver = createBodyResolver(resolveSession, trace, file, StatementFilter.NONE)
         bodyResolver.resolveConstructorParameterDefaultValuesAndAnnotations(DataFlowInfo.EMPTY, trace, klass, constructorDescriptor, scope)
 
         return trace

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

@@ -28,19 +28,16 @@ import org.jetbrains.kotlin.idea.stubindex.JetProbablyNothingPropertyShortNameIn
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.BodyResolveCache
+import org.jetbrains.kotlin.resolve.StatementFilter
 import org.jetbrains.kotlin.resolve.lazy.*
 import org.jetbrains.kotlin.storage.MemoizedFunctionToNotNull
 
 public class ResolveElementCache(resolveSession: ResolveSession, private val project: Project) : ElementResolver(resolveSession), BodyResolveCache {
     // Recreate internal cache after change of modification count
-    private val additionalResolveCache: CachedValue<MemoizedFunctionToNotNull<JetElement, BindingContext>> = CachedValuesManager.getManager(project).createCachedValue(
-            object : CachedValueProvider<MemoizedFunctionToNotNull<JetElement, BindingContext>> {
-                override fun compute(): CachedValueProvider.Result<MemoizedFunctionToNotNull<JetElement, BindingContext>> {
-                    val manager = resolveSession.getStorageManager()
-                    val cacheFunction = manager.createSoftlyRetainedMemoizedFunction<JetElement, BindingContext> { element ->
-                        performElementAdditionalResolve(element, element, BodyResolveMode.FULL).first
-                    }
-                    return CachedValueProvider.Result.create(cacheFunction,
+    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())
                 }
@@ -58,36 +55,53 @@ public class ResolveElementCache(resolveSession: ResolveSession, private val pro
             false)
 
     override fun getElementAdditionalResolve(resolveElement: JetElement, contextElement: JetElement, bodyResolveMode: BodyResolveMode): BindingContext {
-        if (bodyResolveMode != BodyResolveMode.FULL && !hasElementAdditionalResolveCached(resolveElement) && resolveElement is JetDeclaration) {
-            if (bodyResolveMode == BodyResolveMode.PARTIAL) {
+        val fullResolveMap = fullResolveCache.getValue()
+        fullResolveMap[resolveElement]?.let { return it } // check if full additional resolve already performed
+
+        when (bodyResolveMode) {
+            BodyResolveMode.FULL -> {
+                val bindingContext = performElementAdditionalResolve(resolveElement, resolveElement, BodyResolveMode.FULL).first
+                fullResolveMap[resolveElement] = bindingContext
+                return bindingContext
+            }
+
+            BodyResolveMode.PARTIAL -> {
+                if (resolveElement !is JetDeclaration) {
+                    return getElementAdditionalResolve(resolveElement, contextElement, BodyResolveMode.FULL)
+                }
+
                 val statementToResolve = PartialBodyResolveFilter.findStatementToResolve(contextElement, resolveElement)
-                val map = partialBodyResolveCache.getValue()
-                map[statementToResolve ?: resolveElement]?.let { return it }
+                val partialResolveMap = partialBodyResolveCache.getValue()
+                partialResolveMap[statementToResolve ?: resolveElement]?.let { return it } // partial resolve is already cached for this statement
 
                 val (bindingContext, statementFilter) = performElementAdditionalResolve(resolveElement, contextElement, BodyResolveMode.PARTIAL)
 
-                if (statementFilter is PartialBodyResolveFilter) {
-                    for (statement in statementFilter.allStatementsToResolve) {
-                        if (!map.containsKey(statement) && bindingContext[BindingContext.PROCESSED, statement] == true) {
-                            map[statement] = bindingContext
-                        }
+                if (statementFilter == StatementFilter.NONE) { // partial resolve is not supported for the given declaration - full resolve performed instead
+                    fullResolveMap[resolveElement] = bindingContext
+                    return bindingContext
+                }
+
+                for (statement in (statementFilter as PartialBodyResolveFilter).allStatementsToResolve) {
+                    if (!partialResolveMap.containsKey(statement) && bindingContext[BindingContext.PROCESSED, statement] == true) {
+                        partialResolveMap[statement] = bindingContext
                     }
                 }
-                map[resolveElement] = bindingContext // we use the whole declaration key in the map to obtain resolve not inside any block (e.g. default parameter values)
+                partialResolveMap[resolveElement] = bindingContext // we use the whole declaration key in the map to obtain resolve not inside any block (e.g. default parameter values)
 
                 return bindingContext
             }
-            else {
+
+            BodyResolveMode.PARTIAL_FOR_COMPLETION -> {
+                if (resolveElement !is JetDeclaration) {
+                    return getElementAdditionalResolve(resolveElement, contextElement, BodyResolveMode.FULL)
+                }
+
+                // not cached
                 return performElementAdditionalResolve(resolveElement, contextElement, bodyResolveMode).first
             }
         }
-
-        return additionalResolveCache.getValue().invoke(resolveElement)
     }
 
-    private fun hasElementAdditionalResolveCached(element: JetElement)
-            = additionalResolveCache.hasUpToDateValue() && additionalResolveCache.getValue().isComputed(element)
-
     override fun createAdditionalCheckerProvider(file: JetFile, module: ModuleDescriptor)
             = TargetPlatformDetector.getPlatform(file).createAdditionalCheckerProvider(module)
 

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

@@ -28,7 +28,7 @@ public class ResolveElementCacheTest : JetLightCodeInsightFixtureTestCase() {
 
     private val FILE_TEXT =
 """
-class C {
+class C(param1: String = "", param2: Int = 0) {
     fun a(p: Int = 0) {
         b(1, 2)
         val x = c()
@@ -45,6 +45,7 @@ class C {
 
     private data class Data(
             val file: JetFile,
+            val klass: JetClass,
             val members: List<JetDeclaration>,
             val statements: List<JetExpression>,
             val factory: JetPsiFactory
@@ -57,7 +58,7 @@ class C {
         val function = members.first() as JetNamedFunction
         val statements = (function.getBodyExpression() as JetBlockExpression).getStatements()
         myFixture.getProject().executeWriteCommand("") {
-            Data(file, members, statements, JetPsiFactory(getProject())).handler()
+            Data(file, klass, members, statements, JetPsiFactory(getProject())).handler()
         }
     }
 
@@ -156,4 +157,14 @@ class C {
             assert(bindingContext3 !== bindingContext2)
         }
     }
+
+    public fun testFullResolvedCachedWhenPartialForConstructorInvoked() {
+        doTest {
+            val defaultValue1 = klass.getPrimaryConstructorParameters()[0].getDefaultValue()!!
+            val defaultValue2 = klass.getPrimaryConstructorParameters()[1].getDefaultValue()!!
+            val bindingContext1 = defaultValue1.analyze(BodyResolveMode.PARTIAL)
+            val bindingContext2 = defaultValue2.analyze(BodyResolveMode.FULL)
+            assert(bindingContext1 === bindingContext2)
+        }
+    }
 }