Extend import resolution for library-to-source analysis

Use composite importing scope for references when resolution anchors are enabled.
Composite scope provides additional descriptors from scope of resolution anchor module.
Overriding old importing scope with a new one is not possible as it breaks library
dependencies on other libraries, which are inaccessible in anchor scope (scope for sources).

KT-24309 In Progress

compiler/config/src/org/jetbrains/kotlin/config/AnalysisFlags.kt

@@ -47,4 +47,7 @@ object AnalysisFlags {
 
     @JvmStatic
     val reportErrorsOnIrDependencies by AnalysisFlag.Delegates.Boolean
+
+    @JvmStatic
+    val libraryToSourceAnalysis by AnalysisFlag.Delegates.Boolean
 }

compiler/config/src/org/jetbrains/kotlin/config/CommonConfigurationKeys.kt

@@ -61,3 +61,6 @@ var CompilerConfiguration.languageVersionSettings: LanguageVersionSettings
 
 val LanguageVersionSettings.isTypeRefinementEnabled: Boolean
     get() = getFlag(AnalysisFlags.useTypeRefinement)
+
+val LanguageVersionSettings.isLibraryToSourceAnalysisEnabled: Boolean
+    get() = getFlag(AnalysisFlags.libraryToSourceAnalysis)

compiler/frontend/src/org/jetbrains/kotlin/resolve/QualifiedExpressionResolver.kt

@@ -20,6 +20,7 @@ import com.intellij.psi.impl.source.DummyHolder
 import com.intellij.util.SmartList
 import org.jetbrains.kotlin.config.AnalysisFlags
 import org.jetbrains.kotlin.config.LanguageVersionSettings
+import org.jetbrains.kotlin.config.isLibraryToSourceAnalysisEnabled
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.diagnostics.Errors
 import org.jetbrains.kotlin.incremental.KotlinLookupLocation
@@ -35,6 +36,7 @@ import org.jetbrains.kotlin.resolve.calls.CallExpressionElement
 import org.jetbrains.kotlin.resolve.calls.checkers.UnderscoreUsageChecker
 import org.jetbrains.kotlin.resolve.calls.unrollToLeftMostQualifiedExpression
 import org.jetbrains.kotlin.resolve.descriptorUtil.module
+import org.jetbrains.kotlin.resolve.scopes.CompositePrioritizedImportingScope
 import org.jetbrains.kotlin.resolve.scopes.ImportingScope
 import org.jetbrains.kotlin.resolve.scopes.LexicalScope
 import org.jetbrains.kotlin.resolve.scopes.receivers.*
@@ -221,6 +223,31 @@ class QualifiedExpressionResolver(val languageVersionSettings: LanguageVersionSe
         trace: BindingTrace,
         excludedImportNames: Collection<FqName>,
         packageFragmentForVisibilityCheck: PackageFragmentDescriptor?
+    ): ImportingScope? {
+        fun processReferenceInContextOf(moduleDescriptor: ModuleDescriptor): ImportingScope? =
+            doProcessImportReference(
+                importDirective,
+                moduleDescriptor,
+                trace,
+                excludedImportNames,
+                packageFragmentForVisibilityCheck
+            )
+
+        val primaryImportingScope = processReferenceInContextOf(moduleDescriptor)
+        if (!languageVersionSettings.isLibraryToSourceAnalysisEnabled) return primaryImportingScope
+
+        val resolutionAnchor = moduleDescriptor.getResolutionAnchorIfAny() ?: return primaryImportingScope
+        val secondaryImportingScope = processReferenceInContextOf(resolutionAnchor) ?: return primaryImportingScope
+        if (primaryImportingScope == null) return secondaryImportingScope
+        return CompositePrioritizedImportingScope(primaryImportingScope, secondaryImportingScope)
+    }
+
+    private fun doProcessImportReference(
+        importDirective: KtImportInfo,
+        moduleDescriptor: ModuleDescriptor,
+        trace: BindingTrace,
+        excludedImportNames: Collection<FqName>,
+        packageFragmentForVisibilityCheck: PackageFragmentDescriptor?
     ): ImportingScope? { // null if some error happened
         ProgressIndicatorAndCompilationCanceledStatus.checkCanceled()
 

compiler/resolution/src/org/jetbrains/kotlin/resolve/scopes/Scopes.kt

@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.incremental.components.LookupLocation
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.utils.Printer
+import org.jetbrains.kotlin.utils.addToStdlib.ifNotEmpty
 
 // see utils/ScopeUtils.kt
 
@@ -169,3 +170,51 @@ abstract class BaseImportingScope(parent: ImportingScope?) : BaseHierarchicalSco
         changeNamesForAliased: Boolean
     ): Collection<DeclarationDescriptor> = emptyList()
 }
+
+class CompositePrioritizedImportingScope(
+    private val primaryScope: ImportingScope,
+    private val secondaryScope: ImportingScope,
+) : ImportingScope {
+    override val parent: ImportingScope?
+        get() = primaryScope.parent ?: secondaryScope.parent
+
+    override fun getContributedPackage(name: Name): PackageViewDescriptor? {
+        return primaryScope.getContributedPackage(name) ?: secondaryScope.getContributedPackage(name)
+    }
+
+    override fun getContributedDescriptors(
+        kindFilter: DescriptorKindFilter,
+        nameFilter: (Name) -> Boolean,
+        changeNamesForAliased: Boolean
+    ): Collection<DeclarationDescriptor> {
+        return primaryScope.getContributedDescriptors(kindFilter, nameFilter, changeNamesForAliased)
+            .ifEmpty { secondaryScope.getContributedDescriptors(kindFilter, nameFilter, changeNamesForAliased) }
+    }
+
+    override fun computeImportedNames(): Set<Name>? {
+        val primaryNames = primaryScope.computeImportedNames()
+        val secondaryNames = secondaryScope.computeImportedNames()
+        return primaryNames?.union(secondaryNames.orEmpty()) ?: secondaryNames
+    }
+
+    override fun printStructure(p: Printer) {
+        p.println(primaryScope::class.java.simpleName)
+    }
+
+    override fun getContributedClassifier(name: Name, location: LookupLocation): ClassifierDescriptor? {
+        return primaryScope.getContributedClassifier(name, location) ?: secondaryScope.getContributedClassifier(name, location)
+    }
+
+    override fun getContributedVariables(name: Name, location: LookupLocation): Collection<VariableDescriptor> {
+        return primaryScope.getContributedVariables(name, location).ifEmpty {
+            secondaryScope.getContributedVariables(name, location)
+        }
+    }
+
+    override fun getContributedFunctions(name: Name, location: LookupLocation): Collection<FunctionDescriptor> {
+        return primaryScope.getContributedFunctions(name, location).ifEmpty {
+            secondaryScope.getContributedFunctions(name, location)
+        }
+    }
+
+}

core/descriptors/src/org/jetbrains/kotlin/descriptors/findClassInModule.kt

@@ -19,6 +19,7 @@ package org.jetbrains.kotlin.descriptors
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.resolve.RESOLUTION_ANCHOR_PROVIDER_CAPABILITY
+import org.jetbrains.kotlin.resolve.getResolutionAnchorIfAny
 
 fun ModuleDescriptor.findClassifierAcrossModuleDependencies(classId: ClassId): ClassifierDescriptor? = withAnchorFallback {
     val packageViewDescriptor = getPackage(classId.packageFqName)
@@ -40,7 +41,7 @@ fun ModuleDescriptor.findClassifierAcrossModuleDependencies(classId: ClassId): C
 private inline fun ModuleDescriptor.withAnchorFallback(
     crossinline doSearch: ModuleDescriptor.() -> ClassifierDescriptor? 
 ): ClassifierDescriptor? {
-    val anchor = getCapability(RESOLUTION_ANCHOR_PROVIDER_CAPABILITY)?.getResolutionAnchor(this)
+    val anchor = getResolutionAnchorIfAny()
     return if (anchor == null) doSearch() else doSearch() ?: anchor.doSearch()
 }
 

core/descriptors/src/org/jetbrains/kotlin/resolve/ResolutionAnchorProvider.kt

@@ -19,3 +19,6 @@ interface ResolutionAnchorProvider {
 }
 
 val RESOLUTION_ANCHOR_PROVIDER_CAPABILITY = ModuleCapability<ResolutionAnchorProvider>("ResolutionAnchorProvider")
+
+fun ModuleDescriptor.getResolutionAnchorIfAny(): ModuleDescriptor? =
+    getCapability(RESOLUTION_ANCHOR_PROVIDER_CAPABILITY)?.getResolutionAnchor(this)

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

@@ -79,8 +79,11 @@ class KotlinIdeResolutionAnchorService(
 
     override fun getResolutionAnchor(moduleDescriptor: ModuleDescriptor): ModuleDescriptor? {
         if (!project.libraryToSourceAnalysisEnabled) return null
+
         val moduleInfo = moduleDescriptor.moduleInfo ?: return null
-        val mapped = moduleToAnchor[moduleInfo] ?: return null
+        val keyModuleInfo = if (moduleInfo is SourceForBinaryModuleInfo) moduleInfo.binariesModuleInfo else moduleInfo
+        val mapped = moduleToAnchor[keyModuleInfo] ?: return null
+
         return KotlinCacheService.getInstance(project)
             .getResolutionFacadeByModuleInfo(mapped, mapped.platform)
             ?.moduleDescriptor

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

@@ -250,6 +250,9 @@ private fun MutableMap<AnalysisFlag<*>, Any>.initIDESpecificAnalysisSettings(pro
     if (KotlinMultiplatformAnalysisModeComponent.getMode(project) == KotlinMultiplatformAnalysisModeComponent.Mode.COMPOSITE) {
         put(AnalysisFlags.useTypeRefinement, true)
     }
+    if (KotlinLibraryToSourceAnalysisComponent.isEnabled(project)) {
+        put(AnalysisFlags.libraryToSourceAnalysis, true)
+    }
     put(AnalysisFlags.ideMode, true)
 }