backend: implement tailrec lowering

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/common/descriptors/DescriptorUtils.kt

@@ -3,24 +3,48 @@ package org.jetbrains.kotlin.backend.common.descriptors
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.ConstructorDescriptor
 import org.jetbrains.kotlin.descriptors.ParameterDescriptor
+import org.jetbrains.kotlin.descriptors.ValueParameterDescriptor
 
 /**
  * @return naturally-ordered list of all parameters available inside the function body.
  */
 internal val CallableDescriptor.allParameters: List<ParameterDescriptor>
+    get() = if (this is ConstructorDescriptor) {
+        listOf(this.constructedClass.thisAsReceiverParameter) + explicitParameters
+    } else {
+        explicitParameters
+    }
+
+/**
+ * @return naturally-ordered list of the parameters that can have values specified at call site.
+ */
+internal val CallableDescriptor.explicitParameters: List<ParameterDescriptor>
     get() {
-        val receivers = mutableListOf<ParameterDescriptor>()
+        val result = ArrayList<ParameterDescriptor>(valueParameters.size + 2)
 
-        if (this is ConstructorDescriptor)
+        this.dispatchReceiverParameter?.let {
-            receivers.add(this.constructedClass.thisAsReceiverParameter)
+            result.add(it)
+        }
 
-        val dispatchReceiverParameter = this.dispatchReceiverParameter
+        this.extensionReceiverParameter?.let {
-        if (dispatchReceiverParameter != null)
+            result.add(it)
-            receivers.add(dispatchReceiverParameter)
+        }
 
-        val extensionReceiverParameter = this.extensionReceiverParameter
+        result.addAll(valueParameters)
-        if (extensionReceiverParameter != null)
-            receivers.add(extensionReceiverParameter)
 
-        return receivers + this.valueParameters
+        return result
-    }
+    }
+
+/**
+ * Returns the parameter in the original function corresponding to given parameter of this function.
+ *
+ * Note: `parameter.original` doesn't seem to be always the parameter of `this.original`.
+ *
+ * @param parameter must be declared in this function
+ */
+fun CallableDescriptor.getOriginalParameter(parameter: ParameterDescriptor): ParameterDescriptor = when (parameter) {
+    is ValueParameterDescriptor -> this.original.valueParameters[parameter.index]
+    this.dispatchReceiverParameter -> this.original.dispatchReceiverParameter!!
+    this.extensionReceiverParameter -> this.original.extensionReceiverParameter!!
+    else -> TODO("$parameter in $this")
+}

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

@@ -0,0 +1,157 @@
+package org.jetbrains.kotlin.backend.common.lower
+
+import org.jetbrains.kotlin.backend.common.BackendContext
+import org.jetbrains.kotlin.backend.common.DeepCopyIrTreeWithDeclarations
+import org.jetbrains.kotlin.backend.common.FunctionLoweringPass
+import org.jetbrains.kotlin.backend.common.collectTailRecursionCalls
+import org.jetbrains.kotlin.backend.common.descriptors.explicitParameters
+import org.jetbrains.kotlin.backend.common.descriptors.getOriginalParameter
+import org.jetbrains.kotlin.descriptors.ParameterDescriptor
+import org.jetbrains.kotlin.descriptors.ValueDescriptor
+import org.jetbrains.kotlin.descriptors.ValueParameterDescriptor
+import org.jetbrains.kotlin.descriptors.VariableDescriptor
+import org.jetbrains.kotlin.ir.builders.*
+import org.jetbrains.kotlin.ir.declarations.IrFunction
+import org.jetbrains.kotlin.ir.expressions.*
+import org.jetbrains.kotlin.ir.util.getArguments
+import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
+import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
+
+/**
+ * This pass lowers tail recursion calls in `tailrec` functions.
+ *
+ * Note: it currently can't handle local functions and classes declared in default arguments.
+ * See [DeepCopyIrTreeWithDeclarations].
+ */
+internal class TailrecLowering(val context: BackendContext) : FunctionLoweringPass {
+    override fun lower(irFunction: IrFunction) {
+        lowerTailRecursionCalls(context, irFunction)
+    }
+}
+
+private fun lowerTailRecursionCalls(context: BackendContext, irFunction: IrFunction) {
+    val tailRecursionCalls = collectTailRecursionCalls(irFunction)
+    if (tailRecursionCalls.isEmpty()) {
+        return
+    }
+
+    val descriptor = irFunction.descriptor
+    val oldBody = irFunction.body as IrBlockBody
+    val builder = context.createIrBuilder(descriptor).at(oldBody)
+
+    val parameters = descriptor.explicitParameters
+
+    irFunction.body = builder.irBlockBody {
+        // Define variables containing current values of parameters:
+        val parameterToVariable = parameters.associate {
+            it to defineTemporaryVar(irGet(it), nameHint = it.suggestVariableName())
+        }
+        // (these variables are to be updated on any tail call).
+
+        +irWhile().apply {
+            val loop = this
+            condition = irTrue()
+
+            body = irBlock(startOffset, endOffset, resultType = context.builtIns.unitType) {
+                // Read variables containing current values of parameters:
+                val parameterToNew = parameters.associate {
+                    val variable = parameterToVariable[it]!!
+                    it to defineTemporary(irGet(variable), nameHint = it.suggestVariableName())
+                }
+
+                val transformer = BodyTransformer(builder, irFunction, loop,
+                        parameterToNew, parameterToVariable, tailRecursionCalls)
+
+                oldBody.statements.forEach {
+                    +it.transform(transformer, null)
+                }
+
+                +irBreak(loop)
+            }
+        }
+    }
+}
+
+private class BodyTransformer(
+        val builder: IrBuilderWithScope,
+        val irFunction: IrFunction,
+        val loop: IrLoop,
+        val parameterToNew: Map<ParameterDescriptor, ValueDescriptor>,
+        val parameterToVariable: Map<ParameterDescriptor, VariableDescriptor>,
+        val tailRecursionCalls: Set<IrCall>
+) : IrElementTransformerVoid() {
+
+    val parameters = irFunction.descriptor.explicitParameters
+
+    override fun visitGetValue(expression: IrGetValue): IrExpression {
+        expression.transformChildrenVoid(this)
+        val value = parameterToNew[expression.descriptor] ?: return expression
+        return builder.at(expression).irGet(value)
+    }
+
+    override fun visitCall(expression: IrCall): IrExpression {
+        expression.transformChildrenVoid(this)
+        if (expression !in tailRecursionCalls) {
+            return expression
+        }
+
+        return builder.at(expression).genTailCall(expression)
+    }
+
+    private fun IrBuilderWithScope.genTailCall(expression: IrCall) = this.irBlock(expression) {
+        // Get all specified arguments:
+        val parameterToArgument = expression.getArguments().map { (parameter, argument) ->
+            expression.descriptor.getOriginalParameter(parameter) to argument
+        }
+
+        // For each specified argument set the corresponding variable to it in the correct order:
+        parameterToArgument.forEach { (parameter, argument) ->
+            at(argument)
+            // Note that argument can use values of parameters, so it is important that
+            // references to parameters are mapped using `parameterToNew`, not `parameterToVariable`.
+            +irSetVar(parameterToVariable[parameter]!!, argument)
+        }
+
+        val specifiedParameters = parameterToArgument.map { (parameter, _) -> parameter }.toSet()
+
+        // For each unspecified argument set the corresponding variable to default:
+        parameters.filter { it !in specifiedParameters }.forEach { parameter ->
+
+            val originalDefaultValue = irFunction.getDefaultArgumentExpression(parameter) ?:
+                    throw Error("no argument specified for $parameter")
+
+            // Copy default value, mapping parameters to variables containing freshly computed arguments:
+            val defaultValue = originalDefaultValue.transform(object : DeepCopyIrTreeWithDeclarations() {
+                override fun mapValueReference(descriptor: ValueDescriptor): ValueDescriptor {
+                    return parameterToVariable[descriptor] ?: super.mapValueReference(descriptor)
+                }
+            }, data = null)
+
+            +irSetVar(parameterToVariable[parameter]!!, defaultValue)
+        }
+
+        // Jump to the entry:
+        +irContinue(loop)
+    }
+}
+
+private fun IrFunction.getDefaultArgumentExpression(parameter: ParameterDescriptor): IrExpression? {
+    if (parameter !is ValueParameterDescriptor) {
+        return null
+    }
+
+    val body = this.getDefault(parameter) ?: return null
+
+    if (body !is IrExpressionBody) {
+        throw Error("unexpected default argument body: $body")
+    }
+
+    return body.expression
+}
+
+private fun ParameterDescriptor.suggestVariableName(): String = if (name.isSpecial) {
+    val oldNameStr = name.asString()
+    "$" + oldNameStr.substring(1, oldNameStr.length - 1)
+} else {
+    name.identifier
+}

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/KonanLower.kt

@@ -25,13 +25,6 @@ internal class KonanLower(val context: Context) {
         phaser.phase(KonanPhase.LOWER_ENUMS) {
             EnumClassLowering(context).run(irFile)
         }
-        phaser.phase(KonanPhase.LOWER_DEFAULT_PARAMETER_EXTENT) {
-            DefaultArgumentStubGenerator(context).runOnFilePostfix(irFile)
-            DefaultParameterInjector(context).runOnFilePostfix(irFile)
-        }
-        phaser.phase(KonanPhase.LOWER_BUILTIN_OPERATORS) {
-            BuiltinOperatorLowering(context).runOnFilePostfix(irFile)
-        }
         phaser.phase(KonanPhase.LOWER_SHARED_VARIABLES) {
             SharedVariablesLowering(context).runOnFilePostfix(irFile)
         }
@@ -48,6 +41,16 @@ internal class KonanLower(val context: Context) {
         phaser.phase(KonanPhase.LOWER_LOCAL_FUNCTIONS) {
             LocalDeclarationsLowering(context).runOnFilePostfix(irFile)
         }
+        phaser.phase(KonanPhase.LOWER_TAILREC) {
+            TailrecLowering(context).runOnFilePostfix(irFile)
+        }
+        phaser.phase(KonanPhase.LOWER_DEFAULT_PARAMETER_EXTENT) {
+            DefaultArgumentStubGenerator(context).runOnFilePostfix(irFile)
+            DefaultParameterInjector(context).runOnFilePostfix(irFile)
+        }
+        phaser.phase(KonanPhase.LOWER_BUILTIN_OPERATORS) {
+            BuiltinOperatorLowering(context).runOnFilePostfix(irFile)
+        }
         phaser.phase(KonanPhase.LOWER_INNER_CLASSES) {
             InnerClassLowering(context).runOnFilePostfix(irFile)
         }

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/KonanPhases.kt

@@ -26,6 +26,7 @@ enum class KonanPhase(val description: String,
     /* ... ... */ LOWER_INITIALIZERS("Initializers lowering"),
     /* ... ... */ BRIDGES_BUILDING("Bridges building"),
     /* ... ... */ LOWER_DELEGATION("Delegation lowering"),
+    /* ... ... */ LOWER_TAILREC("tailrec lowering"),
     /* ... */ BITCODE("LLVM BitCode Generation"),
     /* ... ... */ RTTI("RTTI Generation"),
     /* ... ... */ CODEGEN("Code Generation"),

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

@@ -1,6 +1,7 @@
 package org.jetbrains.kotlin.backend.konan.lower
 
 import org.jetbrains.kotlin.backend.common.DeclarationContainerLoweringPass
+import org.jetbrains.kotlin.backend.common.descriptors.explicitParameters
 import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
 import org.jetbrains.kotlin.backend.konan.KonanBackendContext
 import org.jetbrains.kotlin.backend.konan.descriptors.getKonanInternalClass
@@ -109,9 +110,7 @@ private class CallableReferencesUnbinder(val lower: CallableReferenceLowering,
         val boundArgs = expression.getArguments()
         val boundParams = boundArgs.map { it.first }
 
-        val allParams = with (descriptor) {
+        val allParams = descriptor.explicitParameters
-            listOf(dispatchReceiverParameter, extensionReceiverParameter).filterNotNull() + valueParameters
-        }
         val unboundParams = allParams - boundParams
 
         val startOffset = expression.startOffset

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/ir/builders/IrBuilders.kt

@@ -1,8 +1,11 @@
 package org.jetbrains.kotlin.ir.builders
 
+import org.jetbrains.kotlin.descriptors.ValueDescriptor
 import org.jetbrains.kotlin.descriptors.VariableDescriptor
 import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.expressions.IrLoop
 import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
+import org.jetbrains.kotlin.ir.expressions.impl.*
 import org.jetbrains.kotlin.types.KotlinType
 
 inline fun IrBuilderWithScope.irLetSequence(
@@ -16,4 +19,18 @@ inline fun IrBuilderWithScope.irLetSequence(
 ): IrExpression = irBlock(startOffset, endOffset, origin, resultType) {
     val irTemporary = defineTemporary(value, nameHint)
     this.body(irTemporary)
-}
+}
+
+fun IrBuilderWithScope.irWhile(origin: IrStatementOrigin? = null) =
+        IrWhileLoopImpl(startOffset, endOffset, context.builtIns.unitType, origin)
+
+fun IrBuilderWithScope.irBreak(loop: IrLoop) =
+        IrBreakImpl(startOffset, endOffset, context.builtIns.nothingType, loop)
+
+fun IrBuilderWithScope.irContinue(loop: IrLoop) =
+        IrContinueImpl(startOffset, endOffset, context.builtIns.nothingType, loop)
+
+fun IrBuilderWithScope.irTrue() = IrConstImpl.boolean(startOffset, endOffset, context.builtIns.booleanType, true)
+
+fun IrBuilderWithScope.irGet(value: ValueDescriptor) =
+        IrGetValueImpl(startOffset, endOffset, value)