Implemented ability to analyze to multiple elements at once

compiler/frontend/src/org/jetbrains/kotlin/resolve/BindingContext.java

@@ -161,7 +161,7 @@ public interface BindingContext {
     /**
      * Has type of current expression has been already resolved
      */
-    WritableSlice<KtExpression, Boolean> PROCESSED = Slices.createSimpleSetSlice();
+    WritableSlice<KtExpression, Boolean> PROCESSED = Slices.createSimpleSlice();
     WritableSlice<KtElement, Boolean> USED_AS_EXPRESSION = Slices.createSimpleSetSlice();
     WritableSlice<KtElement, Boolean> USED_AS_RESULT_OF_LAMBDA = Slices.createSimpleSetSlice();
     WritableSlice<KtElement, Boolean> UNREACHABLE_CODE = Slices.createSimpleSetSlice();

compiler/frontend/src/org/jetbrains/kotlin/resolve/BindingContextUtils.java

@@ -181,7 +181,7 @@ public class BindingContextUtils {
     @Nullable
     public static KotlinTypeInfo getRecordedTypeInfo(@NotNull KtExpression expression, @NotNull BindingContext context) {
         // noinspection ConstantConditions
-        if (!context.get(BindingContext.PROCESSED, expression)) return null;
+        if (context.get(BindingContext.PROCESSED, expression) != Boolean.TRUE) return null;
         // NB: should never return null if expression is already processed
         KotlinTypeInfo result = context.get(BindingContext.EXPRESSION_TYPE_INFO, expression);
         return result != null ? result : TypeInfoFactoryKt.noTypeInfo(DataFlowInfoFactory.EMPTY);

compiler/frontend/src/org/jetbrains/kotlin/types/expressions/ExpressionTypingVisitorDispatcher.java

@@ -182,7 +182,7 @@ public abstract class ExpressionTypingVisitorDispatcher extends KtVisitor<Kotlin
                         result = expression.accept(visitor, context);
                         // Some recursive definitions (object expressions) must put their types in the cache manually:
                         //noinspection ConstantConditions
-                        if (context.trace.get(BindingContext.PROCESSED, expression)) {
+                        if (context.trace.get(BindingContext.PROCESSED, expression) == Boolean.TRUE) {
                             KotlinType type = context.trace.getBindingContext().getType(expression);
                             return result.replaceType(type);
                         }

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

@@ -30,6 +30,7 @@ interface ResolutionFacade {
     val project: Project
 
     fun analyze(element: KtElement, bodyResolveMode: BodyResolveMode = BodyResolveMode.FULL): BindingContext
+    fun analyze(elements: Collection<KtElement>, bodyResolveMode: BodyResolveMode): BindingContext
 
     fun analyzeFullyAndGetResult(elements: Collection<KtElement>): AnalysisResult
 

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

@@ -17,7 +17,6 @@
 package org.jetbrains.kotlin.resolve.lazy
 
 import com.intellij.psi.PsiElement
-import com.intellij.util.SmartList
 import org.jetbrains.kotlin.KtNodeTypes
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
@@ -26,14 +25,12 @@ import org.jetbrains.kotlin.resolve.StatementFilter
 import org.jetbrains.kotlin.util.isProbablyNothing
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 import org.jetbrains.kotlin.utils.addToStdlib.swap
-import java.util.ArrayList
+import java.util.*
-import java.util.HashMap
-import java.util.HashSet
 
 //TODO: do resolve anonymous object's body
 
 class PartialBodyResolveFilter(
-        elementToResolve: KtElement,
+        elementsToResolve: Collection<KtElement>,
         private val declaration: KtDeclaration,
         probablyNothingCallableNames: ProbablyNothingCallableNames,
         forCompletion: Boolean
@@ -50,7 +47,7 @@ class PartialBodyResolveFilter(
         get() = statementMarks.allMarkedStatements()
 
     init {
-        assert(declaration.isAncestor(elementToResolve))
+        elementsToResolve.forEach { assert(declaration.isAncestor(it)) }
         assert(!KtPsiUtil.isLocal(declaration)) { "Should never be invoked on local declaration otherwise we may miss some local declarations with type Nothing" }
 
         declaration.forEachDescendantOfType<KtCallableDeclaration> { declaration ->
@@ -67,7 +64,9 @@ class PartialBodyResolveFilter(
             }
         }
 
-        statementMarks.mark(elementToResolve, if (forCompletion) MarkLevel.NEED_COMPLETION else MarkLevel.NEED_REFERENCE_RESOLVE)
+        elementsToResolve.forEach {
+            statementMarks.mark(it, if (forCompletion) MarkLevel.NEED_COMPLETION else MarkLevel.NEED_REFERENCE_RESOLVE)
+        }
         declaration.forTopLevelBlocksInside { processBlock(it) }
     }
 

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

@@ -48,7 +48,7 @@ class CodeFragmentAnalyzer(
         val codeFragmentElement = codeFragment.getContentElement()
 
         val (scopeForContextElement, dataFlowInfo) = getScopeAndDataFlowForAnalyzeFragment(codeFragment) {
-            resolveElementCache!!.resolveToElement(it, bodyResolveMode)
+            resolveElementCache!!.resolveToElements(listOf(it), bodyResolveMode)
         } ?: return
 
         when (codeFragmentElement) {

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

@@ -46,8 +46,12 @@ internal class ResolutionFacadeImpl(
     fun findModuleDescriptor(ideaModuleInfo: IdeaModuleInfo) = projectFacade.findModuleDescriptor(ideaModuleInfo)
 
     override fun analyze(element: KtElement, bodyResolveMode: BodyResolveMode): BindingContext {
-        val resolveElementCache = getFrontendService(element, ResolveElementCache::class.java)
+        return analyze(listOf(element), bodyResolveMode)
-        return resolveElementCache.resolveToElement(element, bodyResolveMode)
+    }
+
+    override fun analyze(elements: Collection<KtElement>, bodyResolveMode: BodyResolveMode): BindingContext {
+        val resolveElementCache = getFrontendService(elements.first(), ResolveElementCache::class.java)
+        return resolveElementCache.resolveToElements(elements, bodyResolveMode)
     }
 
     override fun analyzeFullyAndGetResult(elements: Collection<KtElement>): AnalysisResult

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

@@ -31,7 +31,7 @@ class IdeaLocalDescriptorResolver(
         private val absentDescriptorHandler: AbsentDescriptorHandler
 ) : LocalDescriptorResolver {
     override fun resolveLocalDeclaration(declaration: KtDeclaration): DeclarationDescriptor {
-        val context = resolveElementCache.resolveToElement(declaration, BodyResolveMode.FULL)
+        val context = resolveElementCache.resolveToElements(listOf(declaration), BodyResolveMode.FULL)
         return context.get(BindingContext.DECLARATION_TO_DESCRIPTOR, declaration)
                ?: absentDescriptorHandler.diagnoseDescriptorNotFound(declaration)
     }

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

@@ -38,6 +38,7 @@ import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.forEachDescendantOfType
 import org.jetbrains.kotlin.psi.psiUtil.getElementTextWithContext
+import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
 import org.jetbrains.kotlin.resolve.*
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo
@@ -45,7 +46,8 @@ import org.jetbrains.kotlin.resolve.lazy.*
 import org.jetbrains.kotlin.resolve.lazy.descriptors.LazyClassDescriptor
 import org.jetbrains.kotlin.resolve.lazy.descriptors.LazyScriptDescriptor
 import org.jetbrains.kotlin.resolve.scopes.LexicalScope
-import org.jetbrains.kotlin.utils.addToStdlib.check
+import org.jetbrains.kotlin.utils.addIfNotNull
+import java.util.*
 
 class ResolveElementCache(
         private val resolveSession: ResolveSession,
@@ -112,13 +114,18 @@ class ResolveElementCache(
     }
 
     override fun resolveFunctionBody(function: KtNamedFunction)
-            = getElementAdditionalResolve(function, function, BodyResolveMode.FULL)
+            = getElementsAdditionalResolve(function, emptyList(), BodyResolveMode.FULL)
 
     fun resolvePrimaryConstructorParametersDefaultValues(ktClass: KtClass): BindingContext {
         return constructorAdditionalResolve(resolveSession, ktClass, ktClass.getContainingKtFile()).bindingContext
     }
 
+    @Deprecated("Use getElementsAdditionalResolve")
     fun getElementAdditionalResolve(resolveElement: KtElement, contextElement: KtElement, bodyResolveMode: BodyResolveMode): BindingContext {
+        return getElementsAdditionalResolve(resolveElement, listOf(contextElement), bodyResolveMode)
+    }
+
+    fun getElementsAdditionalResolve(resolveElement: KtElement, contextElements: Collection<KtElement>, bodyResolveMode: BodyResolveMode): BindingContext {
         // check if full additional resolve already performed and is up-to-date
         val fullResolveMap = fullResolveCache.value
         val cachedFullResolve = fullResolveMap[resolveElement]
@@ -133,23 +140,25 @@ class ResolveElementCache(
 
         when (bodyResolveMode) {
             BodyResolveMode.FULL -> {
-                val bindingContext = performElementAdditionalResolve(resolveElement, resolveElement, BodyResolveMode.FULL).first
+                val bindingContext = performElementAdditionalResolve(resolveElement, emptyList(), BodyResolveMode.FULL).first
                 fullResolveMap[resolveElement] = CachedFullResolve(bindingContext, resolveElement)
                 return bindingContext
             }
 
             BodyResolveMode.PARTIAL -> {
                 if (resolveElement !is KtDeclaration) {
-                    return getElementAdditionalResolve(resolveElement, contextElement, BodyResolveMode.FULL)
+                    return getElementsAdditionalResolve(resolveElement, emptyList(), BodyResolveMode.FULL)
                 }
 
                 val file = resolveElement.getContainingKtFile()
-                val statementToResolve = PartialBodyResolveFilter.findStatementToResolve(contextElement, resolveElement)
+                val statementsToResolve = contextElements.map { PartialBodyResolveFilter.findStatementToResolve(it, resolveElement) }.distinct()
                 val partialResolveMap = partialBodyResolveCache.value
-                partialResolveMap[statementToResolve ?: resolveElement]
+                val cachedResults = statementsToResolve.map { partialResolveMap[it ?: resolveElement] }
-                        ?.check { it.isUpToDate(file) }
+                if (cachedResults.all { it != null && it.isUpToDate(file) }) { // partial resolve is already cached for these statements
-                        ?.let { return it.bindingContext } // partial resolve is already cached for this statement
+                    return CompositeBindingContext.create(cachedResults.map { it!!.bindingContext }.distinct())
-                val (bindingContext, statementFilter) = performElementAdditionalResolve(resolveElement, contextElement, BodyResolveMode.PARTIAL)
+                }
+
+                val (bindingContext, statementFilter) = performElementAdditionalResolve(resolveElement, contextElements, BodyResolveMode.PARTIAL)
 
                 if (statementFilter == StatementFilter.NONE) { // partial resolve is not supported for the given declaration - full resolve performed instead
                     fullResolveMap[resolveElement] = CachedFullResolve(bindingContext, resolveElement)
@@ -172,35 +181,48 @@ class ResolveElementCache(
 
             BodyResolveMode.PARTIAL_FOR_COMPLETION -> {
                 if (resolveElement !is KtDeclaration) {
-                    return getElementAdditionalResolve(resolveElement, contextElement, BodyResolveMode.FULL)
+                    return getElementsAdditionalResolve(resolveElement, emptyList(), BodyResolveMode.FULL)
                 }
 
                 // not cached
-                return performElementAdditionalResolve(resolveElement, contextElement, bodyResolveMode).first
+                return performElementAdditionalResolve(resolveElement, contextElements, bodyResolveMode).first
             }
         }
     }
 
-    fun resolveToElement(element: KtElement, bodyResolveMode: BodyResolveMode = BodyResolveMode.FULL): BindingContext {
+    fun resolveToElements(elements: Collection<KtElement>, bodyResolveMode: BodyResolveMode = BodyResolveMode.FULL): BindingContext {
-        var contextElement = element
+        val elementsByAdditionalResolveElement: Map<KtElement?, List<KtElement>> = elements.groupBy { findElementOfAdditionalResolve(it) }
 
-        val elementOfAdditionalResolve = findElementOfAdditionalResolve(contextElement)
+        val bindingContexts = ArrayList<BindingContext>()
-
+        val declarationsToResolve = ArrayList<KtDeclaration>()
-        if (elementOfAdditionalResolve is KtParameter) {
+        var addResolveSessionBindingContext = false
-            // Parameters for function literal could be met inside other parameters. We can't make resolveToDescriptors for internal elements.
+        for ((elementOfAdditionalResolve, contextElements) in elementsByAdditionalResolveElement) {
-            contextElement = elementOfAdditionalResolve
+            if (elementOfAdditionalResolve != null) {
-        }
+                if (elementOfAdditionalResolve !is KtParameter) {
-        else if (elementOfAdditionalResolve != null) {
+                    val bindingContext = getElementsAdditionalResolve(elementOfAdditionalResolve, contextElements, bodyResolveMode)
-            return getElementAdditionalResolve(elementOfAdditionalResolve, contextElement, bodyResolveMode)
+                    bindingContexts.add(bindingContext)
+                }
+                else {
+                    // Parameters for function literal could be met inside other parameters. We can't make resolveToDescriptors for internal elements.
+                    declarationsToResolve.addIfNotNull(elementOfAdditionalResolve.getNonStrictParentOfType<KtDeclaration>())
+                    addResolveSessionBindingContext = true
+                }
+            }
+            else {
+                contextElements
+                        .map { it.getNonStrictParentOfType<KtDeclaration>() }
+                        .filterNotNull()
+                        .filterTo(declarationsToResolve) { it !is KtAnonymousInitializer }
+                addResolveSessionBindingContext = true
+            }
         }
 
-        val declaration = contextElement.getParentOfType<KtDeclaration>(false)
+        declarationsToResolve.forEach { resolveSession.resolveToDescriptor(it) }
-        if (declaration != null && declaration !is KtAnonymousInitializer) {
+        if (addResolveSessionBindingContext) {
-            // Activate descriptor resolution
+            bindingContexts.add(resolveSession.bindingContext)
-            resolveSession.resolveToDescriptor(declaration)
         }
 
-        return resolveSession.bindingContext
+        return CompositeBindingContext.create(bindingContexts)
     }
 
     private fun findElementOfAdditionalResolve(element: KtElement): KtElement? {
@@ -258,7 +280,7 @@ class ResolveElementCache(
         }
     }
 
-    private fun performElementAdditionalResolve(resolveElement: KtElement, contextElement: KtElement,
+    private fun performElementAdditionalResolve(resolveElement: KtElement, contextElements: Collection<KtElement>,
                                                 bodyResolveMode: BodyResolveMode): Pair<BindingContext, StatementFilter> {
         val file = resolveElement.getContainingKtFile()
 
@@ -268,7 +290,7 @@ class ResolveElementCache(
             assert(resolveElement is KtDeclaration)
             if (bodyResolveMode != BodyResolveMode.FULL) {
                 statementFilterUsed = PartialBodyResolveFilter(
-                        contextElement,
+                        contextElements,
                         resolveElement as KtDeclaration,
                         probablyNothingCallableNames(),
                         bodyResolveMode == BodyResolveMode.PARTIAL_FOR_COMPLETION)

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

@@ -17,9 +17,11 @@
 package org.jetbrains.kotlin.idea
 
 import com.intellij.psi.PsiDocumentManager
+import junit.framework.TestCase
 import org.intellij.lang.annotations.Language
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.caches.resolve.analyzeFullyAndGetResult
+import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade
 import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptor
 import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCase
 import org.jetbrains.kotlin.idea.test.KotlinLightProjectDescriptor
@@ -48,6 +50,7 @@ class C(param1: String = "", param2: Int = 0) {
     }
 
     fun c() {
+        x(2)
     }
 }
 """
@@ -361,4 +364,61 @@ class C(param1: String = "", param2: Int = 0) {
         val annotationArguments = ((file.declarations[0]) as KtClass).getPrimaryConstructor()!!.annotationEntries[0].valueArgumentList!!
         annotationArguments.analyzeFullyAndGetResult()
     }
+
+    fun testFullResolveMultiple() {
+        doTest {
+            val aBody = (members[0] as KtFunction).bodyExpression as KtBlockExpression
+            val statement1InFunA = aBody.statements[0]
+            val statement2InFunA = aBody.statements[1]
+            val statementInFunB = ((members[1] as KtFunction).bodyExpression as KtBlockExpression).statements[0]
+            val statementInFunC = ((members[2] as KtFunction).bodyExpression as KtBlockExpression).statements[0]
+
+            val bindingContext = checkResolveMultiple(BodyResolveMode.FULL, statement1InFunA, statementInFunB)
+
+            TestCase.assertEquals(true, bindingContext[BindingContext.PROCESSED, statement2InFunA])
+            TestCase.assertEquals(null, bindingContext[BindingContext.PROCESSED, statementInFunC])
+        }
+    }
+
+    fun testPartialResolveMultiple() {
+        doTest {
+            val aBody = (members[0] as KtFunction).bodyExpression as KtBlockExpression
+            val statement1InFunA = aBody.statements[0]
+            val statement2InFunA = aBody.statements[1]
+            val statementInFunB = ((members[1] as KtFunction).bodyExpression as KtBlockExpression).statements[0]
+            val constructorParameterDefault = klass.getPrimaryConstructor()!!.valueParameters[1].defaultValue!!
+            val funC = members[2]
+
+            checkResolveMultiple(BodyResolveMode.PARTIAL, statement1InFunA, statement2InFunA, statementInFunB, constructorParameterDefault, funC)
+        }
+    }
+
+    fun testPartialResolveMultipleInOneFunction() {
+        doTest {
+            val aBody = (members[0] as KtFunction).bodyExpression as KtBlockExpression
+            val statement1InFunA = aBody.statements[0]
+            val statement2InFunA = aBody.statements[1]
+
+            val bindingContext = checkResolveMultiple(BodyResolveMode.PARTIAL, statement1InFunA, statement2InFunA)
+
+            val bindingContext1 = statement1InFunA.analyze(BodyResolveMode.PARTIAL)
+            assert(bindingContext1 === bindingContext)
+        }
+    }
+
+    private fun checkResolveMultiple(mode: BodyResolveMode, vararg expressions: KtExpression): BindingContext {
+        val resolutionFacade = expressions.first().getResolutionFacade()
+        val bindingContext = resolutionFacade.analyze(expressions.asList(), mode)
+
+        expressions.forEach {
+            if (it !is KtDeclaration) {
+                TestCase.assertEquals(true, bindingContext[BindingContext.PROCESSED, it])
+            }
+            else {
+                TestCase.assertNotNull(bindingContext[BindingContext.DECLARATION_TO_DESCRIPTOR, it])
+            }
+        }
+
+        return bindingContext
+    }
 }

idea/tests/org/jetbrains/kotlin/idea/codeInsight/AbstractOutOfBlockModificationTest.java

@@ -108,7 +108,7 @@ public abstract class AbstractOutOfBlockModificationTest extends KotlinLightCode
             @SuppressWarnings("ConstantConditions")
             boolean expressionProcessed = context.get(
                     BindingContext.PROCESSED,
-                    ktExpression instanceof KtFunctionLiteral ? (KtLambdaExpression) ktExpression.getParent() : ktExpression);
+                    ktExpression instanceof KtFunctionLiteral ? (KtLambdaExpression) ktExpression.getParent() : ktExpression) == Boolean.TRUE;
 
             assertEquals("Expected out-of-block should result expression analyzed and vise versa", expectedOutOfBlock,
                          expressionProcessed);

idea/tests/org/jetbrains/kotlin/idea/resolve/AbstractAdditionalResolveDescriptorRendererTest.kt

@@ -34,7 +34,7 @@ abstract class AbstractAdditionalResolveDescriptorRendererTest : AbstractDescrip
     override fun getDescriptor(declaration: KtDeclaration, container: ComponentProvider): DeclarationDescriptor {
         if (declaration is KtAnonymousInitializer || KtPsiUtil.isLocal(declaration)) {
             return container.get<ResolveElementCache>()
-                    .resolveToElement(declaration, BodyResolveMode.FULL)
+                    .resolveToElements(listOf(declaration), BodyResolveMode.FULL)
                     .get(BindingContext.DECLARATION_TO_DESCRIPTOR, declaration)!!
         }
         return container.get<ResolveSession>().resolveToDescriptor(declaration)