K2 Scripting: make FirScript a CFG owner

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

@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirModuleData
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.fir.expressions.FirStatement
+import org.jetbrains.kotlin.fir.references.FirControlFlowGraphReference
 import org.jetbrains.kotlin.fir.symbols.impl.FirScriptSymbol
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
@@ -21,12 +22,13 @@ import org.jetbrains.kotlin.name.Name
 /**
  * Generated from: [org.jetbrains.kotlin.fir.tree.generator.FirTreeBuilder.script]
  */
-abstract class FirScript : FirDeclaration() {
+abstract class FirScript : FirDeclaration(), FirControlFlowGraphOwner {
     abstract override val source: KtSourceElement?
     abstract override val annotations: List<FirAnnotation>
     abstract override val moduleData: FirModuleData
     abstract override val origin: FirDeclarationOrigin
     abstract override val attributes: FirDeclarationAttributes
+    abstract override val controlFlowGraphReference: FirControlFlowGraphReference?
     abstract val name: Name
     abstract val statements: List<FirStatement>
     abstract override val symbol: FirScriptSymbol
@@ -43,6 +45,8 @@ abstract class FirScript : FirDeclaration() {
 
     abstract override fun replaceAnnotations(newAnnotations: List<FirAnnotation>)
 
+    abstract override fun replaceControlFlowGraphReference(newControlFlowGraphReference: FirControlFlowGraphReference?)
+
     abstract fun replaceStatements(newStatements: List<FirStatement>)
 
     abstract override fun <D> transformAnnotations(transformer: FirTransformer<D>, data: D): FirScript

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

@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.fir.builder.toMutableOrEmpty
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.fir.expressions.FirStatement
+import org.jetbrains.kotlin.fir.references.FirControlFlowGraphReference
 import org.jetbrains.kotlin.fir.symbols.impl.FirScriptSymbol
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
@@ -38,6 +39,7 @@ internal class FirScriptImpl(
     override var contextReceivers: MutableOrEmptyList<FirContextReceiver>,
     override val resultPropertyName: Name?,
 ) : FirScript() {
+    override var controlFlowGraphReference: FirControlFlowGraphReference? = null
 
     init {
         symbol.bind(this)
@@ -46,6 +48,7 @@ internal class FirScriptImpl(
 
     override fun <R, D> acceptChildren(visitor: FirVisitor<R, D>, data: D) {
         annotations.forEach { it.accept(visitor, data) }
+        controlFlowGraphReference?.accept(visitor, data)
         statements.forEach { it.accept(visitor, data) }
         parameters.forEach { it.accept(visitor, data) }
         contextReceivers.forEach { it.accept(visitor, data) }
@@ -53,6 +56,7 @@ internal class FirScriptImpl(
 
     override fun <D> transformChildren(transformer: FirTransformer<D>, data: D): FirScriptImpl {
         transformAnnotations(transformer, data)
+        controlFlowGraphReference = controlFlowGraphReference?.transform(transformer, data)
         transformStatements(transformer, data)
         parameters.transformInplace(transformer, data)
         contextReceivers.transformInplace(transformer, data)
@@ -73,6 +77,10 @@ internal class FirScriptImpl(
         annotations = newAnnotations.toMutableOrEmpty()
     }
 
+    override fun replaceControlFlowGraphReference(newControlFlowGraphReference: FirControlFlowGraphReference?) {
+        controlFlowGraphReference = newControlFlowGraphReference
+    }
+
     override fun replaceStatements(newStatements: List<FirStatement>) {
         statements = newStatements.toMutableOrEmpty()
     }

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/visitors/FirDefaultVisitor.kt

@@ -59,9 +59,6 @@ abstract class FirDefaultVisitor<out R, in D> : FirVisitor<R, D>() {
     override fun visitRegularClass(regularClass: FirRegularClass, data: D): R =
         visitClass(regularClass, data)
 
-    override fun visitScript(script: FirScript, data: D): R =
-        visitDeclaration(script, data)
-
     override fun visitCodeFragment(codeFragment: FirCodeFragment, data: D): R =
         visitDeclaration(codeFragment, data)
 

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/visitors/FirDefaultVisitorVoid.kt

@@ -46,8 +46,6 @@ abstract class FirDefaultVisitorVoid : FirVisitorVoid() {
 
     override fun visitRegularClass(regularClass: FirRegularClass) = visitClass(regularClass)
 
-    override fun visitScript(script: FirScript) = visitDeclaration(script)
-
     override fun visitCodeFragment(codeFragment: FirCodeFragment) = visitDeclaration(codeFragment)
 
     override fun visitAnonymousFunctionExpression(anonymousFunctionExpression: FirAnonymousFunctionExpression) = visitExpression(anonymousFunctionExpression)

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

@@ -65,7 +65,7 @@ object FirTreeBuilder : AbstractFirTreeBuilder() {
     val backingField by element(Declaration, variable, typeParametersOwner, statement)
     val constructor by element(Declaration, function, typeParameterRefsOwner, contractDescriptionOwner)
     val file by element(Declaration, declaration, controlFlowGraphOwner)
-    val script by element(Declaration, declaration)
+    val script by element(Declaration, declaration, controlFlowGraphOwner)
     val codeFragment by element(Declaration, declaration)
     val packageDirective by element(Other)