K2 Scripting: add support for result field

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/declaration/FirImplicitNothingReturnTypeChecker.kt

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.diagnostics.reportOn
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
 import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
 import org.jetbrains.kotlin.fir.declarations.FirProperty
 import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
 import org.jetbrains.kotlin.fir.declarations.utils.isOverride
@@ -22,6 +23,7 @@ object FirImplicitNothingReturnTypeChecker : FirCallableDeclarationChecker() {
         if (declaration is FirProperty && declaration.isLocal) return
         if (declaration.isOverride) return
         if (declaration.symbol.hasExplicitReturnType) return
+        if (declaration.origin == FirDeclarationOrigin.ScriptCustomization.ResultProperty) return
         if (declaration.returnTypeRef.coneType.isNothing) {
             val factory = when (declaration) {
                 is FirSimpleFunction -> FirErrors.IMPLICIT_NOTHING_RETURN_TYPE

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrDeclarationStorage.kt

@@ -859,6 +859,7 @@ class Fir2IrDeclarationStorage(
             isLateInit -> setter?.visibility ?: status.visibility
             isConst -> status.visibility
             hasJvmFieldAnnotation(session) -> status.visibility
+            origin == FirDeclarationOrigin.ScriptCustomization.ResultProperty -> status.visibility
             else -> Visibilities.Private
         }
 

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrVisitor.kt

@@ -208,6 +208,7 @@ class Fir2IrVisitor(
         return irClass
     }
 
+    @OptIn(UnexpandedTypeCheck::class)
     override fun visitScript(script: FirScript, data: Any?): IrElement {
         return declarationStorage.getCachedIrScript(script)!!.also { irScript ->
             irScript.parent = conversionScope.parentFromStack()
@@ -246,6 +247,17 @@ class Fir2IrVisitor(
                 for (statement in script.statements) {
                     val irStatement = if (statement is FirDeclaration) {
                         when {
+                            statement is FirProperty && statement.origin == FirDeclarationOrigin.ScriptCustomization.ResultProperty -> {
+                                // Generating the result property only for expressions with a meaningful result type
+                                // otherwise skip the property and convert the expression into the statement
+                                if (statement.returnTypeRef.let { (it.isUnit || it.isNothing || it.isNullableNothing) } == true) {
+                                    statement.initializer!!.toIrStatement()
+                                } else {
+                                    (statement.accept(this@Fir2IrVisitor, null) as? IrDeclaration)?.also {
+                                        irScript.resultProperty = (it as? IrProperty)?.symbol
+                                    }
+                                }
+                            }
                             statement is FirVariable && statement.isDestructuringDeclarationContainerVariable == true -> {
                                 statement.convertWithOffsets { startOffset, endOffset ->
                                     IrCompositeImpl(

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/FirScript.kt

@@ -31,6 +31,7 @@ abstract class FirScript : FirDeclaration() {
     abstract override val symbol: FirScriptSymbol
     abstract val parameters: List<FirVariable>
     abstract val contextReceivers: List<FirContextReceiver>
+    abstract val resultPropertyName: Name?
 
     override fun <R, D> accept(visitor: FirVisitor<R, D>, data: D): R = visitor.visitScript(this, data)
 

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/builder/FirScriptBuilder.kt

@@ -48,6 +48,7 @@ class FirScriptBuilder : FirAnnotationContainerBuilder {
     lateinit var symbol: FirScriptSymbol
     val parameters: MutableList<FirVariable> = mutableListOf()
     val contextReceivers: MutableList<FirContextReceiver> = mutableListOf()
+    var resultPropertyName: Name? = null
 
     override fun build(): FirScript {
         return FirScriptImpl(
@@ -62,6 +63,7 @@ class FirScriptBuilder : FirAnnotationContainerBuilder {
             symbol,
             parameters,
             contextReceivers.toMutableOrEmpty(),
+            resultPropertyName,
         )
     }
 
@@ -92,5 +94,6 @@ inline fun buildScriptCopy(original: FirScript, init: FirScriptBuilder.() -> Uni
     copyBuilder.symbol = original.symbol
     copyBuilder.parameters.addAll(original.parameters)
     copyBuilder.contextReceivers.addAll(original.contextReceivers)
+    copyBuilder.resultPropertyName = original.resultPropertyName
     return copyBuilder.apply(init).build()
 }

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/impl/FirScriptImpl.kt

@@ -43,6 +43,7 @@ internal class FirScriptImpl(
     override val symbol: FirScriptSymbol,
     override val parameters: MutableList<FirVariable>,
     override var contextReceivers: MutableOrEmptyList<FirContextReceiver>,
+    override val resultPropertyName: Name?,
 ) : FirScript() {
     init {
         symbol.bind(this)

compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/FirDeclarationOrigin.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.declarations
 
 import org.jetbrains.kotlin.GeneratedDeclarationKey
+import org.jetbrains.kotlin.fir.declarations.utils.FirScriptCustomizationKind
 
 sealed class FirDeclarationOrigin(
     private val displayName: String? = null,
@@ -36,7 +37,10 @@ sealed class FirDeclarationOrigin(
     object Delegated : FirDeclarationOrigin()
     object RenamedForOverride : FirDeclarationOrigin()
     object WrappedIntegerOperator : FirDeclarationOrigin()
-    object ScriptCustomization : FirDeclarationOrigin()
+    sealed class ScriptCustomization(val kind: FirScriptCustomizationKind) : FirDeclarationOrigin() {
+        object Default : ScriptCustomization(FirScriptCustomizationKind.DEFAULT)
+        object ResultProperty : ScriptCustomization(FirScriptCustomizationKind.RESULT_PROPERTY)
+    }
     class Plugin(val key: GeneratedDeclarationKey) : FirDeclarationOrigin(displayName = "Plugin[$key]", generated = true)
 
     override fun toString(): String {

compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/utils/FirScriptUtils.kt

@@ -7,3 +7,8 @@ package org.jetbrains.kotlin.fir.declarations.utils
 
 const val SCRIPT_SPECIAL_NAME_STRING = "<script>"
 const val SCRIPT_RECEIVER_NAME_PREFIX = "\$script_receiver"
+
+enum class FirScriptCustomizationKind {
+    DEFAULT,
+    RESULT_PROPERTY
+}

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/NodeConfigurator.kt

@@ -493,6 +493,7 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
             +symbol("FirScriptSymbol")
             +FieldList("parameters", variable, withReplace = false)
             +fieldList(contextReceiver, useMutableOrEmpty = true)
+            +field("resultPropertyName", nameType, nullable = true)
         }
 
         codeFragment.configure {