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,14 +53,21 @@ 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.
+        declaration.acceptChildrenVoid(ClosureCollectorVisitor())
+        // Second pass - add callees' closures for callers' ones.
+        declaration.acceptChildrenVoid(CallsCollectorVisitor())
+    }
 
-    private fun <E> MutableList<E>.pop() = this.removeAt(size - 1)
+    private abstract class ClosureCollectorVisitorBase : IrElementVisitorVoid {
+        open protected fun recordFunctionClosure(functionDescriptor: FunctionDescriptor, closure: Closure) {}
+        open protected fun recordClassClosure(classDescriptor: ClassDescriptor, closure: Closure) {}
 
-    private fun <E> MutableList<E>.peek(): E? = if (size == 0) null else this[size - 1]
+        protected val closuresStack = mutableListOf<ClosureBuilder>()
+        protected val classClosures = mutableMapOf<ClassDescriptor, Closure>()
 
         override fun visitElement(element: IrElement) {
             element.acceptChildrenVoid(this)
@@ -127,4 +139,31 @@ abstract class AbstractClosureAnnotator : IrElementVisitorVoid {
             closuresStack.peek()?.declareVariable(declaration.descriptor)
             super.visitVariable(declaration)
         }
+    }
+
+    private inner class ClosureCollectorVisitor : ClosureCollectorVisitorBase() {
+        override fun recordFunctionClosure(functionDescriptor: FunctionDescriptor, closure: Closure) {
+            functionClosures[functionDescriptor] = closure
+        }
+
+        override fun recordClassClosure(classDescriptor: ClassDescriptor, closure: Closure) =
+                this@AbstractClosureAnnotator.recordClassClosure(classDescriptor, closure)
+    }
+
+    private inner class CallsCollectorVisitor : ClosureCollectorVisitorBase() {
+        override fun recordFunctionClosure(functionDescriptor: FunctionDescriptor, closure: Closure) =
+                this@AbstractClosureAnnotator.recordFunctionClosure(functionDescriptor, closure)
+
+
+        override fun visitCall(expression: IrCall) {
+            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() {