[WASM] Lazy properties initialization

compiler/cli/cli-js/src/org/jetbrains/kotlin/cli/js/K2JsIrCompiler.kt

@@ -317,13 +317,15 @@ class K2JsIrCompiler : CLICompiler<K2JSCompilerArguments>() {
                 sourceModule
             }
 
+
             if (arguments.wasm) {
                 val res = compileWasm(
                     module,
                     PhaseConfig(wasmPhases),
                     IrFactoryImpl,
                     exportedDeclarations = setOf(FqName("main")),
-                    emitNameSection  = arguments.wasmDebug,
+                    emitNameSection = arguments.wasmDebug,
+                    propertyLazyInitialization = arguments.irPropertyLazyInitialization,
                 )
                 val outputWasmFile = outputFile.withReplacedExtensionOrNull(outputFile.extension, "wasm")!!
                 outputWasmFile.writeBytes(res.wasm)

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

@@ -30,6 +30,7 @@ interface JsCommonBackendContext : CommonBackendContext {
     override val mapping: JsMapping
 
     val reflectionSymbols: ReflectionSymbols
+    val propertyLazyInitialization: PropertyLazyInitialization
 
     override val inlineClassesUtils: JsCommonInlineClassesUtils
 

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

@@ -294,8 +294,6 @@ class JsIntrinsics(private val irBuiltIns: IrBuiltIns, val context: JsIrBackendC
     val doNotIntrinsifyAnnotationSymbol = context.symbolTable.referenceClass(context.getJsInternalClass("DoNotIntrinsify"))
     val jsFunAnnotationSymbol = context.symbolTable.referenceClass(context.getJsInternalClass("JsFun"))
 
-    val jsEagerInitializationAnnotationSymbol = context.symbolTable.referenceClass(context.getJsInternalClass("EagerInitialization"))
-
     // TODO move CharSequence-related stiff to IntrinsifyCallsLowering
     val charSequenceClassSymbol = context.symbolTable.referenceClass(context.getClass(FqName("kotlin.CharSequence")))
     val charSequenceLengthPropertyGetterSymbol by context.lazy2 {

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

@@ -56,7 +56,7 @@ class JsIrBackendContext(
     override val scriptMode: Boolean = false,
     override val es6mode: Boolean = false,
     val dceRuntimeDiagnostic: RuntimeDiagnostic? = null,
-    val propertyLazyInitialization: Boolean = false,
+    propertyLazyInitialization: Boolean = false,
     val baseClassIntoMetadata: Boolean = false,
     val safeExternalBoolean: Boolean = false,
     val safeExternalBooleanDiagnostic: RuntimeDiagnostic? = null,
@@ -65,8 +65,6 @@ class JsIrBackendContext(
     val icCompatibleIr2Js: Boolean = false,
 ) : JsCommonBackendContext {
 
-    val fileToInitializationFuns: MutableMap<IrFile, IrSimpleFunction?> = mutableMapOf()
-    val fileToInitializerPureness: MutableMap<IrFile, Boolean> = mutableMapOf()
     val fieldToInitializer: MutableMap<IrField, IrExpression> = mutableMapOf()
 
     val localClassNames: MutableMap<IrClass, String> = mutableMapOf()
@@ -142,6 +140,11 @@ class JsIrBackendContext(
     val intrinsics: JsIntrinsics = JsIntrinsics(irBuiltIns, this)
     override val reflectionSymbols: ReflectionSymbols get() = intrinsics.reflectionSymbols
 
+    override val propertyLazyInitialization: PropertyLazyInitialization = PropertyLazyInitialization(
+        enabled = propertyLazyInitialization,
+        eagerInitialization = symbolTable.referenceClass(getJsInternalClass("EagerInitialization"))
+    )
+
     override val catchAllThrowableType: IrType
         get() = dynamicType
 
@@ -207,7 +210,8 @@ class JsIrBackendContext(
 
             override val getContinuation = symbolTable.referenceSimpleFunction(getJsInternalFunction("getContinuation"))
 
-            override val coroutineContextGetter = symbolTable.referenceSimpleFunction(context.coroutineSymbols.coroutineContextProperty.getter!!)
+            override val coroutineContextGetter =
+                symbolTable.referenceSimpleFunction(context.coroutineSymbols.coroutineContextProperty.getter!!)
 
             override val suspendCoroutineUninterceptedOrReturn =
                 symbolTable.referenceSimpleFunction(getJsInternalFunction(COROUTINE_SUSPEND_OR_RETURN_JS_NAME))

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

@@ -0,0 +1,16 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.ir.backend.js
+
+import org.jetbrains.kotlin.ir.declarations.IrFile
+import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
+import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
+
+class PropertyLazyInitialization(val enabled: Boolean, eagerInitialization: IrClassSymbol) {
+    val fileToInitializationFuns: MutableMap<IrFile, IrSimpleFunction?> = mutableMapOf()
+    val fileToInitializerPureness: MutableMap<IrFile, Boolean> = mutableMapOf()
+    val eagerInitialization: IrClassSymbol = eagerInitialization
+}

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

@@ -13,8 +13,7 @@ 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.JsCommonBackendContext
-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.ir.builders.declarations.addFunction
@@ -28,25 +27,23 @@ import kotlin.collections.component2
 import kotlin.collections.set
 
 class PropertyLazyInitLowering(
-    private val context: JsIrBackendContext
+    private val context: JsCommonBackendContext
 ) : BodyLoweringPass {
 
     private val irBuiltIns
         get() = context.irBuiltIns
 
-    private val calculator = JsIrArithBuilder(context)
-
     private val irFactory
         get() = context.irFactory
 
     private val fileToInitializationFuns
-        get() = context.fileToInitializationFuns
+        get() = context.propertyLazyInitialization.fileToInitializationFuns
 
     private val fileToInitializerPureness
-        get() = context.fileToInitializerPureness
+        get() = context.propertyLazyInitialization.fileToInitializerPureness
 
     override fun lower(irBody: IrBody, container: IrDeclaration) {
-        if (!context.propertyLazyInitialization) {
+        if (!context.propertyLazyInitialization.enabled) {
             return
         }
 
@@ -132,32 +129,35 @@ class PropertyLazyInitLowering(
         body = irFactory.createBlockBody(
             UNDEFINED_OFFSET,
             UNDEFINED_OFFSET,
-            buildBodyWithIfGuard(initializers, initializedField)
+            listOf(buildBodyWithIfGuard(initializers, initializedField))
         )
     }
 
     private fun buildBodyWithIfGuard(
         initializers: Map<IrField, IrExpression>,
         initializedField: IrField
-    ): List<IrStatement> {
+    ): IrStatement {
-        val statements = initializers
+        val statements = buildList<IrStatement> {
-            .map { (field, expression) ->
+            val upGuard = createIrSetField(
-                createIrSetField(field, expression)
+                initializedField,
+                JsIrBuilder.buildBoolean(context.irBuiltIns.booleanType, true)
+            )
+            add(upGuard)
+            initializers.forEach { (field, expression) ->
+                add(createIrSetField(field, expression))
             }
-
+            add(JsIrBuilder.buildBlock(irBuiltIns.unitType))
-        val upGuard = createIrSetField(
+        }
-            initializedField,
-            JsIrBuilder.buildBoolean(context.irBuiltIns.booleanType, true)
-        )
 
         return JsIrBuilder.buildIfElse(
             type = irBuiltIns.unitType,
-            cond = calculator.not(createIrGetField(initializedField)),
+            cond = createIrGetField(initializedField),
-            thenBranch = JsIrBuilder.buildComposite(
+            thenBranch = JsIrBuilder.buildBlock(irBuiltIns.unitType),
+            elseBranch = JsIrBuilder.buildComposite(
                 type = irBuiltIns.unitType,
-                statements = mutableListOf(upGuard).apply { addAll(statements) }
+                statements = statements
             )
-        ).let { listOf(it) }
+        )
     }
 
     companion object {
@@ -188,14 +188,14 @@ private fun allFieldsInFilePure(fieldToInitializer: Collection<IrExpression>): B
         }
 
 class RemoveInitializersForLazyProperties(
-    private val context: JsIrBackendContext
+    private val context: JsCommonBackendContext
 ) : DeclarationTransformer {
 
     private val fileToInitializerPureness
-        get() = context.fileToInitializerPureness
+        get() = context.propertyLazyInitialization.fileToInitializerPureness
 
     override fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {
-        if (!context.propertyLazyInitialization) {
+        if (!context.propertyLazyInitialization.enabled) {
             return null
         }
 
@@ -235,7 +235,10 @@ class RemoveInitializersForLazyProperties(
     }
 }
 
-private fun calculateFieldToExpression(declarations: Collection<IrDeclaration>, context: JsIrBackendContext): Map<IrField, IrExpression> =
+private fun calculateFieldToExpression(
+    declarations: Collection<IrDeclaration>,
+    context: JsCommonBackendContext
+): Map<IrField, IrExpression> =
     declarations
         .asSequence()
         .filter { it.isCompatibleDeclaration(context) }
@@ -276,9 +279,9 @@ private fun IrDeclaration.propertyWithPersistentSafe(transform: IrDeclaration.()
 private fun <T> IrDeclaration.withPersistentSafe(transform: IrDeclaration.() -> T?): T? =
     transform()
 
-private fun IrDeclaration.isCompatibleDeclaration(context: JsIrBackendContext) =
+private fun IrDeclaration.isCompatibleDeclaration(context: JsCommonBackendContext) =
     correspondingProperty?.let {
-        it.isForLazyInit() && !it.hasAnnotation(context.intrinsics.jsEagerInitializationAnnotationSymbol)
+        it.isForLazyInit() && !it.hasAnnotation(context.propertyLazyInitialization.eagerInitialization)
     } ?: true && withPersistentSafe { origin in compatibleOrigins } == true
 
 private val compatibleOrigins = listOf(

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/WasmBackendContext.kt

@@ -35,6 +35,7 @@ class WasmBackendContext(
     override val irBuiltIns: IrBuiltIns,
     val symbolTable: SymbolTable,
     val irModuleFragment: IrModuleFragment,
+    propertyLazyInitialization: Boolean,
     override val configuration: CompilerConfiguration,
 ) : JsCommonBackendContext {
     override val builtIns = module.builtIns
@@ -104,6 +105,9 @@ class WasmBackendContext(
     val wasmSymbols: WasmSymbols = WasmSymbols(this@WasmBackendContext, symbolTable)
     override val reflectionSymbols: ReflectionSymbols get() = wasmSymbols.reflectionSymbols
 
+    override val propertyLazyInitialization: PropertyLazyInitialization =
+        PropertyLazyInitialization(enabled = propertyLazyInitialization, eagerInitialization = wasmSymbols.eagerInitialization)
+
     override val ir = object : Ir<WasmBackendContext>(this, irModuleFragment) {
         override val symbols: Symbols<WasmBackendContext> = wasmSymbols
         override fun shouldGenerateHandlerParameterForDefaultBodyFun() = true

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/WasmLoweringPhases.kt

@@ -496,6 +496,18 @@ private val forLoopsLoweringPhase = makeWasmModulePhase(
     description = "[Optimization] For loops lowering"
 )
 
+private val propertyLazyInitLoweringPhase = makeWasmModulePhase(
+    ::PropertyLazyInitLowering,
+    name = "PropertyLazyInitLowering",
+    description = "Make property init as lazy"
+)
+
+private val removeInitializersForLazyProperties = makeWasmModulePhase(
+    ::RemoveInitializersForLazyProperties,
+    name = "RemoveInitializersForLazyProperties",
+    description = "Remove property initializers if they was initialized lazily"
+)
+
 private val propertyAccessorInlinerLoweringPhase = makeWasmModulePhase(
     ::PropertyAccessorInlineLowering,
     name = "PropertyAccessorInlineLowering",
@@ -578,6 +590,8 @@ val wasmPhases = NamedCompilerPhase(
             returnableBlockLoweringPhase then
 
             forLoopsLoweringPhase then
+            propertyLazyInitLoweringPhase then
+            removeInitializersForLazyProperties then
             propertyAccessorInlinerLoweringPhase then
             stringConcatenationLowering then
 

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/WasmSymbols.kt

@@ -63,6 +63,8 @@ class WasmSymbols(
 
     internal val reflectionSymbols: WasmReflectionSymbols = WasmReflectionSymbols()
 
+    internal val eagerInitialization: IrClassSymbol = getIrClass(FqName("kotlin.EagerInitialization"))
+
     override val throwNullPointerException = getInternalFunction("THROW_NPE")
     override val throwISE = getInternalFunction("THROW_ISE")
     override val throwTypeCastException = getInternalFunction("THROW_CCE")

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/compiler.kt

@@ -29,6 +29,7 @@ fun compileWasm(
     phaseConfig: PhaseConfig,
     irFactory: IrFactory,
     exportedDeclarations: Set<FqName> = emptySet(),
+    propertyLazyInitialization: Boolean,
     emitNameSection: Boolean = false,
 ): WasmCompilerResult {
     val mainModule = depsDescriptors.mainModule
@@ -46,7 +47,7 @@ fun compileWasm(
     }
 
     val moduleDescriptor = moduleFragment.descriptor
-    val context = WasmBackendContext(moduleDescriptor, irBuiltIns, symbolTable, moduleFragment, configuration)
+    val context = WasmBackendContext(moduleDescriptor, irBuiltIns, symbolTable, moduleFragment, propertyLazyInitialization, configuration)
 
     // Load declarations referenced during `context` initialization
     allModules.forEach {

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

@@ -1,5 +1,4 @@
 // IGNORE_BACKEND: JS
-// IGNORE_BACKEND: WASM
 // PROPERTY_LAZY_INITIALIZATION
 
 // FILE: A.kt

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

@@ -1,6 +1,5 @@
 // IGNORE_BACKEND: JS
 // IGNORE_BACKEND: NATIVE
-// IGNORE_BACKEND: WASM
 // PROPERTY_LAZY_INITIALIZATION
 
 // FILE: A.kt

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

@@ -1,5 +1,4 @@
 // IGNORE_BACKEND: JS
-// IGNORE_BACKEND: WASM
 // SPLIT_PER_MODULE
 // PROPERTY_LAZY_INITIALIZATION
 

js/js.tests/test/org/jetbrains/kotlin/js/testOld/BasicWasmBoxTest.kt

@@ -32,10 +32,8 @@ import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtNamedFunction
 import org.jetbrains.kotlin.psi.KtPsiFactory
 import org.jetbrains.kotlin.resolve.CompilerEnvironment
-import org.jetbrains.kotlin.test.DebugMode
+import org.jetbrains.kotlin.test.*
-import org.jetbrains.kotlin.test.Directives
+import org.jetbrains.kotlin.test.directives.JsEnvironmentConfigurationDirectives
-import org.jetbrains.kotlin.test.KotlinTestWithEnvironment
-import org.jetbrains.kotlin.test.TestFiles
 import org.jetbrains.kotlin.test.util.KtTestUtil
 import java.io.Closeable
 import java.io.File
@@ -178,11 +176,14 @@ abstract class BasicWasmBoxTest(
             AnalyzerWithCompilerReport(config.configuration)
         )
 
+        val directives = KotlinTestUtils.parseDirectives(FileUtil.loadFile(testFile))
+
         val compilerResult = compileWasm(
             sourceModule,
             phaseConfig = phaseConfig,
             irFactory = IrFactoryImpl,
             exportedDeclarations = setOf(FqName.fromSegments(listOfNotNull(testPackage, testFunction))),
+            propertyLazyInitialization = directives.contains(JsEnvironmentConfigurationDirectives.PROPERTY_LAZY_INITIALIZATION.name),
             emitNameSection = true,
         )
 

libraries/stdlib/wasm/src/kotlin/annotations.kt

@@ -3,11 +3,14 @@
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
-package kotlin.js
+package kotlin
 
+/**
+ * Forces a top-level property to be initialized eagerly, opposed to lazily on the first access to file and/or property.
+ */
 @ExperimentalStdlibApi
 @Retention(AnnotationRetention.BINARY)
 @Target(AnnotationTarget.PROPERTY)
 @SinceKotlin("1.6")
 @Deprecated("This annotation is a temporal migration assistance and may be removed in the future releases, please consider filing an issue about the case where it is needed")
-public annotation class EagerInitialization
+public annotation class EagerInitialization