Support selective filtering of implicits for extensions resolution in REPL

plugins/scripting/scripting-compiler-impl/src/org/jetbrains/kotlin/scripting/resolve/resolveCompilationOptions.kt

@@ -0,0 +1,23 @@
+/*
+ * 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.scripting.resolve
+
+import kotlin.script.experimental.api.KotlinType
+import kotlin.script.experimental.api.ScriptCompilationConfigurationKeys
+import kotlin.script.experimental.util.PropertiesCollection
+
+/**
+ * Classes for which instances extensions should not be resolved if they are used in implicit context
+ */
+val ScriptCompilationConfigurationKeys.skipExtensionsResolutionForImplicits
+        by PropertiesCollection.key<Collection<KotlinType>>(emptyList())
+
+/**
+ * Extensions resolution for these classes in implicit context of their instances will be done only for innermost instance
+ * in scopes chain. Instances of each class in collection are handled separately.
+ */
+val ScriptCompilationConfigurationKeys.skipExtensionsResolutionForImplicitsExceptInnermost
+        by PropertiesCollection.key<Collection<KotlinType>>(emptyList())

plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/impl/KJvmReplCompilerBase.kt

@@ -18,10 +18,15 @@ import org.jetbrains.kotlin.codegen.KotlinCodegenFacade
 import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.descriptors.ScriptDescriptor
 import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.resolve.calls.tower.ImplicitsExtensionsResolutionFilter
+import org.jetbrains.kotlin.scripting.compiler.plugin.repl.ReplImplicitsExtensionsResolutionFilter
 import org.jetbrains.kotlin.scripting.compiler.plugin.repl.JvmReplCompilerStageHistory
 import org.jetbrains.kotlin.scripting.compiler.plugin.repl.JvmReplCompilerState
 import org.jetbrains.kotlin.scripting.compiler.plugin.repl.ReplCodeAnalyzerBase
 import org.jetbrains.kotlin.scripting.definitions.ScriptDependenciesProvider
+import org.jetbrains.kotlin.scripting.resolve.skipExtensionsResolutionForImplicits
+import org.jetbrains.kotlin.scripting.resolve.skipExtensionsResolutionForImplicitsExceptInnermost
+import org.jetbrains.kotlin.utils.addToStdlib.min
 import java.util.concurrent.atomic.AtomicInteger
 import kotlin.script.experimental.api.*
 import kotlin.script.experimental.host.ScriptingHostConfiguration
@@ -33,18 +38,19 @@ import kotlin.script.experimental.util.add
 
 open class KJvmReplCompilerBase<AnalyzerT : ReplCodeAnalyzerBase> protected constructor(
     protected val hostConfiguration: ScriptingHostConfiguration = defaultJvmScriptingHostConfiguration,
-    val initAnalyzer: (SharedScriptCompilationContext) -> AnalyzerT
+    val initAnalyzer: (SharedScriptCompilationContext, ImplicitsExtensionsResolutionFilter) -> AnalyzerT
 ) : ReplCompiler<KJvmCompiledScript>, ScriptCompiler {
     val state = JvmReplCompilerState({ createReplCompilationState(it, initAnalyzer) })
     val history = JvmReplCompilerStageHistory(state)
     protected val scriptPriority = AtomicInteger()
+    private val resolutionFilter = ReplImplicitsExtensionsResolutionFilter()
 
     override var lastCompiledSnippet: LinkedSnippetImpl<KJvmCompiledScript>? = null
         protected set
 
     fun createReplCompilationState(
         scriptCompilationConfiguration: ScriptCompilationConfiguration,
-        initAnalyzer: (SharedScriptCompilationContext) -> AnalyzerT /* = { ReplCodeAnalyzer1(it.environment) } */
+        initAnalyzer: (SharedScriptCompilationContext, ImplicitsExtensionsResolutionFilter) -> AnalyzerT /* = { ReplCodeAnalyzer1(it.environment) } */
     ): ReplCompilationState<AnalyzerT> {
         val context = withMessageCollectorAndDisposable(disposeOnSuccess = false) { messageCollector, disposable ->
             createIsolatedCompilationContext(
@@ -54,7 +60,10 @@ open class KJvmReplCompilerBase<AnalyzerT : ReplCodeAnalyzerBase> protected cons
                 disposable
             ).asSuccess()
         }.valueOr { throw IllegalStateException("Unable to initialize repl compiler:\n  ${it.reports.joinToString("\n  ")}") }
-        return ReplCompilationState(context, initAnalyzer)
+
+        updateResolutionFilter(scriptCompilationConfiguration)
+
+        return ReplCompilationState(context, initAnalyzer, resolutionFilter)
     }
 
     override suspend fun compile(
@@ -63,6 +72,8 @@ open class KJvmReplCompilerBase<AnalyzerT : ReplCodeAnalyzerBase> protected cons
     ): ResultWithDiagnostics<LinkedSnippet<KJvmCompiledScript>> =
         snippets.map { snippet ->
             withMessageCollector(snippet) { messageCollector ->
+                updateResolutionFilter(configuration)
+
                 val initialConfiguration = configuration.refineBeforeParsing(snippet).valueOr {
                     return it
                 }
@@ -225,10 +236,44 @@ open class KJvmReplCompilerBase<AnalyzerT : ReplCodeAnalyzerBase> protected cons
             ).asSuccess()
         }
 
+    protected fun updateResolutionFilter(configuration: ScriptCompilationConfiguration) {
+        val updatedConfiguration = updateConfigurationWithPreviousScripts(configuration)
+
+        val classesToSkip =
+            updatedConfiguration[ScriptCompilationConfiguration.skipExtensionsResolutionForImplicits]!!
+        val classesToSkipAfterFirstTime =
+            updatedConfiguration[ScriptCompilationConfiguration.skipExtensionsResolutionForImplicitsExceptInnermost]!!
+
+        resolutionFilter.update(classesToSkip, classesToSkipAfterFirstTime)
+    }
+
+    private fun updateConfigurationWithPreviousScripts(
+        configuration: ScriptCompilationConfiguration
+    ): ScriptCompilationConfiguration {
+        val allPreviousLines =
+            generateSequence(lastCompiledSnippet) { it.previous }
+                .map { KotlinType(it.get().scriptClassFQName) }
+                .toList()
+
+        val skipFirstTime = allPreviousLines.subList(0, min(1, allPreviousLines.size))
+        val skipAlways =
+            if (allPreviousLines.isEmpty()) emptyList()
+            else allPreviousLines.subList(1, allPreviousLines.size)
+
+        return ScriptCompilationConfiguration(configuration) {
+            skipExtensionsResolutionForImplicits.update {
+                it?.also { it.toMutableList().addAll(skipAlways) } ?: skipAlways
+            }
+            skipExtensionsResolutionForImplicitsExceptInnermost.update {
+                it?.also { it.toMutableList().addAll(skipFirstTime) } ?: skipFirstTime
+            }
+        }
+    }
+
     companion object {
         fun create(hostConfiguration: ScriptingHostConfiguration = defaultJvmScriptingHostConfiguration) =
-            KJvmReplCompilerBase(hostConfiguration) {
+            KJvmReplCompilerBase(hostConfiguration) { context, resolutionFilter ->
-                ReplCodeAnalyzerBase(it.environment)
+                ReplCodeAnalyzerBase(context.environment, implicitsResolutionFilter = resolutionFilter)
             }
     }
 
@@ -236,13 +281,14 @@ open class KJvmReplCompilerBase<AnalyzerT : ReplCodeAnalyzerBase> protected cons
 
 class ReplCompilationState<AnalyzerT : ReplCodeAnalyzerBase>(
     val context: SharedScriptCompilationContext,
-    val analyzerInit: (context: SharedScriptCompilationContext) -> AnalyzerT
+    val analyzerInit: (context: SharedScriptCompilationContext, implicitsResolutionFilter: ImplicitsExtensionsResolutionFilter) -> AnalyzerT,
+    override val implicitsResolutionFilter: ImplicitsExtensionsResolutionFilter
 ) : JvmReplCompilerState.Compilation {
     override val disposable: Disposable? get() = context.disposable
     override val baseScriptCompilationConfiguration: ScriptCompilationConfiguration get() = context.baseScriptCompilationConfiguration
     override val environment: KotlinCoreEnvironment get() = context.environment
     override val analyzerEngine: AnalyzerT by lazy {
         // ReplCodeAnalyzer1(context.environment)
-        analyzerInit(context)
+        analyzerInit(context, implicitsResolutionFilter)
     }
 }

plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/repl/ReplCodeAnalyzer.kt

@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.resolve.LazyTopDownAnalyzer
 import org.jetbrains.kotlin.resolve.TopDownAnalysisContext
 import org.jetbrains.kotlin.resolve.TopDownAnalysisMode
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfoFactory
+import org.jetbrains.kotlin.resolve.calls.tower.ImplicitsExtensionsResolutionFilter
 import org.jetbrains.kotlin.resolve.diagnostics.Diagnostics
 import org.jetbrains.kotlin.resolve.extensions.ExtraImportsProviderExtension
 import org.jetbrains.kotlin.resolve.lazy.*
@@ -40,13 +41,14 @@ import kotlin.script.experimental.jvm.util.SnippetsHistory
 
 open class ReplCodeAnalyzerBase(
     environment: KotlinCoreEnvironment,
-    val trace: BindingTraceContext = NoScopeRecordCliBindingTrace()
+    val trace: BindingTraceContext = NoScopeRecordCliBindingTrace(),
+    implicitsResolutionFilter: ImplicitsExtensionsResolutionFilter? = null
 ) {
     protected val scriptDeclarationFactory: ScriptMutableDeclarationProviderFactory
 
     protected val container: ComponentProvider
     protected val topDownAnalysisContext: TopDownAnalysisContext
-    protected val topDownAnalyzer: LazyTopDownAnalyzer
+    private val topDownAnalyzer: LazyTopDownAnalyzer
     protected val resolveSession: ResolveSession
     protected val replState = ResettableAnalyzerState()
 
@@ -62,7 +64,8 @@ open class ReplCodeAnalyzerBase(
             trace,
             environment.configuration,
             environment::createPackagePartProvider,
-            { _, _ -> ScriptMutableDeclarationProviderFactory() }
+            { _, _ -> ScriptMutableDeclarationProviderFactory() },
+            implicitsResolutionFilter = implicitsResolutionFilter
         )
 
         this.module = container.get()
@@ -117,13 +120,17 @@ open class ReplCodeAnalyzerBase(
         return doAnalyze(psiFile, importedScripts, codeLine.addNo(priority))
     }
 
+    protected fun runAnalyzer(linePsi: KtFile, importedScripts: List<KtFile>): TopDownAnalysisContext {
+        return topDownAnalyzer.analyzeDeclarations(topDownAnalysisContext.topDownAnalysisMode, listOf(linePsi) + importedScripts)
+    }
+
     private fun doAnalyze(linePsi: KtFile, importedScripts: List<KtFile>, codeLine: SourceCodeByReplLine): ReplLineAnalysisResult {
         scriptDeclarationFactory.setDelegateFactory(
             FileBasedDeclarationProviderFactory(resolveSession.storageManager, listOf(linePsi) + importedScripts)
         )
         replState.submitLine(linePsi)
 
-        val context = topDownAnalyzer.analyzeDeclarations(topDownAnalysisContext.topDownAnalysisMode, listOf(linePsi) + importedScripts)
+        val context = runAnalyzer(linePsi, importedScripts)
 
         val diagnostics = trace.bindingContext.diagnostics
         val hasErrors = diagnostics.any { it.severity == Severity.ERROR }

plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/repl/ReplImplicitsExtensionsResolutionFilter.kt

@@ -0,0 +1,52 @@
+/*
+ * 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.scripting.compiler.plugin.repl
+
+import org.jetbrains.kotlin.resolve.calls.tower.ImplicitsExtensionsResolutionFilter
+import org.jetbrains.kotlin.resolve.calls.tower.ScopeWithImplicitsExtensionsResolutionInfo
+import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
+import org.jetbrains.kotlin.resolve.scopes.HierarchicalScope
+import org.jetbrains.kotlin.resolve.scopes.LexicalScope
+import org.jetbrains.kotlin.resolve.scopes.receivers.ImplicitClassReceiver
+import java.util.concurrent.locks.ReentrantReadWriteLock
+import kotlin.concurrent.read
+import kotlin.concurrent.write
+import kotlin.script.experimental.api.KotlinType
+
+class ReplImplicitsExtensionsResolutionFilter : ImplicitsExtensionsResolutionFilter {
+    private val lock = ReentrantReadWriteLock()
+    private var classesToSkipNames: Set<String> = emptySet()
+    private var classesToSkipFirstTimeNames: Set<String> = emptySet()
+
+    fun update(
+        classesToSkip: Collection<KotlinType> = emptyList(),
+        classesToSkipAfterFirstTime: Collection<KotlinType> = emptyList()
+    ) = lock.write {
+        classesToSkipNames = classesToSkip.mapTo(hashSetOf()) { it.typeName }
+        classesToSkipFirstTimeNames = classesToSkipAfterFirstTime.mapTo(hashSetOf()) { it.typeName }
+    }
+
+    override fun getScopesWithInfo(
+        scopes: Sequence<HierarchicalScope>
+    ): Sequence<ScopeWithImplicitsExtensionsResolutionInfo> {
+        val processedReceivers = mutableSetOf<String>()
+        return scopes.map { scope ->
+            val receiver = (scope as? LexicalScope)?.implicitReceiver?.value
+            val keep = receiver?.let {
+                lock.read {
+                    when (val descriptorFqName = (it as? ImplicitClassReceiver)?.declarationDescriptor?.fqNameSafe?.asString()) {
+                        null -> true
+                        in classesToSkipNames -> false
+                        in classesToSkipFirstTimeNames -> processedReceivers.add(descriptorFqName)
+                        else -> true
+                    }
+                }
+            } ?: true
+
+            ScopeWithImplicitsExtensionsResolutionInfo(scope, keep)
+        }
+    }
+}

plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/repl/jvmReplCompilation.kt

@@ -10,6 +10,7 @@ import com.intellij.openapi.util.Disposer
 import org.jetbrains.kotlin.cli.common.repl.*
 import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
 import org.jetbrains.kotlin.descriptors.ScriptDescriptor
+import org.jetbrains.kotlin.resolve.calls.tower.ImplicitsExtensionsResolutionFilter
 import java.util.concurrent.locks.ReentrantReadWriteLock
 import kotlin.concurrent.write
 import kotlin.script.experimental.api.*
@@ -60,5 +61,6 @@ class JvmReplCompilerState<CompilationT : JvmReplCompilerState.Compilation>(
         val baseScriptCompilationConfiguration: ScriptCompilationConfiguration
         val environment: KotlinCoreEnvironment
         val analyzerEngine: ReplCodeAnalyzerBase
+        val implicitsResolutionFilter: ImplicitsExtensionsResolutionFilter
     }
 }

plugins/scripting/scripting-ide-services/src/org/jetbrains/kotlin/scripting/ide_services/compiler/KJvmReplCompilerWithIdeServices.kt

@@ -11,10 +11,7 @@ import org.jetbrains.kotlin.descriptors.PropertyDescriptor
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.renderer.DescriptorRenderer
 import org.jetbrains.kotlin.resolve.BindingContext
-import org.jetbrains.kotlin.scripting.compiler.plugin.impl.KJvmReplCompilerBase
+import org.jetbrains.kotlin.scripting.compiler.plugin.impl.*
-import org.jetbrains.kotlin.scripting.compiler.plugin.impl.ScriptDiagnosticsMessageCollector
-import org.jetbrains.kotlin.scripting.compiler.plugin.impl.failure
-import org.jetbrains.kotlin.scripting.compiler.plugin.impl.withMessageCollector
 import org.jetbrains.kotlin.scripting.ide_services.compiler.impl.*
 import kotlin.script.experimental.api.*
 import kotlin.script.experimental.host.ScriptingHostConfiguration
@@ -22,8 +19,8 @@ import kotlin.script.experimental.jvm.defaultJvmScriptingHostConfiguration
 import kotlin.script.experimental.jvm.util.calcAbsolute
 
 class KJvmReplCompilerWithIdeServices(hostConfiguration: ScriptingHostConfiguration = defaultJvmScriptingHostConfiguration) :
-    KJvmReplCompilerBase<IdeLikeReplCodeAnalyzer>(hostConfiguration, {
+    KJvmReplCompilerBase<IdeLikeReplCodeAnalyzer>(hostConfiguration, { sharedScriptCompilationContext, scopeProcessor ->
-        IdeLikeReplCodeAnalyzer(it.environment)
+        IdeLikeReplCodeAnalyzer(sharedScriptCompilationContext.environment, scopeProcessor)
     }),
     ReplCompleter, ReplCodeAnalyzer {
 
@@ -99,6 +96,8 @@ class KJvmReplCompilerWithIdeServices(hostConfiguration: ScriptingHostConfigurat
         cursor: SourceCode.Position? = null,
         getNewSnippet: (SourceCode, Int) -> SourceCode = { code, _ -> code }
     ): ResultWithDiagnostics<AnalyzeWithCursorResult> {
+        updateResolutionFilter(configuration)
+
         val initialConfiguration = configuration.refineBeforeParsing(snippet).valueOr {
             return it
         }

plugins/scripting/scripting-ide-services/src/org/jetbrains/kotlin/scripting/ide_services/compiler/impl/IdeLikeReplCodeAnalyzer.kt

@@ -14,12 +14,16 @@ import org.jetbrains.kotlin.descriptors.PropertyDescriptor
 import org.jetbrains.kotlin.descriptors.ScriptDescriptor
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.calls.tower.ImplicitsExtensionsResolutionFilter
 import org.jetbrains.kotlin.resolve.diagnostics.Diagnostics
 import org.jetbrains.kotlin.resolve.lazy.declarations.FileBasedDeclarationProviderFactory
 import org.jetbrains.kotlin.scripting.compiler.plugin.repl.ReplCodeAnalyzerBase
 import org.jetbrains.kotlin.scripting.definitions.ScriptPriorities
 
-class IdeLikeReplCodeAnalyzer(private val environment: KotlinCoreEnvironment) : ReplCodeAnalyzerBase(environment, CliBindingTrace()) {
+class IdeLikeReplCodeAnalyzer(
+    private val environment: KotlinCoreEnvironment,
+    implicitsResolutionFilter: ImplicitsExtensionsResolutionFilter
+) : ReplCodeAnalyzerBase(environment, CliBindingTrace(), implicitsResolutionFilter) {
     interface ReplLineAnalysisResultWithStateless : ReplLineAnalysisResult {
         // Result of stateless analyse, which may be used for reporting errors
         // without code generation
@@ -54,7 +58,7 @@ class IdeLikeReplCodeAnalyzer(private val environment: KotlinCoreEnvironment) :
         )
         replState.submitLine(linePsi)
 
-        val context = topDownAnalyzer.analyzeDeclarations(topDownAnalysisContext.topDownAnalysisMode, listOf(linePsi) + importedScripts)
+        val context = runAnalyzer(linePsi, importedScripts)
 
         val resultPropertyDescriptor = (context.scripts[linePsi.script] as? ScriptDescriptor)?.resultValue