Add snippet result type to analysis output of REPL IDE services

libraries/scripting/common/src/kotlin/script/experimental/api/replAnalysis.kt

@@ -5,16 +5,13 @@
 
 package kotlin.script.experimental.api
 
+import kotlin.script.experimental.util.PropertiesCollection
+
 /**
  * Type for [ReplCompleter.complete] return value
  */
 typealias ReplCompletionResult = Sequence<SourceCodeCompletionVariant>
 
-/**
- * Type for [ReplCodeAnalyzer.analyze] return value
- */
-typealias ReplAnalyzerResult = Sequence<ScriptDiagnostic>
-
 /**
  * Single code completion result
  */
@@ -25,6 +22,33 @@ data class SourceCodeCompletionVariant(
     val icon: String
 )
 
+interface ReplAnalyzerResultKeys
+
+/**
+ * The container for analysis results data returned from the [ReplCodeAnalyzer.analyze]
+ */
+class ReplAnalyzerResult(baseConfigurations: Iterable<ReplAnalyzerResult>, body: Builder.() -> Unit = {}) :
+    PropertiesCollection(Builder(baseConfigurations).apply(body).data) {
+
+    constructor(body: Builder.() -> Unit = {}) : this(emptyList(), body)
+
+    class Builder internal constructor(baseConfigurations: Iterable<ReplAnalyzerResult>) :
+        ReplAnalyzerResultKeys,
+        PropertiesCollection.Builder(baseConfigurations)
+
+    companion object : ReplAnalyzerResultKeys
+}
+
+/**
+ * Script compile-time errors and warning with their locations
+ */
+val ReplAnalyzerResultKeys.analysisDiagnostics by PropertiesCollection.key<Sequence<ScriptDiagnostic>>(emptySequence())
+
+/**
+ * String representing snippet return value, or null, if script returns nothing
+ */
+val ReplAnalyzerResultKeys.renderedResultType by PropertiesCollection.key<String?>()
+
 /**
  * Interface for REPL context completion
  */

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

@@ -6,12 +6,16 @@
 package org.jetbrains.kotlin.scripting.ide_services.compiler
 
 import org.jetbrains.kotlin.cli.common.messages.AnalyzerWithCompilerReport
+import org.jetbrains.kotlin.descriptors.ModuleDescriptor
+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.ide_services.compiler.impl.IdeLikeReplCodeAnalyzer
+import org.jetbrains.kotlin.scripting.ide_services.compiler.impl.*
-import org.jetbrains.kotlin.scripting.ide_services.compiler.impl.getKJvmCompletion
-import org.jetbrains.kotlin.scripting.ide_services.compiler.impl.prepareCodeForCompletion
 import kotlin.script.experimental.api.*
 import kotlin.script.experimental.host.ScriptingHostConfiguration
 import kotlin.script.experimental.jvm.defaultJvmScriptingHostConfiguration
@@ -29,58 +33,33 @@ class KJvmReplCompilerWithIdeServices(hostConfiguration: ScriptingHostConfigurat
         configuration: ScriptCompilationConfiguration
     ): ResultWithDiagnostics<ReplCompletionResult> =
         withMessageCollector(snippet) { messageCollector ->
-            val initialConfiguration = configuration.refineBeforeParsing(snippet).valueOr {
+            val analyzeResult = analyzeWithCursor(
-                return it
+                messageCollector, snippet, configuration, cursor
-            }
+            ) { snippet, cursorAbs ->
-
+                val newText =
-            val cursorAbs = cursor.calcAbsolute(snippet)
+                    prepareCodeForCompletion(snippet.text, cursorAbs)
-
+                object : SourceCode {
-            val newText =
+                    override val text: String
-                prepareCodeForCompletion(snippet.text, cursorAbs)
+                        get() = newText
-            val newSnippet = object : SourceCode {
+                    override val name: String?
-                override val text: String
+                        get() = snippet.name
-                    get() = newText
+                    override val locationId: String?
-                override val name: String?
+                        get() = snippet.locationId
-                    get() = snippet.name
-                override val locationId: String?
-                    get() = snippet.locationId
-
-            }
-
-            val compilationState = state.getCompilationState(initialConfiguration)
-
-            val (_, errorHolder, snippetKtFile) = prepareForAnalyze(
-                newSnippet,
-                messageCollector,
-                compilationState,
-                checkSyntaxErrors = false
-            ).valueOr { return@withMessageCollector it }
-
-            val analysisResult =
-                compilationState.analyzerEngine.statelessAnalyzeWithImportedScripts(snippetKtFile, emptyList(), scriptPriority.get() + 1)
-            AnalyzerWithCompilerReport.reportDiagnostics(analysisResult.diagnostics, errorHolder)
-
-            val (_, bindingContext, resolutionFacade, moduleDescriptor) = when (analysisResult) {
-                is IdeLikeReplCodeAnalyzer.ReplLineAnalysisResultWithStateless.Stateless -> {
-                    analysisResult
                 }
-                else -> return failure(
-                    newSnippet,
-                    messageCollector,
-                    "Unexpected result ${analysisResult::class.java}"
-                )
             }
 
-            return getKJvmCompletion(
+            with(analyzeResult.valueOr { return it }) {
-                snippetKtFile,
+                return getKJvmCompletion(
-                bindingContext,
+                    ktScript,
-                resolutionFacade,
+                    bindingContext,
-                moduleDescriptor,
+                    resolutionFacade,
-                cursorAbs
+                    moduleDescriptor,
-            ).asSuccess(messageCollector.diagnostics)
+                    cursorAbs
+                ).asSuccess(messageCollector.diagnostics)
+            }
         }
 
-    private fun List<ScriptDiagnostic>.toAnalyzeResult() = (filter {
+    private fun List<ScriptDiagnostic>.toAnalyzeResultSequence() = (filter {
         when (it.severity) {
             ScriptDiagnostic.Severity.FATAL,
             ScriptDiagnostic.Severity.ERROR,
@@ -96,24 +75,73 @@ class KJvmReplCompilerWithIdeServices(hostConfiguration: ScriptingHostConfigurat
         configuration: ScriptCompilationConfiguration
     ): ResultWithDiagnostics<ReplAnalyzerResult> {
         return withMessageCollector(snippet) { messageCollector ->
-            val initialConfiguration = configuration.refineBeforeParsing(snippet).valueOr {
+            val analyzeResult = analyzeWithCursor(
-                return it
+                messageCollector, snippet, configuration
+            )
+
+            with(analyzeResult.valueOr { return it }) {
+                val resultRenderedType = resultProperty?.let {
+                    DescriptorRenderer.SHORT_NAMES_IN_TYPES.renderType(it.type)
+                }
+                return ReplAnalyzerResult {
+                    analysisDiagnostics(messageCollector.diagnostics.toAnalyzeResultSequence())
+                    renderedResultType(resultRenderedType)
+                }.asSuccess()
             }
-
-            val compilationState = state.getCompilationState(initialConfiguration)
-
-            val (_, errorHolder, snippetKtFile) = prepareForAnalyze(
-                snippet,
-                messageCollector,
-                compilationState,
-                checkSyntaxErrors = true
-            ).valueOr { return@withMessageCollector messageCollector.diagnostics.toAnalyzeResult().asSuccess() }
-
-            val analysisResult =
-                compilationState.analyzerEngine.statelessAnalyzeWithImportedScripts(snippetKtFile, emptyList(), scriptPriority.get() + 1)
-            AnalyzerWithCompilerReport.reportDiagnostics(analysisResult.diagnostics, errorHolder)
-
-            messageCollector.diagnostics.toAnalyzeResult().asSuccess()
         }
     }
+
+    private fun analyzeWithCursor(
+        messageCollector: ScriptDiagnosticsMessageCollector,
+        snippet: SourceCode,
+        configuration: ScriptCompilationConfiguration,
+        cursor: SourceCode.Position? = null,
+        getNewSnippet: (SourceCode, Int) -> SourceCode = { code, _ -> code }
+    ): ResultWithDiagnostics<AnalyzeWithCursorResult> {
+        val initialConfiguration = configuration.refineBeforeParsing(snippet).valueOr {
+            return it
+        }
+
+        val cursorAbs = cursor?.calcAbsolute(snippet) ?: -1
+        val newSnippet = if (cursorAbs == -1) snippet else getNewSnippet(snippet, cursorAbs)
+
+        val compilationState = state.getCompilationState(initialConfiguration)
+
+        val (_, errorHolder, snippetKtFile) = prepareForAnalyze(
+            newSnippet,
+            messageCollector,
+            compilationState,
+            checkSyntaxErrors = false
+        ).valueOr { return it }
+
+        val analysisResult =
+            compilationState.analyzerEngine.statelessAnalyzeWithImportedScripts(snippetKtFile, emptyList(), scriptPriority.get() + 1)
+        AnalyzerWithCompilerReport.reportDiagnostics(analysisResult.diagnostics, errorHolder)
+
+        val (_, bindingContext, resolutionFacade, moduleDescriptor, resultProperty) = when (analysisResult) {
+            is IdeLikeReplCodeAnalyzer.ReplLineAnalysisResultWithStateless.Stateless -> {
+                analysisResult
+            }
+            else -> return failure(
+                newSnippet,
+                messageCollector,
+                "Unexpected result ${analysisResult::class.java}"
+            )
+        }
+
+        return AnalyzeWithCursorResult(
+            snippetKtFile, bindingContext, resolutionFacade, moduleDescriptor, cursorAbs, resultProperty
+        ).asSuccess()
+    }
+
+    companion object {
+        data class AnalyzeWithCursorResult(
+            val ktScript: KtFile,
+            val bindingContext: BindingContext,
+            val resolutionFacade: KotlinResolutionFacadeForRepl,
+            val moduleDescriptor: ModuleDescriptor,
+            val cursorAbs: Int,
+            val resultProperty: PropertyDescriptor?,
+        )
+    }
 }

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

@@ -10,6 +10,8 @@ import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
 import org.jetbrains.kotlin.container.getService
 import org.jetbrains.kotlin.descriptors.ClassDescriptorWithResolutionScopes
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor
+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.diagnostics.Diagnostics
@@ -25,7 +27,8 @@ class IdeLikeReplCodeAnalyzer(private val environment: KotlinCoreEnvironment) :
             override val diagnostics: Diagnostics,
             val bindingContext: BindingContext,
             val resolutionFacade: KotlinResolutionFacadeForRepl,
-            val moduleDescriptor: ModuleDescriptor
+            val moduleDescriptor: ModuleDescriptor,
+            val resultProperty: PropertyDescriptor?,
         ) :
             ReplLineAnalysisResultWithStateless {
             override val scriptDescriptor: ClassDescriptorWithResolutionScopes? get() = null
@@ -51,7 +54,9 @@ class IdeLikeReplCodeAnalyzer(private val environment: KotlinCoreEnvironment) :
         )
         replState.submitLine(linePsi)
 
-        topDownAnalyzer.analyzeDeclarations(topDownAnalysisContext.topDownAnalysisMode, listOf(linePsi) + importedScripts)
+        val context = topDownAnalyzer.analyzeDeclarations(topDownAnalysisContext.topDownAnalysisMode, listOf(linePsi) + importedScripts)
+
+        val resultPropertyDescriptor = (context.scripts[linePsi.script] as? ScriptDescriptor)?.resultValue
 
         val moduleDescriptor = container.getService(ModuleDescriptor::class.java)
         val resolutionFacade =
@@ -61,7 +66,8 @@ class IdeLikeReplCodeAnalyzer(private val environment: KotlinCoreEnvironment) :
             diagnostics,
             trace.bindingContext,
             resolutionFacade,
-            moduleDescriptor
+            moduleDescriptor,
+            resultPropertyDescriptor,
         )
     }