[IR BE] Make lateinit fields and variables nullable

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

@@ -18,29 +18,46 @@ package org.jetbrains.kotlin.backend.common.lower
 
 import org.jetbrains.kotlin.backend.common.CommonBackendContext
 import org.jetbrains.kotlin.backend.common.FileLoweringPass
+import org.jetbrains.kotlin.backend.common.descriptors.WrappedVariableDescriptor
 import org.jetbrains.kotlin.backend.common.ir.Symbols
 import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.builders.*
+import org.jetbrains.kotlin.ir.builders.declarations.buildField
 import org.jetbrains.kotlin.ir.declarations.*
-import org.jetbrains.kotlin.ir.expressions.IrCall
-import org.jetbrains.kotlin.ir.expressions.IrExpression
-import org.jetbrains.kotlin.ir.expressions.IrGetValue
-import org.jetbrains.kotlin.ir.expressions.IrPropertyReference
-import org.jetbrains.kotlin.ir.expressions.impl.IrBlockBodyImpl
-import org.jetbrains.kotlin.ir.expressions.impl.IrConstImpl
+import org.jetbrains.kotlin.ir.declarations.impl.IrVariableImpl
+import org.jetbrains.kotlin.ir.expressions.*
+import org.jetbrains.kotlin.ir.expressions.impl.*
+import org.jetbrains.kotlin.ir.symbols.impl.IrVariableSymbolImpl
 import org.jetbrains.kotlin.ir.types.isPrimitiveType
+import org.jetbrains.kotlin.ir.types.makeNullable
 import org.jetbrains.kotlin.ir.util.resolveFakeOverride
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 
 class LateinitLowering(val context: CommonBackendContext) : FileLoweringPass {
     override fun lower(irFile: IrFile) {
+        val nullableFields = mutableMapOf<IrField, IrField>()
+        val nullableVariables = mutableMapOf<IrVariable, IrVariable>()
+
+        // Transform declarations
         irFile.transformChildrenVoid(object : IrElementTransformerVoid() {
             override fun visitProperty(declaration: IrProperty): IrStatement {
                 declaration.transformChildrenVoid(this)
                 if (declaration.isLateinit && declaration.origin != IrDeclarationOrigin.FAKE_OVERRIDE) {
-                    transformGetter(declaration.backingField!!, declaration.getter!!)
+                    val oldField = declaration.backingField!!
+                    val newField = buildField {
+                        updateFrom(oldField)
+                        type = oldField.type.makeNullable()
+                        name = oldField.name
+                    }.also { newField ->
+                        newField.parent = oldField.parent
+                        declaration.backingField = newField
+                    }
+
+                    nullableFields[oldField] = newField
+
+                    transformGetter(newField, declaration.getter!!)
                 }
                 return declaration
             }
@@ -50,41 +67,27 @@ class LateinitLowering(val context: CommonBackendContext) : FileLoweringPass {
 
                 if (!declaration.isLateinit) return declaration
 
-                declaration.run { initializer = IrConstImpl.constNull(startOffset, endOffset, type) }
-
-                return declaration
-            }
-
-            override fun visitGetValue(expression: IrGetValue): IrExpression {
-                val irVar = expression.symbol.owner as? IrVariable ?: return expression
-
-                if (!irVar.isLateinit) return expression
-
-                val parent = irVar.parent as IrSymbolOwner
-
-                val irBuilder = context.createIrBuilder(parent.symbol, expression.startOffset, expression.endOffset)
-
-                return irBuilder.run {
-                    irIfThenElse(
-                        expression.type, irEqualsNull(irGet(irVar)),
-                        throwUninitializedPropertyAccessException(irVar.name.asString()),
-                        irGet(irVar)
-                    )
+                val descriptor = WrappedVariableDescriptor()
+                val type = declaration.type.makeNullable()
+                val newVar = IrVariableImpl(
+                    declaration.startOffset,
+                    declaration.endOffset,
+                    declaration.origin,
+                    IrVariableSymbolImpl(descriptor),
+                    declaration.name,
+                    type,
+                    true,
+                    false,
+                    true
+                ).also {
+                    descriptor.bind(it)
+                    it.parent = declaration.parent
+                    it.initializer = IrConstImpl.constNull(declaration.startOffset, declaration.endOffset, type)
                 }
-            }
 
-            override fun visitCall(expression: IrCall): IrExpression {
-                expression.transformChildrenVoid(this)
+                nullableVariables[declaration] = newVar
 
-                if (!Symbols.isLateinitIsInitializedPropertyGetter(expression.symbol)) return expression
-
-                val receiver = expression.extensionReceiver as IrPropertyReference
-
-                val property = receiver.getter?.owner?.resolveFakeOverride()?.correspondingProperty!!.also { assert(it.isLateinit) }
-
-                return expression.run { context.createIrBuilder(symbol, startOffset, endOffset) }.run {
-                    irNotEquals(irGetField(receiver.dispatchReceiver, property.backingField!!), irNull())
-                }
+                return newVar
             }
 
             private fun transformGetter(backingField: IrField, getter: IrFunction) {
@@ -111,6 +114,62 @@ class LateinitLowering(val context: CommonBackendContext) : FileLoweringPass {
                 }
             }
         })
+
+        // Transform usages
+        irFile.transformChildrenVoid(object : IrElementTransformerVoid() {
+            override fun visitGetValue(expression: IrGetValue): IrExpression {
+                val irVar = nullableVariables[expression.symbol.owner] ?: return expression
+
+                val parent = irVar.parent as IrSymbolOwner
+                val irBuilder = context.createIrBuilder(parent.symbol, expression.startOffset, expression.endOffset)
+
+                return irBuilder.run {
+                    irIfThenElse(
+                        expression.type, irEqualsNull(irGet(irVar)),
+                        throwUninitializedPropertyAccessException(irVar.name.asString()),
+                        irGet(irVar)
+                    )
+                }
+            }
+
+            override fun visitSetVariable(expression: IrSetVariable): IrExpression {
+                expression.transformChildrenVoid(this)
+                val newVar = nullableVariables[expression.symbol.owner] ?: return expression
+                return with(expression) {
+                    IrSetVariableImpl(startOffset, endOffset, type, newVar.symbol, value, origin)
+                }
+            }
+
+            override fun visitGetField(expression: IrGetField): IrExpression {
+                expression.transformChildrenVoid(this)
+                val newField = nullableFields[expression.symbol.owner] ?: return expression
+                return with(expression) {
+                    IrGetFieldImpl(startOffset, endOffset, newField.symbol, newField.type, receiver, origin, superQualifierSymbol)
+                }
+            }
+
+            override fun visitSetField(expression: IrSetField): IrExpression {
+                expression.transformChildrenVoid(this)
+                val newField = nullableFields[expression.symbol.owner] ?: return expression
+                return with(expression) {
+                    IrSetFieldImpl(startOffset, endOffset, newField.symbol, receiver, value, type, origin, superQualifierSymbol)
+                }
+            }
+
+            override fun visitCall(expression: IrCall): IrExpression {
+                expression.transformChildrenVoid(this)
+
+                if (!Symbols.isLateinitIsInitializedPropertyGetter(expression.symbol)) return expression
+
+                val receiver = expression.extensionReceiver as IrPropertyReference
+
+                val property = receiver.getter?.owner?.resolveFakeOverride()?.correspondingProperty!!.also { assert(it.isLateinit) }
+
+                return expression.run { context.createIrBuilder(symbol, startOffset, endOffset) }.run {
+                    irNotEquals(irGetField(receiver.dispatchReceiver, property.backingField!!), irNull())
+                }
+            }
+        })
     }
 
     private fun IrBuilderWithScope.throwUninitializedPropertyAccessException(name: String) =