backend: Add closures of calls in function closure

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

@@ -5,7 +5,6 @@ import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.IrCall
 import org.jetbrains.kotlin.ir.expressions.IrValueAccessExpression
-import org.jetbrains.kotlin.ir.expressions.impl.IrSetterCallImpl
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
 import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
 import org.jetbrains.kotlin.ir.visitors.acceptVoid
@@ -15,7 +14,13 @@ import org.jetbrains.kotlin.resolve.descriptorUtil.getSuperClassOrAny
 // TODO: synchronize with JVM BE
 class Closure(val capturedValues: List<ValueDescriptor>)
 
-abstract class AbstractClosureAnnotator : IrElementVisitorVoid {
+private fun <E> MutableList<E>.push(element: E) = this.add(element)
+
+private fun <E> MutableList<E>.pop() = this.removeAt(size - 1)
+
+private fun <E> MutableList<E>.peek(): E? = if (size == 0) null else this[size - 1]
+
+abstract class AbstractClosureAnnotator  {
     protected abstract fun recordFunctionClosure(functionDescriptor: FunctionDescriptor, closure: Closure)
     protected abstract fun recordClassClosure(classDescriptor: ClassDescriptor, closure: Closure)
 
@@ -48,83 +53,117 @@ abstract class AbstractClosureAnnotator : IrElementVisitorVoid {
         }
     }
 
-    private val classClosures = mutableMapOf<ClassDescriptor, Closure>()
+    val functionClosures = mutableMapOf<FunctionDescriptor, Closure>()
-    private val closuresStack = mutableListOf<ClosureBuilder>()
 
-    private fun <E> MutableList<E>.push(element: E) = this.add(element)
+    fun annotate(declaration: IrDeclaration) {
-
+        // First pass - collect all closures for classes and functions.
-    private fun <E> MutableList<E>.pop() = this.removeAt(size - 1)
+        declaration.acceptChildrenVoid(ClosureCollectorVisitor())
-
+        // Second pass - add callees' closures for callers' ones.
-    private fun <E> MutableList<E>.peek(): E? = if (size == 0) null else this[size - 1]
+        declaration.acceptChildrenVoid(CallsCollectorVisitor())
-
-    override fun visitElement(element: IrElement) {
-        element.acceptChildrenVoid(this)
     }
 
-    override fun visitClass(declaration: IrClass) {
+    private abstract class ClosureCollectorVisitorBase : IrElementVisitorVoid {
-        val classDescriptor = declaration.descriptor
+        open protected fun recordFunctionClosure(functionDescriptor: FunctionDescriptor, closure: Closure) {}
-        val closureBuilder = ClosureBuilder()
+        open protected fun recordClassClosure(classDescriptor: ClassDescriptor, closure: Closure) {}
 
-        closureBuilder.declareVariable(classDescriptor.thisAsReceiverParameter)
+        protected val closuresStack = mutableListOf<ClosureBuilder>()
-        if (classDescriptor.isInner)
+        protected val classClosures = mutableMapOf<ClassDescriptor, Closure>()
-            closureBuilder.declareVariable((classDescriptor.containingDeclaration as ClassDescriptor).thisAsReceiverParameter)
 
-        closuresStack.push(closureBuilder)
+        override fun visitElement(element: IrElement) {
-        declaration.acceptChildrenVoid(this)
+            element.acceptChildrenVoid(this)
-        closuresStack.pop()
-
-        val superClassClosure = classClosures[classDescriptor.getSuperClassOrAny()]
-        if (superClassClosure != null) {
-            // Capture all values from the super class since we need to call constructor of super class
-            // with its captured values.
-            closureBuilder.addNested(superClassClosure)
         }
 
-        val closure = closureBuilder.buildClosure()
+        override fun visitClass(declaration: IrClass) {
+            val classDescriptor = declaration.descriptor
+            val closureBuilder = ClosureBuilder()
 
-        if (DescriptorUtils.isLocal(classDescriptor)) {
+            closureBuilder.declareVariable(classDescriptor.thisAsReceiverParameter)
-            recordClassClosure(classDescriptor, closure)
+            if (classDescriptor.isInner)
-            classClosures[classDescriptor] = closure
+                closureBuilder.declareVariable((classDescriptor.containingDeclaration as ClassDescriptor).thisAsReceiverParameter)
+
+            closuresStack.push(closureBuilder)
+            declaration.acceptChildrenVoid(this)
+            closuresStack.pop()
+
+            val superClassClosure = classClosures[classDescriptor.getSuperClassOrAny()]
+            if (superClassClosure != null) {
+                // Capture all values from the super class since we need to call constructor of super class
+                // with its captured values.
+                closureBuilder.addNested(superClassClosure)
+            }
+
+            val closure = closureBuilder.buildClosure()
+
+            if (DescriptorUtils.isLocal(classDescriptor)) {
+                recordClassClosure(classDescriptor, closure)
+                classClosures[classDescriptor] = closure
+            }
+
+            closuresStack.peek()?.addNested(closure)
         }
 
-        closuresStack.peek()?.addNested(closure)
+        override fun visitFunction(declaration: IrFunction) {
-    }
+            val functionDescriptor = declaration.descriptor
+            val closureBuilder = ClosureBuilder()
 
-    override fun visitFunction(declaration: IrFunction) {
+            functionDescriptor.valueParameters.forEach { closureBuilder.declareVariable(it) }
-        val functionDescriptor = declaration.descriptor
+            closureBuilder.declareVariable(functionDescriptor.dispatchReceiverParameter)
-        val closureBuilder = ClosureBuilder()
+            closureBuilder.declareVariable(functionDescriptor.extensionReceiverParameter)
+            if (functionDescriptor is ConstructorDescriptor)
+                closureBuilder.declareVariable(functionDescriptor.constructedClass.thisAsReceiverParameter)
 
-        functionDescriptor.valueParameters.forEach { closureBuilder.declareVariable(it) }
+            closuresStack.push(closureBuilder)
-        closureBuilder.declareVariable(functionDescriptor.dispatchReceiverParameter)
+            declaration.acceptChildrenVoid(this)
-        closureBuilder.declareVariable(functionDescriptor.extensionReceiverParameter)
+            closuresStack.pop()
-        if (functionDescriptor is ConstructorDescriptor)
-            closureBuilder.declareVariable(functionDescriptor.constructedClass.thisAsReceiverParameter)
 
-        closuresStack.push(closureBuilder)
+            val closure = closureBuilder.buildClosure()
-        declaration.acceptChildrenVoid(this)
-        closuresStack.pop()
 
-        val closure = closureBuilder.buildClosure()
+            if (DescriptorUtils.isLocal(functionDescriptor)) {
+                recordFunctionClosure(functionDescriptor, closure)
+            }
 
-        if (DescriptorUtils.isLocal(functionDescriptor)) {
+            closuresStack.peek()?.addNested(closure)
-            recordFunctionClosure(functionDescriptor, closure)
         }
 
-        closuresStack.peek()?.addNested(closure)
+        override fun visitLocalDelegatedProperty(declaration: IrLocalDelegatedProperty) {
+            // Getter and setter of local delegated properties are special generated functions and don't have closure.
+            declaration.delegate.initializer?.acceptVoid(this)
+        }
+
+        override fun visitVariableAccess(expression: IrValueAccessExpression) {
+            closuresStack.peek()?.seeVariable(expression.descriptor)
+            super.visitVariableAccess(expression)
+        }
+
+        override fun visitVariable(declaration: IrVariable) {
+            closuresStack.peek()?.declareVariable(declaration.descriptor)
+            super.visitVariable(declaration)
+        }
     }
 
-    override fun visitLocalDelegatedProperty(declaration: IrLocalDelegatedProperty) {
+    private inner class ClosureCollectorVisitor : ClosureCollectorVisitorBase() {
-        // Getter and setter of local delegated properties are special generated functions and don't have closure.
+        override fun recordFunctionClosure(functionDescriptor: FunctionDescriptor, closure: Closure) {
-        declaration.delegate.initializer?.acceptVoid(this)
+            functionClosures[functionDescriptor] = closure
+        }
+
+        override fun recordClassClosure(classDescriptor: ClassDescriptor, closure: Closure) =
+                this@AbstractClosureAnnotator.recordClassClosure(classDescriptor, closure)
     }
 
-    override fun visitVariableAccess(expression: IrValueAccessExpression) {
+    private inner class CallsCollectorVisitor : ClosureCollectorVisitorBase() {
-        closuresStack.peek()?.seeVariable(expression.descriptor)
+        override fun recordFunctionClosure(functionDescriptor: FunctionDescriptor, closure: Closure) =
-        super.visitVariableAccess(expression)
+                this@AbstractClosureAnnotator.recordFunctionClosure(functionDescriptor, closure)
-    }
 
-    override fun visitVariable(declaration: IrVariable) {
+
-        closuresStack.peek()?.declareVariable(declaration.descriptor)
+        override fun visitCall(expression: IrCall) {
-        super.visitVariable(declaration)
+            expression.acceptChildrenVoid(this)
+            val descriptor = expression.descriptor
+            if (descriptor is FunctionDescriptor && DescriptorUtils.isLocal(descriptor)) {
+                val callClosure = functionClosures[descriptor]
+                callClosure?.let {
+                    closuresStack.peek()?.addNested(callClosure)
+                }
+            }
+        }
     }
 }

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

@@ -623,7 +623,7 @@ class LocalDeclarationsLowering(val context: BackendContext) : DeclarationContai
 
 
         private fun collectClosures() {
-            memberDeclaration.acceptChildrenVoid(object : AbstractClosureAnnotator() {
+            object : AbstractClosureAnnotator() {
                 override fun recordFunctionClosure(functionDescriptor: FunctionDescriptor, closure: Closure) {
                     localFunctions[functionDescriptor]?.closure = closure
                 }
@@ -631,7 +631,7 @@ class LocalDeclarationsLowering(val context: BackendContext) : DeclarationContai
                 override fun recordClassClosure(classDescriptor: ClassDescriptor, closure: Closure) {
                     localClasses[classDescriptor]?.closure = closure
                 }
-            })
+            }.annotate(memberDeclaration)
         }
 
         private fun collectLocalDeclarations() {