Refactor default arguments lowering

2019-12-24 17:30:09 +03:00
parent 9327776706
commit 64a5039ebe
7 changed files with 209 additions and 114 deletions
@@ -154,8 +154,13 @@ fun IrValueParameter.copyTo(
        WrappedValueParameterDescriptor(this.descriptor.annotations, this.descriptor.source)
    }
    val symbol = IrValueParameterSymbolImpl(descriptor)
-    val defaultValueCopy = defaultValue?.deepCopyWithVariables()
-    defaultValueCopy?.patchDeclarationParents(irFunction)
+    val defaultValueCopy = defaultValue?.let { originalDefault ->
+        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
@@ -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
-) : DeclarationContainerLoweringPass {
+    private val skipExternalMethods: Boolean = false,
+    private val forceSetOverrideSymbols: Boolean = true
+) : DeclarationTransformer {

-    override fun lower(irDeclarationContainer: IrDeclarationContainer) {
-        irDeclarationContainer.transformDeclarationsFlat { memberDeclaration ->
-            if (memberDeclaration is IrFunction)
-                lower(memberDeclaration)
-            else
-                null
-        }
+    override fun lower(irFile: IrFile) {
+        runPostfix(true).toFileLoweringPass().lower(irFile)
    }

-    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)
-            ?: return listOf(irFunction)
+        val newIrFunction = irFunction.generateDefaultsFunction(context, skipInlineMethods, skipExternalMethods, forceSetOverrideSymbols, visibility)
+            ?: return null
        if (newIrFunction.origin == IrDeclarationOrigin.FAKE_OVERRIDE) {
            return listOf(irFunction, newIrFunction)
        }
@@ -59,80 +63,77 @@ open class DefaultArgumentStubGenerator(
        log { "$irFunction -> $newIrFunction" }
        val builder = context.createIrBuilder(newIrFunction.symbol)

-        newIrFunction.body = builder.irBlockBody(newIrFunction) {
-            val params = mutableListOf<IrValueDeclaration>()
-            val variables = mutableMapOf<IrValueDeclaration, IrValueDeclaration>()
+        newIrFunction.body = IrBlockBodyImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET) {
+            statements += builder.irBlockBody(newIrFunction) {
+                val params = mutableListOf<IrValueDeclaration>()
+                val variables = mutableMapOf<IrValueDeclaration, IrValueDeclaration>()

-            irFunction.dispatchReceiverParameter?.let {
-                variables[it] = newIrFunction.dispatchReceiverParameter!!
-            }
-
-            irFunction.extensionReceiverParameter?.let {
-                variables[it] = newIrFunction.extensionReceiverParameter!!
-            }
-
-            // In order to deal with forward references in default value lambdas,
-            // accesses to the parameter before it has been determined if there is
-            // a default value or not is redirected to the actual parameter of the
-            // $default function. This is to ensure that examples such as:
-            //
-            // fun f(f1: () -> String = { f2() },
-            //       f2: () -> String = { "OK" }) = f1()
-            //
-            // works correctly so that `f() { "OK" }` returns "OK" and
-            // `f()` throws a NullPointerException.
-            irFunction.valueParameters.associateWithTo(variables) {
-                newIrFunction.valueParameters[it.index]
-            }
-
-            var sourceParameterIndex = -1
-            for (valueParameter in irFunction.valueParameters) {
-                if (!valueParameter.isMovedReceiver()) {
-                    ++sourceParameterIndex
+                irFunction.dispatchReceiverParameter?.let {
+                    variables[it] = newIrFunction.dispatchReceiverParameter!!
                }
-                val parameter = newIrFunction.valueParameters[valueParameter.index]
-                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 expressionBody = valueParameter.defaultValue!!
-                    expressionBody.patchDeclarationParents(newIrFunction)
-                    expressionBody.transformChildrenVoid(object : IrElementTransformerVoid() {
-                        override fun visitGetValue(expression: IrGetValue): IrExpression {
-                            log { "GetValue: ${expression.symbol.owner}" }
-                            val valueSymbol = variables[expression.symbol.owner] ?: return expression
-                            return irGet(valueSymbol)
-                        }
-                    })
-
-                    selectArgumentOrDefault(defaultFlag, parameter, expressionBody.expression)
-                } else {
-                    parameter
+                irFunction.extensionReceiverParameter?.let {
+                    variables[it] = newIrFunction.extensionReceiverParameter!!
                }
-                params.add(remapped)
-                variables[valueParameter] = remapped
-            }

-            when (irFunction) {
-                is IrConstructor -> +irDelegatingConstructorCall(irFunction).apply {
-                    passTypeArgumentsFrom(newIrFunction.parentAsClass)
-                    // This is for Kotlin/Native, which differs from the other backends in that constructors
-                    // apparently do have dispatch receivers (though *probably* not type arguments, but copy
-                    // those as well just in case):
-                    passTypeArgumentsFrom(newIrFunction, offset = newIrFunction.parentAsClass.typeParameters.size)
-                    dispatchReceiver = newIrFunction.dispatchReceiverParameter?.let { irGet(it) }
-                    params.forEachIndexed { i, variable -> putValueArgument(i, irGet(variable)) }
+                // In order to deal with forward references in default value lambdas,
+                // accesses to the parameter before it has been determined if there is
+                // a default value or not is redirected to the actual parameter of the
+                // $default function. This is to ensure that examples such as:
+                //
+                // fun f(f1: () -> String = { f2() },
+                //       f2: () -> String = { "OK" }) = f1()
+                //
+                // works correctly so that `f() { "OK" }` returns "OK" and
+                // `f()` throws a NullPointerException.
+                irFunction.valueParameters.associateWithTo(variables) {
+                    newIrFunction.valueParameters[it.index]
                }
-                is IrSimpleFunction -> +irReturn(dispatchToImplementation(irFunction, newIrFunction, params))
-                else -> error("Unknown function declaration")
-            }
-        }
-        // Remove default argument initializers.
-        irFunction.valueParameters.forEach {
-            if (it.defaultValue != null) {
-                it.defaultValue = IrExpressionBodyImpl(IrErrorExpressionImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, it.type, "Default Stub"))
-            }
+
+                var sourceParameterIndex = -1
+                for (valueParameter in irFunction.valueParameters) {
+                    if (!valueParameter.isMovedReceiver()) {
+                        ++sourceParameterIndex
+                    }
+                    val parameter = newIrFunction.valueParameters[valueParameter.index]
+                    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 expressionBody = valueParameter.defaultValue!!
+                        expressionBody.patchDeclarationParents(newIrFunction)
+                        expressionBody.transformChildrenVoid(object : IrElementTransformerVoid() {
+                            override fun visitGetValue(expression: IrGetValue): IrExpression {
+                                log { "GetValue: ${expression.symbol.owner}" }
+                                val valueSymbol = variables[expression.symbol.owner] ?: return expression
+                                return irGet(valueSymbol)
+                            }
+                        })
+
+                        selectArgumentOrDefault(defaultFlag, parameter, expressionBody.expression)
+                    } else {
+                        parameter
+                    }
+                    params.add(remapped)
+                    variables[valueParameter] = remapped
+                }
+
+                when (irFunction) {
+                    is IrConstructor -> +irDelegatingConstructorCall(irFunction).apply {
+                        passTypeArgumentsFrom(newIrFunction.parentAsClass)
+                        // This is for Kotlin/Native, which differs from the other backends in that constructors
+                        // apparently do have dispatch receivers (though *probably* not type arguments, but copy
+                        // those as well just in case):
+                        passTypeArgumentsFrom(newIrFunction, offset = newIrFunction.parentAsClass.typeParameters.size)
+                        dispatchReceiver = newIrFunction.dispatchReceiverParameter?.let { irGet(it) }
+                        params.forEachIndexed { i, variable -> putValueArgument(i, irGet(variable)) }
+                    }
+                    is IrSimpleFunction -> +irReturn(dispatchToImplementation(irFunction, newIrFunction, params))
+                    else -> error("Unknown function declaration")
+                }
+            }.statements
        }
        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
-) : IrElementTransformerVoid(), BodyLoweringPass, FileLoweringPass {
-
-    override fun lower(irFile: IrFile) {
-        irFile.transformChildrenVoid(this)
-    }
+    private val skipExternalMethods: Boolean = false,
+    private val forceSetOverrideSymbols: Boolean = true
+) : IrElementTransformerVoid(), BodyLoweringPass {

    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 {
-    override fun lower(irFunction: IrFunction) {
-        if (!context.scriptMode) {
-            irFunction.valueParameters.forEach { it.defaultValue = null }
+// Remove default argument initializers.
+class DefaultParameterCleaner(
+    val context: CommonBackendContext,
+    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,27 +406,32 @@ 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 {
-            it.owner.generateDefaultsFunction(
-                context,
-                skipInlineMethods,
-                skipExternalMethods,
-                visibility
-            )?.symbol as IrSimpleFunctionSymbol?
-        }
-        if (overriddenStubs.isNotEmpty()) {
+        if (overriddenSymbols.any { it.owner.findBaseFunctionWithDefaultArguments(skipInlineMethods, skipExternalMethods) != null })
            return generateDefaultsFunctionImpl(context, IrDeclarationOrigin.FAKE_OVERRIDE, visibility).also {
-                (it as IrSimpleFunction).overriddenSymbols += overriddenStubs
-                context.ir.defaultParameterDeclarationsCache[this] = it
+                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?
+                    }
+                }
            }
-        }
    }
    // Note: this is intentionally done *after* checking for overrides. While normally `override fun`s
    // have no default parameters, there is an exception in case of interface delegation:
@@ -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
@@ -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
@@ -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

@@ -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(
@@ -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
@@ -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