Make ExpectDeclarationsRemoveLowering a DeclarationTransformer

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

@@ -6,20 +6,22 @@
 package org.jetbrains.kotlin.backend.common.lower
 
 import org.jetbrains.kotlin.backend.common.BackendContext
-import org.jetbrains.kotlin.backend.common.FileLoweringPass
+import org.jetbrains.kotlin.backend.common.DeclarationTransformer
-import org.jetbrains.kotlin.ir.declarations.IrFile
+import org.jetbrains.kotlin.ir.declarations.*
-import org.jetbrains.kotlin.ir.util.ExpectDeclarationRemover
+import org.jetbrains.kotlin.ir.util.*
-import org.jetbrains.kotlin.ir.util.patchDeclarationParents
-import org.jetbrains.kotlin.ir.visitors.acceptVoid
 
 /**
  * This pass removes all declarations with `isExpect == true`.
  */
-class ExpectDeclarationsRemoveLowering(val context: BackendContext, keepOptionalAnnotations: Boolean = false) : FileLoweringPass {
+class ExpectDeclarationsRemoveLowering(context: BackendContext, keepOptionalAnnotations: Boolean = false) : DeclarationTransformer {
 
-    val visitor = ExpectDeclarationRemover(context.ir.symbols.externalSymbolTable, doRemove = true, keepOptionalAnnotations)
+    private val remover = ExpectDeclarationRemover(
+        symbolTable = context.ir.symbols.externalSymbolTable,
+        doRemove = true,
+        keepOptionalAnnotations = keepOptionalAnnotations
+    )
 
-    override fun lower(irFile: IrFile) {
+    override fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {
-        irFile.acceptVoid(visitor)
+        return remover.transformFlat(declaration)
     }
 }

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/util/ExpectDeclarationRemover.kt

@@ -11,13 +11,13 @@ import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.IrGetValue
+import org.jetbrains.kotlin.ir.expressions.impl.IrExpressionBodyImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrGetValueImpl
 import org.jetbrains.kotlin.ir.symbols.IrValueParameterSymbol
 import org.jetbrains.kotlin.ir.symbols.IrValueSymbol
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
 import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
-import org.jetbrains.kotlin.ir.visitors.acceptVoid
 import org.jetbrains.kotlin.resolve.checkers.ExpectedActualDeclarationChecker
 import org.jetbrains.kotlin.resolve.descriptorUtil.module
 import org.jetbrains.kotlin.resolve.multiplatform.ExpectedActualResolver
@@ -33,67 +33,81 @@ class ExpectDeclarationRemover(
     }
 
     override fun visitFile(declaration: IrFile) {
-        // All declarations with `isExpect == true` are nested into a top-level declaration with `isExpect == true`.
+        super.visitFile(declaration)
         declaration.declarations.removeAll {
-            // TODO: rewrite findCompatibleActualForExpected using IR structures instead of descriptors
+            shouldRemoveTopLevelDeclaration(it)
-            val descriptor = it.descriptor
+        }
-            if (descriptor !is MemberDescriptor || !descriptor.isExpect ||
+    }
-                (keepOptionalAnnotations && descriptor is ClassDescriptor &&
+
-                        ExpectedActualDeclarationChecker.shouldGenerateExpectClass(descriptor))
+    override fun visitValueParameter(declaration: IrValueParameter) {
-            ) {
+        super.visitValueParameter(declaration)
-                false
+        tryCopyDefaultArguments(declaration)
-            } else {
+    }
-                copyDefaultArgumentsFromExpectToActual(it)
+
-                doRemove
+    fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {
+
+        if (declaration.isTopLevelDeclaration && shouldRemoveTopLevelDeclaration(declaration)) {
+            return emptyList()
+        }
+
+        if (declaration is IrValueParameter) {
+            tryCopyDefaultArguments(declaration)
+        }
+
+        return null
+    }
+
+    private fun shouldRemoveTopLevelDeclaration(declaration: IrDeclaration): Boolean {
+        // TODO: rewrite findCompatibleActualForExpected using IR structures instead of descriptors
+        val descriptor = declaration.descriptor
+        return doRemove && descriptor is MemberDescriptor && descriptor.isExpect &&
+                !(keepOptionalAnnotations && descriptor is ClassDescriptor && ExpectedActualDeclarationChecker.shouldGenerateExpectClass(descriptor))
+    }
+
+    private fun tryCopyDefaultArguments(declaration: IrValueParameter) {
+        // Keep actual default value if present. They are generally not allowed but can be suppressed with
+        // @Suppress("ACTUAL_FUNCTION_WITH_DEFAULT_ARGUMENTS")
+        if (declaration.defaultValue != null) {
+            return
+        }
+
+        val function = declaration.parent as? IrFunction ?: return
+
+        if (function is IrConstructor &&
+            ExpectedActualDeclarationChecker.isOptionalAnnotationClass(function.descriptor.constructedClass)
+        ) {
+            return
+        }
+
+        if (!function.descriptor.isActual) return
+
+        val index = declaration.index
+
+        if (index < 0) return
+
+        assert(function.valueParameters[index] == declaration)
+
+        // If the containing declaration is an `expect class` that matches an `actual typealias`,
+        // the `actual fun` or `actual constructor` for this may be in a different module.
+        // Nothing we can do with those.
+        // TODO they may not actually have the defaults though -- may be a frontend bug.
+        val expectParameter = function.findExpectForActual()?.valueParameters?.get(index) ?: return
+
+        val defaultValue = expectParameter.defaultValue ?: return
+
+        defaultValue.let { originalDefault ->
+            declaration.defaultValue = IrExpressionBodyImpl(originalDefault.startOffset, originalDefault.endOffset) {
+                expression = originalDefault.expression.deepCopyWithSymbols(function).remapExpectValueSymbols()
             }
         }
     }
 
-    private fun copyDefaultArgumentsFromExpectToActual(declaration: IrDeclaration) {
-        declaration.acceptVoid(object : IrElementVisitorVoid {
-            override fun visitElement(element: IrElement) {
-                element.acceptChildrenVoid(this)
-            }
-
-            override fun visitValueParameter(declaration: IrValueParameter) {
-                super.visitValueParameter(declaration)
-
-                val defaultValue = declaration.defaultValue ?: return
-                val function = declaration.parent as IrFunction
-
-                val index = declaration.index
-                assert(function.valueParameters[index] == declaration)
-
-                if (function is IrConstructor &&
-                    ExpectedActualDeclarationChecker.isOptionalAnnotationClass(function.descriptor.constructedClass)
-                ) {
-                    return
-                }
-
-                // If the containing declaration is an `expect class` that matches an `actual typealias`,
-                // the `actual fun` or `actual constructor` for this may be in a different module.
-                // Nothing we can do with those.
-                // TODO they may not actually have the defaults though -- may be a frontend bug.
-                val actualFunction = function.findActualForExpected() ?: return
-                val actualParameter = actualFunction.valueParameters[index]
-
-                // Keep actual default value if present. They are generally not allowed but can be suppressed with
-                // @Suppress("ACTUAL_FUNCTION_WITH_DEFAULT_ARGUMENTS")
-                if (actualParameter.defaultValue != null)
-                    return
-
-                actualParameter.defaultValue = defaultValue.also {
-                    it.expression = it.expression.remapExpectValueSymbols()
-                    // Default value might have some declarations inside. Patching parents.
-                    it.expression.patchDeclarationParents(actualParameter.parent)
-                }
-            }
-        })
-    }
-
     private fun IrFunction.findActualForExpected(): IrFunction? =
         descriptor.findActualForExpect()?.let { symbolTable.referenceFunction(it).owner }
 
+    private fun IrFunction.findExpectForActual(): IrFunction? =
+        descriptor.findExpectForActual()?.let { symbolTable.referenceFunction(it).owner }
+
     private fun IrClass.findActualForExpected(): IrClass? =
         descriptor.findActualForExpect()?.let { symbolTable.referenceClass(it).owner }
 
@@ -105,6 +119,14 @@ class ExpectDeclarationRemover(
         findCompatibleActualForExpected(descriptor.module).singleOrNull()
     } as T?
 
+    private inline fun <reified T : MemberDescriptor> T.findExpectForActual() = with(ExpectedActualResolver) {
+        val descriptor = this@findExpectForActual
+
+        if (!descriptor.isActual) error(this) else {
+            findCompatibleExpectedForActual(descriptor.module).singleOrNull()
+        }
+    } as T?
+
     private fun IrExpression.remapExpectValueSymbols(): IrExpression {
         return this.transform(object : IrElementTransformerVoid() {