Make InlineClassDeclarationLowering a DeclarationTransformer

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

@@ -5,14 +5,14 @@
 
 package org.jetbrains.kotlin.backend.common
 
-import org.jetbrains.kotlin.ir.declarations.IrDeclaration
+import org.jetbrains.kotlin.ir.declarations.*
-import org.jetbrains.kotlin.ir.declarations.IrFunction
 import kotlin.reflect.KMutableProperty0
 import kotlin.reflect.KProperty
 
 interface Mapping {
     val defaultArgumentsDispatchFunction: Delegate<IrFunction, IrFunction>
     val defaultArgumentsOriginalFunction: Delegate<IrFunction, IrFunction>
+    val inlineClassMemberToStatic: Delegate<IrFunction, IrSimpleFunction>
 
     abstract class Delegate<K : IrDeclaration, V> {
         abstract operator fun get(key: K): V?
@@ -31,6 +31,7 @@ open class DefaultMapping : Mapping {
 
     override val defaultArgumentsDispatchFunction: Mapping.Delegate<IrFunction, IrFunction> = newMapping()
     override val defaultArgumentsOriginalFunction: Mapping.Delegate<IrFunction, IrFunction> = newMapping()
+    override val inlineClassMemberToStatic: Mapping.Delegate<IrFunction, IrSimpleFunction> = newMapping()
 
     protected fun <K : IrDeclaration, V> newMapping() = object : Mapping.Delegate<K, V>() {
         private val map: MutableMap<K, V> = mutableMapOf()

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

@@ -6,14 +6,16 @@
 package org.jetbrains.kotlin.backend.common.lower
 
 import org.jetbrains.kotlin.backend.common.BodyLoweringPass
-import org.jetbrains.kotlin.backend.common.ClassLoweringPass
 import org.jetbrains.kotlin.backend.common.CommonBackendContext
-import org.jetbrains.kotlin.backend.common.FileLoweringPass
+import org.jetbrains.kotlin.backend.common.DeclarationTransformer
+import org.jetbrains.kotlin.backend.common.getOrPut
 import org.jetbrains.kotlin.backend.common.ir.createStaticFunctionWithReceivers
 import org.jetbrains.kotlin.ir.IrStatement
+import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.builders.*
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.*
+import org.jetbrains.kotlin.ir.expressions.impl.IrBlockBodyImpl
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
@@ -22,40 +24,63 @@ import org.jetbrains.kotlin.name.Name
 
 private const val INLINE_CLASS_IMPL_SUFFIX = "-impl"
 
-// TODO: Support incremental compilation
 class InlineClassLowering(val context: CommonBackendContext) {
-    private val transformedFunction = if (context.scriptMode) context.transformedFunction else mutableMapOf()
+    private val transformedFunction = context.mapping.inlineClassMemberToStatic
 
-    val inlineClassDeclarationLowering = object : ClassLoweringPass {
+    val inlineClassDeclarationLowering = object : DeclarationTransformer {
-        override fun lower(irClass: IrClass) {
-            if (!irClass.isInline) return
 
-            irClass.transformDeclarationsFlat { declaration ->
+        override fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {
-                when (declaration) {
+            val irClass = declaration.parent as? IrClass ?: return null
-                    is IrConstructor -> listOf(transformConstructor(declaration))
+            if (!irClass.isInline) return null
+
+            return when (declaration) {
+                is IrConstructor -> transformConstructor(declaration)
                 is IrSimpleFunction -> transformMethodFlat(declaration)
-                    is IrProperty -> listOf(declaration)  // Getters and setters should be flattened
+                else -> null
-                    is IrField -> listOf(declaration)
-                    is IrClass -> listOf(declaration)
-                    else -> error("Unexpected declaration: $declaration")
-                }
             }
         }
 
-        private fun transformConstructor(irConstructor: IrConstructor): IrDeclaration {
+        private fun transformConstructor(irConstructor: IrConstructor): List<IrDeclaration>? {
-            if (irConstructor.isPrimary) return irConstructor
+            if (irConstructor.isPrimary) return null
 
             // Secondary constructors are lowered into static function
-            val result = transformedFunction.getOrPut(irConstructor.symbol) { createStaticBodilessMethod(irConstructor).symbol }.owner
+            val result = getOrCreateStaticMethod(irConstructor)
+
+            transformConstructorBody(irConstructor, result)
+
+            return listOf(result)
+        }
+
+
+        private fun transformMethodFlat(function: IrSimpleFunction): List<IrDeclaration>? {
+            // TODO: Support fake-overridden methods without boxing
+            if (function.isStaticMethodOfClass || !function.isReal)
+                return null
+
+            val staticMethod = getOrCreateStaticMethod(function)
+
+            transformMethodBodyFlat(function, staticMethod)
+            function.body = delegateToStaticMethod(function, staticMethod)
+
+            if (function.overriddenSymbols.isEmpty())  // Function is used only in unboxed context
+                return listOf(staticMethod)
+
+            return listOf(function, staticMethod)
+        }
+
+        private fun transformConstructorBody(irConstructor: IrConstructor, staticMethod: IrSimpleFunction) {
+            if (irConstructor.isPrimary) return // TODO error() maybe?
+
             val irClass = irConstructor.parentAsClass
 
             // Copied and adapted from Kotlin/Native InlineClassTransformer
-            result.body = context.createIrBuilder(result.symbol).irBlockBody(result) {
+            staticMethod.body = IrBlockBodyImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET) {
+                statements += context.createIrBuilder(staticMethod.symbol).irBlockBody(staticMethod) {
 
                     // Secondary ctors of inline class must delegate to some other constructors.
                     // Use these delegating call later to initialize this variable.
                     lateinit var thisVar: IrVariable
-                val parameterMapping = result.valueParameters.associateBy {
+                    val parameterMapping = staticMethod.valueParameters.associateBy {
                         irConstructor.valueParameters[it.index].symbol
                     }
 
@@ -68,7 +93,7 @@ class InlineClassLowering(val context: CommonBackendContext) {
                                         expression,
                                         irType = irClass.defaultType
                                     )
-                                thisVar.parent = result
+                                    thisVar.parent = staticMethod
                                 }
                             }
 
@@ -85,7 +110,7 @@ class InlineClassLowering(val context: CommonBackendContext) {
                             override fun visitDeclaration(declaration: IrDeclaration): IrStatement {
                                 declaration.transformChildrenVoid(this)
                                 if (declaration.parent == irConstructor)
-                                declaration.parent = result
+                                    declaration.parent = staticMethod
                                 return declaration
                             }
 
@@ -104,22 +129,21 @@ class InlineClassLowering(val context: CommonBackendContext) {
                         }, null)
                     }
                     +irReturn(irGet(thisVar))
+                }.statements
+            }
         }
 
-            return result
+        private fun transformMethodBodyFlat(function: IrSimpleFunction, staticMethod: IrSimpleFunction) {
-        }
-
-
-        private fun transformMethodFlat(function: IrSimpleFunction): List<IrDeclaration> {
             // TODO: Support fake-overridden methods without boxing
-            if (function.isStaticMethodOfClass || !function.isReal)
+            if (function.isStaticMethodOfClass || !function.isReal) return // TODO error()
-                return listOf(function)
 
-            val staticMethod = createStaticBodilessMethod(function)
+            val functionBody = function.body
-            transformedFunction[function.symbol] = staticMethod.symbol
 
             // Move function body to static method, transforming value parameters and nested declarations
-            function.body!!.transformChildrenVoid(object : IrElementTransformerVoid() {
+            staticMethod.body = IrBlockBodyImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET) {
+                statements.addAll((functionBody as IrBlockBody).statements)
+
+                transformChildrenVoid(object : IrElementTransformerVoid() {
                     override fun visitDeclaration(declaration: IrDeclaration): IrStatement {
                         declaration.transformChildrenVoid(this)
                         if (declaration.parent == function)
@@ -149,14 +173,13 @@ class InlineClassLowering(val context: CommonBackendContext) {
                         )
                     }
                 })
+            }
+        }
 
-            staticMethod.body = function.body
+        private fun delegateToStaticMethod(function: IrSimpleFunction, staticMethod: IrSimpleFunction): IrBlockBody {
-
-            if (function.overriddenSymbols.isEmpty())  // Function is used only in unboxed context
-                return listOf(staticMethod)
-
             // Delegate original function to static implementation
-            function.body = context.createIrBuilder(function.symbol).irBlockBody {
+            return IrBlockBodyImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET) {
+                statements += context.createIrBuilder(function.symbol).irBlockBody {
                     +irReturn(
                         irCall(staticMethod).apply {
                             val parameters =
@@ -170,10 +193,15 @@ class InlineClassLowering(val context: CommonBackendContext) {
                             }
                         }
                     )
+                }.statements
+            }
         }
 
-            return listOf(function, staticMethod)
     }
+
+    private fun getOrCreateStaticMethod(function: IrFunction): IrSimpleFunction =
+        transformedFunction.getOrPut(function) {
+            createStaticBodilessMethod(function)
         }
 
     val inlineClassUsageLowering = object : BodyLoweringPass {
@@ -220,13 +248,6 @@ class InlineClassLowering(val context: CommonBackendContext) {
                         else -> irCall(expression, getOrCreateStaticMethod(function))
                     }
                 }
-
-                private fun getOrCreateStaticMethod(function: IrFunction): IrSimpleFunctionSymbol =
-                    transformedFunction.getOrPut(function.symbol) {
-                        createStaticBodilessMethod(function).also {
-                            function.parentAsClass.declarations.add(it)
-                        }.symbol
-                    }
             })
         }
     }

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

@@ -47,7 +47,9 @@ class JsIrBackendContext(
     override val configuration: CompilerConfiguration, // TODO: remove configuration from backend context
     override val scriptMode: Boolean = false
 ) : JsCommonBackendContext {
-    override val transformedFunction = mutableMapOf<IrFunctionSymbol, IrSimpleFunctionSymbol>()
+    override val transformedFunction
+        get() = error("Use Mapping.inlineClassMemberToStatic instead")
+
     override val lateinitNullableFields = mutableMapOf<IrField, IrField>()
 
     val memberMap = mutableMapOf<IrSimpleFunctionSymbol, IrSimpleFunction>()

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

@@ -327,15 +327,16 @@ private val secondaryFactoryInjectorLoweringPhase = makeJsModulePhase(
     prerequisite = setOf(innerClassesLoweringPhase)
 )
 
-private val inlineClassLoweringPhase = makeCustomJsModulePhase(
+private val inlineClassDeclarationLoweringPhase = makeJsModulePhase(
-    { context, module ->
+    { InlineClassLowering(it).inlineClassDeclarationLowering },
-        InlineClassLowering(context).run {
+    name = "InlineClassDeclarationLowering",
-            inlineClassDeclarationLowering.runOnFilesPostfix(module)
+    description = "Handle inline class declarations"
-            inlineClassUsageLowering.lower(module)
+)
-        }
+
-    },
+private val inlineClassUsageLoweringPhase = makeJsModulePhase(
-    name = "InlineClassLowering",
+    { InlineClassLowering(it).inlineClassUsageLowering },
-    description = "Handle inline classes"
+    name = "InlineClassUsageLowering",
+    description = "Handle inline class usages"
 )
 
 private val autoboxingTransformerPhase = makeJsModulePhase(
@@ -451,7 +452,8 @@ val jsPhases = namedIrModulePhase(
             secondaryConstructorLoweringPhase then
             secondaryFactoryInjectorLoweringPhase then
             classReferenceLoweringPhase then
-            inlineClassLoweringPhase then
+            inlineClassDeclarationLoweringPhase then
+            inlineClassUsageLoweringPhase then
             autoboxingTransformerPhase then
             blockDecomposerLoweringPhase then
             constLoweringPhase then

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

@@ -229,15 +229,16 @@ private val bridgesConstructionPhase = makeWasmModulePhase(
     description = "Generate bridges"
 )
 
-private val inlineClassLoweringPhase = makeCustomWasmModulePhase(
+private val inlineClassDeclarationLoweringPhase = makeWasmModulePhase(
-    { context, module ->
+    { InlineClassLowering(it).inlineClassDeclarationLowering },
-        InlineClassLowering(context).run {
+    name = "InlineClassDeclarationLowering",
-            inlineClassDeclarationLowering.runOnFilesPostfix(module)
+    description = "Handle inline class declarations"
-            inlineClassUsageLowering.lower(module)
+)
-        }
+
-    },
+private val inlineClassUsageLoweringPhase = makeWasmModulePhase(
-    name = "InlineClassLowering",
+    { InlineClassLowering(it).inlineClassUsageLowering },
-    description = "Handle inline classes"
+    name = "InlineClassUsageLowering",
+    description = "Handle inline class usages"
 )
 
 //private val autoboxingTransformerPhase = makeJsModulePhase(
@@ -384,7 +385,8 @@ val wasmPhases = namedIrModulePhase<WasmBackendContext>(
 //            TODO: Reimplement
 //            classReferenceLoweringPhase then
 
-            inlineClassLoweringPhase then
+            inlineClassDeclarationLoweringPhase then
+            inlineClassUsageLoweringPhase then
 
 //            TODO: Commonize box/unbox intrinsics
 //            autoboxingTransformerPhase then