Refactor LateinitLowering into DeclarationTransformer's and a BodyLoweringPass

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

@@ -12,6 +12,7 @@ import kotlin.reflect.KProperty
 interface Mapping {
     val defaultArgumentsDispatchFunction: Delegate<IrFunction, IrFunction>
     val defaultArgumentsOriginalFunction: Delegate<IrFunction, IrFunction>
+    val lateInitFieldToNullableField: Delegate<IrField, IrField>
     val inlineClassMemberToStatic: Delegate<IrFunction, IrSimpleFunction>
 
     abstract class Delegate<K : IrDeclaration, V> {
@@ -31,6 +32,7 @@ open class DefaultMapping : Mapping {
 
     override val defaultArgumentsDispatchFunction: Mapping.Delegate<IrFunction, IrFunction> = newMapping()
     override val defaultArgumentsOriginalFunction: Mapping.Delegate<IrFunction, IrFunction> = newMapping()
+    override val lateInitFieldToNullableField: Mapping.Delegate<IrField, IrField> = newMapping()
     override val inlineClassMemberToStatic: Mapping.Delegate<IrFunction, IrSimpleFunction> = newMapping()
 
     protected fun <K : IrDeclaration, V> newMapping() = object : Mapping.Delegate<K, V>() {

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

@@ -16,8 +16,7 @@
 
 package org.jetbrains.kotlin.backend.common.lower
 
-import org.jetbrains.kotlin.backend.common.CommonBackendContext
+import org.jetbrains.kotlin.backend.common.*
-import org.jetbrains.kotlin.backend.common.FileLoweringPass
 import org.jetbrains.kotlin.backend.common.ir.Symbols
 import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.builders.*
@@ -35,43 +34,77 @@ import org.jetbrains.kotlin.ir.util.resolveFakeOverride
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 
-class LateinitLowering(val backendContext: CommonBackendContext) : FileLoweringPass {
+class NullableFieldsForLateinitCreationLowering(val backendContext: CommonBackendContext) : DeclarationTransformer {
 
-    private val nullableFields = backendContext.lateinitNullableFields
+    override fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {
-    private fun buildOrGetNullableField(originalField: IrField): IrField {
+        if (declaration is IrField) {
-        if (originalField.type.isMarkedNullable()) return originalField
+            declaration.correspondingPropertySymbol?.owner?.let { property ->
-        return nullableFields.getOrPut(originalField) {
+                if (property.isRealLateinit) {
-            buildField {
+                    val newField = backendContext.buildOrGetNullableField(declaration)
-                updateFrom(originalField)
+                    if (declaration != newField && declaration.parent != property.parent) return listOf(newField)
-                type = originalField.type.makeNullable()
+                }
-                name = originalField.name
+            }
-            }.apply {
+        }
-                parent = originalField.parent
+        return null
-                correspondingPropertySymbol = originalField.correspondingPropertySymbol
+    }
-                annotations += originalField.annotations
+}
+
+// Transform declarations
+class NullableFieldsDeclarationLowering(val backendContext: CommonBackendContext) : DeclarationTransformer {
+
+    override fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {
+        when (declaration) {
+            is IrProperty -> {
+                if (declaration.isRealLateinit) {
+                    declaration.backingField = backendContext.buildOrGetNullableField(declaration.backingField!!)
+                }
+            }
+
+            is IrSimpleFunction -> {
+                declaration.correspondingPropertySymbol?.owner?.let { property ->
+                    if (declaration == property.getter && property.isRealLateinit) {
+                        // f = buildOrGetNullableField is idempotent, i.e. f(f(x)) == f(x)
+                        transformGetter(backendContext.buildOrGetNullableField(property.backingField!!), declaration)
+                    }
+                }
+            }
+        }
+
+        return null
+    }
+
+    private fun transformGetter(backingField: IrField, getter: IrFunction) {
+        val type = backingField.type
+        assert(!type.isPrimitiveType()) { "'lateinit' modifier is not allowed on primitive types" }
+        val startOffset = getter.startOffset
+        val endOffset = getter.endOffset
+        getter.body = IrBlockBodyImpl(startOffset, endOffset) {
+            val irBuilder = backendContext.createIrBuilder(getter.symbol, startOffset, endOffset)
+            irBuilder.run {
+                val resultVar = scope.createTemporaryVariable(
+                    irGetField(getter.dispatchReceiverParameter?.let { irGet(it) }, backingField)
+                )
+                resultVar.parent = getter
+                statements.add(resultVar)
+                val throwIfNull = irIfThenElse(
+                    context.irBuiltIns.nothingType,
+                    irNotEquals(irGet(resultVar), irNull()),
+                    irReturn(irGet(resultVar)),
+                    backendContext.throwUninitializedPropertyAccessException(this, backingField.name.asString())
+                )
+                statements.add(throwIfNull)
             }
         }
     }
+}
 
-    override fun lower(irFile: IrFile) {
+// Transform usages
+class LateinitUsageLowering(val backendContext: CommonBackendContext) : BodyLoweringPass {
 
+    override fun lower(irBody: IrBody, container: IrDeclaration) {
         val nullableVariables = mutableMapOf<IrVariable, IrVariable>()
 
-        // Transform declarations
+        irBody.transformChildrenVoid(object : IrElementTransformerVoid() {
-        irFile.transformChildrenVoid(object : IrElementTransformerVoid() {
-            override fun visitProperty(declaration: IrProperty): IrStatement {
-                declaration.transformChildrenVoid(this)
-                if (declaration.isLateinit && declaration.origin != IrDeclarationOrigin.FAKE_OVERRIDE) {
-                    val oldField = declaration.backingField!!
-                    val newField = buildOrGetNullableField(oldField)
-
-                    declaration.backingField = newField
-
-                    transformGetter(newField, declaration.getter!!)
-                }
-                return declaration
-            }
-
             override fun visitVariable(declaration: IrVariable): IrStatement {
                 declaration.transformChildrenVoid(this)
 
@@ -92,41 +125,17 @@ class LateinitLowering(val backendContext: CommonBackendContext) : FileLoweringP
                 ).also {
                     descriptor.bind(it)
                     it.parent = declaration.parent
-                    it.initializer = IrConstImpl.constNull(declaration.startOffset, declaration.endOffset, backendContext.irBuiltIns.nothingNType)
+                    it.initializer =
+                        IrConstImpl.constNull(declaration.startOffset, declaration.endOffset, backendContext.irBuiltIns.nothingNType)
                 }
 
                 nullableVariables[declaration] = newVar
 
                 return newVar
             }
-
-            private fun transformGetter(backingField: IrField, getter: IrFunction) {
-                val type = backingField.type
-                assert(!type.isPrimitiveType()) { "'lateinit' modifier is not allowed on primitive types" }
-                val startOffset = getter.startOffset
-                val endOffset = getter.endOffset
-                val irBuilder = backendContext.createIrBuilder(getter.symbol, startOffset, endOffset)
-                irBuilder.run {
-                    val body = IrBlockBodyImpl(startOffset, endOffset)
-                    val resultVar = scope.createTemporaryVariable(
-                        irGetField(getter.dispatchReceiverParameter?.let { irGet(it) }, backingField)
-                    )
-                    resultVar.parent = getter
-                    body.statements.add(resultVar)
-                    val throwIfNull = irIfThenElse(
-                        context.irBuiltIns.nothingType,
-                        irNotEquals(irGet(resultVar), irNull()),
-                        irReturn(irGet(resultVar)),
-                        backendContext.throwUninitializedPropertyAccessException(this, backingField.name.asString())
-                    )
-                    body.statements.add(throwIfNull)
-                    getter.body = body
-                }
-            }
         })
 
-        // Transform usages
+        irBody.transformChildrenVoid(object : IrElementTransformerVoid() {
-        irFile.transformChildrenVoid(object : IrElementTransformerVoid() {
             override fun visitGetValue(expression: IrGetValue): IrExpression {
                 val irVar = nullableVariables[expression.symbol.owner] ?: return expression
 
@@ -152,7 +161,7 @@ class LateinitLowering(val backendContext: CommonBackendContext) : FileLoweringP
 
             override fun visitGetField(expression: IrGetField): IrExpression {
                 expression.transformChildrenVoid(this)
-                val newField = nullableFields[expression.symbol.owner] ?: return expression
+                val newField = backendContext.mapping.lateInitFieldToNullableField[expression.symbol.owner] ?: return expression
                 return with(expression) {
                     IrGetFieldImpl(startOffset, endOffset, newField.symbol, newField.type, receiver, origin, superQualifierSymbol)
                 }
@@ -160,7 +169,7 @@ class LateinitLowering(val backendContext: CommonBackendContext) : FileLoweringP
 
             override fun visitSetField(expression: IrSetField): IrExpression {
                 expression.transformChildrenVoid(this)
-                val newField = nullableFields[expression.symbol.owner] ?: return expression
+                val newField = backendContext.mapping.lateInitFieldToNullableField[expression.symbol.owner] ?: return expression
                 return with(expression) {
                     IrSetFieldImpl(startOffset, endOffset, newField.symbol, receiver, value, type, origin, superQualifierSymbol)
                 }
@@ -177,7 +186,9 @@ class LateinitLowering(val backendContext: CommonBackendContext) : FileLoweringP
                     receiver.getter?.owner?.resolveFakeOverride()?.correspondingPropertySymbol!!.owner.also { assert(it.isLateinit) }
 
                 val nullableField =
-                    buildOrGetNullableField(property.backingField ?: error("Lateinit property is supposed to have backing field"))
+                    backendContext.buildOrGetNullableField(
+                        property.backingField ?: error("Lateinit property is supposed to have backing field")
+                    )
 
                 return expression.run { backendContext.createIrBuilder(symbol, startOffset, endOffset) }.run {
                     irNotEquals(irGetField(receiver.dispatchReceiver, nullableField), irNull())
@@ -186,3 +197,20 @@ class LateinitLowering(val backendContext: CommonBackendContext) : FileLoweringP
         })
     }
 }
+
+private fun CommonBackendContext.buildOrGetNullableField(originalField: IrField): IrField {
+    if (originalField.type.isMarkedNullable()) return originalField
+    return mapping.lateInitFieldToNullableField.getOrPut(originalField) {
+        buildField {
+            updateFrom(originalField)
+            type = originalField.type.makeNullable()
+            name = originalField.name
+        }.apply {
+            parent = originalField.parent
+            correspondingPropertySymbol = originalField.correspondingPropertySymbol
+            annotations += originalField.annotations
+        }
+    }
+}
+
+private val IrProperty.isRealLateinit get() = isLateinit && origin != IrDeclarationOrigin.FAKE_OVERRIDE

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

@@ -50,7 +50,8 @@ class JsIrBackendContext(
     override val transformedFunction
         get() = error("Use Mapping.inlineClassMemberToStatic instead")
 
-    override val lateinitNullableFields = mutableMapOf<IrField, IrField>()
+    override val lateinitNullableFields
+        get() = error("Use Mapping.lateInitFieldToNullableField instead")
 
     val memberMap = mutableMapOf<IrSimpleFunctionSymbol, IrSimpleFunction>()
 

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

@@ -73,9 +73,21 @@ private val expectDeclarationsRemovingPhase = makeJsModulePhase(
     description = "Remove expect declaration from module fragment"
 )
 
-private val lateinitLoweringPhase = makeJsModulePhase(
+private val lateinitNullableFieldsPhase = makeJsModulePhase(
-    ::LateinitLowering,
+    ::NullableFieldsForLateinitCreationLowering,
-    name = "LateinitLowering",
+    name = "LateinitNullableFields",
+    description = "Create nullable fields for lateinit properties"
+)
+
+private val lateinitDeclarationLoweringPhase = makeJsModulePhase(
+    ::NullableFieldsDeclarationLowering,
+    name = "LateinitDeclarations",
+    description = "Reference nullable fields from properties and getters + insert checks"
+)
+
+private val lateinitUsageLoweringPhase = makeJsModulePhase(
+    ::LateinitUsageLowering,
+    name = "LateinitUsage",
     description = "Insert checks for lateinit field references"
 )
 
@@ -422,7 +434,9 @@ val jsPhases = namedIrModulePhase(
             functionInliningPhase then
             createScriptFunctionsPhase then
             provisionalFunctionExpressionPhase then
-            lateinitLoweringPhase then
+            lateinitNullableFieldsPhase then
+            lateinitDeclarationLoweringPhase then
+            lateinitUsageLoweringPhase then
             tailrecLoweringPhase then
             enumClassConstructorLoweringPhase then
             sharedVariablesLoweringPhase then

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmLower.kt

@@ -79,9 +79,21 @@ private val expectDeclarationsRemovingPhase = makeIrModulePhase<JvmBackendContex
     description = "Remove expect declaration from module fragment"
 )
 
-private val lateinitPhase = makeIrFilePhase(
+private val lateinitNullableFieldsPhase = makeIrFilePhase(
-    ::LateinitLowering,
+    ::NullableFieldsForLateinitCreationLowering,
-    name = "Lateinit",
+    name = "LateinitNullableFields",
+    description = "Create nullable fields for lateinit properties"
+)
+
+private val lateinitDeclarationLoweringPhase = makeIrFilePhase(
+    ::NullableFieldsDeclarationLowering,
+    name = "LateinitDeclarations",
+    description = "Reference nullable fields from properties and getters + insert checks"
+)
+
+private val lateinitUsageLoweringPhase = makeIrFilePhase(
+    ::LateinitUsageLowering,
+    name = "LateinitUsage",
     description = "Insert checks for lateinit field references"
 )
 
@@ -241,7 +253,9 @@ private val jvmFilePhases =
         arrayConstructorPhase then
         checkNotNullPhase then
 
-        lateinitPhase then
+        lateinitNullableFieldsPhase then
+        lateinitDeclarationLoweringPhase then
+        lateinitUsageLoweringPhase then
 
         moveOrCopyCompanionObjectFieldsPhase then
         inlineCallableReferenceToLambdaPhase then

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

@@ -71,12 +71,25 @@ private val expectDeclarationsRemovingPhase = makeWasmModulePhase(
     description = "Remove expect declaration from module fragment"
 )
 
-private val lateinitLoweringPhase = makeWasmModulePhase(
+private val lateinitNullableFieldsPhase = makeWasmModulePhase(
-    ::LateinitLowering,
+    ::NullableFieldsForLateinitCreationLowering,
-    name = "LateinitLowering",
+    name = "LateinitNullableFields",
+    description = "Create nullable fields for lateinit properties"
+)
+
+private val lateinitDeclarationLoweringPhase = makeWasmModulePhase(
+    ::NullableFieldsDeclarationLowering,
+    name = "LateinitDeclarations",
+    description = "Reference nullable fields from properties and getters + insert checks"
+)
+
+private val lateinitUsageLoweringPhase = makeWasmModulePhase(
+    ::LateinitUsageLowering,
+    name = "LateinitUsage",
     description = "Insert checks for lateinit field references"
 )
 
+
 // TODO make all lambda-related stuff work with IrFunctionExpression and drop this phase
 private val provisionalFunctionExpressionPhase = makeWasmModulePhase(
     { ProvisionalFunctionExpressionLowering() },
@@ -329,7 +342,9 @@ val wasmPhases = namedIrModulePhase<WasmBackendContext>(
             // arrayConstructorPhase then
 
             functionInliningPhase then
-            lateinitLoweringPhase then
+            lateinitNullableFieldsPhase then
+            lateinitDeclarationLoweringPhase then
+            lateinitUsageLoweringPhase then
             tailrecLoweringPhase then
 
             enumClassConstructorLoweringPhase then