From 1eec2acf43bd1d510119f133e46163811fd1ea04 Mon Sep 17 00:00:00 2001
From: Anton Bannykh <Anton.Bannykh@jetbrains.com>
Date: Thu, 19 Dec 2019 19:02:27 +0300
Subject: [PATCH] Make LocalClassPopupLowering a BodyLoweringPass

---
 .../common/lower/LocalClassPopupLowering.kt   | 41 ++++++++++---------
 1 file changed, 21 insertions(+), 20 deletions(-)

diff --git a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/LocalClassPopupLowering.kt b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/LocalClassPopupLowering.kt
index e4c55f0dbdf..07cbf61684f 100644
--- a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/LocalClassPopupLowering.kt
+++ b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/LocalClassPopupLowering.kt
@@ -5,38 +5,39 @@
 
 package org.jetbrains.kotlin.backend.common.lower
 
-import org.jetbrains.kotlin.backend.common.BackendContext
-import org.jetbrains.kotlin.backend.common.FileLoweringPass
-import org.jetbrains.kotlin.backend.common.IrElementTransformerVoidWithContext
-import org.jetbrains.kotlin.backend.common.ScopeWithIr
+import org.jetbrains.kotlin.backend.common.*
 import org.jetbrains.kotlin.backend.common.ir.addChild
 import org.jetbrains.kotlin.ir.IrStatement
-import org.jetbrains.kotlin.ir.declarations.IrClass
-import org.jetbrains.kotlin.ir.declarations.IrDeclarationContainer
-import org.jetbrains.kotlin.ir.declarations.IrFile
-import org.jetbrains.kotlin.ir.declarations.IrScript
+import org.jetbrains.kotlin.ir.declarations.*
+import org.jetbrains.kotlin.ir.expressions.IrBody
 import org.jetbrains.kotlin.ir.expressions.impl.IrCompositeImpl
 
 //This lower takes part of old LocalDeclarationLowering job to pop up local classes from functions
-open class LocalClassPopupLowering(val context: BackendContext) : FileLoweringPass {
+open class LocalClassPopupLowering(val context: BackendContext) : BodyLoweringPass {
     override fun lower(irFile: IrFile) {
+        runOnFilePostfix(irFile, withLocalDeclarations = true)
+    }
+
+    override fun lower(irBody: IrBody, container: IrDeclaration) {
         val extractedLocalClasses = arrayListOf<Pair<IrClass, IrDeclarationContainer>>()
 
-        irFile.transform(object : IrElementTransformerVoidWithContext() {
+        irBody.transform(object : IrElementTransformerVoidWithContext() {
 
             override fun visitClassNew(declaration: IrClass): IrStatement {
-                val newDeclaration = super.visitClassNew(declaration)
-                if (newDeclaration !is IrClass) return newDeclaration
+                val currentScope =
+                    if (allScopes.size > 1) allScopes[allScopes.lastIndex - 1] else createScope(container as IrSymbolOwner)
+                if (!shouldPopUp(declaration, currentScope)) return declaration
 
-                val currentScope = allScopes[allScopes.lastIndex - 1]
-                if (!shouldPopUp(declaration, currentScope)) return newDeclaration
+                val newContainer = run {
+                    var currentParent = declaration.parent
+                    while (currentParent is IrDeclaration && currentParent !is IrClass && currentParent !is IrScript) {
+                        currentParent = currentParent.parent
+                    }
 
-                val newContainer = allScopes.asReversed().drop(1/*skip self*/).firstOrNull {
-                    //find first class local or not;
-                    // to reproduce original LocalDeclarationLowering behaviour add: '&& !it.irElement.isLocal' condition
-                    it.irElement is IrClass || it.irElement is IrScript
-                }?.irElement?.let { it as? IrClass ?: it as? IrScript } ?: currentFile
-                extractedLocalClasses.add(newDeclaration to newContainer)
+                    currentParent as IrDeclarationContainer // IrClass or IrScript or IrPackageFragment
+                }
+
+                extractedLocalClasses.add(declaration to newContainer)
                 return IrCompositeImpl(declaration.startOffset, declaration.endOffset, context.irBuiltIns.unitType)
             }
         }, null)