Allow skipping extensions resolution for implicit receivers

compiler/cli/src/org/jetbrains/kotlin/cli/jvm/compiler/TopDownAnalyzerFacadeForJVM.kt

@@ -66,6 +66,7 @@ import org.jetbrains.kotlin.platform.TargetPlatform
 import org.jetbrains.kotlin.platform.jvm.JvmPlatforms
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.resolve.*
+import org.jetbrains.kotlin.resolve.calls.tower.ImplicitsExtensionsResolutionFilter
 import org.jetbrains.kotlin.resolve.jvm.JavaDescriptorResolver
 import org.jetbrains.kotlin.resolve.jvm.extensions.AnalysisHandlerExtension
 import org.jetbrains.kotlin.resolve.jvm.extensions.PackageFragmentProviderExtension
@@ -135,7 +136,8 @@ object TopDownAnalyzerFacadeForJVM {
         declarationProviderFactory: (StorageManager, Collection<KtFile>) -> DeclarationProviderFactory,
         targetEnvironment: TargetEnvironment = CompilerEnvironment,
         sourceModuleSearchScope: GlobalSearchScope = newModuleSearchScope(project, files),
-        klibList: List<KotlinLibrary> = emptyList()
+        klibList: List<KotlinLibrary> = emptyList(),
+        implicitsResolutionFilter: ImplicitsExtensionsResolutionFilter? = null
     ): ComponentProvider {
         val jvmTarget = configuration.get(JVMConfigurationKeys.JVM_TARGET, JvmTarget.DEFAULT)
         val languageVersionSettings = configuration.languageVersionSettings
@@ -186,7 +188,8 @@ object TopDownAnalyzerFacadeForJVM {
                 targetEnvironment, lookupTracker, expectActualTracker,
                 packagePartProvider(dependencyScope), languageVersionSettings,
                 useBuiltInsProvider = true,
-                configureJavaClassFinder = configureJavaClassFinder
+                configureJavaClassFinder = configureJavaClassFinder,
+                implicitsResolutionFilter = implicitsResolutionFilter
             )
 
             moduleClassResolver.compiledCodeResolver = dependenciesContainer.get()
@@ -220,7 +223,8 @@ object TopDownAnalyzerFacadeForJVM {
             partProvider, languageVersionSettings,
             useBuiltInsProvider = true,
             configureJavaClassFinder = configureJavaClassFinder,
-            javaClassTracker = configuration[JVMConfigurationKeys.JAVA_CLASSES_TRACKER]
+            javaClassTracker = configuration[JVMConfigurationKeys.JAVA_CLASSES_TRACKER],
+            implicitsResolutionFilter = implicitsResolutionFilter
         ).apply {
             initJvmBuiltInsForTopDownAnalysis()
             (partProvider as? IncrementalPackagePartProvider)?.deserializationConfiguration = get()

compiler/frontend.java/src/org/jetbrains/kotlin/frontend/java/di/injection.kt

@@ -42,6 +42,7 @@ import org.jetbrains.kotlin.load.kotlin.PackagePartProvider
 import org.jetbrains.kotlin.load.kotlin.VirtualFileFinderFactory
 import org.jetbrains.kotlin.platform.TargetPlatform
 import org.jetbrains.kotlin.resolve.*
+import org.jetbrains.kotlin.resolve.calls.tower.ImplicitsExtensionsResolutionFilter
 import org.jetbrains.kotlin.resolve.jvm.JavaDescriptorResolver
 import org.jetbrains.kotlin.resolve.jvm.JvmDiagnosticComponents
 import org.jetbrains.kotlin.resolve.jvm.multiplatform.OptionalAnnotationPackageFragmentProvider
@@ -63,7 +64,8 @@ fun createContainerForLazyResolveWithJava(
     languageVersionSettings: LanguageVersionSettings,
     useBuiltInsProvider: Boolean,
     configureJavaClassFinder: (StorageComponentContainer.() -> Unit)? = null,
-    javaClassTracker: JavaClassesTracker? = null
+    javaClassTracker: JavaClassesTracker? = null,
+    implicitsResolutionFilter: ImplicitsExtensionsResolutionFilter? = null,
 ): StorageComponentContainer = createContainer("LazyResolveWithJava", JvmPlatformAnalyzerServices) {
     configureModule(moduleContext, jvmPlatform, JvmPlatformAnalyzerServices, bindingTrace, languageVersionSettings)
 
@@ -73,6 +75,7 @@ fun createContainerForLazyResolveWithJava(
     useInstance(moduleContentScope)
     useInstance(packagePartProvider)
     useInstance(declarationProviderFactory)
+    useInstanceIfNotNull(implicitsResolutionFilter)
 
     useInstance(VirtualFileFinderFactory.getInstance(moduleContext.project).create(moduleContentScope))
 

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/NewResolutionOldInference.kt

@@ -74,6 +74,7 @@ class NewResolutionOldInference(
     private val coroutineInferenceSupport: CoroutineInferenceSupport,
     private val deprecationResolver: DeprecationResolver,
     private val typeApproximator: TypeApproximator,
+    private val implicitsResolutionFilter: ImplicitsExtensionsResolutionFilter,
     private val callResolver: CallResolver,
     private val candidateInterceptor: CandidateInterceptor
 ) {
@@ -174,7 +175,7 @@ class NewResolutionOldInference(
 
         val dynamicScope = dynamicCallableDescriptors.createDynamicDescriptorScope(context.call, context.scope.ownerDescriptor)
         val scopeTower = ImplicitScopeTowerImpl(
-            context, dynamicScope, syntheticScopes, context.call.createLookupLocation(), typeApproximator, callResolver, candidateInterceptor
+            context, dynamicScope, syntheticScopes, context.call.createLookupLocation(), typeApproximator, implicitsResolutionFilter, callResolver, candidateInterceptor
         )
 
         val shouldUseOperatorRem = languageVersionSettings.supportsFeature(LanguageFeature.OperatorRem)
@@ -362,6 +363,7 @@ class NewResolutionOldInference(
         override val syntheticScopes: SyntheticScopes,
         override val location: LookupLocation,
         override val typeApproximator: TypeApproximator,
+        override val implicitsResolutionFilter: ImplicitsExtensionsResolutionFilter,
         val callResolver: CallResolver,
         val candidateInterceptor: CandidateInterceptor
     ) : ImplicitScopeTower {

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/PSICallResolver.kt

@@ -61,6 +61,7 @@ class PSICallResolver(
     private val kotlinToResolvedCallTransformer: KotlinToResolvedCallTransformer,
     private val kotlinCallResolver: KotlinCallResolver,
     private val typeApproximator: TypeApproximator,
+    private val implicitsResolutionFilter: ImplicitsExtensionsResolutionFilter,
     private val argumentTypeResolver: ArgumentTypeResolver,
     private val effectSystem: EffectSystem,
     private val constantExpressionEvaluator: ConstantExpressionEvaluator,
@@ -393,6 +394,7 @@ class PSICallResolver(
         override val isNewInferenceEnabled: Boolean get() = context.languageVersionSettings.supportsFeature(LanguageFeature.NewInference)
         override val lexicalScope: LexicalScope get() = context.scope
         override val typeApproximator: TypeApproximator get() = this@PSICallResolver.typeApproximator
+        override val implicitsResolutionFilter: ImplicitsExtensionsResolutionFilter get() = this@PSICallResolver.implicitsResolutionFilter
         private val cache = HashMap<ReceiverParameterDescriptor, ReceiverValueWithSmartCastInfo>()
 
         override fun getImplicitReceiver(scope: LexicalScope): ReceiverValueWithSmartCastInfo? {

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/tower/ImplicitScopeTower.kt

@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.resolve.calls.model.KotlinCallDiagnostic
 import org.jetbrains.kotlin.resolve.calls.tower.ResolutionCandidateApplicability.*
 import org.jetbrains.kotlin.resolve.scopes.*
 import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValueWithSmartCastInfo
+import org.jetbrains.kotlin.resolve.scopes.utils.parentsWithSelf
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.TypeApproximator
 
@@ -44,6 +45,11 @@ interface ImplicitScopeTower {
 
     val typeApproximator: TypeApproximator
 
+    val implicitsResolutionFilter: ImplicitsExtensionsResolutionFilter
+
+    fun allScopesWithImplicitsResolutionInfo(): Sequence<ScopeWithImplicitsExtensionsResolutionInfo> =
+        implicitsResolutionFilter.getScopesWithInfo(lexicalScope.parentsWithSelf)
+
     fun interceptFunctionCandidates(
         resolutionScope: ResolutionScope,
         name: Name,

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/tower/ImplicitsExtensionsResolutionFilter.kt

@@ -0,0 +1,29 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.resolve.calls.tower
+
+import org.jetbrains.kotlin.container.DefaultImplementation
+import org.jetbrains.kotlin.resolve.scopes.HierarchicalScope
+
+@DefaultImplementation(ImplicitsExtensionsResolutionFilter.Default::class)
+interface ImplicitsExtensionsResolutionFilter {
+    fun getScopesWithInfo(
+        scopes: Sequence<HierarchicalScope>
+    ): Sequence<ScopeWithImplicitsExtensionsResolutionInfo>
+
+    object Default : ImplicitsExtensionsResolutionFilter {
+        override fun getScopesWithInfo(
+            scopes: Sequence<HierarchicalScope>
+        ): Sequence<ScopeWithImplicitsExtensionsResolutionInfo> = scopes.map { scope ->
+            ScopeWithImplicitsExtensionsResolutionInfo(scope, true)
+        }
+    }
+}
+
+class ScopeWithImplicitsExtensionsResolutionInfo(
+    val scope: HierarchicalScope,
+    val resolveExtensionsForImplicitReceiver: Boolean,
+)

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/tower/TowerResolver.kt

@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.progress.ProgressIndicatorAndCompilationCanceledStatus
 import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind
 import org.jetbrains.kotlin.resolve.descriptorUtil.HIDES_MEMBERS_NAME_LIST
+import org.jetbrains.kotlin.resolve.scopes.HierarchicalScope
 import org.jetbrains.kotlin.resolve.scopes.ImportingScope
 import org.jetbrains.kotlin.resolve.scopes.LexicalScope
 import org.jetbrains.kotlin.resolve.scopes.ResolutionScope
@@ -185,7 +186,7 @@ class TowerResolver {
                 TowerData.TowerLevel(localLevel).process()?.let { return it }
             }
 
-            for (scope in implicitScopeTower.lexicalScope.parentsWithSelf) {
+            fun processScope(scope: HierarchicalScope, resolveExtensionsForImplicitReceiver: Boolean): Collection<C>? {
                 if (scope is LexicalScope) {
                     // statics
                     if (!scope.kind.withLocalDescriptors) {
@@ -194,12 +195,23 @@ class TowerResolver {
                     }
 
                     implicitScopeTower.getImplicitReceiver(scope)
-                        ?.let(this::processImplicitReceiver)
+                        ?.let { processImplicitReceiver(it, resolveExtensionsForImplicitReceiver) }
                         ?.let { return it }
                 } else {
                     TowerData.TowerLevel(ImportingScopeBasedTowerLevel(implicitScopeTower, scope as ImportingScope))
                         .process(scope.mayFitForName(name))?.let { return it }
                 }
+                return null
+            }
+
+            if (implicitScopeTower.implicitsResolutionFilter === ImplicitsExtensionsResolutionFilter.Default) {
+                for (scope in implicitScopeTower.lexicalScope.parentsWithSelf) {
+                    processScope(scope, true)?.let { return it }
+                }
+            } else {
+                for (scopeInfo in implicitScopeTower.allScopesWithImplicitsResolutionInfo()) {
+                    processScope(scopeInfo.scope, scopeInfo.resolveExtensionsForImplicitReceiver)?.let { return it }
+                }
             }
 
             recordLookups()
@@ -207,7 +219,7 @@ class TowerResolver {
             return resultCollector.getFinalCandidates()
         }
 
-        private fun processImplicitReceiver(implicitReceiver: ReceiverValueWithSmartCastInfo): Collection<C>? {
+        private fun processImplicitReceiver(implicitReceiver: ReceiverValueWithSmartCastInfo, resolveExtensions: Boolean): Collection<C>? {
             if (isNameForHidesMember) {
                 // hides members extensions
                 TowerData.BothTowerLevelAndImplicitReceiver(hidesMembersLevel, implicitReceiver).process()?.let { return it }
@@ -220,17 +232,19 @@ class TowerResolver {
             // synthetic properties
             TowerData.BothTowerLevelAndImplicitReceiver(syntheticLevel, implicitReceiver).process()?.let { return it }
 
-            // invokeExtension on local variable
+            if (resolveExtensions) {
-            TowerData.OnlyImplicitReceiver(implicitReceiver).process()?.let { return it }
+                // invokeExtension on local variable
+                TowerData.OnlyImplicitReceiver(implicitReceiver).process()?.let { return it }
 
-            // local extensions for implicit receiver
+                // local extensions for implicit receiver
-            for (localLevel in localLevels) {
+                for (localLevel in localLevels) {
-                TowerData.BothTowerLevelAndImplicitReceiver(localLevel, implicitReceiver).process()?.let { return it }
+                    TowerData.BothTowerLevelAndImplicitReceiver(localLevel, implicitReceiver).process()?.let { return it }
-            }
+                }
 
-            // extension for implicit receiver
+                // extension for implicit receiver
-            for (nonLocalLevel in nonLocalLevels) {
+                for (nonLocalLevel in nonLocalLevels) {
-                TowerData.BothTowerLevelAndImplicitReceiver(nonLocalLevel, implicitReceiver).process()?.let { return it }
+                    TowerData.BothTowerLevelAndImplicitReceiver(nonLocalLevel, implicitReceiver).process()?.let { return it }
+                }
             }
 
             return null