[JS IR] Migrate on body lowering pass and declaration transformer

^KT-43222 fixed

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/DeclarationOrigins.kt

@@ -14,5 +14,5 @@ object JsLoweredDeclarationOrigin : IrDeclarationOrigin {
     object JS_CLOSURE_BOX_CLASS : IrStatementOriginImpl("JS_CLOSURE_BOX_CLASS")
     object JS_CLOSURE_BOX_CLASS_DECLARATION : IrDeclarationOriginImpl("JS_CLOSURE_BOX_CLASS_DECLARATION")
     object BRIDGE_TO_EXTERNAL_FUNCTION : IrDeclarationOriginImpl("BRIDGE_TO_EXTERNAL_FUNCTION")
-    object OBJECT_GET_INSTANCE_FUNCTION : IrDeclarationOriginImpl("BRIDGE_TO_EXTERNAL_FUNCTION")
+    object OBJECT_GET_INSTANCE_FUNCTION : IrDeclarationOriginImpl("OBJECT_GET_INSTANCE_FUNCTION")
 }

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/JsIrBackendContext.kt

@@ -45,6 +45,8 @@ class JsIrBackendContext(
     override val scriptMode: Boolean = false,
     override val es6mode: Boolean = false
 ) : JsCommonBackendContext {
+    val fileToInitialisationFuns: MutableMap<IrFile, IrSimpleFunction> = mutableMapOf()
+
     override val extractedLocalClasses: MutableSet<IrClass> = hashSetOf()
 
     override val builtIns = module.builtIns

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/JsLoweringPhases.kt

@@ -369,12 +369,18 @@ private val forLoopsLoweringPhase = makeBodyLoweringPhase(
     description = "[Optimization] For loops lowering"
 )
 
-private val propertyLazyInitLoweringPhase = makeFileLoweringPhase(
+private val propertyLazyInitLoweringPhase = makeBodyLoweringPhase(
     ::PropertyLazyInitLowering,
     name = "PropertyLazyInitLowering",
     description = "Make property init as lazy"
 )
 
+private val removeInitializersForLazyProperties = makeDeclarationTransformerPhase(
+    ::RemoveInitializersForLazyProperties,
+    name = "RemoveInitializersForLazyProperties",
+    description = "Make property init as lazy"
+)
+
 private val propertyAccessorInlinerLoweringPhase = makeBodyLoweringPhase(
     ::PropertyAccessorInlineLowering,
     name = "PropertyAccessorInlineLowering",
@@ -751,6 +757,7 @@ val loweringList = listOf<Lowering>(
     forLoopsLoweringPhase,
     primitiveCompanionLoweringPhase,
     propertyLazyInitLoweringPhase,
+    removeInitializersForLazyProperties,
     propertyAccessorInlinerLoweringPhase,
     foldConstantLoweringPhase,
     privateMembersLoweringPhase,

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/JsMapping.kt

@@ -8,6 +8,8 @@ package org.jetbrains.kotlin.ir.backend.js
 import org.jetbrains.kotlin.backend.common.DefaultMapping
 import org.jetbrains.kotlin.backend.common.Mapping
 import org.jetbrains.kotlin.ir.declarations.*
+import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.expressions.IrExpressionBody
 
 class JsMapping : DefaultMapping() {
     val outerThisFieldSymbols = newMapping<IrClass, IrField>()
@@ -30,6 +32,8 @@ class JsMapping : DefaultMapping() {
     val enumEntryToCorrespondingField = newMapping<IrEnumEntry, IrField>()
     val enumClassToInitEntryInstancesFun = newMapping<IrClass, IrSimpleFunction>()
 
+    val lazyInitialisedFields = newMapping<IrField, IrExpression>()
+
     // Triggers `StageController.lazyLower` on access
     override fun <K : IrDeclaration, V> newMapping(): Mapping.Delegate<K, V> = object : Mapping.Delegate<K, V>() {
         private val map: MutableMap<K, V> = mutableMapOf()

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/PropertyLazyInitLowering.kt

@@ -5,27 +5,31 @@
 
 package org.jetbrains.kotlin.ir.backend.js.lower
 
-import org.jetbrains.kotlin.backend.common.FileLoweringPass
+import org.jetbrains.kotlin.backend.common.BodyLoweringPass
+import org.jetbrains.kotlin.backend.common.DeclarationTransformer
 import org.jetbrains.kotlin.backend.common.ir.isTopLevel
 import org.jetbrains.kotlin.descriptors.DescriptorVisibilities.INTERNAL
 import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.backend.js.JsIrBackendContext
+import org.jetbrains.kotlin.ir.backend.js.JsLoweredDeclarationOrigin
 import org.jetbrains.kotlin.ir.backend.js.ir.JsIrArithBuilder
 import org.jetbrains.kotlin.ir.backend.js.ir.JsIrBuilder
 import org.jetbrains.kotlin.ir.backend.js.utils.isPure
 import org.jetbrains.kotlin.ir.builders.declarations.addFunction
 import org.jetbrains.kotlin.ir.builders.declarations.buildField
 import org.jetbrains.kotlin.ir.declarations.*
+import org.jetbrains.kotlin.ir.declarations.persistent.PersistentIrElementBase
+import org.jetbrains.kotlin.ir.declarations.persistent.carriers.DeclarationCarrier
 import org.jetbrains.kotlin.ir.expressions.*
-import org.jetbrains.kotlin.ir.util.statements
-import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
-import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 import org.jetbrains.kotlin.name.Name
 import kotlin.collections.component1
 import kotlin.collections.component2
 
-class PropertyLazyInitLowering(private val context: JsIrBackendContext) : FileLoweringPass {
+class PropertyLazyInitLowering(
+    private val context: JsIrBackendContext
+) : BodyLoweringPass {
+
     private val irBuiltIns
         get() = context.irBuiltIns
 
@@ -34,31 +38,65 @@ class PropertyLazyInitLowering(private val context: JsIrBackendContext) : FileLo
     private val irFactory
         get() = context.irFactory
 
-    override fun lower(irFile: IrFile) {
-        val functions = TopLevelFunsSearcher()
-            .search(irFile)
+    val fileToInitialisationFuns
+        get() = context.fileToInitialisationFuns
 
-        val fieldToInitializer = calculateFieldToExpression(
-            functions
+    override fun lower(irBody: IrBody, container: IrDeclaration) {
+        val property = container.correspondingProperty ?: return
+
+        val topLevelProperty = property
+            .takeIf { it.isForLazyInit() }
+            ?.takeIf { it.backingField?.initializer != null }
+            ?: return
+
+        val file = topLevelProperty.parent as? IrFile
+            ?: return
+
+        val initFun = fileToInitialisationFuns[file]
+            ?: createInitialisationFunction(file)?.also { fileToInitialisationFuns[file] = it }
+            ?: return
+
+        val initialisationCall = JsIrBuilder.buildCall(
+            target = initFun.symbol,
+            type = initFun.returnType
         )
 
-        val allPropertyInitializersPure = fieldToInitializer
-            .all { it.value.isPure(anyVariable = true) }
+        when (container) {
+            is IrSimpleFunction ->
+                irBody.addInitialisation(initialisationCall, container)
+            is IrField -> {
+                property
+                    .let { listOf(it.getter, it.setter) }
+                    .filterNotNull()
+                    .forEach {
+                        irBody.addInitialisation(initialisationCall, it)
+                    }
+            }
+        }
+    }
 
-        if (allPropertyInitializersPure) return
+    private fun createInitialisationFunction(
+        file: IrFile
+    ): IrSimpleFunction? {
+        val fileName = file.name
 
-        fieldToInitializer.onEach { it.key.initializer = null }
+        val declarations = ArrayList(file.declarations)
+
+        val fieldToInitializer = calculateFieldToExpression(
+            declarations
+        )
 
-        if (fieldToInitializer.isEmpty()) return
+//        if (allFieldsInFilePure(fieldToInitializer.values)) {
+//            return null
+//        }
 
-        val fileName = irFile.name
         val initialisedField = irFactory.createInitialisationField(fileName)
             .apply {
-                irFile.declarations.add(this)
-                parent = irFile
+                file.declarations.add(this)
+                parent = file
             }
 
-        val initialisationFun = irFactory.addFunction(irFile) {
+        return irFactory.addFunction(file) {
             name = Name.identifier("init properties $fileName")
             returnType = irBuiltIns.unitType
             visibility = INTERNAL
@@ -69,16 +107,6 @@ class PropertyLazyInitLowering(private val context: JsIrBackendContext) : FileLo
                 initialisedField
             )
         }
-
-        functions
-            .asSequence()
-            .filterNotNull()
-            .forEach { function ->
-                val newBody = function.body?.let { body ->
-                    irFactory.bodyWithFunctionCall(body, initialisationFun)
-                }
-                function.body = newBody
-            }
     }
 
     private fun IrFactory.createInitialisationField(fileName: String): IrField =
@@ -94,7 +122,6 @@ class PropertyLazyInitLowering(private val context: JsIrBackendContext) : FileLo
         initializers: Map<IrField, IrExpression>,
         initialisedField: IrField
     ) {
-
         body = irFactory.createBlockBody(
             UNDEFINED_OFFSET,
             UNDEFINED_OFFSET,
@@ -105,7 +132,7 @@ class PropertyLazyInitLowering(private val context: JsIrBackendContext) : FileLo
     private fun buildBodyWithIfGuard(
         initializers: Map<IrField, IrExpression>,
         initialisedField: IrField
-    ): List<IrWhen> {
+    ): List<IrStatement> {
         val statements = initializers
             .map { (field, expression) ->
                 createIrSetField(field, expression)
@@ -127,18 +154,28 @@ class PropertyLazyInitLowering(private val context: JsIrBackendContext) : FileLo
     }
 }
 
-private fun calculateFieldToExpression(functions: Collection<IrSimpleFunction>): Map<IrField, IrExpression> =
-    functions
-        .asSequence()
-        .mapNotNull { it.correspondingPropertySymbol }
-        .map { it.owner }
-        .filter { it.isTopLevel }
-        .filterNot { it.isConst }
-        .distinct()
-        .mapNotNull { it.backingField }
-        .filter { it.initializer != null }
-        .map { it to it.initializer!!.expression }
-        .toMap()
+private fun IrBody.addInitialisation(
+    initCall: IrCall,
+    container: IrSimpleFunction
+) {
+    when (this) {
+        is IrExpressionBody -> {
+            expression = JsIrBuilder.buildComposite(
+                type = container.returnType,
+                statements = listOf(
+                    initCall,
+                    expression
+                )
+            )
+        }
+        is IrBlockBody -> {
+            statements.add(
+                0,
+                initCall
+            )
+        }
+    }
+}
 
 private fun createIrGetField(field: IrField): IrGetField {
     return JsIrBuilder.buildGetField(
@@ -156,35 +193,70 @@ private fun createIrSetField(field: IrField, expression: IrExpression): IrSetFie
     )
 }
 
-private fun IrFactory.bodyWithFunctionCall(
-    body: IrBody,
-    functionToCall: IrSimpleFunction
-): IrBody = createBlockBody(
-    body.startOffset,
-    body.endOffset,
-    mutableListOf<IrStatement>(
-        JsIrBuilder.buildCall(
-            target = functionToCall.symbol,
-            type = functionToCall.returnType
-        )
-    ).apply { addAll(body.statements) }
-)
+private fun allFieldsInFilePure(fieldToInitializer: Collection<IrExpression>) =
+    fieldToInitializer.all { it.isPure(anyVariable = true) }
 
-private class TopLevelFunsSearcher : IrElementTransformerVoid() {
+class RemoveInitializersForLazyProperties(
+    private val context: JsIrBackendContext
+) : DeclarationTransformer {
 
-    private val topLevelFuns = mutableSetOf<IrSimpleFunction>()
+    val fileToInitialisationFuns
+        get() = context.fileToInitialisationFuns
 
-    fun search(irFile: IrFile): Set<IrSimpleFunction> {
-        irFile.transformChildrenVoid(this)
-        return topLevelFuns
+    override fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {
+//        val file = declaration.parent as? IrFile ?: return null
+//        val declarations = ArrayList(file.declarations)
+//
+//        val fields = calculateFieldToExpression(declarations).values
+//
+//        if (allFieldsInFilePure(fields)) {
+//            return null
+//        }
+
+        declaration.correspondingProperty
+            ?.takeIf { it.isForLazyInit() }
+            ?.backingField
+            ?.let { it.initializer = null }
+
+        return null
     }
+}
+
+private fun calculateFieldToExpression(declarations: Collection<IrDeclaration>): Map<IrField, IrExpression> =
+    declarations
+        .asSequence()
+        .map { it.correspondingProperty }
+        .filterNotNull()
+        .filter { it.isForLazyInit() }
+        .distinct()
+        .mapNotNull { it.backingField }
+        .filter { it.initializer != null }
+        .map { it to it.initializer!!.expression }
+        .toMap()
 
-    override fun visitSimpleFunction(declaration: IrSimpleFunction): IrStatement {
+private fun IrProperty.isForLazyInit() = isTopLevel && !isConst
 
-        if (declaration.isTopLevel) {
-            topLevelFuns.add(declaration)
-        }
+private val IrDeclaration.correspondingProperty: IrProperty?
+    get() {
+        if (this !is IrSimpleFunction && this !is IrField && this !is IrProperty)
+            return null
 
-        return super.visitSimpleFunction(declaration)
+        // Objects are in fact fields and we need to ignore them
+        val originField = (this as? DeclarationCarrier)?.originField
+        if (originField == JsLoweredDeclarationOrigin.OBJECT_GET_INSTANCE_FUNCTION || originField == IrDeclarationOrigin.FIELD_FOR_OBJECT_INSTANCE)
+            return null
+
+        return when (this) {
+            is IrProperty -> this
+            is IrSimpleFunction -> propertyWithPersistentSafe {
+                correspondingPropertySymbol?.owner
+            }
+            is IrField -> propertyWithPersistentSafe {
+                correspondingPropertySymbol?.owner
+            }
+            else -> error("Can be only IrProperty, IrSimpleFunction or IrField")
+        }
     }
-}
+
+private fun IrDeclaration.propertyWithPersistentSafe(transform: IrDeclaration.() -> IrProperty?): IrProperty? =
+    transform()

compiler/testData/codegen/box/properties/lazyInitializationPure.kt

@@ -7,12 +7,25 @@ val a = "A"
 // FILE: B.kt
 val b = "B".apply {}
 
-val c = "C"
+val c = b
+
+// FILE: C.kt
+val d = "D".apply {}
+
+val e = d
 
 // FILE: main.kt
 
 fun box(): String {
-    return if (js("a") == "A" && js("typeof b") == "undefined" && js("typeof c") == "undefined")
+    d
+    e
+    return if (
+        js("a") === "A" &&
+        js("typeof b") == "undefined" &&
+        js("typeof c") == "undefined" &&
+        js("d") === "D" &&
+        js("e") === "D"
+    )
         "OK"
-    else "fail"
+    else "a = ${js("a")}; typeof b = ${js("typeof b")}; typeof c = ${js("typeof c")}; d = ${js("d")}; e = ${js("e")}"
 }