Support selective filtering of implicits for extensions resolution in REPL

2026-03-15 08:31:27 +00:00 · 2020-06-04 13:43:47 +03:00
parent 017f640f26
commit 94de114894
7 changed files with 152 additions and 19 deletions
--- a/plugins/scripting/scripting-compiler-impl/src/org/jetbrains/kotlin/scripting/resolve/resolveCompilationOptions.kt
+++ b/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())
--- a/plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/impl/KJvmReplCompilerBase.kt
+++ b/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) {
-                ReplCodeAnalyzerBase(it.environment)
+            KJvmReplCompilerBase(hostConfiguration) { context, resolutionFilter ->
+                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)
    }
 }
--- a/plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/repl/ReplCodeAnalyzer.kt
+++ b/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 }
--- a/plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/repl/ReplImplicitsExtensionsResolutionFilter.kt
+++ b/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)
+        }
+    }
+}
--- a/plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/repl/jvmReplCompilation.kt
+++ b/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
    }
 }
--- a/plugins/scripting/scripting-ide-services/src/org/jetbrains/kotlin/scripting/ide_services/compiler/KJvmReplCompilerWithIdeServices.kt
+++ b/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.ScriptDiagnosticsMessageCollector
-import org.jetbrains.kotlin.scripting.compiler.plugin.impl.failure
-import org.jetbrains.kotlin.scripting.compiler.plugin.impl.withMessageCollector
+import org.jetbrains.kotlin.scripting.compiler.plugin.impl.*
 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, {
-        IdeLikeReplCodeAnalyzer(it.environment)
+    KJvmReplCompilerBase<IdeLikeReplCodeAnalyzer>(hostConfiguration, { sharedScriptCompilationContext, scopeProcessor ->
+        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
        }
--- a/plugins/scripting/scripting-ide-services/src/org/jetbrains/kotlin/scripting/ide_services/compiler/impl/IdeLikeReplCodeAnalyzer.kt
+++ b/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