JVM_IR: replace calls to expect declarations in ExprectDeclarationRemover

References to expect declarations still remain in IrTypes;
we will need to remove those as well sooner or later.

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

@@ -5,46 +5,172 @@
 
 package org.jetbrains.kotlin.backend.common.lower
 
-import org.jetbrains.kotlin.descriptors.ClassDescriptor
-import org.jetbrains.kotlin.descriptors.ClassKind
-import org.jetbrains.kotlin.descriptors.FunctionDescriptor
-import org.jetbrains.kotlin.descriptors.MemberDescriptor
+import org.jetbrains.kotlin.backend.common.BackendContext
+import org.jetbrains.kotlin.backend.common.FileLoweringPass
+import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.ir.IrElement
+import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.ObsoleteDescriptorBasedAPI
 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.IrGetValueImpl
+import org.jetbrains.kotlin.ir.expressions.*
+import org.jetbrains.kotlin.ir.expressions.impl.*
+import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
 import org.jetbrains.kotlin.ir.symbols.IrValueParameterSymbol
 import org.jetbrains.kotlin.ir.symbols.IrValueSymbol
 import org.jetbrains.kotlin.ir.types.extractTypeParameters
 import org.jetbrains.kotlin.ir.util.*
 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.resolve.descriptorUtil.module
 import org.jetbrains.kotlin.resolve.multiplatform.*
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 // `doRemove` means should expect-declaration be removed from IR
 @OptIn(ObsoleteDescriptorBasedAPI::class)
-class ExpectDeclarationRemover(val symbolTable: ReferenceSymbolTable, private val doRemove: Boolean) : IrElementVisitorVoid {
+class ExpectDeclarationRemover(val symbolTable: ReferenceSymbolTable, private val doRemove: Boolean)
+    : IrElementTransformerVoid(), FileLoweringPass {
+
+    constructor(context: BackendContext) : this(context.ir.symbols.externalSymbolTable, true)
+
     private val typeParameterSubstitutionMap = mutableMapOf<Pair<IrFunction, IrFunction>, Map<IrTypeParameter, IrTypeParameter>>()
 
-    override fun visitElement(element: IrElement) {
-        element.acceptChildrenVoid(this)
+    override fun lower(irFile: IrFile) {
+        visitFile(irFile)
     }
 
-    override fun visitFile(declaration: IrFile) {
-        super.visitFile(declaration)
+    override fun visitElement(element: IrElement): IrElement {
+        element.transformChildrenVoid()
+        return element
+    }
+
+    override fun visitFile(declaration: IrFile): IrFile {
         declaration.declarations.removeAll {
             shouldRemoveTopLevelDeclaration(it)
         }
+        return super.visitFile(declaration)
     }
 
-    override fun visitValueParameter(declaration: IrValueParameter) {
-        super.visitValueParameter(declaration)
+    override fun visitValueParameter(declaration: IrValueParameter): IrStatement {
         tryCopyDefaultArguments(declaration)
+        return super.visitValueParameter(declaration)
+    }
+
+    override fun visitConstructorCall(expression: IrConstructorCall): IrExpression {
+        val nExpression = super.visitConstructorCall(expression) as IrConstructorCall
+        if (!nExpression.symbol.owner.isExpect) return nExpression
+
+        val newCallee = symbolTable.referenceConstructor(
+            nExpression.symbol.descriptor.findActualForExpect() as? ClassConstructorDescriptor ?: return nExpression
+        )
+        with(nExpression) {
+            return IrConstructorCallImpl(
+                startOffset, endOffset, type, newCallee, typeArgumentsCount, constructorTypeArgumentsCount, valueArgumentsCount, origin
+            ).also {
+                it.attributeOwnerId = attributeOwnerId
+                it.copyTypeAndValueArgumentsFrom(nExpression)
+            }
+        }
+    }
+
+    override fun visitDelegatingConstructorCall(expression: IrDelegatingConstructorCall): IrExpression {
+        val nExpression = super.visitDelegatingConstructorCall(expression) as IrDelegatingConstructorCall
+        if (!nExpression.symbol.owner.isExpect) return nExpression
+
+        val newCallee = symbolTable.referenceConstructor(
+            nExpression.symbol.descriptor.findActualForExpect() as? ClassConstructorDescriptor ?: return nExpression
+        )
+        with(nExpression) {
+            return IrDelegatingConstructorCallImpl(
+                startOffset, endOffset, type, newCallee, typeArgumentsCount, valueArgumentsCount
+            ).also {
+                it.attributeOwnerId = attributeOwnerId
+                it.copyTypeAndValueArgumentsFrom(nExpression)
+            }
+        }
+    }
+
+    override fun visitEnumConstructorCall(expression: IrEnumConstructorCall): IrExpression {
+        val nExpression = super.visitEnumConstructorCall(expression) as IrEnumConstructorCall
+        if (!nExpression.symbol.owner.isExpect) return nExpression
+
+        val newCallee = symbolTable.referenceConstructor(
+            nExpression.symbol.descriptor.findActualForExpect() as? ClassConstructorDescriptor ?: return nExpression
+        )
+        with(nExpression) {
+            return IrEnumConstructorCallImpl(
+                startOffset, endOffset, type, newCallee, typeArgumentsCount, valueArgumentsCount
+            ).also {
+                it.attributeOwnerId = attributeOwnerId
+                it.copyTypeAndValueArgumentsFrom(nExpression)
+            }
+        }
+    }
+
+    override fun visitCall(expression: IrCall): IrExpression {
+        val nExpression = super.visitCall(expression) as IrCall
+        if (!nExpression.symbol.owner.isExpect) return nExpression
+
+        val newCallee = symbolTable.referenceSimpleFunction(
+            nExpression.symbol.descriptor.findActualForExpect() as? FunctionDescriptor ?: return nExpression
+        )
+        return irCall(nExpression, newCallee).also {
+            it.attributeOwnerId = nExpression.attributeOwnerId
+        }
+    }
+
+    override fun visitPropertyReference(expression: IrPropertyReference): IrExpression {
+        val nExpression = super.visitPropertyReference(expression) as IrPropertyReference
+        if (!nExpression.symbol.owner.isExpect) return nExpression
+
+        val newSymbol = symbolTable.referenceProperty(
+            nExpression.symbol.descriptor.findActualForExpect() as? PropertyDescriptor ?: return nExpression
+        )
+        val newGetter = newSymbol.descriptor.getter?.let { symbolTable.referenceSimpleFunction(it) }
+        val newSetter = newSymbol.descriptor.setter?.let { symbolTable.referenceSimpleFunction(it) }
+        with(nExpression) {
+            return IrPropertyReferenceImpl(
+                startOffset, endOffset, type,
+                newSymbol, typeArgumentsCount,
+                field, newGetter, newSetter,
+                origin
+            ).also {
+                it.attributeOwnerId = attributeOwnerId
+                copyTypeArgumentsFrom(nExpression)
+                it.dispatchReceiver = dispatchReceiver
+                it.extensionReceiver = extensionReceiver
+            }
+        }
+    }
+
+    override fun visitFunctionReference(expression: IrFunctionReference): IrExpression {
+        val nExpression = super.visitFunctionReference(expression) as IrFunctionReference
+        if (!nExpression.symbol.owner.isExpect) return nExpression
+
+        val newCallee = symbolTable.referenceSimpleFunction(
+            nExpression.symbol.descriptor.findActualForExpect() as? FunctionDescriptor ?: return nExpression
+        )
+        with(nExpression) {
+            return IrFunctionReferenceImpl(
+                startOffset, endOffset, type, newCallee, typeArgumentsCount, valueArgumentsCount, reflectionTarget, origin
+            ).also {
+                it.attributeOwnerId = attributeOwnerId
+                it.copyTypeArgumentsFrom(nExpression)
+                it.dispatchReceiver = dispatchReceiver
+                it.extensionReceiver = extensionReceiver
+            }
+        }
+    }
+
+    override fun visitClassReference(expression: IrClassReference): IrExpression {
+        val nExpression = super.visitClassReference(expression) as IrClassReference
+        val oldSymbol = nExpression.symbol as? IrClassSymbol ?: return nExpression
+        if (!oldSymbol.owner.isExpect) return nExpression
+        
+        val newSymbol = symbolTable.referenceClass(
+            oldSymbol.descriptor.findActualForExpect() as? ClassDescriptor ?: return nExpression
+        )
+        with(nExpression) {
+            return IrClassReferenceImpl(startOffset, endOffset, type, newSymbol, classType)
+        }
     }
 
     fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {

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

@@ -65,8 +65,8 @@ private val arrayConstructorPhase = makeIrFilePhase(
     description = "Transform `Array(size) { index -> value }` into a loop"
 )
 
-internal val expectDeclarationsRemovingPhase = makeIrModulePhase(
-    ::ExpectDeclarationsRemoveLowering,
+internal val expectDeclarationsRemovingPhase = makeIrModulePhase<CommonBackendContext>(
+    ::ExpectDeclarationRemover,
     name = "ExpectDeclarationsRemoving",
     description = "Remove expect declaration from module fragment"
 )

compiler/ir/serialization.js/src/org/jetbrains/kotlin/ir/backend/js/klib.kt

@@ -181,7 +181,7 @@ fun generateIrForKlibSerialization(
     }
 
     if (!configuration.expectActualLinker) {
-        moduleFragment.acceptVoid(ExpectDeclarationRemover(psi2IrContext.symbolTable, false))
+        moduleFragment.transform(ExpectDeclarationRemover(psi2IrContext.symbolTable, false), null)
     }
 
     return moduleFragment