Make LocalDeclarationsLowering a BodyLoweringPass

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/ClosureAnnotator.kt

@@ -17,9 +17,6 @@
 package org.jetbrains.kotlin.backend.common.lower
 
 import org.jetbrains.kotlin.backend.common.ir.ir2string
-import org.jetbrains.kotlin.backend.common.peek
-import org.jetbrains.kotlin.backend.common.pop
-import org.jetbrains.kotlin.backend.common.push
 import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.declarations.*
@@ -29,19 +26,21 @@ import org.jetbrains.kotlin.ir.symbols.IrValueSymbol
 import org.jetbrains.kotlin.ir.types.IrSimpleType
 import org.jetbrains.kotlin.ir.types.IrType
 import org.jetbrains.kotlin.ir.types.IrTypeProjection
+import org.jetbrains.kotlin.ir.util.constructedClass
 import org.jetbrains.kotlin.ir.util.isLocal
-import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
+import org.jetbrains.kotlin.ir.visitors.IrElementVisitor
-import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
-import kotlin.collections.set
 
 data class Closure(val capturedValues: List<IrValueSymbol>, val capturedTypeParameters: List<IrTypeParameter>)
 
-class ClosureAnnotator(irFile: IrFile) {
+class ClosureAnnotator(body: IrBody, declaration: IrDeclaration) {
     private val closureBuilders = mutableMapOf<IrDeclaration, ClosureBuilder>()
 
     init {
         // Collect all closures for classes and functions. Collect call graph
-        irFile.acceptChildrenVoid(ClosureCollectorVisitor())
+        if (declaration is IrSimpleFunction && declaration.name.asString() == "arrayIterator") {
+            1
+        }
+        body.accept(ClosureCollectorVisitor(), declaration.closureBuilderOrNull ?: declaration.parentClosureBuilder)
     }
 
     fun getFunctionClosure(declaration: IrFunction) = getClosure(declaration)
@@ -126,129 +125,150 @@ class ClosureAnnotator(irFile: IrFile) {
         }
     }
 
-    private inner class ClosureCollectorVisitor : IrElementVisitorVoid {
+    private fun includeInParent(builder: ClosureBuilder) {
-
+        // We don't include functions or classes in a parent function when they are declared.
-        val closuresStack = mutableListOf<ClosureBuilder>()
+        // Instead we will include them when are is used (use = call for a function or constructor call for a class).
-
+        val parentBuilder = builder.owner.parentClosureBuilder
-        fun includeInParent(builder: ClosureBuilder) {
+        if (parentBuilder != null && parentBuilder.owner !is IrFunction) {
-            // We don't include functions or classes in a parent function when they are declared.
+            parentBuilder.include(builder)
-            // Instead we will include them when are is used (use = call for a function or constructor call for a class).
-            val parentBuilder = closuresStack.peek()
-            if (parentBuilder != null && parentBuilder.owner !is IrFunction) {
-                parentBuilder.include(builder)
-            }
         }
+    }
 
-        override fun visitElement(element: IrElement) {
+    private val IrClass.closureBuilder: ClosureBuilder
-            element.acceptChildrenVoid(this)
+        get() = closureBuilders.getOrPut(this) {
-        }
+            val closureBuilder = ClosureBuilder(this)
-
-        override fun visitClass(declaration: IrClass) {
-            val closureBuilder = ClosureBuilder(declaration)
-            closureBuilders[declaration] = closureBuilder
 
             collectPotentiallyCapturedTypeParameters(closureBuilder)
 
-            closureBuilder.declareVariable(declaration.thisReceiver)
+            closureBuilder.declareVariable(this.thisReceiver)
-            if (declaration.isInner) {
+            if (this.isInner) {
-                closureBuilder.declareVariable((declaration.parent as IrClass).thisReceiver)
+                closureBuilder.declareVariable((this.parent as IrClass).thisReceiver)
                 includeInParent(closureBuilder)
             }
 
-            declaration.declarations.firstOrNull { it is IrConstructor && it.isPrimary }?.let {
+            this.declarations.firstOrNull { it is IrConstructor && it.isPrimary }?.let {
                 val constructor = it as IrConstructor
                 constructor.valueParameters.forEach { v -> closureBuilder.declareVariable(v) }
             }
 
-            closuresStack.push(closureBuilder)
+            closureBuilder
-            declaration.acceptChildrenVoid(this)
-            closuresStack.pop()
         }
 
-        override fun visitFunction(declaration: IrFunction) {
+    private val IrFunction.closureBuilder: ClosureBuilder
-            val closureBuilder = ClosureBuilder(declaration)
+        get() = closureBuilders.getOrPut(this) {
-            closureBuilders[declaration] = closureBuilder
+            val closureBuilder = ClosureBuilder(this)
 
             collectPotentiallyCapturedTypeParameters(closureBuilder)
 
-            declaration.valueParameters.forEach { closureBuilder.declareVariable(it) }
+            this.valueParameters.forEach { closureBuilder.declareVariable(it) }
-            closureBuilder.declareVariable(declaration.dispatchReceiverParameter)
+            closureBuilder.declareVariable(this.dispatchReceiverParameter)
-            closureBuilder.declareVariable(declaration.extensionReceiverParameter)
+            closureBuilder.declareVariable(this.extensionReceiverParameter)
 
-            if (declaration is IrConstructor) {
+            if (this is IrConstructor) {
-                val constructedClass = (declaration.parent as IrClass)
+                val constructedClass = (this.parent as IrClass)
                 closureBuilder.declareVariable(constructedClass.thisReceiver)
 
                 // Include closure of the class in the constructor closure.
-                val classBuilder = closuresStack.peek()
+                val classBuilder = constructedClass.closureBuilder
-                classBuilder?.let {
+                closureBuilder.include(classBuilder)
-                    assert(classBuilder.owner == constructedClass)
-                    closureBuilder.include(classBuilder)
-                }
             }
 
-            closuresStack.push(closureBuilder)
+            closureBuilder
-            declaration.acceptChildrenVoid(this)
+        }
-            closuresStack.pop()
+
+    private fun collectPotentiallyCapturedTypeParameters(closureBuilder: ClosureBuilder) {
+
+        fun ClosureBuilder.doCollect() {
+            if (owner !is IrClass) {
+                (owner as? IrTypeParametersContainer)?.let { container ->
+                    for (tp in container.typeParameters) {
+                        closureBuilder.addPotentiallyCapturedTypeParameter(tp)
+                    }
+                }
+
+                owner.parentClosureBuilder?.doCollect()
+            }
+        }
+
+        closureBuilder.owner.parentClosureBuilder?.doCollect()
+    }
+
+    private val IrDeclaration.parentClosureBuilder: ClosureBuilder?
+        get() = when (val p = parent) {
+            is IrClass -> p.closureBuilder
+            is IrFunction -> p.closureBuilder
+            is IrDeclaration -> p.parentClosureBuilder
+            else -> null
+        }
+
+    private val IrDeclaration.closureBuilderOrNull: ClosureBuilder?
+        get() = when (this) {
+            is IrClass -> closureBuilder
+            is IrFunction -> closureBuilder
+            else -> null
+        }
+
+    private inner class ClosureCollectorVisitor() : IrElementVisitor<Unit, ClosureBuilder?> {
+
+        override fun visitElement(element: IrElement, data: ClosureBuilder?) {
+            element.acceptChildren(this, data)
+        }
+
+        override fun visitClass(declaration: IrClass, data: ClosureBuilder?) {
+            declaration.acceptChildren(this, declaration.closureBuilder)
+        }
+
+        override fun visitFunction(declaration: IrFunction, data: ClosureBuilder?) {
+            val closureBuilder = declaration.closureBuilder
+
+            declaration.acceptChildren(this, closureBuilder)
 
             includeInParent(closureBuilder)
         }
 
-        override fun visitVariableAccess(expression: IrValueAccessExpression) {
+        override fun visitValueAccess(expression: IrValueAccessExpression, data: ClosureBuilder?) {
-            closuresStack.peek()?.seeVariable(expression.symbol)
+            data?.seeVariable(expression.symbol)
-            super.visitVariableAccess(expression)
+            super.visitValueAccess(expression, data)
         }
 
-        override fun visitVariable(declaration: IrVariable) {
+        override fun visitVariable(declaration: IrVariable, data: ClosureBuilder?) {
-            closuresStack.peek()?.declareVariable(declaration)
+            data?.declareVariable(declaration)
-            super.visitVariable(declaration)
+            super.visitVariable(declaration, data)
         }
 
-        override fun visitFunctionAccess(expression: IrFunctionAccessExpression) {
+        override fun visitFunctionAccess(expression: IrFunctionAccessExpression, data: ClosureBuilder?) {
-            super.visitFunctionAccess(expression)
+            super.visitFunctionAccess(expression, data)
-            processMemberAccess(expression.symbol.owner)
+            processMemberAccess(expression.symbol.owner, data)
         }
 
-        override fun visitFunctionReference(expression: IrFunctionReference) {
+        override fun visitFunctionReference(expression: IrFunctionReference, data: ClosureBuilder?) {
-            super.visitFunctionReference(expression)
+            super.visitFunctionReference(expression, data)
-            processMemberAccess(expression.symbol.owner)
+            processMemberAccess(expression.symbol.owner, data)
         }
 
-        override fun visitPropertyReference(expression: IrPropertyReference) {
+        override fun visitPropertyReference(expression: IrPropertyReference, data: ClosureBuilder?) {
-            super.visitPropertyReference(expression)
+            super.visitPropertyReference(expression, data)
-            expression.getter?.let { processMemberAccess(it.owner) }
+            expression.getter?.let { processMemberAccess(it.owner, data) }
-            expression.setter?.let { processMemberAccess(it.owner) }
+            expression.setter?.let { processMemberAccess(it.owner, data) }
         }
 
-        override fun visitExpression(expression: IrExpression) {
+        override fun visitExpression(expression: IrExpression, data: ClosureBuilder?) {
-            super.visitExpression(expression)
+            super.visitExpression(expression, data)
-            val typeParameterContainerScopeBuilder = closuresStack.peek()?.let {
+            val typeParameterContainerScopeBuilder = data?.let {
-                if (it.owner is IrConstructor) {
+                (it.owner as? IrConstructor)?.closureBuilder ?: it
-                    closuresStack[closuresStack.size - 2]
-                } else it
             }
             typeParameterContainerScopeBuilder?.seeType(expression.type)
         }
 
-        private fun processMemberAccess(declaration: IrDeclaration) {
+        private fun processMemberAccess(declaration: IrDeclaration, parentClosure: ClosureBuilder?) {
             if (declaration.isLocal) {
                 if (declaration is IrSimpleFunction && declaration.visibility != Visibilities.LOCAL) {
                     return
                 }
 
-                val builder = closureBuilders[declaration]
+                val builder = declaration.closureBuilderOrNull
                 builder?.let {
-                    closuresStack.peek()?.include(builder)
+                    parentClosure?.include(builder)
-                }
-            }
-        }
-
-        private fun collectPotentiallyCapturedTypeParameters(closureBuilder: ClosureBuilder) {
-            closuresStack.takeLastWhile { it.owner !is IrClass }.forEach {
-                (it.owner as? IrTypeParametersContainer)?.let { container ->
-                    for (tp in container.typeParameters) {
-                        closureBuilder.addPotentiallyCapturedTypeParameter(tp)
-                    }
                 }
             }
         }

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/LocalDeclarationsLowering.kt

@@ -5,10 +5,7 @@
 
 package org.jetbrains.kotlin.backend.common.lower
 
-import org.jetbrains.kotlin.backend.common.BackendContext
+import org.jetbrains.kotlin.backend.common.*
-import org.jetbrains.kotlin.backend.common.FileLoweringPass
-import org.jetbrains.kotlin.backend.common.IrElementVisitorVoidWithContext
-import org.jetbrains.kotlin.backend.common.ScopeWithIr
 import org.jetbrains.kotlin.backend.common.descriptors.synthesizedName
 import org.jetbrains.kotlin.backend.common.ir.*
 import org.jetbrains.kotlin.descriptors.Modality
@@ -43,6 +40,7 @@ import org.jetbrains.kotlin.ir.types.impl.IrUninitializedType
 import org.jetbrains.kotlin.ir.types.impl.makeTypeProjection
 import org.jetbrains.kotlin.ir.util.constructedClass
 import org.jetbrains.kotlin.ir.util.dump
+import org.jetbrains.kotlin.ir.util.file
 import org.jetbrains.kotlin.ir.util.parentAsClass
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
@@ -95,7 +93,7 @@ class LocalDeclarationsLowering(
     val localNameProvider: LocalNameProvider = LocalNameProvider.DEFAULT,
     val visibilityPolicy: VisibilityPolicy = VisibilityPolicy.DEFAULT
 ) :
-    FileLoweringPass {
+    BodyLoweringPass {
 
     object DECLARATION_ORIGIN_FIELD_FOR_CAPTURED_VALUE :
         IrDeclarationOriginImpl("FIELD_FOR_CAPTURED_VALUE", isSynthetic = true)
@@ -106,10 +104,23 @@ class LocalDeclarationsLowering(
     private object STATEMENT_ORIGIN_INITIALIZER_OF_FIELD_FOR_CAPTURED_VALUE :
         IrStatementOriginImpl("INITIALIZER_OF_FIELD_FOR_CAPTURED_VALUE")
 
-    override fun lower(irFile: IrFile) {
+    override fun lower(irBody: IrBody, container: IrDeclaration) {
-        LocalDeclarationsTransformer(irFile).lowerLocalDeclarations()
+        LocalDeclarationsTransformer(irBody, container).lowerLocalDeclarations()
     }
 
+    private class ScopeWithCounter(scope: Scope, irElement: IrElement) : ScopeWithIr(scope, irElement) {
+        // Continuous numbering across all declarations in the container.
+        var counter: Int = 0
+    }
+
+    private val scopeMap: MutableMap<IrSymbolOwner, ScopeWithCounter> = mutableMapOf()
+    // Need to keep LocalFunctionContext.index
+    private val IrSymbolOwner.scopeWithCounter: ScopeWithCounter
+        get() = scopeMap.getOrPut(this) {
+            ScopeWithCounter(Scope(symbol), this)
+        }
+
+
     private abstract class LocalContext {
         val capturedTypeParameterToTypeParameter: MutableMap<IrTypeParameter, IrTypeParameter> = mutableMapOf()
 
@@ -203,7 +214,7 @@ class LocalDeclarationsLowering(
             abbreviation.annotations
         )
 
-    private inner class LocalDeclarationsTransformer(val irFile: IrFile) {
+    private inner class LocalDeclarationsTransformer(val irBody: IrBody, val container: IrDeclaration) {
         val localFunctions: MutableMap<IrFunction, LocalFunctionContext> = LinkedHashMap()
         val localClasses: MutableMap<IrClass, LocalClassContext> = LinkedHashMap()
         val localClassConstructors: MutableMap<IrConstructor, LocalClassConstructorContext> = LinkedHashMap()
@@ -481,7 +492,7 @@ class LocalDeclarationsLowering(
                 rewriteClassMembers(it.declaration, it)
             }
 
-            rewriteFunctionBody(irFile, null)
+            rewriteFunctionBody(container, null)
         }
 
         private fun createNewCall(oldCall: IrCall, newCallee: IrFunction) =
@@ -794,7 +805,7 @@ class LocalDeclarationsLowering(
 
         private fun collectClosureForLocalDeclarations() {
             //TODO: maybe use for granular declarations
-            val annotator = ClosureAnnotator(irFile)
+            val annotator = ClosureAnnotator(irBody, container)
 
             localFunctions.forEach { (declaration, context) ->
                 context.closure = annotator.getFunctionClosure(declaration)
@@ -806,19 +817,25 @@ class LocalDeclarationsLowering(
         }
 
         private fun collectLocalDeclarations() {
-            class ScopeWithCounter(scope: Scope, irElement: IrElement) : ScopeWithIr(scope, irElement) {
+            val enclosingFileScope = container.file.scopeWithCounter
-                // Continuous numbering across all declarations in the container.
-                var counter: Int = 0
-            }
 
-            irFile.acceptVoid(object : IrElementVisitorVoidWithContext() {
+            val enclosingClassScope = run {
+                var currentParent = container as? IrClass ?: container.parent
+                while (currentParent is IrDeclaration && currentParent !is IrClass) {
+                    currentParent = currentParent.parent
+                }
+
+                currentParent as? IrClass
+            }?.scopeWithCounter
+
+            irBody.acceptVoid(object : IrElementVisitorVoidWithContext() {
 
                 override fun visitElement(element: IrElement) {
                     element.acceptChildrenVoid(this)
                 }
 
                 override fun createScope(declaration: IrSymbolOwner): ScopeWithIr {
-                    return ScopeWithCounter(Scope(declaration.symbol), declaration)
+                    return ScopeWithCounter(Scope(declaration.symbol), declaration) // Don't cache local declarations
                 }
 
                 override fun visitSimpleFunction(declaration: IrSimpleFunction) {
@@ -826,7 +843,7 @@ class LocalDeclarationsLowering(
 
                     if (declaration.visibility == Visibilities.LOCAL) {
                         val scopeWithIr =
-                            (currentClass ?: currentFile /*file is required for K/N cause file declarations are not split by classes*/
+                            (currentClass ?: enclosingClassScope ?: enclosingFileScope /*file is required for K/N cause file declarations are not split by classes*/
                             ?: error("No scope for ${declaration.dump()}"))
                         localFunctions[declaration] =
                             LocalFunctionContext(
@@ -858,7 +875,8 @@ class LocalDeclarationsLowering(
                 }
 
                 private val inInlineFunctionScope: Boolean
-                    get() = allScopes.any { scope -> (scope.irElement as? IrFunction)?.isInline ?: false }
+                    get() = allScopes.any { scope -> (scope.irElement as? IrFunction)?.isInline ?: false } ||
+                            generateSequence(container) { it.parent as? IrDeclaration }.any { it is IrFunction && it.isInline }
             })
         }
     }