K2 Scripting: add proper CFG processing for scripts

#KT-59683 fixed
#KT-63328 fixed

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/FirLazyBodiesCalculator.kt

@@ -858,4 +858,13 @@ private abstract class FirLazyBodiesCalculatorTransformer : FirTransformer<Persi
 
         return codeFragment
     }
+
+    override fun transformScript(script: FirScript, data: PersistentList<FirRegularClass>): FirScript {
+        for (declaration in script.declarations) {
+            if (declaration !is FirAnonymousInitializer) continue
+            declaration.accept(this, data)
+        }
+
+        return script
+    }
 }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirAbstractBodyTargetResolver.kt

@@ -10,15 +10,13 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirResolveT
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.LLFirReturnTypeCalculatorWithJump
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.LLFirLockProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.FirLazyBodiesCalculator
-import org.jetbrains.kotlin.analysis.low.level.api.fir.util.isElementWhichShouldBeResolvedAsPartOfScript
-import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirElementWithResolveState
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirAbstractBodyResolveTransformerDispatcher
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirResolveContextCollector
-import org.jetbrains.kotlin.fir.visitors.FirTransformer
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.ImplicitBodyResolveComputationSession
 import org.jetbrains.kotlin.fir.visitors.transformSingle
 
 internal abstract class LLFirAbstractBodyTargetResolver(
@@ -81,22 +79,4 @@ internal abstract class LLFirAbstractBodyTargetResolver(
         target.transformSingle(transformer, ResolutionMode.ContextIndependent)
     }
 
-    protected fun resolveScript(script: FirScript) {
-        transformer.declarationsTransformer?.withScript(script) {
-            script.parameters.forEach { it.transformSingle(transformer, ResolutionMode.ContextIndependent) }
-            script.transformDeclarations(
-                transformer = object : FirTransformer<Any?>() {
-                    override fun <E : FirElement> transformElement(element: E, data: Any?): E {
-                        if (element !is FirDeclaration || !element.isElementWhichShouldBeResolvedAsPartOfScript) return element
-
-                        transformer.firResolveContextCollector?.addDeclarationContext(element, transformer.context)
-                        return element.transformSingle(transformer, ResolutionMode.ContextIndependent)
-                    }
-                },
-                data = null,
-            )
-
-            script
-        }
-    }
 }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirBodyLazyResolver.kt

@@ -18,9 +18,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.llFirMo
 import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLFirResolvableResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.*
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.checkDelegatedConstructorIsResolved
-import org.jetbrains.kotlin.fir.FirElementWithResolveState
+import org.jetbrains.kotlin.fir.*
-import org.jetbrains.kotlin.fir.FirFileAnnotationsContainer
-import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.utils.getExplicitBackingField
 import org.jetbrains.kotlin.fir.expressions.*
@@ -28,20 +26,17 @@ import org.jetbrains.kotlin.fir.expressions.builder.buildLazyDelegatedConstructo
 import org.jetbrains.kotlin.fir.expressions.builder.buildMultiDelegatedConstructorCall
 import org.jetbrains.kotlin.fir.expressions.impl.FirContractCallBlock
 import org.jetbrains.kotlin.fir.expressions.impl.FirLazyDelegatedConstructorCall
-import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.references.FirSuperReference
 import org.jetbrains.kotlin.fir.references.FirThisReference
 import org.jetbrains.kotlin.fir.references.builder.buildExplicitSuperReference
 import org.jetbrains.kotlin.fir.references.builder.buildExplicitThisReference
-import org.jetbrains.kotlin.fir.resolve.FirCodeFragmentContext
+import org.jetbrains.kotlin.fir.resolve.*
-import org.jetbrains.kotlin.fir.resolve.ScopeSession
-import org.jetbrains.kotlin.fir.resolve.SessionHolderImpl
-import org.jetbrains.kotlin.fir.resolve.codeFragmentContext
 import org.jetbrains.kotlin.fir.resolve.dfa.FirControlFlowGraphReferenceImpl
 import org.jetbrains.kotlin.fir.resolve.dfa.RealVariable
 import org.jetbrains.kotlin.fir.resolve.dfa.cfg.isUsedInControlFlowGraphBuilderForClass
 import org.jetbrains.kotlin.fir.resolve.dfa.cfg.isUsedInControlFlowGraphBuilderForFile
+import org.jetbrains.kotlin.fir.resolve.dfa.cfg.isUsedInControlFlowGraphBuilderForScript
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirBodyResolveTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirResolveContextCollector
 import org.jetbrains.kotlin.fir.resolve.transformers.contracts.FirContractsDslNames
@@ -50,6 +45,8 @@ import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhase
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.isResolved
 import org.jetbrains.kotlin.fir.utils.exceptions.withFirEntry
+import org.jetbrains.kotlin.fir.visitors.FirTransformer
+import org.jetbrains.kotlin.fir.visitors.transformSingle
 import org.jetbrains.kotlin.psi.KtCodeFragment
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.utils.exceptions.checkWithAttachment
@@ -105,6 +102,7 @@ private class LLFirBodyTargetResolver(
         firResolveContextCollector = firResolveContextCollector,
     ) {
         override val preserveCFGForClasses: Boolean get() = false
+        override val buildCfgForScripts: Boolean get() = false
         override val buildCfgForFiles: Boolean get() = false
 
         /**
@@ -153,6 +151,12 @@ private class LLFirBodyTargetResolver(
 
                 return true
             }
+            is FirScript -> {
+                if (target.resolvePhase >= resolverPhase) return true
+                // resolve properties so they are available for CFG building in resolveScript
+                resolveMembersForControlFlowGraph(target)
+                return false
+            }
             is FirCodeFragment -> {
                 resolveCodeFragmentContext(target)
                 performCustomResolveUnderLock(target) {
@@ -224,6 +228,35 @@ private class LLFirBodyTargetResolver(
         }
     }
 
+    private fun calculateControlFlowGraph(target: FirScript) {
+        checkWithAttachment(
+            target.controlFlowGraphReference == null,
+            { "'controlFlowGraphReference' should be 'null' if the script phase < $resolverPhase)" },
+        ) {
+            withFirEntry("firScript", target)
+        }
+
+        val dataFlowAnalyzer = transformer.declarationsTransformer.dataFlowAnalyzer
+        dataFlowAnalyzer.enterScript(target, buildGraph = true)
+        val controlFlowGraph = dataFlowAnalyzer.exitScript()
+            ?: errorWithAttachment("CFG should not be 'null' as 'buildGraph' is specified") {
+                withFirEntry("firScript", target)
+            }
+
+        target.replaceControlFlowGraphReference(FirControlFlowGraphReferenceImpl(controlFlowGraph))
+    }
+
+    private fun resolveMembersForControlFlowGraph(target: FirScript) {
+        withScript(target) {
+            for (member in target.declarations) {
+                if (member is FirControlFlowGraphOwner && member.isUsedInControlFlowGraphBuilderForScript && member !is FirAnonymousInitializer) {
+                    member.lazyResolveToPhase(resolverPhase.previous)
+                    performResolve(member)
+                }
+            }
+        }
+    }
+
     private fun resolveCodeFragmentContext(firCodeFragment: FirCodeFragment) {
         val ktCodeFragment = firCodeFragment.psi as? KtCodeFragment
             ?: errorWithAttachment("Code fragment source not found") {
@@ -281,12 +314,35 @@ private class LLFirBodyTargetResolver(
 
     override fun rawResolve(target: FirElementWithResolveState) {
         when (target) {
-            is FirScript -> target.let(::resolveScript)
+            is FirScript -> {
+                resolveScript(target)
+                calculateControlFlowGraph(target)
+            }
+
             else -> super.rawResolve(target)
         }
 
         LLFirDeclarationModificationService.bodyResolved(target, resolverPhase)
     }
+
+    protected fun resolveScript(script: FirScript) {
+        transformer.declarationsTransformer.withScript(script) {
+            script.parameters.forEach { it.transformSingle(transformer, ResolutionMode.ContextIndependent) }
+            script.transformDeclarations(
+                transformer = object : FirTransformer<Any?>() {
+                    override fun <E : FirElement> transformElement(element: E, data: Any?): E {
+                        if (element !is FirDeclaration || !element.isElementWhichShouldBeResolvedAsPartOfScript) return element
+
+                        transformer.firResolveContextCollector?.addDeclarationContext(element, transformer.context)
+                        return element.transformSingle(transformer, ResolutionMode.ContextIndependent)
+                    }
+                },
+                data = null,
+            )
+
+            script
+        }
+    }
 }
 
 internal object BodyStateKeepers {
@@ -294,27 +350,19 @@ internal object BodyStateKeepers {
         val oldDeclarations = script.declarations
         if (oldDeclarations.none { it.isElementWhichShouldBeResolvedAsPartOfScript }) return@stateKeeper
 
-        add(RESULT_PROPERTY, designation)
+        val lastProperty = oldDeclarations.lastOrNull()
-        add(FirScript::declarations, FirScript::replaceDeclarations) {
+        if (lastProperty is FirProperty &&
+            lastProperty.origin == FirDeclarationOrigin.ScriptCustomization.ResultProperty &&
+            lastProperty.bodyResolveState != FirPropertyBodyResolveState.ALL_BODIES_RESOLVED
+        ) {
+            add(RESULT_PROPERTY, designation)
             val recreatedDeclarations = FirLazyBodiesCalculator.createDeclarationsForScript(script)
             requireSameSize(oldDeclarations, recreatedDeclarations)
-
+            lastProperty.replaceInitializer((recreatedDeclarations.last() as FirProperty).initializer)
-            ArrayList<FirDeclaration>(oldDeclarations.size).apply {
-                oldDeclarations.zip(recreatedDeclarations).mapTo(this) { (old, new) ->
-                    when {
-                        !old.isElementWhichShouldBeResolvedAsPartOfScript -> old
-                        old is FirProperty && old.origin == FirDeclarationOrigin.ScriptCustomization.ResultProperty -> {
-                            if (old.bodyResolveState != FirPropertyBodyResolveState.ALL_BODIES_RESOLVED) {
-                                old.replaceInitializer((new as FirProperty).initializer)
-                            }
-                            old
-                        }
-
-                        else -> new
-                    }
-                }
-            }
         }
+
+        entityList(oldDeclarations.mapNotNull { it as? FirAnonymousInitializer }, ANONYMOUS_INITIALIZER, designation)
+        add(FirScript::controlFlowGraphReference, FirScript::replaceControlFlowGraphReference)
     }
 
     private val RESULT_PROPERTY: StateKeeper<FirScript, FirDesignationWithFile> = stateKeeper { script, _ ->

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirImplicitTypesLazyResolver.kt

@@ -10,12 +10,14 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.throwUnexpectedFirEle
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.LLFirLockProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.LLFirDeclarationModificationService
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.checkReturnTypeRefIsResolved
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.forEachDependentDeclaration
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.isScriptDependentDeclaration
 import org.jetbrains.kotlin.fir.FirElementWithResolveState
 import org.jetbrains.kotlin.fir.FirFileAnnotationsContainer
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
+import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirImplicitAwareBodyResolveTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirResolveContextCollector
@@ -27,6 +29,7 @@ import org.jetbrains.kotlin.fir.symbols.impl.FirValueParameterSymbol
 import org.jetbrains.kotlin.fir.types.FirImplicitTypeRef
 import org.jetbrains.kotlin.fir.util.setMultimapOf
 import org.jetbrains.kotlin.fir.utils.exceptions.withFirSymbolEntry
+import org.jetbrains.kotlin.fir.visitors.transformSingle
 import org.jetbrains.kotlin.utils.exceptions.errorWithAttachment
 
 internal object LLFirImplicitTypesLazyResolver : LLFirLazyResolver(FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE) {
@@ -140,6 +143,7 @@ internal class LLFirImplicitBodyTargetResolver(
         returnTypeCalculator = createReturnTypeCalculator(firResolveContextCollector = firResolveContextCollector),
     ) {
         override val preserveCFGForClasses: Boolean get() = false
+        override val buildCfgForScripts: Boolean get() = false
         override val buildCfgForFiles: Boolean get() = false
         override fun transformForeignAnnotationCall(symbol: FirBasedSymbol<*>, annotationCall: FirAnnotationCall): FirAnnotationCall {
             llImplicitBodyResolveComputationSession.postponeForeignAnnotationResolution(symbol)
@@ -202,7 +206,16 @@ internal class LLFirImplicitBodyTargetResolver(
 
     override fun rawResolve(target: FirElementWithResolveState) {
         when {
-            target is FirScript -> resolveScript(target)
+            target is FirScript -> {
+                transformer.declarationsTransformer.withScript(target) {
+                    target.forEachDependentDeclaration {
+                        it.transformSingle(transformer, ResolutionMode.ContextIndependent)
+                    }
+
+                    target
+                }
+            }
+
             target is FirCallableDeclaration && target.attributes.callableCopySubstitutionForTypeUpdater != null -> {
                 transformer.returnTypeCalculator.callableCopyTypeCalculator.computeReturnType(target)
                 Unit

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt

@@ -251,14 +251,14 @@ abstract class FirDataFlowAnalyzer(
 
     // ----------------------------------- Scripts ------------------------------------------
 
-    fun enterScript(script: FirScript) {
+    fun enterScript(script: FirScript, buildGraph: Boolean) {
-        graphBuilder.enterScript(script).mergeIncomingFlow()
+        graphBuilder.enterScript(script, buildGraph)?.mergeIncomingFlow()
     }
 
-    fun exitScript(): ControlFlowGraph {
+    fun exitScript(): ControlFlowGraph? {
         val (node, graph) = graphBuilder.exitScript()
-        node.mergeIncomingFlow()
+        node?.mergeIncomingFlow()
-        graph.completePostponedNodes()
+        graph?.completePostponedNodes()
         return graph
     }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphBuilder.kt

@@ -22,7 +22,9 @@ import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.util.ListMultimap
 import org.jetbrains.kotlin.fir.util.listMultimapOf
+import org.jetbrains.kotlin.fir.utils.exceptions.withFirEntry
 import org.jetbrains.kotlin.utils.addToStdlib.runIf
+import org.jetbrains.kotlin.utils.exceptions.requireWithAttachment
 import org.jetbrains.kotlin.utils.getOrPutNullable
 
 data class FirAnonymousFunctionReturnExpressionInfo(val expression: FirExpression, val isExplicit: Boolean)
@@ -625,14 +627,60 @@ class ControlFlowGraphBuilder {
         }
     }
 
-    fun enterScript(script: FirScript): ScriptEnterNode {
+    // ----------------------------------- Scripts -----------------------------------
-        return enterGraph(script, "SCRIPT_GRAPH", ControlFlowGraph.Kind.Function) {
+
+    fun enterScript(script: FirScript, buildGraph: Boolean): ScriptEnterNode? {
+        if (!buildGraph) {
+            graphs.push(ControlFlowGraph(declaration = null, "<discarded script graph>", ControlFlowGraph.Kind.Script))
+            return null
+        }
+        return enterGraph(script, script.name.asString(), ControlFlowGraph.Kind.Script) {
             createScriptEnterNode(it) to createScriptExitNode(it)
         }
     }
 
-    fun exitScript(): Pair<ScriptExitNode, ControlFlowGraph> {
+    fun exitScript(): Pair<ScriptExitNode?, ControlFlowGraph?> {
-        return exitGraph()
+        require(currentGraph.kind == ControlFlowGraph.Kind.Script)
+        if (currentGraph.declaration == null) {
+            graphs.pop() // Discard empty script graph.
+            return null to null
+        }
+
+        val enterNode = lastNodes.pop() as ScriptEnterNode
+        val exitNode = currentGraph.exitNode as ScriptExitNode
+
+        val script = enterNode.fir
+        requireWithAttachment(
+            (script as FirControlFlowGraphOwner).controlFlowGraphReference == null,
+            { "Unexpected state: script already has a CFG attached" }
+        ) {
+            withFirEntry("script", script)
+        }
+
+        val calledInPlace = mutableListOf<ControlFlowGraph>()
+
+        script.declarations.forEachGraphOwner {
+            val graph = it.controlFlowGraphReference?.controlFlowGraph ?: return@forEachGraphOwner
+            if (it.isUsedInControlFlowGraphBuilderForScript) {
+                calledInPlace.add(graph)
+            }
+        }
+
+        val lastNode = calledInPlace.fold<_, CFGNode<*>>(enterNode) { lastNode, graph ->
+            if (lastNode !== enterNode || lastNode.previousNodes.isEmpty()) {
+                addEdgeToSubGraph(lastNode, graph.enterNode)
+            }
+            graph.exitNode
+        }
+
+        addEdge(lastNode, exitNode, preferredKind = EdgeKind.CfgForward, propagateDeadness = false)
+        if (calledInPlace.isNotEmpty()) {
+            // Fake edge to enforce ordering.
+            addEdge(enterNode, exitNode, preferredKind = EdgeKind.DeadForward, propagateDeadness = false)
+        }
+
+        enterNode.subGraphs = calledInPlace
+        return exitNode to popGraph()
     }
 
     fun enterCodeFragment(codeFragment: FirCodeFragment): CodeFragmentEnterNode {
@@ -1495,6 +1543,15 @@ val FirControlFlowGraphOwner.isUsedInControlFlowGraphBuilderForFile: Boolean
         else -> false
     }
 
+/**
+ * @return true for [FirControlFlowGraphOwner] which, as a script statement, should be part of the script
+ */
+val FirControlFlowGraphOwner.isUsedInControlFlowGraphBuilderForScript: Boolean
+    get() = when (this) {
+        is FirProperty, is FirField, is FirAnonymousInitializer -> memberShouldHaveGraph
+        else -> false
+    }
+
 // TODO Refactor annotation arguments phase to not build CFG so that we can use resolvedType instead, see KT-61834
 @OptIn(UnresolvedExpressionTypeAccess::class)
 private val FirExpression.hasNothingType: Boolean

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

@@ -32,6 +32,7 @@ abstract class FirAbstractBodyResolveTransformerDispatcher(
 ) : FirAbstractBodyResolveTransformer(phase) {
 
     open val preserveCFGForClasses: Boolean get() = !implicitTypeOnly
+    open val buildCfgForScripts: Boolean get() = !implicitTypeOnly
     open val buildCfgForFiles: Boolean get() = !implicitTypeOnly
 
     final override val context: BodyResolveContext =

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

@@ -139,7 +139,9 @@ open class FirDeclarationsResolveTransformer(
     override fun transformProperty(property: FirProperty, data: ResolutionMode): FirProperty = whileAnalysing(session, property) {
         require(property !is FirSyntheticProperty) { "Synthetic properties should not be processed by body transformers" }
 
-        if (property.isLocal) {
+        // script top level destructuring declaration container variables should be treated as properties here
+        // to avoid CFG/DFA complications
+        if (property.isLocal && property.origin != FirDeclarationOrigin.Synthetic.ScriptTopLevelDestructuringDeclarationContainer) {
             prepareSignatureForBodyResolve(property)
             property.transformStatus(this, property.resolveStatus().mode())
             property.getter?.let { it.transformStatus(this, it.resolveStatus(containingProperty = property).mode()) }
@@ -551,7 +553,10 @@ open class FirDeclarationsResolveTransformer(
             .transformOtherChildren(transformer, ResolutionMode.ContextIndependent)
 
         context.storeVariable(variable, session)
-        dataFlowAnalyzer.exitLocalVariableDeclaration(variable, hadExplicitType)
+        if (variable.origin != FirDeclarationOrigin.ScriptCustomization.Parameter) {
+            // script parameters should not be added to CFG to avoid graph building compilations
+            dataFlowAnalyzer.exitLocalVariableDeclaration(variable, hadExplicitType)
+        }
         return variable
     }
 
@@ -685,17 +690,23 @@ open class FirDeclarationsResolveTransformer(
         }
     }
 
-    fun withScript(script: FirScript, action: () -> FirScript): FirScript {
+    open fun withScript(script: FirScript, action: () -> FirScript): FirScript {
-        dataFlowAnalyzer.enterScript(script)
         val result = context.withScript(script, components) {
+            // see todo in withFile
+            dataFlowAnalyzer.enterScript(script, buildGraph = transformer.buildCfgForScripts)
             action()
         }
-        dataFlowAnalyzer.exitScript() // TODO: FirScript should be a FirControlFlowGraphOwner, KT-59683
+        val controlFlowGraph = dataFlowAnalyzer.exitScript()
+        if (controlFlowGraph != null) {
+            result.replaceControlFlowGraphReference(FirControlFlowGraphReferenceImpl(controlFlowGraph))
+        }
         return result
     }
 
-    override fun transformScript(script: FirScript, data: ResolutionMode): FirScript = withScript(script) {
+    override fun transformScript(script: FirScript, data: ResolutionMode): FirScript = whileAnalysing(session, script) {
-        transformDeclarationContent(script, data) as FirScript
+        withScript(script) {
+            transformDeclarationContent(script, data) as FirScript
+        }
     }
 
     override fun transformCodeFragment(codeFragment: FirCodeFragment, data: ResolutionMode): FirCodeFragment {
@@ -740,6 +751,7 @@ open class FirDeclarationsResolveTransformer(
         val result = context.withFile(file, components) {
             // TODO Must be done within 'withFile' as the context - any the analyzer - is cleared as the first step.
             //  yuk. maybe the clear shouldn't happen for `enterFile`? or at maybe separately?
+            // also check whether it is applicable in withScript`
             dataFlowAnalyzer.enterFile(file, buildGraph = transformer.buildCfgForFiles)
 
             action()

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/contracts/FirAbstractContractResolveTransformerDispatcher.kt

@@ -96,10 +96,6 @@ abstract class FirAbstractContractResolveTransformerDispatcher(
         }
 
 
-        override fun transformScript(script: FirScript, data: ResolutionMode): FirScript {
-            return script
-        }
-
         override fun transformProperty(property: FirProperty, data: ResolutionMode): FirProperty {
             if (
                 property.getter?.hasContractToResolve != true && property.setter?.hasContractToResolve != true ||
@@ -283,6 +279,16 @@ abstract class FirAbstractContractResolveTransformerDispatcher(
             }
         }
 
+        override fun withScript(script: FirScript, action: () -> FirScript): FirScript {
+            return context.withScript(script, components) {
+                action()
+            }
+        }
+
+        override fun transformScript(script: FirScript, data: ResolutionMode): FirScript {
+            return script
+        }
+
         override fun transformAnonymousObject(
             anonymousObject: FirAnonymousObject,
             data: ResolutionMode

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirAnnotationArgumentsTransformer.kt

@@ -202,6 +202,12 @@ private class FirDeclarationsResolveTransformerForAnnotationArguments(
         }
     }
 
+    override fun withScript(script: FirScript, action: () -> FirScript): FirScript {
+        return context.withScript(script, components) {
+            action()
+        }
+    }
+
     override fun transformAnonymousInitializer(
         anonymousInitializer: FirAnonymousInitializer,
         data: ResolutionMode

compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/CFGNode.kt

@@ -341,7 +341,12 @@ class AnonymousObjectExpressionExitNode(owner: ControlFlowGraph, override val fi
 
 // ----------------------------------- Scripts ------------------------------------------
 
-class ScriptEnterNode(owner: ControlFlowGraph, override val fir: FirScript, level: Int) : CFGNode<FirScript>(owner, level), GraphEnterNodeMarker {
+class ScriptEnterNode(owner: ControlFlowGraph, override val fir: FirScript, level: Int)
+    : CFGNodeWithSubgraphs<FirScript>(owner, level), GraphEnterNodeMarker
+{
+    @set:CfgInternals
+    override lateinit var subGraphs: List<ControlFlowGraph>
+
     override fun <R, D> accept(visitor: ControlFlowGraphVisitor<R, D>, data: D): R {
         return visitor.visitScriptEnterNode(this, data)
     }

compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraph.kt

@@ -36,6 +36,7 @@ class ControlFlowGraph(val declaration: FirDeclaration?, val name: String, val k
         Class,
         Constructor,
         Function,
+        Script,
         LocalFunction,
         AnonymousFunction,
         AnonymousFunctionCalledInPlace,