rra/ilgonmic/lazy-init-with-js-fun

[JS IR] Use initializer of field to use in js fun

[JS IR] Use string folding to consider const-like string vals as pure

Merge-request: KT-MR-4677

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

@@ -26,6 +26,7 @@ import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.declarations.impl.IrExternalPackageFragmentImpl
 import org.jetbrains.kotlin.ir.declarations.impl.IrFileImpl
 import org.jetbrains.kotlin.ir.expressions.IrCall
+import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
 import org.jetbrains.kotlin.ir.symbols.IrClassifierSymbol
 import org.jetbrains.kotlin.ir.symbols.IrFileSymbol
@@ -61,6 +62,7 @@ class JsIrBackendContext(
 ) : JsCommonBackendContext {
     val fileToInitializationFuns: MutableMap<IrFile, IrSimpleFunction?> = mutableMapOf()
     val fileToInitializerPureness: MutableMap<IrFile, Boolean> = mutableMapOf()
+    val fieldToInitializer: MutableMap<IrField, IrExpression> = mutableMapOf()
 
     val extractedLocalClasses: MutableSet<IrClass> = hashSetOf()
 

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

@@ -134,7 +134,7 @@ private class JsCodeOutlineTransformer(
             return null
 
         val jsCodeArg = expression.getValueArgument(0) ?: error("Expected js code string")
-        val jsStatements = translateJsCodeIntoStatementList(jsCodeArg) ?: return null
+        val jsStatements = translateJsCodeIntoStatementList(jsCodeArg, backendContext) ?: return null
 
         // Collect used Kotlin local variables and parameters.
         val kotlinLocalsUsedInJs = mutableListOf<IrValueDeclaration>()

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

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.ir.backend.js.lower
 
 import org.jetbrains.kotlin.backend.common.BodyLoweringPass
 import org.jetbrains.kotlin.backend.common.DeclarationTransformer
+import org.jetbrains.kotlin.backend.common.ir.isPure
 import org.jetbrains.kotlin.backend.common.ir.isTopLevel
 import org.jetbrains.kotlin.descriptors.DescriptorVisibilities.INTERNAL
 import org.jetbrains.kotlin.ir.IrStatement
@@ -15,7 +16,6 @@ import org.jetbrains.kotlin.ir.backend.js.JsIrBackendContext
 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.prependFunctionCall
-import org.jetbrains.kotlin.backend.common.ir.isPure
 import org.jetbrains.kotlin.ir.builders.declarations.addFunction
 import org.jetbrains.kotlin.ir.builders.declarations.buildField
 import org.jetbrains.kotlin.ir.declarations.*
@@ -25,6 +25,7 @@ import org.jetbrains.kotlin.ir.util.hasAnnotation
 import org.jetbrains.kotlin.name.Name
 import kotlin.collections.component1
 import kotlin.collections.component2
+import kotlin.collections.set
 
 class PropertyLazyInitLowering(
     private val context: JsIrBackendContext
@@ -71,24 +72,11 @@ class PropertyLazyInitLowering(
 
         val initializationCall = JsIrBuilder.buildCall(
             target = initFun.symbol,
-            type = initFun.returnType
+            type = initFun.returnType,
+            origin = PROPERTY_INIT_FUN_CALL
         )
 
-        when (container) {
+        if (container is IrSimpleFunction) irBody.prependFunctionCall(initializationCall)
-            is IrSimpleFunction ->
-                irBody.prependFunctionCall(initializationCall)
-            is IrField -> {
-                container
-                    .correspondingProperty
-                    ?.takeIf { it.isForLazyInit() }
-                    ?.takeIf { it.backingField?.initializer != null }
-                    ?.let { listOf(it.getter, it.setter) }
-                    ?.filterNotNull()
-                    ?.forEach {
-                        irBody.prependFunctionCall(initializationCall)
-                    }
-            }
-        }
     }
 
     private fun createInitializationFunction(
@@ -173,6 +161,10 @@ class PropertyLazyInitLowering(
             )
         ).let { listOf(it) }
     }
+
+    companion object {
+        object PROPERTY_INIT_FUN_CALL : IrStatementOriginImpl("PROPERTY_INIT_FUN_CALL")
+    }
 }
 
 private fun createIrGetField(field: IrField): IrGetField {
@@ -191,8 +183,11 @@ private fun createIrSetField(field: IrField, expression: IrExpression): IrSetFie
     )
 }
 
-private fun allFieldsInFilePure(fieldToInitializer: Collection<IrExpression>) =
+private fun allFieldsInFilePure(fieldToInitializer: Collection<IrExpression>): Boolean =
-    fieldToInitializer.all { it.isPure(anyVariable = true) }
+    fieldToInitializer
+        .all { expression ->
+            expression.isPure(anyVariable = true)
+        }
 
 class RemoveInitializersForLazyProperties(
     private val context: JsIrBackendContext

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/transformers/irToJs/IrElementToJsExpressionTransformer.kt

@@ -192,8 +192,10 @@ class IrElementToJsExpressionTransformer : BaseIrElementToJsNodeTransformer<JsEx
 
     override fun visitCall(expression: IrCall, context: JsGenerationContext): JsExpression {
         if (context.checkIfJsCode(expression.symbol)) {
-            val statements = translateJsCodeIntoStatementList(expression.getValueArgument(0) ?: error("JsCode is expected"))
+            val statements = translateJsCodeIntoStatementList(
-                ?: error("Cannot compute js code for ${expression.render()}")
+                expression.getValueArgument(0) ?: error("JsCode is expected"),
+                context.staticContext.backendContext
+            ) ?: error("Cannot compute js code for ${expression.render()}")
 
             if (statements.isEmpty()) return JsPrefixOperation(JsUnaryOperator.VOID, JsIntLiteral(3)) // TODO: report warning or even error
 

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/transformers/irToJs/IrElementToJsStatementTransformer.kt

@@ -133,8 +133,10 @@ class IrElementToJsStatementTransformer : BaseIrElementToJsNodeTransformer<JsSta
 
     override fun visitCall(expression: IrCall, data: JsGenerationContext): JsStatement {
         if (data.checkIfJsCode(expression.symbol)) {
-            val statements = translateJsCodeIntoStatementList(expression.getValueArgument(0) ?: error("JsCode is expected"))
+            val statements = translateJsCodeIntoStatementList(
-                ?: error("Cannot compute js code for ${expression.render()}")
+                expression.getValueArgument(0) ?: error("JsCode is expected"),
+                data.staticContext.backendContext
+            ) ?: error("Cannot compute js code for ${expression.render()}")
             return when (statements.size) {
                 0 -> JsEmpty
                 1 -> statements.single().withSource(expression, data)

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/transformers/irToJs/jsCode.kt

@@ -8,6 +8,8 @@ package org.jetbrains.kotlin.ir.backend.js.transformers.irToJs
 import com.google.gwt.dev.js.ThrowExceptionOnErrorReporter
 import com.google.gwt.dev.js.rhino.CodePosition
 import org.jetbrains.kotlin.ir.IrElement
+import org.jetbrains.kotlin.ir.backend.js.JsIrBackendContext
+import org.jetbrains.kotlin.ir.backend.js.lower.PropertyLazyInitLowering
 import org.jetbrains.kotlin.ir.declarations.IrVariable
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
@@ -20,55 +22,10 @@ import org.jetbrains.kotlin.js.backend.ast.JsStatement
 import org.jetbrains.kotlin.js.parser.parseExpressionOrStatement
 
 // Returns null if constant expression could not be parsed
-fun translateJsCodeIntoStatementList(code: IrExpression): List<JsStatement>? {
+fun translateJsCodeIntoStatementList(code: IrExpression, context: JsIrBackendContext): List<JsStatement>? {
     // TODO: support proper symbol linkage and label clash resolution
 
-    fun foldString(expression: IrExpression): String? {
+    return parseJsCode(foldString(code, context) ?: return null)
-        val builder = StringBuilder()
-        var foldingFailed = false
-        expression.acceptVoid(object : IrElementVisitorVoid {
-            override fun visitElement(element: IrElement) {
-                foldingFailed = true
-            }
-
-            override fun visitGetValue(expression: IrGetValue) {
-                expression.symbol.owner.acceptVoid(this)
-            }
-
-            override fun visitVariable(declaration: IrVariable) {
-                declaration.initializer?.let {
-                    it.acceptVoid(this)
-                    return
-                }
-
-                super.visitVariable(declaration)
-            }
-
-            override fun visitGetField(expression: IrGetField) {
-                expression.symbol.owner.initializer?.expression?.acceptVoid(this)
-            }
-
-            override fun visitCall(expression: IrCall) {
-                if (expression.origin != IrStatementOrigin.PLUS) {
-                    return super.visitCall(expression)
-                }
-
-                expression.acceptChildrenVoid(this)
-            }
-
-            override fun <T> visitConst(expression: IrConst<T>) {
-                builder.append(expression.kind.valueOf(expression))
-            }
-
-            override fun visitStringConcatenation(expression: IrStringConcatenation) = expression.acceptChildrenVoid(this)
-        })
-
-        if (foldingFailed) return null
-
-        return builder.toString()
-    }
-
-    return parseJsCode(foldString(code) ?: return null)
 }
 
 fun parseJsCode(jsCode: String): List<JsStatement>? {
@@ -82,4 +39,74 @@ fun parseJsCode(jsCode: String): List<JsStatement>? {
     // TODO: write debug info, see how it's done in CallExpressionTranslator.parseJsCode
 
     return parseExpressionOrStatement(jsCode, ThrowExceptionOnErrorReporter, currentScope, CodePosition(0, 0), "<js-code>")
+}
+
+fun foldString(expression: IrExpression, context: JsIrBackendContext): String? {
+    val builder = StringBuilder()
+    var foldingFailed = false
+    expression.acceptVoid(object : IrElementVisitorVoid {
+        override fun visitElement(element: IrElement) {
+            foldingFailed = true
+        }
+
+        override fun visitGetValue(expression: IrGetValue) {
+            expression.symbol.owner.acceptVoid(this)
+        }
+
+        override fun visitVariable(declaration: IrVariable) {
+            declaration.initializer?.let {
+                it.acceptVoid(this)
+                return
+            }
+
+            super.visitVariable(declaration)
+        }
+
+        override fun visitGetField(expression: IrGetField) {
+            val owner = expression.symbol.owner
+            owner.initializer?.expression?.acceptVoid(this)
+                ?: context.fieldToInitializer[owner]?.acceptVoid(this)
+        }
+
+        override fun visitCall(expression: IrCall) {
+            val owner = expression.symbol.owner
+            return when {
+                expression.origin == IrStatementOrigin.PLUS ->
+                    expression.acceptChildrenVoid(this)
+                expression.origin == PropertyLazyInitLowering.Companion.PROPERTY_INIT_FUN_CALL -> {
+                    owner.body?.acceptChildrenVoid(InitFunVisitor(context))
+                    expression.acceptChildrenVoid(this)
+                }
+                owner == owner.correspondingPropertySymbol?.owner?.getter -> {
+                    owner.body?.acceptChildrenVoid(this)
+                    expression.acceptChildrenVoid(this)
+                }
+                else -> super.visitCall(expression)
+            }
+        }
+
+        override fun visitReturn(expression: IrReturn) {
+            expression.acceptChildrenVoid(this)
+        }
+
+        override fun <T> visitConst(expression: IrConst<T>) {
+            builder.append(expression.kind.valueOf(expression))
+        }
+
+        override fun visitStringConcatenation(expression: IrStringConcatenation) = expression.acceptChildrenVoid(this)
+    })
+
+    if (foldingFailed) return null
+
+    return builder.toString()
+}
+
+private class InitFunVisitor(private val context: JsIrBackendContext) : IrElementVisitorVoid {
+    override fun visitElement(element: IrElement) {
+        element.acceptChildrenVoid(this)
+    }
+
+    override fun visitSetField(expression: IrSetField) {
+        context.fieldToInitializer[expression.symbol.owner] = expression.value
+    }
 }