Make IrScript a statement container first, update JS support

The IrScript is now the IrStatementContainer with significant
statements order, and the code is adapted accordingly

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/Lower.kt

@@ -17,9 +17,12 @@
 package org.jetbrains.kotlin.backend.common
 
 import org.jetbrains.kotlin.ir.IrElement
+import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.IrBody
+import org.jetbrains.kotlin.ir.expressions.IrStatementContainer
 import org.jetbrains.kotlin.ir.util.transformFlat
+import org.jetbrains.kotlin.ir.util.transformSubsetFlat
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitor
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
 import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
@@ -111,11 +114,6 @@ private class DeclarationContainerLoweringVisitor(
         declaration.acceptChildrenVoid(this)
         loweringPass.lower(declaration)
     }
-
-    override fun visitScript(declaration: IrScript) {
-        declaration.acceptChildrenVoid(this)
-        loweringPass.lower(declaration)
-    }
 }
 
 fun BodyLoweringPass.runOnFilePostfix(
@@ -160,11 +158,8 @@ private class BodyLoweringVisitor(
     }
 
     override fun visitScript(declaration: IrScript, data: IrDeclaration?) {
-        ArrayList(declaration.declarations).forEach { it.accept(this, declaration) }
-        if (withLocalDeclarations) {
-            declaration.statements.forEach { it.accept(this, null) }
-        }
         declaration.thisReceiver.accept(this, declaration)
+        ArrayList(declaration.statements).forEach { it.accept(this, declaration) }
     }
 }
 
@@ -214,22 +209,26 @@ interface DeclarationTransformer : FileLoweringPass {
 
             val visitor = this
 
+            fun IrDeclaration.replaceInContainer(container: MutableList<in IrDeclaration>, result: List<IrDeclaration>): Boolean {
+                var index = container.indexOf(this)
+                if (index == -1) {
+                    index = container.indexOf(declaration)
+                } else {
+                    container.removeAt(index)
+                    --index
+                }
+                return container.addAll(index + 1, result)
+            }
+
             fun IrDeclaration.transform() {
 
                 acceptVoid(visitor)
 
                 val result = transformer.transformFlatRestricted(this)
                 if (result != null) {
-                    (parent as? IrDeclarationContainer)?.let {
-                        var index = it.declarations.indexOf(this)
-                        if (index == -1) {
-                            index = it.declarations.indexOf(declaration)
-                        } else {
-                            it.declarations.removeAt(index)
-                            --index
-                        }
-
-                        it.declarations.addAll(index + 1, result)
+                    when (val parentCopy = parent) {
+                        is IrDeclarationContainer -> replaceInContainer(parentCopy.declarations, result)
+                        is IrStatementContainer -> replaceInContainer(parentCopy.statements, result)
                     }
                 }
             }
@@ -248,12 +247,12 @@ interface DeclarationTransformer : FileLoweringPass {
         }
 
         override fun visitScript(declaration: IrScript) {
-            ArrayList(declaration.declarations).forEach { it.accept(this, null) }
-            declaration.declarations.transformFlat(transformer::transformFlatRestricted)
-
-            if (transformer.withLocalDeclarations) {
-                declaration.statements.forEach { it.accept(this, null) }
+            ArrayList(declaration.statements).forEach {
+                if (transformer.withLocalDeclarations || it is IrDeclaration) {
+                    it.accept(this, null)
+                }
             }
+            declaration.statements.transformSubsetFlat(transformer::transformFlatRestricted)
 
             declaration.thisReceiver.accept(this, null)
         }

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

@@ -7,10 +7,12 @@ package org.jetbrains.kotlin.backend.common.lower
 
 import org.jetbrains.kotlin.backend.common.*
 import org.jetbrains.kotlin.backend.common.ir.addChild
+import org.jetbrains.kotlin.backend.common.ir.setDeclarationsParent
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.IrBody
+import org.jetbrains.kotlin.ir.expressions.IrStatementContainer
 import org.jetbrains.kotlin.ir.expressions.impl.IrCompositeImpl
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
 import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
@@ -22,8 +24,12 @@ open class LocalClassPopupLowering(val context: BackendContext) : BodyLoweringPa
         runOnFilePostfix(irFile, withLocalDeclarations = true, allowDeclarationModification = true)
     }
 
+    private data class ExtractedLocalClass(
+        val local: IrClass, val newContainer: IrDeclarationParent, val extractedUnder: IrStatement?
+    )
+
     override fun lower(irBody: IrBody, container: IrDeclaration) {
-        val extractedLocalClasses = arrayListOf<Pair<IrClass, IrDeclarationContainer>>()
+        val extractedLocalClasses = arrayListOf<ExtractedLocalClass>()
 
         irBody.transform(object : IrElementTransformerVoidWithContext() {
 
@@ -32,22 +38,46 @@ open class LocalClassPopupLowering(val context: BackendContext) : BodyLoweringPa
                     if (allScopes.size > 1) allScopes[allScopes.lastIndex - 1] else createScope(container as IrSymbolOwner)
                 if (!shouldPopUp(declaration, currentScope)) return declaration
 
-                val newContainer = run {
-                    var currentParent = declaration.parent
-                    while (currentParent is IrDeclaration && currentParent !is IrClass && currentParent !is IrScript) {
-                        currentParent = currentParent.parent
+                var extractedUnder: IrStatement? = declaration
+                var newContainer = declaration.parent
+                while (newContainer is IrDeclaration && newContainer !is IrClass && newContainer !is IrScript) {
+                    extractedUnder = newContainer
+                    newContainer = newContainer.parent
+                }
+                when (newContainer) {
+                    is IrStatementContainer -> {
+                        // TODO: check if it is the correct behavior
+                        if (extractedUnder == declaration) {
+                            extractedUnder = (newContainer.statements.indexOf(extractedUnder) + 1)
+                                .takeIf { it > 0 && it < newContainer.statements.size }
+                                ?.let { newContainer.statements[it] }
+                        }
+                        extractedLocalClasses.add(ExtractedLocalClass(declaration, newContainer, extractedUnder))
                     }
-
-                    currentParent as IrDeclarationContainer // IrClass or IrScript or IrPackageFragment
+                    is IrDeclarationContainer -> extractedLocalClasses.add(ExtractedLocalClass(declaration, newContainer, extractedUnder))
+                    else -> error("Inexpected container type $newContainer")
                 }
 
-                extractedLocalClasses.add(declaration to newContainer)
                 return IrCompositeImpl(declaration.startOffset, declaration.endOffset, context.irBuiltIns.unitType)
             }
         }, null)
 
-        for ((local, newContainer) in extractedLocalClasses) {
-            newContainer.addChild(local)
+        for ((local, newContainer, extractedUnder) in extractedLocalClasses) {
+            when (newContainer) {
+                is IrStatementContainer -> {
+                    val insertIndex = extractedUnder?.let { newContainer.statements.indexOf(it) } ?: -1
+                    if (insertIndex >= 0) {
+                        newContainer.statements.add(insertIndex, local)
+                    } else {
+                        newContainer.statements.add(local)
+                    }
+                    local.setDeclarationsParent(newContainer)
+                }
+                is IrDeclarationContainer -> {
+                    newContainer.addChild(local)
+                }
+                else -> error("Inexpected container type $newContainer")
+            }
 
             local.acceptVoid(object : IrElementVisitorVoid {
                 override fun visitElement(element: IrElement) {

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

@@ -19,9 +19,9 @@ package org.jetbrains.kotlin.backend.common.lower
 import org.jetbrains.kotlin.backend.common.BackendContext
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.descriptors.impl.ValueParameterDescriptorImpl
-import org.jetbrains.kotlin.ir.ObsoleteDescriptorBasedAPI
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.IrStatement
+import org.jetbrains.kotlin.ir.ObsoleteDescriptorBasedAPI
 import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.builders.*
 import org.jetbrains.kotlin.ir.declarations.*
@@ -162,6 +162,12 @@ open class IrBuildingTransformer(private val context: BackendContext) : IrElemen
             return super.visitEnumEntry(declaration)
         }
     }
+
+    override fun visitScript(declaration: IrScript): IrStatement {
+        withBuilder(declaration.symbol) {
+            return super.visitScript(declaration)
+        }
+    }
 }
 
 fun IrConstructor.callsSuper(irBuiltIns: IrBuiltIns): Boolean {