IR: optimize f(); return in tailrec fun f(): Unit

It's equivalent to `return f()`.

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

@@ -16,15 +16,15 @@
 
 package org.jetbrains.kotlin.backend.common
 
-import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.declarations.IrClass
 import org.jetbrains.kotlin.ir.declarations.IrFunction
 import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
 import org.jetbrains.kotlin.ir.expressions.*
-import org.jetbrains.kotlin.ir.types.IrSimpleType
 import org.jetbrains.kotlin.ir.types.classOrNull
+import org.jetbrains.kotlin.ir.types.IdSignatureValues
 import org.jetbrains.kotlin.ir.types.isUnit
+import org.jetbrains.kotlin.ir.util.parentClassOrNull
 import org.jetbrains.kotlin.ir.util.usesDefaultArguments
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitor
 
@@ -42,8 +42,9 @@ fun collectTailRecursionCalls(irFunction: IrFunction): Set<IrCall> {
         return emptySet()
     }
 
-    val result = mutableSetOf<IrCall>()
+    val isUnitReturn = irFunction.returnType.isUnit()
 
+    val result = mutableSetOf<IrCall>()
     val visitor = object : IrElementVisitor<Unit, ElementKind> {
 
         override fun visitElement(element: IrElement, data: ElementKind) {
@@ -73,18 +74,39 @@ fun collectTailRecursionCalls(irFunction: IrFunction): Set<IrCall> {
             expression.value.accept(this, valueKind)
         }
 
-        override fun visitContainerExpression(expression: IrContainerExpression, data: ElementKind) {
+        override fun visitExpressionBody(body: IrExpressionBody, data: ElementKind) =
+            body.acceptChildren(this, data)
+
+        override fun visitBlockBody(body: IrBlockBody, data: ElementKind) =
+            visitStatementContainer(body, data)
+
+        override fun visitContainerExpression(expression: IrContainerExpression, data: ElementKind) =
+            visitStatementContainer(expression, data)
+
+        private fun visitStatementContainer(expression: IrStatementContainer, data: ElementKind) {
             expression.statements.forEachIndexed { index, irStatement ->
-                val statementKind = if (index == expression.statements.lastIndex) {
+                val statementKind = when {
                     // The last statement defines the result of the container expression, so it has the same kind.
-                    data
-                } else {
-                    ElementKind.NOT_SURE
+                    index == expression.statements.lastIndex -> data
+                    // In a Unit-returning function, any statement directly followed by a `return` is a tail statement.
+                    isUnitReturn && expression.statements[index + 1].let {
+                        it is IrReturn && it.returnTargetSymbol == irFunction.symbol && it.value.isUnitRead()
+                    } -> ElementKind.TAIL_STATEMENT
+                    else -> ElementKind.NOT_SURE
                 }
                 irStatement.accept(this, statementKind)
             }
         }
 
+        private fun IrExpression.isUnitRead(): Boolean =
+            when (this) {
+                is IrGetObjectValue -> symbol
+                // On the JVM, if SingletonReferencesLowering has already finished, a `Unit` reference
+                // is now an IrGetField to the INSTANCE field.
+                is IrGetField -> symbol.owner.parentClassOrNull?.symbol
+                else -> null
+            }?.signature == IdSignatureValues.unit
+
         override fun visitWhen(expression: IrWhen, data: ElementKind) {
             expression.branches.forEach {
                 it.condition.accept(this, ElementKind.NOT_SURE)
@@ -131,25 +153,10 @@ fun collectTailRecursionCalls(irFunction: IrFunction): Set<IrCall> {
             }
 
             result.add(expression)
-
         }
-
-    }
-
-    val body = irFunction.body
-    if (body !is IrBlockBody) {
-        return emptySet() // TODO: should an assert be here instead?
-    }
-
-    body.statements.forEachIndexed { index, irStatement ->
-        val kind = if (index == body.statements.lastIndex && irFunction.returnType.isUnit()) {
-            ElementKind.TAIL_STATEMENT
-        } else {
-            ElementKind.NOT_SURE
-        }
-        irStatement.accept(visitor, kind)
     }
 
+    irFunction.body?.accept(visitor, ElementKind.TAIL_STATEMENT)
     return result
 }
 

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

@@ -52,18 +52,25 @@ open class TailrecLowering(val context: BackendContext) : BodyLoweringPass {
 
                 override fun visitFunction(declaration: IrFunction) {
                     declaration.acceptChildrenVoid(this)
-                    lowerTailRecursionCalls(context, declaration, useProperComputationOrderOfTailrecDefaultParameters())
+                    lowerTailRecursionCalls(declaration)
                 }
             })
 
-            lowerTailRecursionCalls(context, container, useProperComputationOrderOfTailrecDefaultParameters())
+            lowerTailRecursionCalls(container)
         }
     }
 
+    private fun lowerTailRecursionCalls(function: IrFunction) =
+        lowerTailRecursionCalls(context, function, useProperComputationOrderOfTailrecDefaultParameters())
+
     open fun useProperComputationOrderOfTailrecDefaultParameters() = true
 }
 
-private fun lowerTailRecursionCalls(context: BackendContext, irFunction: IrFunction, properComputationOrderOfTailrecDefaultParameters: Boolean) {
+private fun lowerTailRecursionCalls(
+    context: BackendContext,
+    irFunction: IrFunction,
+    properComputationOrderOfTailrecDefaultParameters: Boolean
+) {
     val tailRecursionCalls = collectTailRecursionCalls(irFunction)
     if (tailRecursionCalls.isEmpty()) {
         return
@@ -78,8 +85,8 @@ private fun lowerTailRecursionCalls(context: BackendContext, irFunction: IrFunct
     oldBody.statements.clear()
     oldBody.statements += builder.irBlockBody {
         // Define variables containing current values of parameters:
-        val parameterToVariable = parameters.associate {
-            it to createTmpVariable(irGet(it), nameHint = it.symbol.suggestVariableName(), isMutable = true)
+        val parameterToVariable = parameters.associateWith {
+            createTmpVariable(irGet(it), nameHint = it.symbol.suggestVariableName(), isMutable = true)
         }
         // (these variables are to be updated on any tail call).
 
@@ -89,9 +96,8 @@ private fun lowerTailRecursionCalls(context: BackendContext, irFunction: IrFunct
 
             body = irBlock(startOffset, endOffset, resultType = context.irBuiltIns.unitType) {
                 // Read variables containing current values of parameters:
-                val parameterToNew = parameters.associate {
-                    val variable = parameterToVariable[it]!!
-                    it to createTmpVariable(irGet(variable), nameHint = it.symbol.suggestVariableName())
+                val parameterToNew = parameters.associateWith {
+                    createTmpVariable(irGet(parameterToVariable[it]!!), nameHint = it.symbol.suggestVariableName())
                 }
 
                 val transformer = BodyTransformer(

compiler/testData/codegen/box/diagnostics/functions/tailRecursion/unitBlocks.kt

@@ -7,14 +7,14 @@
 // IGNORE_BACKEND: JS
 
 tailrec fun test(x : Int) : Unit {
-    if (x == 1) {
+    if (x > 800000) {
         test(x - 1)
-    } else if (x == 2) {
+    } else if (x > 600000) {
         test(x - 1)
         return
-    } else if (x == 3) {
-        <!NON_TAIL_RECURSIVE_CALL!>test<!>(x - 1)
-        if (x == 3) {
+    } else if (x > 400000) {
+        <!NON_TAIL_RECURSIVE_CALL!>test<!>(1)
+        if (x > 200000) {
             test(x - 1)
         }
         return