K2 Scripting: add proper CFG processing for scripts

#KT-59683 fixed
#KT-63328 fixed

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

@@ -251,14 +251,14 @@ abstract class FirDataFlowAnalyzer(
 
     // ----------------------------------- Scripts ------------------------------------------
 
-    fun enterScript(script: FirScript) {
-        graphBuilder.enterScript(script).mergeIncomingFlow()
+    fun enterScript(script: FirScript, buildGraph: Boolean) {
+        graphBuilder.enterScript(script, buildGraph)?.mergeIncomingFlow()
     }
 
-    fun exitScript(): ControlFlowGraph {
+    fun exitScript(): ControlFlowGraph? {
         val (node, graph) = graphBuilder.exitScript()
-        node.mergeIncomingFlow()
-        graph.completePostponedNodes()
+        node?.mergeIncomingFlow()
+        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 {
-        return enterGraph(script, "SCRIPT_GRAPH", ControlFlowGraph.Kind.Function) {
+    // ----------------------------------- Scripts -----------------------------------
+
+    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> {
-        return exitGraph()
+    fun exitScript(): Pair<ScriptExitNode?, ControlFlowGraph?> {
+        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 {
-        dataFlowAnalyzer.enterScript(script)
+    open fun withScript(script: FirScript, action: () -> FirScript): FirScript {
         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) {
-        transformDeclarationContent(script, data) as FirScript
+    override fun transformScript(script: FirScript, data: ResolutionMode): FirScript = whileAnalysing(session, script) {
+        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,