K2 Scripting: add support for script args to FIR part

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/KtFirSymbol.kt

@@ -97,6 +97,7 @@ internal tailrec fun FirDeclaration.ktSymbolOrigin(): KtSymbolOrigin = when (ori
     FirDeclarationOrigin.DynamicScope -> buildErrorWithAttachment("Invalid FirDeclarationOrigin ${origin::class.simpleName}") {
         withFirEntry("firToGetOrigin", this@ktSymbolOrigin)
     }
+    FirDeclarationOrigin.ScriptCustomization -> KtSymbolOrigin.PLUGIN
 }
 
 internal fun KtClassLikeSymbol.getSymbolKind(firResolveSession: LLFirResolveSession): KtSymbolKind {

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

@@ -357,7 +357,15 @@ class Fir2IrConverter(
             }
             is FirScript -> {
                 assert(parent is IrFile)
-                declarationStorage.getOrCreateIrScript(declaration)
+                declarationStorage.getOrCreateIrScript(declaration).also { irScript ->
+                    declarationStorage.enterScope(irScript)
+                    for (scriptStatement in declaration.statements) {
+                        if (scriptStatement is FirDeclaration) {
+                            processMemberDeclaration(scriptStatement, null, irScript)
+                        }
+                    }
+                    declarationStorage.leaveScope(irScript)
+                }
             }
             is FirSimpleFunction -> {
                 declarationStorage.getOrCreateIrFunction(

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

@@ -194,7 +194,8 @@ class Fir2IrDeclarationStorage(
             declaration is IrConstructor ||
             declaration is IrAnonymousInitializer ||
             declaration is IrProperty ||
-            declaration is IrEnumEntry
+            declaration is IrEnumEntry ||
+            declaration is IrScript
         ) {
             localStorage.enterCallable()
         }
@@ -205,7 +206,8 @@ class Fir2IrDeclarationStorage(
             declaration is IrConstructor ||
             declaration is IrAnonymousInitializer ||
             declaration is IrProperty ||
-            declaration is IrEnumEntry
+            declaration is IrEnumEntry ||
+            declaration is IrScript
         ) {
             localStorage.leaveCallable()
         }
@@ -768,7 +770,7 @@ class Fir2IrDeclarationStorage(
                 name, inferredType,
                 visibility, isFinal = isFinal,
                 isExternal = property.isExternal,
-                isStatic = property.isStatic || parent !is IrClass,
+                isStatic = property.isStatic || !(parent is IrClass || parent is IrScript),
             ).also {
                 it.correspondingPropertySymbol = this@createBackingField.symbol
             }.apply {
@@ -1579,7 +1581,6 @@ class Fir2IrDeclarationStorage(
                 IrScriptImpl(symbol, script.name, irFactory, startOffset, endOffset).also { irScript ->
                     irScript.origin = SCRIPT_K2_ORIGIN
                     irScript.metadata = FirMetadataSource.Script(script)
-                    irScript.explicitCallParameters = emptyList()
                     irScript.implicitReceiversParameters = emptyList()
                     irScript.providedProperties = emptyList()
                     irScript.providedPropertiesParameters = emptyList()

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

@@ -193,7 +193,10 @@ class Fir2IrVisitor(
     override fun visitScript(script: FirScript, data: Any?): IrElement {
         return declarationStorage.getCachedIrScript(script)!!.also { irScript ->
             irScript.parent = conversionScope.parentFromStack()
-            symbolTable.enterScope(irScript)
+            declarationStorage.enterScope(irScript)
+            irScript.explicitCallParameters = script.parameters.map { parameter ->
+                declarationStorage.createIrVariable(parameter, irScript)
+            }
             conversionScope.withParent(irScript) {
                 for (statement in script.statements) {
                     if (statement is FirDeclaration) {
@@ -205,7 +208,7 @@ class Fir2IrVisitor(
                     }
                 }
             }
-            symbolTable.leaveScope(irScript)
+            declarationStorage.leaveScope(irScript)
         }
     }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/ResolveUtils.kt

@@ -561,6 +561,7 @@ fun FirFunction.getAsForbiddenNamedArgumentsTarget(
         FirDeclarationOrigin.DynamicScope -> null
         FirDeclarationOrigin.RenamedForOverride -> null
         FirDeclarationOrigin.WrappedIntegerOperator -> null
+        FirDeclarationOrigin.ScriptCustomization -> null
         is FirDeclarationOrigin.Plugin -> null // TODO: figure out what to do with plugin generated functions
     }
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt

@@ -502,8 +502,13 @@ class BodyResolveContext(
         val statics = base
             .addNonLocalScopeIfNotNull(towerElementsForScript.staticScope)
 
+        val parameterScope = owner.parameters.fold(FirLocalScope(holder.session)) { scope, parameter ->
+            scope.storeVariable(parameter, holder.session)
+        }
+
         val forMembersResolution =
             statics
+                .addLocalScope(parameterScope)
                 .addContextReceiverGroup(towerElementsForScript.implicitReceivers)
 
         val newContexts = FirRegularTowerDataContexts(

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

@@ -29,6 +29,7 @@ abstract class FirScript : FirDeclaration() {
     abstract val name: Name
     abstract val statements: List<FirStatement>
     abstract override val symbol: FirScriptSymbol
+    abstract val parameters: List<FirVariable>
     abstract val contextReceivers: List<FirContextReceiver>
 
     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

@@ -18,6 +18,7 @@ import org.jetbrains.kotlin.fir.declarations.FirDeclarationAttributes
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.declarations.FirScript
+import org.jetbrains.kotlin.fir.declarations.FirVariable
 import org.jetbrains.kotlin.fir.declarations.impl.FirScriptImpl
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.fir.expressions.FirStatement
@@ -41,6 +42,7 @@ class FirScriptBuilder : FirAnnotationContainerBuilder {
     lateinit var name: Name
     val statements: MutableList<FirStatement> = mutableListOf()
     lateinit var symbol: FirScriptSymbol
+    val parameters: MutableList<FirVariable> = mutableListOf()
     val contextReceivers: MutableList<FirContextReceiver> = mutableListOf()
 
     override fun build(): FirScript {
@@ -54,6 +56,7 @@ class FirScriptBuilder : FirAnnotationContainerBuilder {
             name,
             statements,
             symbol,
+            parameters,
             contextReceivers.toMutableOrEmpty(),
         )
     }

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

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.fir.declarations.FirDeclarationAttributes
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.declarations.FirScript
+import org.jetbrains.kotlin.fir.declarations.FirVariable
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.symbols.impl.FirScriptSymbol
@@ -38,6 +39,7 @@ internal class FirScriptImpl(
     override val name: Name,
     override val statements: MutableList<FirStatement>,
     override val symbol: FirScriptSymbol,
+    override val parameters: MutableList<FirVariable>,
     override var contextReceivers: MutableOrEmptyList<FirContextReceiver>,
 ) : FirScript() {
     init {
@@ -47,12 +49,14 @@ internal class FirScriptImpl(
     override fun <R, D> acceptChildren(visitor: FirVisitor<R, D>, data: D) {
         annotations.forEach { it.accept(visitor, data) }
         statements.forEach { it.accept(visitor, data) }
+        parameters.forEach { it.accept(visitor, data) }
         contextReceivers.forEach { it.accept(visitor, data) }
     }
 
     override fun <D> transformChildren(transformer: FirTransformer<D>, data: D): FirScriptImpl {
         transformAnnotations(transformer, data)
         transformStatements(transformer, data)
+        parameters.transformInplace(transformer, data)
         contextReceivers.transformInplace(transformer, data)
         return this
     }

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

@@ -32,7 +32,7 @@ sealed class FirDeclarationOrigin(
     object Delegated : FirDeclarationOrigin()
     object RenamedForOverride : FirDeclarationOrigin()
     object WrappedIntegerOperator : FirDeclarationOrigin()
-
+    object ScriptCustomization : FirDeclarationOrigin()
     class Plugin(val key: GeneratedDeclarationKey) : FirDeclarationOrigin(displayName = "Plugin[$key]", generated = true)
 
     override fun toString(): String {

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

@@ -471,6 +471,7 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
             +name
             +fieldList(statement).withTransform()
             +symbol("FirScriptSymbol")
+            +FieldList("parameters", variable, withReplace = false)
             +fieldList(contextReceiver, useMutableOrEmpty = true)
         }
 

compiler/ir/backend.jvm/lower/src/org/jetbrains/kotlin/backend/jvm/lower/ScriptLowering.kt

@@ -21,10 +21,7 @@ import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.builders.*
 import org.jetbrains.kotlin.ir.builders.declarations.*
 import org.jetbrains.kotlin.ir.declarations.*
-import org.jetbrains.kotlin.ir.declarations.impl.IrAnonymousInitializerImpl
-import org.jetbrains.kotlin.ir.declarations.impl.IrClassImpl
-import org.jetbrains.kotlin.ir.declarations.impl.IrExternalPackageFragmentImpl
-import org.jetbrains.kotlin.ir.declarations.impl.SCRIPT_K2_ORIGIN
+import org.jetbrains.kotlin.ir.declarations.impl.*
 import org.jetbrains.kotlin.ir.descriptors.toIrBasedKotlinType
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.IrClassReferenceImpl
@@ -334,7 +331,15 @@ private class ScriptsToClassesLowering(val context: JvmBackendContext, val inner
         }.also { irConstructor ->
             irConstructor.valueParameters = buildList {
                 addIfNotNull(irScript.earlierScriptsParameter)
-                addAll(irScript.explicitCallParameters)
+                addAll(irScript.explicitCallParameters.map {
+                    IrValueParameterImpl(
+                        it.startOffset, it.endOffset,
+                        IrDeclarationOrigin.SCRIPT_CALL_PARAMETER, IrValueParameterSymbolImpl(),
+                        it.name, index = 0,
+                        type = it.type, varargElementType = null,
+                        isCrossinline = false, isNoinline = false, isHidden = false, isAssignable = false
+                    ).also { it.parent = irScript }
+                })
                 addAll(irScript.implicitReceiversParameters)
                 addAll(irScript.providedPropertiesParameters)
             }
@@ -806,7 +811,21 @@ private class ScriptToClassTransformer(
     }
 
     override fun visitGetValue(expression: IrGetValue, data: ScriptToClassTransformerContext): IrExpression {
-        if (irScript.needsReceiverProcessing) {
+        val getVar = expression.symbol.owner as? IrVariable
+        if (getVar != null) {
+            if (irScript.explicitCallParameters.contains(getVar)) {
+                val correspondingParam = irScriptClass.constructors.single().valueParameters.find {
+                    it.origin == IrDeclarationOrigin.SCRIPT_CALL_PARAMETER && it.name == getVar.name
+                } ?: error("script explicit parameter ${getVar.name.asString()} not found")
+                val newExpression =
+                    IrGetValueImpl(
+                        expression.startOffset, expression.endOffset,
+                        correspondingParam.type, correspondingParam.symbol,
+                        expression.origin
+                    )
+                return super.visitExpression(newExpression, data)
+            }
+        } else if (irScript.needsReceiverProcessing) {
             val getValueParameter = expression.symbol.owner as? IrValueParameter
             if (getValueParameter != null && getValueParameter.name == SpecialNames.THIS) {
                 val newExpression = getDispatchReceiverExpression(

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/ScriptGenerator.kt

@@ -15,12 +15,14 @@ import org.jetbrains.kotlin.ir.declarations.DescriptorMetadataSource
 import org.jetbrains.kotlin.ir.declarations.IrDeclaration
 import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
 import org.jetbrains.kotlin.ir.declarations.IrValueParameter
+import org.jetbrains.kotlin.ir.declarations.impl.IrVariableImpl
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
 import org.jetbrains.kotlin.ir.expressions.impl.IrCompositeImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrGetValueImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrSetFieldImpl
 import org.jetbrains.kotlin.ir.symbols.impl.IrValueParameterSymbolImpl
+import org.jetbrains.kotlin.ir.symbols.impl.IrVariableSymbolImpl
 import org.jetbrains.kotlin.ir.types.IrSimpleType
 import org.jetbrains.kotlin.ir.util.indexOrMinusOne
 import org.jetbrains.kotlin.ir.util.isCrossinline
@@ -37,7 +39,6 @@ import org.jetbrains.kotlin.psi2ir.intermediate.createTemporaryVariableInBlock
 import org.jetbrains.kotlin.psi2ir.intermediate.setExplicitReceiverValue
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.calls.util.isSingleUnderscore
-import org.jetbrains.kotlin.types.typeUtil.makeNullable
 import org.jetbrains.kotlin.utils.addIfNotNull
 
 class ScriptGenerator(declarationGenerator: DeclarationGenerator) : DeclarationGeneratorExtension(declarationGenerator) {
@@ -104,7 +105,17 @@ class ScriptGenerator(declarationGenerator: DeclarationGenerator) : DeclarationG
                 irScript.earlierScriptsParameter = descriptor.earlierScriptsConstructorParameter?.let(::createValueParameter)
             }
 
-            irScript.explicitCallParameters = descriptor.explicitConstructorParameters.map(::createValueParameter)
+            val explicitCallParams = descriptor.explicitConstructorParameters.map(::createValueParameter)
+
+            irScript.explicitCallParameters = descriptor.explicitConstructorParameters.map {
+                IrVariableImpl(
+                    UNDEFINED_OFFSET, UNDEFINED_OFFSET,
+                    IrDeclarationOrigin.SCRIPT_CALL_PARAMETER, IrVariableSymbolImpl(),
+                    it.name,
+                    it.type.toIrType(),
+                    isVar = false, isConst = false, isLateinit = false
+                ).also { it.parent = irScript }
+            }
 
             irScript.implicitReceiversParameters = descriptor.implicitReceivers.map {
                 makeParameter(it.thisAsReceiverParameter, IrDeclarationOrigin.SCRIPT_IMPLICIT_RECEIVER, parametersIndex++)
@@ -155,7 +166,7 @@ class ScriptGenerator(declarationGenerator: DeclarationGenerator) : DeclarationG
             }.also { irConstructor ->
                 irConstructor.valueParameters = buildList {
                     addIfNotNull(irScript.earlierScriptsParameter)
-                    addAll(irScript.explicitCallParameters)
+                    addAll(explicitCallParams)
                     addAll(irScript.implicitReceiversParameters)
                     addAll(irScript.providedPropertiesParameters)
                 }

compiler/ir/ir.tree/gen/org/jetbrains/kotlin/ir/declarations/IrScript.kt

@@ -30,7 +30,7 @@ abstract class IrScript : IrDeclarationBase(), IrDeclarationWithName, IrDeclarat
 
     abstract var baseClass: IrType?
 
-    abstract var explicitCallParameters: List<IrValueParameter>
+    abstract var explicitCallParameters: List<IrVariable>
 
     abstract var implicitReceiversParameters: List<IrValueParameter>
 

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/declarations/impl/IrScriptImpl.kt

@@ -47,7 +47,7 @@ class IrScriptImpl(
     override var thisReceiver: IrValueParameter? = null
     override var baseClass: IrType? = null
 
-    override lateinit var explicitCallParameters: List<IrValueParameter>
+    override lateinit var explicitCallParameters: List<IrVariable>
     override lateinit var implicitReceiversParameters: List<IrValueParameter>
     override lateinit var providedProperties: List<IrPropertySymbol>
     override lateinit var providedPropertiesParameters: List<IrValueParameter>

compiler/ir/ir.tree/tree-generator/src/org/jetbrains/kotlin/ir/generator/IrTree.kt

@@ -400,7 +400,7 @@ object IrTree : AbstractTreeBuilder() {
         // TODO: consider removing from here and handle appropriately in the lowering
         +field("thisReceiver", valueParameter, mutable = true, isChild = true, nullable = true) // K1
         +field("baseClass", irTypeType, mutable = true, nullable = true) // K1
-        +listField("explicitCallParameters", valueParameter, mutability = Var, isChild = true)
+        +listField("explicitCallParameters", variable, mutability = Var, isChild = true)
         +listField("implicitReceiversParameters", valueParameter, mutability = Var, isChild = true)
         +listField("providedProperties", propertySymbolType, mutability = Var)
         +listField("providedPropertiesParameters", valueParameter, mutability = Var, isChild = true)