[optmz] DataFlowIR: Added scopes for loops

2020-07-09 17:35:02 +05:00
parent 2f35bbec41
commit dc0b4ad238
6 changed files with 222 additions and 101 deletions
@@ -107,7 +107,7 @@ internal class CallGraphBuilder(val context: Context,
    private val arraySet = symbols.arraySet[symbols.array]!!.owner

    private inline fun DataFlowIR.FunctionBody.forEachCallSite(block: (DataFlowIR.Node.Call) -> Unit) =
-            nodes.forEach { node ->
+            forEachNonScopeNode { node ->
                when (node) {
                    is DataFlowIR.Node.Call -> block(node)

@@ -18,6 +18,7 @@ import org.jetbrains.kotlin.backend.konan.ir.*
 import org.jetbrains.kotlin.backend.konan.llvm.functionName
 import org.jetbrains.kotlin.backend.konan.llvm.localHash
 import org.jetbrains.kotlin.ir.IrElement
+import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.declarations.lazy.IrLazyClass
 import org.jetbrains.kotlin.ir.expressions.*
@@ -50,22 +51,22 @@ private fun IrTypeOperator.isCast() =


 private class VariableValues {
-    val elementData = HashMap<IrVariable, MutableSet<IrExpression>>()
+    data class Variable(val loop: IrLoop?, val values: MutableSet<IrExpression>)

-    fun addEmpty(variable: IrVariable) =
-            elementData.getOrPut(variable, { mutableSetOf() })
+    val elementData = HashMap<IrVariable, Variable>()
+
+    fun addEmpty(variable: IrVariable, loop: IrLoop?) {
+        elementData[variable] = Variable(loop, mutableSetOf())
+    }

    fun add(variable: IrVariable, element: IrExpression) =
-            elementData[variable]?.add(element)
+            elementData[variable]!!.values.add(element)

-    fun add(variable: IrVariable, elements: Set<IrExpression>) =
-            elementData[variable]?.addAll(elements)
-
-    fun get(variable: IrVariable): Set<IrExpression>? =
-            elementData[variable]
+    private fun add(variable: IrVariable, elements: Set<IrExpression>) =
+            elementData[variable]?.values?.addAll(elements)

    fun computeClosure() {
-        elementData.forEach { key, _ ->
+        elementData.forEach { (key, _) ->
            add(key, computeValueClosure(key))
        }
    }
@@ -80,8 +81,7 @@ private class VariableValues {

    private fun dfs(value: IrVariable, seen: MutableSet<IrVariable>, result: MutableSet<IrExpression>) {
        seen += value
-        val elements = elementData[value]
-                ?: return
+        val elements = elementData[value]?.values ?: return
        for (element in elements) {
            if (element !is IrGetValue)
                result += element
@@ -98,6 +98,12 @@ private class ExpressionValuesExtractor(val context: Context,
                                        val returnableBlockValues: Map<IrReturnableBlock, List<IrExpression>>,
                                        val suspendableExpressionValues: Map<IrSuspendableExpression, List<IrSuspensionPoint>>) {

+    val unit = IrGetObjectValueImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET,
+            context.irBuiltIns.unitType, context.ir.symbols.unit)
+
+    val nothing = IrGetObjectValueImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET,
+            context.irBuiltIns.nothingType, context.ir.symbols.nothing)
+
    fun forEachValue(expression: IrExpression, block: (IrExpression) -> Unit) {
        when (expression) {
            is IrReturnableBlock -> returnableBlockValues[expression]!!.forEach { forEachValue(it, block) }
@@ -113,9 +119,7 @@ private class ExpressionValuesExtractor(val context: Context,
            is IrContainerExpression -> {
                if (expression.statements.isNotEmpty())
                    forEachValue(
-                            expression = (expression.statements.last() as? IrExpression)
-                                    ?: IrGetObjectValueImpl(expression.startOffset, expression.endOffset,
-                                            context.irBuiltIns.unitType, context.ir.symbols.unit),
+                            expression = (expression.statements.last() as? IrExpression) ?: unit,
                            block      = block
                    )
            }
@@ -156,15 +160,10 @@ private class ExpressionValuesExtractor(val context: Context,

            is IrSetField -> block(expression)

-            else -> {
-                val classSymbol = when {
-                    expression.type.isUnit() -> context.ir.symbols.unit
-                    expression.type.isNothing() -> context.ir.symbols.nothing
-                    else -> TODO(ir2stringWhole(expression))
-                }
-
-                block(IrGetObjectValueImpl(expression.startOffset, expression.endOffset,
-                        expression.type, classSymbol))
+            else -> when {
+                expression.type.isUnit() -> unit
+                expression.type.isNothing() -> nothing
+                else -> TODO(ir2stringWhole(expression))
            }
        }
    }
@@ -250,13 +249,13 @@ internal class ModuleDFGBuilder(val context: Context, val irModule: IrModuleFrag
                DEBUG_OUTPUT(0) {
                    println("FIRST PHASE")
                    visitor.variableValues.elementData.forEach { t, u ->
-                        println("VAR $t:")
-                        u.forEach {
+                        println("VAR $t [LOOP ${u.loop}]:")
+                        u.values.forEach {
                            println("    ${ir2stringWhole(it)}")
                        }
                    }
                    visitor.expressions.forEach { t ->
-                        println("EXP ${ir2stringWhole(t)}")
+                        println("EXP [LOOP ${t.value}] ${ir2stringWhole(t.key)}")
                    }
                }

@@ -266,15 +265,16 @@ internal class ModuleDFGBuilder(val context: Context, val irModule: IrModuleFrag
                DEBUG_OUTPUT(0) {
                    println("SECOND PHASE")
                    visitor.variableValues.elementData.forEach { t, u ->
-                        println("VAR $t:")
-                        u.forEach {
+                        println("VAR $t [LOOP ${u.loop}]:")
+                        u.values.forEach {
                            println("    ${ir2stringWhole(it)}")
                        }
                    }
                }

                val function = FunctionDFGBuilder(expressionValuesExtractor, visitor.variableValues,
-                        declaration, visitor.expressions, visitor.returnValues, visitor.thrownValues, visitor.catchParameters).build()
+                        declaration, visitor.expressions, visitor.parentLoops, visitor.returnValues,
+                        visitor.thrownValues, visitor.catchParameters).build()

                DEBUG_OUTPUT(0) {
                    function.debugOutput()
@@ -304,14 +304,16 @@ internal class ModuleDFGBuilder(val context: Context, val irModule: IrModuleFrag
    }

    private inner class ElementFinderVisitor : IrElementVisitorVoid {
-
-        val expressions = mutableListOf<IrExpression>()
+        val expressions = mutableMapOf<IrExpression, IrLoop?>()
+        val parentLoops = mutableMapOf<IrLoop, IrLoop?>()
        val variableValues = VariableValues()
        val returnValues = mutableListOf<IrExpression>()
        val thrownValues = mutableListOf<IrExpression>()
        val catchParameters = mutableSetOf<IrVariable>()

        private val suspendableExpressionStack = mutableListOf<IrSuspendableExpression>()
+        private val loopStack = mutableListOf<IrLoop>()
+        private val currentLoop get() = loopStack.peek()

        override fun visitElement(element: IrElement) {
            element.acceptChildrenVoid(this)
@@ -331,7 +333,7 @@ internal class ModuleDFGBuilder(val context: Context, val irModule: IrModuleFrag
                is IrVararg,
                is IrConst<*>,
                is IrTypeOperatorCall ->
-                    expressions += expression
+                    expressions += expression to currentLoop
            }

            if (expression is IrCall && expression.symbol == executeImplSymbol) {
@@ -340,6 +342,9 @@ internal class ModuleDFGBuilder(val context: Context, val irModule: IrModuleFrag
                        executeImplProducerInvoke.returnType,
                        executeImplProducerInvoke.symbol)
                producerInvocation.dispatchReceiver = expression.getValueArgument(2)
+
+                expressions += producerInvocation to currentLoop
+
                val jobFunctionReference = expression.getValueArgument(3) as? IrFunctionReference
                        ?: error("A function reference expected")
                val jobInvocation = IrCallImpl(expression.startOffset, expression.endOffset,
@@ -347,7 +352,7 @@ internal class ModuleDFGBuilder(val context: Context, val irModule: IrModuleFrag
                        jobFunctionReference.symbol as IrSimpleFunctionSymbol)
                jobInvocation.putValueArgument(0, producerInvocation)

-                expressions += jobInvocation
+                expressions += jobInvocation to currentLoop
            }

            if (expression is IrReturnableBlock) {
@@ -359,13 +364,21 @@ internal class ModuleDFGBuilder(val context: Context, val irModule: IrModuleFrag
            }
            if (expression is IrSuspensionPoint)
                suspendableExpressionValues[suspendableExpressionStack.peek()!!]!!.add(expression)
+            if (expression is IrLoop) {
+                parentLoops[expression] = currentLoop
+                loopStack.push(expression)
+            }
+
            super.visitExpression(expression)
+
+            if (expression is IrLoop)
+                loopStack.pop()
            if (expression is IrSuspendableExpression)
                suspendableExpressionStack.pop()
        }

        override fun visitSetField(expression: IrSetField) {
-            expressions += expression
+            expressions += expression to currentLoop
            super.visitSetField(expression)
        }

@@ -396,7 +409,7 @@ internal class ModuleDFGBuilder(val context: Context, val irModule: IrModuleFrag
        }

        override fun visitVariable(declaration: IrVariable) {
-            variableValues.addEmpty(declaration)
+            variableValues.addEmpty(declaration, currentLoop)
            super.visitVariable(declaration)
            declaration.initializer?.let { assignVariable(declaration, it) }
        }
@@ -420,21 +433,27 @@ internal class ModuleDFGBuilder(val context: Context, val irModule: IrModuleFrag
            .single { it.name == OperatorNameConventions.INVOKE }
    private val reinterpret = symbols.reinterpret

+    private class Scoped<out T : Any>(val value: T, val scope: DataFlowIR.Node.Scope)
+
    private inner class FunctionDFGBuilder(val expressionValuesExtractor: ExpressionValuesExtractor,
                                           val variableValues: VariableValues,
                                           val declaration: IrDeclaration,
-                                           val expressions: List<IrExpression>,
+                                           val expressions: Map<IrExpression, IrLoop?>,
+                                           val parentLoops: Map<IrLoop, IrLoop?>,
                                           val returnValues: List<IrExpression>,
                                           val thrownValues: List<IrExpression>,
                                           val catchParameters: Set<IrVariable>) {

+        private val rootScope = DataFlowIR.Node.Scope(0, emptyList())
        private val allParameters = (declaration as? IrFunction)?.allParameters ?: emptyList()
-        private val templateParameters = allParameters.withIndex().associateBy({ it.value }, { DataFlowIR.Node.Parameter(it.index) })
+        private val templateParameters = allParameters.withIndex().associateBy({ it.value },
+                { Scoped(DataFlowIR.Node.Parameter(it.index), rootScope) }
+        )

        private val continuationParameter = when {
            declaration !is IrSimpleFunction -> null

-            declaration.isSuspend -> DataFlowIR.Node.Parameter(allParameters.size)
+            declaration.isSuspend -> Scoped(DataFlowIR.Node.Parameter(allParameters.size), rootScope)

            declaration.overrides(invokeSuspendFunctionSymbol.owner) ->           // <this> is a ContinuationImpl inheritor.
                templateParameters[declaration.dispatchReceiverParameter!!]       // It is its own continuation.
@@ -442,23 +461,51 @@ internal class ModuleDFGBuilder(val context: Context, val irModule: IrModuleFrag
            else -> null
        }

-        private fun getContinuation() = continuationParameter ?: error("Function ${declaration.descriptor} has no continuation parameter")
+        private fun getContinuation() = continuationParameter
+                ?: error("Function ${declaration.descriptor} has no continuation parameter")

-        private val nodes = mutableMapOf<IrExpression, DataFlowIR.Node>()
-        private val variables = variableValues.elementData.keys.associate {
-                it to DataFlowIR.Node.Variable(
-                        values = mutableListOf(),
-                        type   = symbolTable.mapType(it.type),
-                        kind   = if (catchParameters.contains(it))
-                                     DataFlowIR.VariableKind.CatchParameter
-                                 else DataFlowIR.VariableKind.Ordinary
-                )
-        }
+        private val nodes = mutableMapOf<IrExpression, Scoped<DataFlowIR.Node>>()
+        private val variables = mutableMapOf<IrVariable, Scoped<DataFlowIR.Node.Variable>>()
+        private val expressionsScopes = mutableMapOf<IrExpression, DataFlowIR.Node.Scope>()

        fun build(): DataFlowIR.Function {
            val isSuspend = declaration is IrSimpleFunction && declaration.isSuspend

-            expressions.forEach { getNode(it) }
+            val scopes = mutableMapOf<IrLoop, DataFlowIR.Node.Scope>()
+            fun transformLoop(loop: IrLoop, parentLoop: IrLoop?): DataFlowIR.Node.Scope {
+                scopes[loop]?.let { return it }
+                val parentScope =
+                        if (parentLoop == null)
+                            rootScope
+                        else transformLoop(parentLoop, parentLoops[parentLoop])
+                val scope = DataFlowIR.Node.Scope(parentScope.depth + 1, emptyList())
+                parentScope.nodes += scope
+                scopes[loop] = scope
+                return scope
+            }
+            parentLoops.forEach { (loop, parentLoop) -> transformLoop(loop, parentLoop) }
+            expressions.forEach { (expression, loop) ->
+                val scope = if (loop == null) rootScope else scopes[loop]!!
+                expressionsScopes[expression] = scope
+            }
+            expressionsScopes[expressionValuesExtractor.unit] = rootScope
+            expressionsScopes[expressionValuesExtractor.nothing] = rootScope
+
+            variableValues.elementData.forEach { (irVariable, variable) ->
+                val loop = variable.loop
+                val scope = if (loop == null) rootScope else scopes[loop]!!
+                val node = DataFlowIR.Node.Variable(
+                        values = mutableListOf(),
+                        type   = symbolTable.mapType(irVariable.type),
+                        kind   = if (catchParameters.contains(irVariable))
+                                     DataFlowIR.VariableKind.CatchParameter
+                                 else DataFlowIR.VariableKind.Ordinary
+                )
+                scope.nodes += node
+                variables[irVariable] = Scoped(node, scope)
+            }
+
+            expressions.forEach { getNode(it.key) }

            val returnNodeType = when (declaration) {
                is IrField -> declaration.type
@@ -476,17 +523,21 @@ internal class ModuleDFGBuilder(val context: Context, val irModule: IrModuleFrag
                    type   = symbolTable.mapClassReferenceType(symbols.throwable.owner),
                    kind   = DataFlowIR.VariableKind.Temporary
            )
-            variables.forEach { variable, node ->
-                variableValues.elementData[variable]!!.forEach {
-                    node.values += expressionToEdge(it)
-                }
+            variables.forEach { (irVariable, node) ->
+                val values = variableValues.elementData[irVariable]!!.values
+                values.forEach { node.value.values += expressionToEdge(it) }
            }
-            val allNodes = nodes.values + variables.values + templateParameters.values + returnsNode + throwsNode +
-                    (if (isSuspend) listOf(continuationParameter!!) else emptyList())
+
+            rootScope.nodes += templateParameters.values.map { it.value }
+            rootScope.nodes += returnsNode
+            rootScope.nodes += throwsNode
+            if (isSuspend)
+                rootScope.nodes += continuationParameter!!.value

            return DataFlowIR.Function(
-                    symbol         = symbolTable.mapFunction(declaration),
-                    body           = DataFlowIR.FunctionBody(allNodes.distinct().toList(), returnsNode, throwsNode)
+                    symbol = symbolTable.mapFunction(declaration),
+                    body   = DataFlowIR.FunctionBody(
+                            rootScope, listOf(rootScope) + scopes.values, returnsNode, throwsNode)
            )
        }

@@ -511,13 +562,25 @@ internal class ModuleDFGBuilder(val context: Context, val irModule: IrModuleFrag
            }
        }

-        private fun expressionToEdge(expression: IrExpression) =
+        private fun expressionToEdge(expression: IrExpression) = expressionToScopedEdge(expression).value
+
+        private fun expressionToScopedEdge(expression: IrExpression) =
                if (expression is IrTypeOperatorCall && expression.operator.isCast())
-                    DataFlowIR.Edge(
-                            getNode(expression.argument),
-                            symbolTable.mapClassReferenceType(expression.typeOperand.erasure().getClass()!!)
-                    )
-                else DataFlowIR.Edge(getNode(expression), null)
+                    getNode(expression.argument).let {
+                        Scoped(
+                                DataFlowIR.Edge(
+                                        it.value,
+                                        symbolTable.mapClassReferenceType(expression.typeOperand.erasure().getClass()!!)
+                                ), it.scope)
+                    }
+                else {
+                    getNode(expression).let {
+                        Scoped(
+                                DataFlowIR.Edge(it.value, null),
+                                it.scope
+                        )
+                    }
+                }

        private fun mapReturnType(actualType: IrType, returnType: IrType): DataFlowIR.Type {
            val returnedInlinedClass = returnType.getInlinedClassNative()
@@ -530,7 +593,7 @@ internal class ModuleDFGBuilder(val context: Context, val irModule: IrModuleFrag
            }
        }

-        private fun getNode(expression: IrExpression): DataFlowIR.Node {
+        private fun getNode(expression: IrExpression): Scoped<DataFlowIR.Node> {
            if (expression is IrGetValue) {
                val valueDeclaration = expression.symbol.owner
                if (valueDeclaration is IrValueParameter)
@@ -543,17 +606,33 @@ internal class ModuleDFGBuilder(val context: Context, val irModule: IrModuleFrag
                    println(ir2stringWhole(expression))
                }
                val values = mutableListOf<IrExpression>()
-                expressionValuesExtractor.forEachValue(expression) { values += it }
-                if (values.size != 1) {
+                val edges = mutableListOf<DataFlowIR.Edge>()
+                var highestScope: DataFlowIR.Node.Scope? = null
+                expressionValuesExtractor.forEachValue(expression) {
+                    values += it
+                    if (it != expression || values.size > 1) {
+                        val edge = expressionToScopedEdge(it)
+                        val scope = edge.scope
+                        if (highestScope == null || highestScope!!.depth > scope.depth)
+                            highestScope = scope
+                        edges += edge.value
+                    }
+                }
+                if (values.size == 1 && values[0] == expression) {
+                    highestScope = expressionsScopes[expression] ?: error("Unknown expression: ${expression.dump()}")
+                }
+                if (values.size == 0)
+                    highestScope = rootScope
+                val node = if (values.size != 1) {
                    DataFlowIR.Node.Variable(
-                            values = values.map { expressionToEdge(it) },
+                            values = edges,
                            type   = symbolTable.mapType(expression.type),
                            kind   = DataFlowIR.VariableKind.Temporary
                    )
                } else {
                    val value = values[0]
                    if (value != expression) {
-                        val edge = expressionToEdge(value)
+                        val edge = edges[0]
                        if (edge.castToType == null)
                            edge.node
                        else
@@ -564,7 +643,7 @@ internal class ModuleDFGBuilder(val context: Context, val irModule: IrModuleFrag
                            )
                    } else {
                        when (value) {
-                            is IrGetValue -> getNode(value)
+                            is IrGetValue -> getNode(value).value

                            is IrVararg -> DataFlowIR.Node.Const(symbolTable.mapType(value.type))

@@ -610,7 +689,7 @@ internal class ModuleDFGBuilder(val context: Context, val irModule: IrModuleFrag
                            }

                            is IrCall -> when (value.symbol) {
-                                getContinuationSymbol -> getContinuation()
+                                getContinuationSymbol -> getContinuation().value

                                in arrayGetSymbols -> {
                                    val callee = value.symbol.owner
@@ -642,7 +721,7 @@ internal class ModuleDFGBuilder(val context: Context, val irModule: IrModuleFrag
                                            value.getTypeArgument(0)!!.getClass()!!
                                    ))

-                                reinterpret -> getNode(value.extensionReceiver!!)
+                                reinterpret -> getNode(value.extensionReceiver!!).value

                                initInstanceSymbol -> {
                                    val thiz = expressionToEdge(value.getValueArgument(0)!!)
@@ -664,7 +743,7 @@ internal class ModuleDFGBuilder(val context: Context, val irModule: IrModuleFrag
                                            .map { expressionToEdge(it.second) }
                                            .let {
                                                if (callee.isSuspend)
-                                                    it + DataFlowIR.Edge(getContinuation(), null)
+                                                    it + DataFlowIR.Edge(getContinuation().value, null)
                                                else
                                                    it
                                            }
@@ -785,6 +864,9 @@ internal class ModuleDFGBuilder(val context: Context, val irModule: IrModuleFrag
                        }
                    }
                }
+
+                highestScope!!.nodes += node
+                Scoped(node, highestScope!!)
            }
        }
    }
@@ -266,17 +266,31 @@ internal object DataFlowIR {
        class Variable(values: List<Edge>, val type: Type, val kind: VariableKind) : Node() {
            val values = mutableListOf<Edge>().also { it += values }
        }
+
+        class Scope(val depth: Int, nodes: List<Node>) : Node() {
+            val nodes = mutableSetOf<Node>().also { it += nodes }
+        }
    }

-    class FunctionBody(val nodes: List<Node>, val returns: Node.Variable, val throws: Node.Variable)
+    // Note: scopes form a tree.
+    class FunctionBody(val rootScope: Node.Scope, val allScopes: List<Node.Scope>,
+                       val returns: Node.Variable, val throws: Node.Variable) {
+        inline fun forEachNonScopeNode(block: (Node) -> Unit) {
+            for (scope in allScopes)
+                for (node in scope.nodes)
+                    if (node !is Node.Scope)
+                        block(node)
+        }
+    }

    class Function(val symbol: FunctionSymbol, val body: FunctionBody) {

        fun debugOutput() {
            println("FUNCTION $symbol")
            println("Params: ${symbol.parameters.contentToString()}")
-            val ids = body.nodes.withIndex().associateBy({ it.value }, { it.index })
-            body.nodes.forEach {
+            val nodes = listOf(body.rootScope) + body.allScopes.flatMap { it.nodes }
+            val ids = nodes.withIndex().associateBy({ it.value }, { it.index })
+            nodes.forEach {
                println("    NODE #${ids[it]!!}")
                printNode(it, ids)
            }
@@ -440,6 +454,15 @@ internal object DataFlowIR {
                    }
                }

+                is Node.Scope -> {
+                    val result = StringBuilder()
+                    result.appendLine("       SCOPE ${node.depth}")
+                    node.nodes.forEach {
+                        result.appendLine("            SUBNODE #${ids[it]!!}")
+                    }
+                    result.toString()
+                }
+
                else -> {
                    "        UNKNOWN: ${node::class.java}\n"
                }
@@ -319,11 +319,15 @@ internal object Devirtualization {
                        }
                }
            }
-            var cur: Node = source
-            do {
-                println("    #${cur.id}")
-                cur = prev[cur]!!
-            } while (cur != node)
+            try {
+                var cur: Node = source
+                do {
+                    println("    #${cur.id}")
+                    cur = prev[cur]!!
+                } while (cur != node)
+            } catch (t: Throwable) {
+                println("Unable to print path")
+            }
        }

        private inner class Condensation(val multiNodes: IntArray, val topologicalOrder: IntArray) {
@@ -600,8 +604,8 @@ internal object Devirtualization {
            val nothing = symbolTable.classMap[context.ir.symbols.nothing.owner]
            for (function in functions.values) {
                if (!constraintGraph.functions.containsKey(function.symbol)) continue
-                for (node in function.body.nodes) {
-                    val virtualCall = node as? DataFlowIR.Node.VirtualCall ?: continue
+                function.body.forEachNonScopeNode { node ->
+                    val virtualCall = node as? DataFlowIR.Node.VirtualCall ?: return@forEachNonScopeNode
                    assert(nodesMap[virtualCall] != null) { "Node for virtual call $virtualCall has not been built" }
                    val receiverNode = constraintGraph.virtualCallSiteReceivers[virtualCall]
                            ?: error("virtualCallSiteReceivers were not built for virtual call $virtualCall")
@@ -615,7 +619,7 @@ internal object Devirtualization {
                            printPathToType(reversedEdges, receiverNode, VIRTUAL_TYPE_ID)
                        }

-                        continue
+                        return@forEachNonScopeNode
                    }

                    DEBUG_OUTPUT(0) {
@@ -780,13 +784,17 @@ internal object Devirtualization {
            private val preliminaryNumberOfNodes =
                    allTypes.size + // A possible source node for each type.
                            functions.size * 2 + // <returns> and <throws> nodes for each function.
-                            functions.values.sumBy { it.body.nodes.size } + // A node for each DataFlowIR.Node.
+                            functions.values.sumBy {
+                                it.body.allScopes.sumBy { it.nodes.size } // A node for each DataFlowIR.Node.
+                            } +
                            functions.values
                                    .sumBy { function ->
-                                        function.body.nodes.count { node ->
-                                            // A cast if types are different.
-                                            node is DataFlowIR.Node.Call
-                                                    && node.returnType.resolved() != node.callee.returnParameter.type.resolved()
+                                        function.body.allScopes.sumBy {
+                                            it.nodes.count { node ->
+                                                // A cast if types are different.
+                                                node is DataFlowIR.Node.Call
+                                                        && node.returnType.resolved() != node.callee.returnParameter.type.resolved()
+                                            }
                                        }
                                    }

@@ -903,17 +911,17 @@ internal object Devirtualization {
                    val body = function.body
                    val functionConstraintGraph = constraintGraph.functions[symbol]!!

-                    body.nodes.forEach { dfgNodeToConstraintNode(functionConstraintGraph, it) }
+                    body.forEachNonScopeNode { dfgNodeToConstraintNode(functionConstraintGraph, it) }
                    addEdge(functionNodesMap[body.returns]!!, functionConstraintGraph.returns)
                    addEdge(functionNodesMap[body.throws]!!, functionConstraintGraph.throws)

                    DEBUG_OUTPUT(0) {
                        println("CONSTRAINT GRAPH FOR $symbol")
-                        val ids = function.body.nodes.asSequence().withIndex().associateBy({ it.value }, { it.index })
-                        for (node in function.body.nodes) {
+                        val ids = function.body.allScopes.flatMap { it.nodes }.withIndex().associateBy({ it.value }, { it.index })
+                        function.body.forEachNonScopeNode { node ->
                            println("FT NODE #${ids[node]}")
                            DataFlowIR.Function.printNode(node, ids)
-                            val constraintNode = functionNodesMap[node] ?: variables[node] ?: break
+                            val constraintNode = functionNodesMap[node] ?: variables[node] ?: return@forEachNonScopeNode
                            println("       CG NODE #${constraintNode.id}: ${constraintNode.toString(allTypes)}")
                            println()
                        }
@@ -99,7 +99,8 @@ internal object EscapeAnalysis {
            return callGraph.nodes.associateBy({ it.symbol }) {
                val function = functions[it.symbol]!!
                val body = function.body
-                val nodesRoles = body.nodes.associate { it to Roles() }
+                val nodesRoles = mutableMapOf<DataFlowIR.Node, Roles>()
+                body.forEachNonScopeNode { nodesRoles[it] = Roles() }

                fun assignRole(node: DataFlowIR.Node, role: Role, infoEntry: RoleInfoEntry?) {
                    nodesRoles[node]!!.add(role, infoEntry)
@@ -107,7 +108,7 @@ internal object EscapeAnalysis {

                body.returns.values.forEach { assignRole(it.node, Role.RETURN_VALUE, null /* TODO */) }
                body.throws.values.forEach  { assignRole(it.node, Role.THROW_VALUE,  null /* TODO */) }
-                for (node in body.nodes) {
+                body.forEachNonScopeNode { node ->
                    when (node) {
                        is DataFlowIR.Node.FieldWrite -> {
                            val receiver = node.receiver
@@ -467,7 +468,12 @@ internal object EscapeAnalysis {
            val functionAnalysisResult = intraproceduralAnalysisResult[functionSymbol]!!
            val nodes = mutableMapOf<DataFlowIR.Node, PointsToGraphNode>()

-            val ids = if (DEBUG > 0) functionAnalysisResult.function.body.nodes.withIndex().associateBy({ it.value }, { it.index }) else null
+            val ids = if (DEBUG > 0)
+                (listOf(functionAnalysisResult.function.body.rootScope)
+                        + functionAnalysisResult.function.body.allScopes.flatMap { it.nodes }
+                        )
+                        .withIndex().associateBy({ it.value }, { it.index })
+            else null

            fun lifetimeOf(node: DataFlowIR.Node) = nodes[node]!!.let {
                when (it.kind) {
@@ -563,7 +569,7 @@ internal object EscapeAnalysis {

            fun print_digraph() {
                println("digraph {")
-                val ids = ids ?: functionAnalysisResult.function.body.nodes.withIndex().associateBy({ it.value }, { it.index })
+                val ids = ids!!
                nodes.forEach { t, u ->
                    u.edges.forEach {
                        println("    ${ids[t]} -> ${ids[it]};")
@@ -630,7 +636,9 @@ internal object EscapeAnalysis {
            }

            fun buildClosure(): FunctionEscapeAnalysisResult {
-                val parameters = functionAnalysisResult.function.body.nodes.filterIsInstance<DataFlowIR.Node.Parameter>()
+                // Parameters are declared in the root scope.
+                val parameters = functionAnalysisResult.function.body.rootScope.nodes
+                        .filterIsInstance<DataFlowIR.Node.Parameter>()
                val reachabilities = mutableListOf<IntArray>()

                DEBUG_OUTPUT(0) {
@@ -168,7 +168,7 @@ internal object LocalEscapeAnalysis {
        }

        fun analyze(lifetimes: MutableMap<IrElement, Lifetime>) {
-            function.body.nodes.forEach { node ->
+            function.body.forEachNonScopeNode { node ->
                evaluateEscapeState(node)
            }
            function.body.returns.escapeState = EscapeState.ARG_ESCAPE