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,160 +24,186 @@ 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
-                    is IrSimpleFunction -> transformMethodFlat(declaration)
+
-                    is IrProperty -> listOf(declaration)  // Getters and setters should be flattened
+            return when (declaration) {
-                    is IrField -> listOf(declaration)
+                is IrConstructor -> transformConstructor(declaration)
-                    is IrClass -> listOf(declaration)
+                is IrSimpleFunction -> transformMethodFlat(declaration)
-                    else -> error("Unexpected declaration: $declaration")
+                else -> null
-                }
             }
         }
 
-        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)
-            val irClass = irConstructor.parentAsClass
 
-            // Copied and adapted from Kotlin/Native InlineClassTransformer
+            transformConstructorBody(irConstructor, result)
-            result.body = context.createIrBuilder(result.symbol).irBlockBody(result) {
 
-                // Secondary ctors of inline class must delegate to some other constructors.
+            return listOf(result)
-                // Use these delegating call later to initialize this variable.
-                lateinit var thisVar: IrVariable
-                val parameterMapping = result.valueParameters.associateBy {
-                    irConstructor.valueParameters[it.index].symbol
-                }
-
-                (irConstructor.body as IrBlockBody).statements.forEach { statement ->
-                    +statement.transform(object : IrElementTransformerVoid() {
-                        override fun visitDelegatingConstructorCall(expression: IrDelegatingConstructorCall): IrExpression {
-                            expression.transformChildrenVoid()
-                            return irBlock(expression) {
-                                thisVar = createTmpVariable(
-                                    expression,
-                                    irType = irClass.defaultType
-                                )
-                                thisVar.parent = result
-                            }
-                        }
-
-                        override fun visitGetValue(expression: IrGetValue): IrExpression {
-                            expression.transformChildrenVoid()
-                            if (expression.symbol == irClass.thisReceiver?.symbol) {
-                                return irGet(thisVar)
-                            }
-
-                            parameterMapping[expression.symbol]?.let { return irGet(it) }
-                            return expression
-                        }
-
-                        override fun visitDeclaration(declaration: IrDeclaration): IrStatement {
-                            declaration.transformChildrenVoid(this)
-                            if (declaration.parent == irConstructor)
-                                declaration.parent = result
-                            return declaration
-                        }
-
-                        override fun visitReturn(expression: IrReturn): IrExpression {
-                            expression.transformChildrenVoid()
-                            if (expression.returnTargetSymbol == irConstructor.symbol) {
-                                return irReturn(irBlock(expression.startOffset, expression.endOffset) {
-                                    +expression.value
-                                    +irGet(thisVar)
-                                })
-                            }
-
-                            return expression
-                        }
-
-                    }, null)
-                }
-                +irReturn(irGet(thisVar))
-            }
-
-            return result
         }
 
 
-        private fun transformMethodFlat(function: IrSimpleFunction): List<IrDeclaration> {
+        private fun transformMethodFlat(function: IrSimpleFunction): List<IrDeclaration>? {
             // TODO: Support fake-overridden methods without boxing
             if (function.isStaticMethodOfClass || !function.isReal)
-                return listOf(function)
+                return null
 
-            val staticMethod = createStaticBodilessMethod(function)
+            val staticMethod = getOrCreateStaticMethod(function)
-            transformedFunction[function.symbol] = staticMethod.symbol
 
-            // Move function body to static method, transforming value parameters and nested declarations
+            transformMethodBodyFlat(function, staticMethod)
-            function.body!!.transformChildrenVoid(object : IrElementTransformerVoid() {
+            function.body = delegateToStaticMethod(function, staticMethod)
-                override fun visitDeclaration(declaration: IrDeclaration): IrStatement {
-                    declaration.transformChildrenVoid(this)
-                    if (declaration.parent == function)
-                        declaration.parent = staticMethod
-
-                    return declaration
-                }
-
-                override fun visitGetValue(expression: IrGetValue): IrExpression {
-                    val valueDeclaration = expression.symbol.owner as? IrValueParameter ?: return super.visitGetValue(expression)
-
-                    return context.createIrBuilder(staticMethod.symbol).irGet(
-                        when (valueDeclaration) {
-                            function.dispatchReceiverParameter, function.parentAsClass.thisReceiver ->
-                                staticMethod.valueParameters[0]
-
-                            function.extensionReceiverParameter ->
-                                staticMethod.valueParameters[1]
-
-                            in function.valueParameters -> {
-                                val offset = if (function.extensionReceiverParameter != null) 2 else 1
-                                staticMethod.valueParameters[valueDeclaration.index + offset]
-                            }
-
-                            else -> return expression
-                        }
-                    )
-                }
-            })
-
-            staticMethod.body = function.body
 
             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 {
-                +irReturn(
-                    irCall(staticMethod).apply {
-                        val parameters =
-                            listOfNotNull(
-                                function.dispatchReceiverParameter!!,
-                                function.extensionReceiverParameter
-                            ) + function.valueParameters
-
-                        for ((index, valueParameter) in parameters.withIndex()) {
-                            putValueArgument(index, irGet(valueParameter))
-                        }
-                    }
-                )
-            }
-
             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
+            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 = staticMethod.valueParameters.associateBy {
+                        irConstructor.valueParameters[it.index].symbol
+                    }
+
+                    (irConstructor.body as IrBlockBody).statements.forEach { statement ->
+                        +statement.transform(object : IrElementTransformerVoid() {
+                            override fun visitDelegatingConstructorCall(expression: IrDelegatingConstructorCall): IrExpression {
+                                expression.transformChildrenVoid()
+                                return irBlock(expression) {
+                                    thisVar = createTmpVariable(
+                                        expression,
+                                        irType = irClass.defaultType
+                                    )
+                                    thisVar.parent = staticMethod
+                                }
+                            }
+
+                            override fun visitGetValue(expression: IrGetValue): IrExpression {
+                                expression.transformChildrenVoid()
+                                if (expression.symbol == irClass.thisReceiver?.symbol) {
+                                    return irGet(thisVar)
+                                }
+
+                                parameterMapping[expression.symbol]?.let { return irGet(it) }
+                                return expression
+                            }
+
+                            override fun visitDeclaration(declaration: IrDeclaration): IrStatement {
+                                declaration.transformChildrenVoid(this)
+                                if (declaration.parent == irConstructor)
+                                    declaration.parent = staticMethod
+                                return declaration
+                            }
+
+                            override fun visitReturn(expression: IrReturn): IrExpression {
+                                expression.transformChildrenVoid()
+                                if (expression.returnTargetSymbol == irConstructor.symbol) {
+                                    return irReturn(irBlock(expression.startOffset, expression.endOffset) {
+                                        +expression.value
+                                        +irGet(thisVar)
+                                    })
+                                }
+
+                                return expression
+                            }
+
+                        }, null)
+                    }
+                    +irReturn(irGet(thisVar))
+                }.statements
+            }
+        }
+
+        private fun transformMethodBodyFlat(function: IrSimpleFunction, staticMethod: IrSimpleFunction) {
+            // TODO: Support fake-overridden methods without boxing
+            if (function.isStaticMethodOfClass || !function.isReal) return // TODO error()
+
+            val functionBody = function.body
+
+            // Move function body to static method, transforming value parameters and nested declarations
+            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)
+                            declaration.parent = staticMethod
+
+                        return declaration
+                    }
+
+                    override fun visitGetValue(expression: IrGetValue): IrExpression {
+                        val valueDeclaration = expression.symbol.owner as? IrValueParameter ?: return super.visitGetValue(expression)
+
+                        return context.createIrBuilder(staticMethod.symbol).irGet(
+                            when (valueDeclaration) {
+                                function.dispatchReceiverParameter, function.parentAsClass.thisReceiver ->
+                                    staticMethod.valueParameters[0]
+
+                                function.extensionReceiverParameter ->
+                                    staticMethod.valueParameters[1]
+
+                                in function.valueParameters -> {
+                                    val offset = if (function.extensionReceiverParameter != null) 2 else 1
+                                    staticMethod.valueParameters[valueDeclaration.index + offset]
+                                }
+
+                                else -> return expression
+                            }
+                        )
+                    }
+                })
+            }
+        }
+
+        private fun delegateToStaticMethod(function: IrSimpleFunction, staticMethod: IrSimpleFunction): IrBlockBody {
+            // Delegate original function to static implementation
+            return IrBlockBodyImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET) {
+                statements += context.createIrBuilder(function.symbol).irBlockBody {
+                    +irReturn(
+                        irCall(staticMethod).apply {
+                            val parameters =
+                                listOfNotNull(
+                                    function.dispatchReceiverParameter!!,
+                                    function.extensionReceiverParameter
+                                ) + function.valueParameters
+
+                            for ((index, valueParameter) in parameters.withIndex()) {
+                                putValueArgument(index, irGet(valueParameter))
+                            }
+                        }
+                    )
+                }.statements
+            }
+        }
+
     }
 
+    private fun getOrCreateStaticMethod(function: IrFunction): IrSimpleFunction =
+        transformedFunction.getOrPut(function) {
+            createStaticBodilessMethod(function)
+        }
+
     val inlineClassUsageLowering = object : BodyLoweringPass {
 
         override fun lower(irBody: IrBody, container: IrDeclaration) {
@@ -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