Refactor default arguments lowering

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

@@ -154,8 +154,13 @@ fun IrValueParameter.copyTo(
         WrappedValueParameterDescriptor(this.descriptor.annotations, this.descriptor.source)
     }
     val symbol = IrValueParameterSymbolImpl(descriptor)
-    val defaultValueCopy = defaultValue?.deepCopyWithVariables()
+    val defaultValueCopy = defaultValue?.let { originalDefault ->
-    defaultValueCopy?.patchDeclarationParents(irFunction)
+        IrExpressionBodyImpl(originalDefault.startOffset, originalDefault.endOffset) {
+            expression = originalDefault.expression.deepCopyWithVariables().also {
+                it.patchDeclarationParents(irFunction)
+            }
+        }
+    }
     return IrValueParameterImpl(
         startOffset, endOffset, origin, symbol,
         name, index, type, varargElementType, isCrossinline, isNoinline

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

@@ -36,22 +36,26 @@ import org.jetbrains.kotlin.name.Name
 open class DefaultArgumentStubGenerator(
     open val context: CommonBackendContext,
     private val skipInlineMethods: Boolean = true,
-    private val skipExternalMethods: Boolean = false
+    private val skipExternalMethods: Boolean = false,
-) : DeclarationContainerLoweringPass {
+    private val forceSetOverrideSymbols: Boolean = true
+) : DeclarationTransformer {
 
-    override fun lower(irDeclarationContainer: IrDeclarationContainer) {
+    override fun lower(irFile: IrFile) {
-        irDeclarationContainer.transformDeclarationsFlat { memberDeclaration ->
+        runPostfix(true).toFileLoweringPass().lower(irFile)
-            if (memberDeclaration is IrFunction)
-                lower(memberDeclaration)
-            else
-                null
-        }
     }
 
-    private fun lower(irFunction: IrFunction): List<IrFunction> {
+    override fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {
+        if (declaration is IrFunction) {
+            return lower(declaration)
+        }
+
+        return null
+    }
+
+    private fun lower(irFunction: IrFunction): List<IrFunction>? {
         val visibility = defaultArgumentStubVisibility(irFunction)
-        val newIrFunction = irFunction.generateDefaultsFunction(context, skipInlineMethods, skipExternalMethods, visibility)
+        val newIrFunction = irFunction.generateDefaultsFunction(context, skipInlineMethods, skipExternalMethods, forceSetOverrideSymbols, visibility)
-            ?: return listOf(irFunction)
+            ?: return null
         if (newIrFunction.origin == IrDeclarationOrigin.FAKE_OVERRIDE) {
             return listOf(irFunction, newIrFunction)
         }
@@ -59,7 +63,8 @@ open class DefaultArgumentStubGenerator(
         log { "$irFunction -> $newIrFunction" }
         val builder = context.createIrBuilder(newIrFunction.symbol)
 
-        newIrFunction.body = builder.irBlockBody(newIrFunction) {
+        newIrFunction.body = IrBlockBodyImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET) {
+            statements += builder.irBlockBody(newIrFunction) {
                 val params = mutableListOf<IrValueDeclaration>()
                 val variables = mutableMapOf<IrValueDeclaration, IrValueDeclaration>()
 
@@ -94,7 +99,8 @@ open class DefaultArgumentStubGenerator(
                     val remapped = if (valueParameter.defaultValue != null) {
                         val mask = irGet(newIrFunction.valueParameters[irFunction.valueParameters.size + valueParameter.index / 32])
                         val bit = irInt(1 shl (sourceParameterIndex % 32))
-                    val defaultFlag = irCallOp(this@DefaultArgumentStubGenerator.context.ir.symbols.intAnd, context.irBuiltIns.intType, mask, bit)
+                        val defaultFlag =
+                            irCallOp(this@DefaultArgumentStubGenerator.context.ir.symbols.intAnd, context.irBuiltIns.intType, mask, bit)
 
                         val expressionBody = valueParameter.defaultValue!!
                         expressionBody.patchDeclarationParents(newIrFunction)
@@ -127,12 +133,7 @@ open class DefaultArgumentStubGenerator(
                     is IrSimpleFunction -> +irReturn(dispatchToImplementation(irFunction, newIrFunction, params))
                     else -> error("Unknown function declaration")
                 }
-        }
+            }.statements
-        // Remove default argument initializers.
-        irFunction.valueParameters.forEach {
-            if (it.defaultValue != null) {
-                it.defaultValue = IrExpressionBodyImpl(IrErrorExpressionImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, it.type, "Default Stub"))
-            }
         }
         return listOf(irFunction, newIrFunction)
     }
@@ -195,17 +196,39 @@ open class DefaultArgumentStubGenerator(
     private fun log(msg: () -> String) = context.log { "DEFAULT-REPLACER: ${msg()}" }
 }
 
+private fun IrFunction.findBaseFunctionWithDefaultArguments(skipInlineMethods: Boolean, skipExternalMethods: Boolean): IrFunction? {
+
+    val visited = mutableSetOf<IrFunction>()
+
+    fun IrFunction.dfsImpl(): IrFunction? {
+        visited += this
+
+        if (isInline && skipInlineMethods) return null
+        if (skipExternalMethods && isExternalOrInheritedFromExternal()) return null
+
+        if (this is IrSimpleFunction) {
+            overriddenSymbols.forEach {
+                val base = it.owner
+                if (base !in visited) base.dfsImpl()?.let { return it }
+            }
+        }
+
+        if (valueParameters.any { it.defaultValue != null }) return this
+
+        return null
+    }
+
+    return dfsImpl()
+}
+
 val DEFAULT_DISPATCH_CALL = object : IrStatementOriginImpl("DEFAULT_DISPATCH_CALL") {}
 
 open class DefaultParameterInjector(
     open val context: CommonBackendContext,
     private val skipInline: Boolean = true,
-    private val skipExternalMethods: Boolean = false
+    private val skipExternalMethods: Boolean = false,
-) : IrElementTransformerVoid(), BodyLoweringPass, FileLoweringPass {
+    private val forceSetOverrideSymbols: Boolean = true
-
+) : IrElementTransformerVoid(), BodyLoweringPass {
-    override fun lower(irFile: IrFile) {
-        irFile.transformChildrenVoid(this)
-    }
 
     override fun lower(irBody: IrBody, container: IrDeclaration) {
         irBody.transformChildrenVoid(this)
@@ -279,12 +302,15 @@ open class DefaultParameterInjector(
         val endOffset = expression.endOffset
         val declaration = expression.symbol.owner
         val visibility = defaultArgumentStubVisibility(declaration)
-        val stubOverride = declaration.generateDefaultsFunction(context, skipInline, skipExternalMethods, visibility) ?: return null
+
-        // We *have* to resolve the fake override here since on the JVM, a default stub for a function implemented
+        // We *have* to find the actual function here since on the JVM, a default stub for a function implemented
         // in an interface does not leave an abstract method after being moved to DefaultImpls (see InterfaceLowering).
         // Calling the fake override on an implementation of that interface would then result in a call to a method
         // that does not actually exist as DefaultImpls is not part of the inheritance hierarchy.
-        val stubFunction = if (stubOverride is IrSimpleFunction) stubOverride.resolveFakeOverride()!! else stubOverride
+        val stubFunction = declaration.findBaseFunctionWithDefaultArguments(skipInline, skipExternalMethods)
+            ?.generateDefaultsFunction(context, skipInline, skipExternalMethods, forceSetOverrideSymbols, visibility)
+            ?: return null
+
         log { "$declaration -> $stubFunction" }
 
         val realArgumentsNumber = declaration.valueParameters.size
@@ -335,11 +361,44 @@ open class DefaultParameterInjector(
     private fun log(msg: () -> String) = context.log { "DEFAULT-INJECTOR: ${msg()}" }
 }
 
-class DefaultParameterCleaner constructor(val context: CommonBackendContext) : FunctionLoweringPass {
+// Remove default argument initializers.
-    override fun lower(irFunction: IrFunction) {
+class DefaultParameterCleaner(
-        if (!context.scriptMode) {
+    val context: CommonBackendContext,
-            irFunction.valueParameters.forEach { it.defaultValue = null }
+    val replaceDefaultValuesWithStubs: Boolean = false
+) : DeclarationTransformer {
+    override fun lower(irFile: IrFile) {
+        runPostfix(true).toFileLoweringPass().lower(irFile)
     }
+
+    override fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {
+        if (declaration is IrValueParameter && !context.scriptMode && declaration.defaultValue != null) {
+            if (replaceDefaultValuesWithStubs) {
+                if (context.mapping.defaultArgumentsOriginalFunction[declaration.parent as IrFunction] == null) {
+                    declaration.defaultValue =
+                        IrExpressionBodyImpl(IrErrorExpressionImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, declaration.type, "Default Stub"))
+                }
+            } else {
+                declaration.defaultValue = null
+            }
+        }
+        return null
+    }
+}
+
+// Sets overriden symbols. Should be used in case `forceSetOverrideSymbols = false`
+class DefaultParameterPatchOverridenSymbolsLowering(
+    val context: CommonBackendContext
+) : DeclarationTransformer {
+    override fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {
+        if (declaration is IrSimpleFunction) {
+            (context.mapping.defaultArgumentsOriginalFunction[declaration] as? IrSimpleFunction)?.run {
+                declaration.overriddenSymbols += overriddenSymbols.mapNotNull {
+                    (context.mapping.defaultArgumentsDispatchFunction[it.owner] as? IrSimpleFunction)?.symbol
+                }
+            }
+        }
+
+        return null
     }
 }
 
@@ -347,25 +406,30 @@ private fun IrFunction.generateDefaultsFunction(
     context: CommonBackendContext,
     skipInlineMethods: Boolean,
     skipExternalMethods: Boolean,
+    forceSetOverrideSymbols: Boolean,
     visibility: Visibility
 ): IrFunction? {
     if (skipInlineMethods && isInline) return null
     if (skipExternalMethods && isExternalOrInheritedFromExternal()) return null
-    context.ir.defaultParameterDeclarationsCache[this]?.let { return it }
+    if (context.mapping.defaultArgumentsOriginalFunction[this] != null) return null
+    context.mapping.defaultArgumentsDispatchFunction[this]?.let { return it }
     if (this is IrSimpleFunction) {
         // If this is an override of a function with default arguments, produce a fake override of a default stub.
-        val overriddenStubs = overriddenSymbols.mapNotNull {
+        if (overriddenSymbols.any { it.owner.findBaseFunctionWithDefaultArguments(skipInlineMethods, skipExternalMethods) != null })
+            return generateDefaultsFunctionImpl(context, IrDeclarationOrigin.FAKE_OVERRIDE, visibility).also {
+                context.mapping.defaultArgumentsDispatchFunction[this] = it
+                context.mapping.defaultArgumentsOriginalFunction[it] = this
+
+                if (forceSetOverrideSymbols) {
+                    (it as IrSimpleFunction).overriddenSymbols += overriddenSymbols.mapNotNull {
                         it.owner.generateDefaultsFunction(
                             context,
                             skipInlineMethods,
                             skipExternalMethods,
+                            forceSetOverrideSymbols,
                             visibility
                         )?.symbol as IrSimpleFunctionSymbol?
                     }
-        if (overriddenStubs.isNotEmpty()) {
-            return generateDefaultsFunctionImpl(context, IrDeclarationOrigin.FAKE_OVERRIDE, visibility).also {
-                (it as IrSimpleFunction).overriddenSymbols += overriddenStubs
-                context.ir.defaultParameterDeclarationsCache[this] = it
                 }
             }
     }
@@ -381,7 +445,8 @@ private fun IrFunction.generateDefaultsFunction(
     // binaries, it's way too late to fix it. Hence the workaround.
     if (valueParameters.any { it.defaultValue != null }) {
         return generateDefaultsFunctionImpl(context, IrDeclarationOrigin.FUNCTION_FOR_DEFAULT_PARAMETER, visibility).also {
-            context.ir.defaultParameterDeclarationsCache[this] = it
+            context.mapping.defaultArgumentsDispatchFunction[this] = it
+            context.mapping.defaultArgumentsOriginalFunction[it] = this
         }
     }
     return null

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

@@ -260,8 +260,15 @@ private val defaultArgumentStubGeneratorPhase = makeJsModulePhase(
     description = "Generate synthetic stubs for functions with default parameter values"
 )
 
+private val defaultArgumentPatchOverridesPhase = makeJsModulePhase(
+    ::DefaultParameterPatchOverridenSymbolsLowering,
+    name = "DefaultArgumentsPatchOverrides",
+    description = "Patch overrides for fake override dispatch functions",
+    prerequisite = setOf(defaultArgumentStubGeneratorPhase)
+)
+
 private val defaultParameterInjectorPhase = makeJsModulePhase(
-    { context -> DefaultParameterInjector(context, skipExternalMethods = true) },
+    { context -> DefaultParameterInjector(context, skipExternalMethods = true, forceSetOverrideSymbols = false) },
     name = "DefaultParameterInjector",
     description = "Replace callsite with default parameters with corresponding stub function",
     prerequisite = setOf(callableReferenceLoweringPhase, innerClassesLoweringPhase)
@@ -470,6 +477,7 @@ val jsPhases = namedIrModulePhase(
             privateMembersLoweringPhase then
             callableReferenceLoweringPhase then
             defaultArgumentStubGeneratorPhase then
+            defaultArgumentPatchOverridesPhase then
             defaultParameterInjectorPhase then
             defaultParameterCleanerPhase then
             jsDefaultCallbackGeneratorPhase then

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/CallableReferenceLowering.kt

@@ -496,7 +496,7 @@ class CallableReferenceLowering(val context: JsIrBackendContext) : FileLoweringP
 
         closureFunction.valueParameters += unboundParamDeclarations
 
-        val callTarget = context.ir.defaultParameterDeclarationsCache[declaration] ?: declaration
+        val callTarget = context.mapping.defaultArgumentsDispatchFunction[declaration] ?: declaration
 
         val target = callTarget.symbol
 

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/JsDefaultArgumentStubGenerator.kt

@@ -23,7 +23,7 @@ import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 import org.jetbrains.kotlin.name.Name
 
-class JsDefaultArgumentStubGenerator(override val context: JsIrBackendContext) : DefaultArgumentStubGenerator(context, true, true) {
+class JsDefaultArgumentStubGenerator(override val context: JsIrBackendContext) : DefaultArgumentStubGenerator(context, true, true, false) {
 
     override fun needSpecialDispatch(irFunction: IrSimpleFunction) = irFunction.isOverridableOrOverrides
 
@@ -71,7 +71,7 @@ class JsDefaultCallbackGenerator(val context: JsIrBackendContext): BodyLoweringP
 
     private fun buildBoundSuperCall(irCall: IrCall): IrExpression {
 
-        val originalFunction = context.ir.defaultParameterDeclarationsCache.entries.first { it.value == irCall.symbol.owner }.key
+        val originalFunction = context.mapping.defaultArgumentsOriginalFunction[irCall.symbol.owner]!!
 
         val reference = irCall.run {
             IrFunctionReferenceImpl(

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

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.backend.common.lower.*
 import org.jetbrains.kotlin.backend.common.lower.loops.forLoopsPhase
 import org.jetbrains.kotlin.backend.common.lower.optimizations.foldConstantLoweringPhase
 import org.jetbrains.kotlin.backend.common.phaser.*
+import org.jetbrains.kotlin.backend.jvm.ir.getJvmVisibilityOfDefaultArgumentStub
 import org.jetbrains.kotlin.backend.jvm.lower.*
 import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.descriptors.Visibility
@@ -162,11 +163,18 @@ private val defaultArgumentStubPhase = makeIrFilePhase(
     prerequisite = setOf(localDeclarationsPhase)
 )
 
+private val defaultArgumentCleanerPhase = makeIrFilePhase(
+    { context: JvmBackendContext -> DefaultParameterCleaner(context, replaceDefaultValuesWithStubs = true) },
+    name = "DefaultParameterCleaner",
+    description = "Replace default values arguments with stubs",
+    prerequisite = setOf(defaultArgumentStubPhase)
+)
+
 private val defaultArgumentInjectorPhase = makeIrFilePhase(
     ::JvmDefaultParameterInjector,
     name = "DefaultParameterInjector",
     description = "Transform calls with default arguments into calls to stubs",
-    prerequisite = setOf(defaultArgumentStubPhase, callableReferencePhase, inlineCallableReferenceToLambdaPhase)
+    prerequisite = setOf(callableReferencePhase, inlineCallableReferenceToLambdaPhase)
 )
 
 private val interfacePhase = makeIrFilePhase(
@@ -301,6 +309,7 @@ private val jvmFilePhases =
 
         defaultArgumentStubPhase then
         defaultArgumentInjectorPhase then
+        defaultArgumentCleanerPhase then
 
         interfacePhase then
         inheritedDefaultMethodsOnClassesPhase then

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

@@ -184,6 +184,13 @@ private val defaultArgumentStubGeneratorPhase = makeWasmModulePhase(
     description = "Generate synthetic stubs for functions with default parameter values"
 )
 
+private val defaultArgumentPatchOverridesPhase = makeWasmModulePhase(
+    ::DefaultParameterPatchOverridenSymbolsLowering,
+    name = "DefaultArgumentsPatchOverrides",
+    description = "Patch overrides for fake override dispatch functions",
+    prerequisite = setOf(defaultArgumentStubGeneratorPhase)
+)
+
 private val defaultParameterInjectorPhase = makeWasmModulePhase(
     { context -> DefaultParameterInjector(context, skipExternalMethods = true) },
     name = "DefaultParameterInjector",
@@ -389,6 +396,7 @@ val wasmPhases = namedIrModulePhase<WasmBackendContext>(
 //            callableReferenceLoweringPhase then
 
             defaultArgumentStubGeneratorPhase then
+            defaultArgumentPatchOverridesPhase then
             defaultParameterInjectorPhase then
             defaultParameterCleanerPhase then