Don't generate nullability assertions in methods for directly invoked lambdas

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

@@ -295,6 +295,12 @@ public class FirBytecodeTextTestGenerated extends AbstractFirBytecodeTextTest {
         runTest("compiler/testData/codegen/bytecodeText/kt3845.kt");
     }
 
+    @Test
+    @TestMetadata("kt46615.kt")
+    public void testKt46615() throws Exception {
+        runTest("compiler/testData/codegen/bytecodeText/kt46615.kt");
+    }
+
     @Test
     @TestMetadata("kt5016.kt")
     public void testKt5016() throws Exception {

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

@@ -134,6 +134,8 @@ class JvmBackendContext(
 
     val inlineMethodGenerationLock = Any()
 
+    val directInvokedLambdas = mutableListOf<IrAttributeContainer>()
+
     init {
         state.mapInlineClass = { descriptor ->
             typeMapper.mapType(referenceClass(descriptor).defaultType)

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

@@ -38,6 +38,7 @@ import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
 import org.jetbrains.kotlin.diagnostics.Errors
 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.*
@@ -302,10 +303,15 @@ class ExpressionCodegen(
 
         // Do not generate non-null checks for suspend functions. When resumed the arguments
         // will be null and the actual values are taken from the continuation.
-
         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/src/org/jetbrains/kotlin/backend/jvm/lower/FunctionReferenceLowering.kt

@@ -26,6 +26,7 @@ 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.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.*
 import org.jetbrains.kotlin.ir.symbols.IrFunctionSymbol
@@ -153,6 +154,10 @@ internal class FunctionReferenceLowering(private val context: JvmBackendContext)
         }
         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

compiler/testData/codegen/bytecodeText/kt46615.kt

@@ -0,0 +1,11 @@
+// TARGET_BACKEND: JVM_IR
+
+fun box() {
+    val str = "OK"
+    val a = { s: String -> s }("OK")
+    val b = { s: String -> s }(str)
+    val c = { s: String -> s }
+    c.invoke("OK")
+}
+
+// 1 checkNotNullParameter

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

@@ -295,6 +295,12 @@ public class IrBytecodeTextTestGenerated extends AbstractIrBytecodeTextTest {
         runTest("compiler/testData/codegen/bytecodeText/kt3845.kt");
     }
 
+    @Test
+    @TestMetadata("kt46615.kt")
+    public void testKt46615() throws Exception {
+        runTest("compiler/testData/codegen/bytecodeText/kt46615.kt");
+    }
+
     @Test
     @TestMetadata("kt5016.kt")
     public void testKt5016() throws Exception {