JVM IR: Move direct invoke optimization into a separate pass

This also changes the transformation to inline the body of a directly
invoked lambda rather than producing a call to an anonymous local
function. The latter is unsupported in inline functions and problematic
from an ABI perspective, since it results in functions whose name
depends on the entire source code up to this point.

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirLocalVariableTestGenerated.java

@@ -43,6 +43,12 @@ public class FirLocalVariableTestGenerated extends AbstractFirLocalVariableTest
         runTest("compiler/testData/debug/localVariables/copyFunction.kt");
     }
 
+    @Test
+    @TestMetadata("directInvoke.kt")
+    public void testDirectInvoke() throws Exception {
+        runTest("compiler/testData/debug/localVariables/directInvoke.kt");
+    }
+
     @Test
     @TestMetadata("doWhile.kt")
     public void testDoWhile() throws Exception {

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

@@ -30,12 +30,7 @@ class IrInvokable(val invokable: IrValueDeclaration) : IrInlinable()
 class IrInlinableLambda(val function: IrSimpleFunction, val boundReceiver: IrValueDeclaration?) : IrInlinable()
 
 // Return the underlying function for a lambda argument without bound or default parameters or varargs.
-private fun IrExpression.asInlinableLambda(builder: IrStatementsBuilder<*>): IrInlinableLambda? {
-    if (this is IrFunctionExpression) {
-        if (function.valueParameters.any { it.isVararg || it.defaultValue != null })
-            return null
-        return IrInlinableLambda(function, null)
-    }
+fun IrExpression.asInlinableFunctionReference(): IrFunctionReference? {
     // A lambda is represented as a block with a function declaration and a reference to it.
     // Inlinable function references are also a kind of lambda; bound receivers are represented as extension receivers.
     if (this !is IrBlock || statements.size != 2)
@@ -49,7 +44,18 @@ private fun IrExpression.asInlinableLambda(builder: IrStatementsBuilder<*>): IrI
         return null
     if (function.valueParameters.any { it.isVararg || it.defaultValue != null })
         return null
-    return IrInlinableLambda(function, reference.extensionReceiver?.let { builder.irTemporary(it) })
+    return reference
+}
+
+private fun IrExpression.asInlinableLambda(builder: IrStatementsBuilder<*>): IrInlinableLambda? {
+    if (this is IrFunctionExpression) {
+        if (function.valueParameters.any { it.isVararg || it.defaultValue != null })
+            return null
+        return IrInlinableLambda(function, null)
+    }
+    return asInlinableFunctionReference()?.let { reference ->
+        IrInlinableLambda(reference.symbol.owner as IrSimpleFunction, reference.extensionReceiver?.let { builder.irTemporary(it) })
+    }
 }
 
 fun IrExpression.asInlinable(builder: IrStatementsBuilder<*>): IrInlinable =
@@ -98,7 +104,7 @@ private fun IrBody.move(
 
 // TODO use a generic inliner (e.g. JS/Native's FunctionInlining.Inliner)
 // Inline simple function calls without type parameters, default parameters, or varargs.
-private fun IrFunction.inline(target: IrDeclarationParent, arguments: List<IrValueDeclaration> = listOf()): IrReturnableBlock =
+fun IrFunction.inline(target: IrDeclarationParent, arguments: List<IrValueDeclaration> = listOf()): IrReturnableBlock =
     IrReturnableBlockImpl(startOffset, endOffset, returnType, IrReturnableBlockSymbolImpl(), null, symbol).apply {
         statements += body!!.move(this@inline, target, symbol, explicitParameters.zip(arguments).toMap()).statements
     }

compiler/ir/backend.jvm/codegen/src/org/jetbrains/kotlin/backend/jvm/codegen/ExpressionCodegen.kt

@@ -32,7 +32,6 @@ import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
 import org.jetbrains.kotlin.diagnostics.BackendErrors
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.declarations.*
-import org.jetbrains.kotlin.ir.declarations.impl.IrFunctionImpl
 import org.jetbrains.kotlin.ir.descriptors.toIrBasedDescriptor
 import org.jetbrains.kotlin.ir.descriptors.toIrBasedKotlinType
 import org.jetbrains.kotlin.ir.expressions.*
@@ -313,12 +312,6 @@ class ExpressionCodegen(
         if (irFunction.isSuspend)
             return
 
-        // As a small optimization, don't generate nullability assertions in methods for directly invoked lambdas
-        if (irFunction is IrFunctionImpl && irFunction.attributeOwnerId in context.directInvokedLambdas) {
-            context.directInvokedLambdas.remove(irFunction.attributeOwnerId)
-            return
-        }
-
         irFunction.extensionReceiverParameter?.let { generateNonNullAssertion(it) }
 
         // Private operator functions don't have null checks on value parameters,

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

@@ -287,6 +287,7 @@ private val jvmFilePhases = listOf(
     jvmLateinitLowering,
 
     inlineCallableReferenceToLambdaPhase,
+    directInvokeLowering,
     functionReferencePhase,
     suspendLambdaPhase,
     propertyReferenceDelegationPhase,

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

@@ -0,0 +1,135 @@
+/*
+ * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.backend.jvm.lower
+
+import org.jetbrains.kotlin.backend.common.FileLoweringPass
+import org.jetbrains.kotlin.backend.common.IrElementTransformerVoidWithContext
+import org.jetbrains.kotlin.backend.common.ir.asInlinableFunctionReference
+import org.jetbrains.kotlin.backend.common.ir.inline
+import org.jetbrains.kotlin.backend.common.lower.at
+import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
+import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase
+import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
+import org.jetbrains.kotlin.ir.builders.irBlock
+import org.jetbrains.kotlin.ir.declarations.IrConstructor
+import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
+import org.jetbrains.kotlin.ir.declarations.IrFile
+import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
+import org.jetbrains.kotlin.ir.declarations.impl.IrVariableImpl
+import org.jetbrains.kotlin.ir.expressions.*
+import org.jetbrains.kotlin.ir.expressions.impl.IrCallImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrConstructorCallImpl
+import org.jetbrains.kotlin.ir.symbols.impl.IrVariableSymbolImpl
+import org.jetbrains.kotlin.ir.util.dump
+import org.jetbrains.kotlin.ir.util.explicitParameters
+import org.jetbrains.kotlin.ir.util.render
+import org.jetbrains.kotlin.util.OperatorNameConventions
+
+internal val directInvokeLowering = makeIrFilePhase(
+    ::DirectInvokeLowering,
+    name = "DirectInvokes",
+    description = "Inline directly invoked lambdas and replace invoked function references with calls"
+)
+
+private class DirectInvokeLowering(private val context: JvmBackendContext) : FileLoweringPass, IrElementTransformerVoidWithContext() {
+    override fun lower(irFile: IrFile) = irFile.transformChildrenVoid()
+
+    override fun visitCall(expression: IrCall): IrExpression {
+        val function = expression.symbol.owner
+        val receiver = expression.dispatchReceiver
+        if (receiver == null || function.name != OperatorNameConventions.INVOKE)
+            return super.visitCall(expression)
+
+        val result = when {
+            // TODO deal with type parameters somehow?
+            // It seems we can't encounter them in the code written by user,
+            // but this might be important later if we actually perform inlining and optimizations on IR.
+            receiver is IrFunctionReference && receiver.symbol.owner.typeParameters.isEmpty() ->
+                visitFunctionReferenceInvoke(expression, receiver)
+
+            receiver is IrBlock ->
+                receiver.asInlinableFunctionReference()?.takeIf { it.extensionReceiver == null }?.let { reference ->
+                    visitLambdaInvoke(expression, reference)
+                } ?: expression
+
+            else ->
+                expression
+        }
+
+        result.transformChildrenVoid()
+        return result
+    }
+
+    private fun visitLambdaInvoke(expression: IrCall, reference: IrFunctionReference): IrExpression {
+        val scope = currentScope!!.scope
+        val declarationParent = scope.getLocalDeclarationParent()
+        val function = reference.symbol.owner
+        if (expression.valueArgumentsCount == 0) {
+            return function.inline(declarationParent)
+        }
+        return context.createIrBuilder(scope.scopeOwnerSymbol).run {
+            at(expression)
+            irBlock {
+                val arguments = function.explicitParameters.withIndex().map { (index, parameter) ->
+                    val argument = expression.getValueArgument(index)!!
+                    IrVariableImpl(
+                        argument.startOffset, argument.endOffset, IrDeclarationOrigin.DEFINED, IrVariableSymbolImpl(), parameter.name,
+                        parameter.type, isVar = false, isConst = false, isLateinit = false
+                    ).apply {
+                        parent = declarationParent
+                        initializer = argument
+                        +this
+                    }
+                }
+                +function.inline(declarationParent, arguments)
+            }
+        }
+    }
+
+    private fun visitFunctionReferenceInvoke(expression: IrCall, receiver: IrFunctionReference): IrExpression =
+        when (val irFun = receiver.symbol.owner) {
+            is IrSimpleFunction ->
+                IrCallImpl(
+                    expression.startOffset, expression.endOffset, expression.type, irFun.symbol,
+                    typeArgumentsCount = irFun.typeParameters.size, valueArgumentsCount = irFun.valueParameters.size
+                ).apply {
+                    copyReceiverAndValueArgumentsForDirectInvoke(receiver, expression)
+                }
+
+            is IrConstructor ->
+                IrConstructorCallImpl(
+                    expression.startOffset, expression.endOffset, expression.type, irFun.symbol,
+                    typeArgumentsCount = irFun.typeParameters.size,
+                    constructorTypeArgumentsCount = 0,
+                    valueArgumentsCount = irFun.valueParameters.size
+                ).apply {
+                    copyReceiverAndValueArgumentsForDirectInvoke(receiver, expression)
+                }
+
+            else ->
+                throw AssertionError("Simple function or constructor expected: ${irFun.render()}")
+        }
+
+    private fun IrFunctionAccessExpression.copyReceiverAndValueArgumentsForDirectInvoke(
+        irFunRef: IrFunctionReference,
+        irInvokeCall: IrFunctionAccessExpression
+    ) {
+        val irFun = irFunRef.symbol.owner
+        var invokeArgIndex = 0
+        if (irFun.dispatchReceiverParameter != null) {
+            dispatchReceiver = irFunRef.dispatchReceiver ?: irInvokeCall.getValueArgument(invokeArgIndex++)
+        }
+        if (irFun.extensionReceiverParameter != null) {
+            extensionReceiver = irFunRef.extensionReceiver ?: irInvokeCall.getValueArgument(invokeArgIndex++)
+        }
+        if (invokeArgIndex + valueArgumentsCount != irInvokeCall.valueArgumentsCount) {
+            throw AssertionError("Mismatching value arguments: $invokeArgIndex arguments used for receivers\n${irInvokeCall.dump()}")
+        }
+        for (i in 0 until valueArgumentsCount) {
+            putValueArgument(i, irInvokeCall.getValueArgument(invokeArgIndex++))
+        }
+    }
+}

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

@@ -7,10 +7,9 @@ package org.jetbrains.kotlin.backend.jvm.lower
 
 import org.jetbrains.kotlin.backend.common.FileLoweringPass
 import org.jetbrains.kotlin.backend.common.IrElementTransformerVoidWithContext
-import org.jetbrains.kotlin.backend.common.ir.*
+import org.jetbrains.kotlin.backend.common.ir.moveBodyTo
 import org.jetbrains.kotlin.backend.common.lower.SamEqualsHashCodeMethodsGenerator
 import org.jetbrains.kotlin.backend.common.lower.VariableRemapper
-import org.jetbrains.kotlin.backend.common.lower.parents
 import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
@@ -26,7 +25,6 @@ import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.builders.*
 import org.jetbrains.kotlin.ir.builders.declarations.*
 import org.jetbrains.kotlin.ir.declarations.*
-import org.jetbrains.kotlin.ir.declarations.impl.IrFunctionImpl
 import org.jetbrains.kotlin.ir.declarations.impl.IrVariableImpl
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.*
@@ -39,7 +37,6 @@ import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 import org.jetbrains.kotlin.load.java.JvmAnnotationNames
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.name.SpecialNames
-import org.jetbrains.kotlin.util.OperatorNameConventions
 
 internal val functionReferencePhase = makeIrFilePhase(
     ::FunctionReferenceLowering,
@@ -103,108 +100,6 @@ internal class FunctionReferenceLowering(private val context: JvmBackendContext)
         return FunctionReferenceBuilder(reference).build()
     }
 
-    override fun visitCall(expression: IrCall): IrExpression {
-        if (expression.symbol.owner.isInvokeOperator()) {
-            when (val receiver = expression.dispatchReceiver) {
-                is IrFunctionReference -> {
-                    rewriteDirectInvokeToFunctionReference(expression, receiver)?.let {
-                        it.transformChildrenVoid()
-                        return it
-                    }
-                }
-                is IrBlock -> {
-                    val last = receiver.statements.last()
-                    if (last is IrFunctionReference) {
-                        rewriteDirectInvokeToLambda(expression, receiver, last)?.let {
-                            it.transformChildrenVoid()
-                            return it
-                        }
-                    }
-                }
-            }
-        }
-
-        expression.transformChildrenVoid(this)
-        return expression
-    }
-
-    private fun IrFunction.isInvokeOperator(): Boolean {
-        // For now, it's enough to check that the function name is 'invoke',
-        // because later we are looking at the dispatch receiver and check whether it's a function reference
-        // or a block returning a function reference.
-        return name == OperatorNameConventions.INVOKE
-    }
-
-    private fun rewriteDirectInvokeToLambda(irInvokeCall: IrCall, irBlock: IrBlock, lastFunRef: IrFunctionReference): IrExpression? {
-        val callee = lastFunRef.symbol.owner
-        if (callee.origin != IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA && !callee.isAnonymousFunction)
-            return null
-        if (callee.parents.any { it is IrSimpleFunction && it.isInline }) {
-            // TODO investigate why inliner gets confused when we pass it a function with some optimized direct invoke.
-            return null
-        }
-        val irDirectCall = rewriteDirectInvokeToFunctionReference(irInvokeCall, lastFunRef)
-            ?: return null
-
-        // We track instances of IrFunctionImpl corresponding to direct invoked lambdas,
-        // so we can perform optimization on it later in ExpressionCodegen.kt
-        if (callee is IrFunctionImpl) context.directInvokedLambdas.add(callee.attributeOwnerId)
-        val newBlock = IrBlockImpl(irBlock.startOffset, irBlock.endOffset, irDirectCall.type)
-        newBlock.statements.addAll(irBlock.statements)
-        newBlock.statements[newBlock.statements.lastIndex] = irDirectCall
-        return newBlock
-    }
-
-    private fun rewriteDirectInvokeToFunctionReference(irInvokeCall: IrCall, irFunRef: IrFunctionReference): IrExpression? {
-        // TODO deal with type parameters somehow?
-        // It seems we can't encounter them in the code written by user,
-        // but this might be important later if we actually perform inlining and optimizations on IR.
-        return when (val irFun = irFunRef.symbol.owner) {
-            is IrSimpleFunction -> {
-                if (irFun.typeParameters.isNotEmpty()) return null
-                IrCallImpl(
-                    irInvokeCall.startOffset, irInvokeCall.endOffset, irInvokeCall.type,
-                    irFun.symbol,
-                    typeArgumentsCount = irFun.typeParameters.size, valueArgumentsCount = irFun.valueParameters.size
-                ).apply {
-                    copyReceiverAndValueArgumentsForDirectInvoke(irFunRef, irInvokeCall)
-                }
-            }
-            is IrConstructor ->
-                IrConstructorCallImpl(
-                    irInvokeCall.startOffset, irInvokeCall.endOffset, irInvokeCall.type,
-                    irFun.symbol,
-                    typeArgumentsCount = irFun.typeParameters.size,
-                    constructorTypeArgumentsCount = 0,
-                    valueArgumentsCount = irFun.valueParameters.size
-                ).apply {
-                    copyReceiverAndValueArgumentsForDirectInvoke(irFunRef, irInvokeCall)
-                }
-            else ->
-                throw AssertionError("Simple function or constructor expected: ${irFun.render()}")
-        }
-    }
-
-    private fun IrFunctionAccessExpression.copyReceiverAndValueArgumentsForDirectInvoke(
-        irFunRef: IrFunctionReference,
-        irInvokeCall: IrFunctionAccessExpression
-    ) {
-        val irFun = irFunRef.symbol.owner
-        var invokeArgIndex = 0
-        if (irFun.dispatchReceiverParameter != null) {
-            dispatchReceiver = irFunRef.dispatchReceiver ?: irInvokeCall.getValueArgument(invokeArgIndex++)
-        }
-        if (irFun.extensionReceiverParameter != null) {
-            extensionReceiver = irFunRef.extensionReceiver ?: irInvokeCall.getValueArgument(invokeArgIndex++)
-        }
-        if (invokeArgIndex + valueArgumentsCount != irInvokeCall.valueArgumentsCount) {
-            throw AssertionError("Mismatching value arguments: $invokeArgIndex arguments used for receivers\n${irInvokeCall.dump()}")
-        }
-        for (i in 0 until valueArgumentsCount) {
-            putValueArgument(i, irInvokeCall.getValueArgument(invokeArgIndex++))
-        }
-    }
-
     private fun wrapLambdaReferenceWithIndySamConversion(
         expression: IrBlock,
         reference: IrFunctionReference,

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

@@ -149,8 +149,6 @@ class JvmBackendContext(
 
     val inlineMethodGenerationLock = Any()
 
-    val directInvokedLambdas = mutableListOf<IrAttributeContainer>()
-
     val publicAbiSymbols = mutableSetOf<IrClassSymbol>()
 
     init {

compiler/testData/codegen/asmLike/typeAnnotations/implicit.asm.ir.txt

@@ -30,10 +30,6 @@ public final class foo/Kotlin : java/lang/Object {
       @Lfoo/TypeAnn;([name="2"]) : METHOD_RETURN, null
       @Lfoo/TypeAnnBinary;([]) : METHOD_RETURN, null // invisible
 
-    private final static java.lang.String foo4$lambda$0(foo.Kotlin this$0)
-      @Lfoo/TypeAnn;([name="2"]) : METHOD_RETURN, null
-      @Lfoo/TypeAnnBinary;([]) : METHOD_RETURN, null // invisible
-
     public final void foo5()
 }
 

compiler/testData/debug/localVariables/directInvoke.kt

@@ -0,0 +1,18 @@
+// FILE: test.kt
+fun box() {
+    { a: String, b: String->
+        a + b
+    }("O", "K")
+}
+
+// EXPECTATIONS
+// EXPECTATIONS JVM_IR
+// test.kt:5 box:
+// test.kt:4 box: a:java.lang.String="O":java.lang.String, b:java.lang.String="K":java.lang.String
+// EXPECTATIONS JVM
+// test.kt:3 box:
+// test.kt:5 box:
+// test.kt:4 invoke: a:java.lang.String="O":java.lang.String, b:java.lang.String="K":java.lang.String
+// test.kt:5 box:
+// EXPECTATIONS
+// test.kt:6 box:

compiler/testData/debug/stepping/anonymousFunctionDirect.kt

@@ -7,11 +7,11 @@ fun box() {
 }
 
 // EXPECTATIONS
-// test.kt:4 box
-// EXPECTATIONS JVM
-// test.kt:5 invoke
 // EXPECTATIONS JVM_IR
-// test.kt:5 box$lambda$0
-// EXPECTATIONS
+// test.kt:5 box
+// EXPECTATIONS JVM
 // test.kt:4 box
+// test.kt:5 invoke
+// test.kt:4 box
+// EXPECTATIONS
 // test.kt:7 box

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrLocalVariableTestGenerated.java

@@ -43,6 +43,12 @@ public class IrLocalVariableTestGenerated extends AbstractIrLocalVariableTest {
         runTest("compiler/testData/debug/localVariables/copyFunction.kt");
     }
 
+    @Test
+    @TestMetadata("directInvoke.kt")
+    public void testDirectInvoke() throws Exception {
+        runTest("compiler/testData/debug/localVariables/directInvoke.kt");
+    }
+
     @Test
     @TestMetadata("doWhile.kt")
     public void testDoWhile() throws Exception {

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/LocalVariableTestGenerated.java

@@ -43,6 +43,12 @@ public class LocalVariableTestGenerated extends AbstractLocalVariableTest {
         runTest("compiler/testData/debug/localVariables/copyFunction.kt");
     }
 
+    @Test
+    @TestMetadata("directInvoke.kt")
+    public void testDirectInvoke() throws Exception {
+        runTest("compiler/testData/debug/localVariables/directInvoke.kt");
+    }
+
     @Test
     @TestMetadata("doWhile.kt")
     public void testDoWhile() throws Exception {