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

^KT-43222 fixed
2020-11-10 15:53:58 +03:00
parent 99d0740234
commit 06b276f9c3
6 changed files with 170 additions and 72 deletions
@@ -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")
 }
@@ -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
@@ -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,
@@ -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()
@@ -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)

-        if (fieldToInitializer.isEmpty()) return
+        val fieldToInitializer = calculateFieldToExpression(
+            declarations
+        )
+
+//        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
    }
+}

-    override fun visitSimpleFunction(declaration: IrSimpleFunction): IrStatement {
+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()

-        if (declaration.isTopLevel) {
-            topLevelFuns.add(declaration)
+private fun IrProperty.isForLazyInit() = isTopLevel && !isConst
+
+private val IrDeclaration.correspondingProperty: IrProperty?
+    get() {
+        if (this !is IrSimpleFunction && this !is IrField && this !is IrProperty)
+            return null
+
+        // 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")
        }
-
-        return super.visitSimpleFunction(declaration)
    }
-}
+
+private fun IrDeclaration.propertyWithPersistentSafe(transform: IrDeclaration.() -> IrProperty?): IrProperty? =
+    transform()
@@ -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")}"
 }