FIR: rewrite lambda return type inference

* `return` should only be added to the last statement if the return
   type is not Unit

 * If there is a `return` without an argument, then the expected return
   type is Unit and the last expression is not a return argument (unless
   it's an incomplete call, in which case it is inferred to return Unit;
   this behavior is questionable, but inherited from K1)

 * There should be a constraint on return arguments even if the expected
   type is Unit, otherwise errors will be missed

 * When the expected type is known, using the call completion results
   writer is pointless (and probably subtly wrong).

^KT-54742 Fixed

.idea/artifacts/kotlin_dom_api_compat_1_8_255_SNAPSHOT.xml

@@ -0,0 +1,8 @@
+<component name="ArtifactManager">
+  <artifact type="jar" name="kotlin-dom-api-compat-1.8.255-SNAPSHOT">
+    <output-path>$PROJECT_DIR$/libraries/kotlin-dom-api-compat/build/libs</output-path>
+    <root id="archive" name="kotlin-dom-api-compat-1.8.255-SNAPSHOT.jar">
+      <element id="module-output" name="kotlin.kotlin-dom-api-compat.main" />
+    </root>
+  </artifact>
+</component>

compiler/fir/analysis-tests/testData/resolve/expresssions/invoke/threeReceiversCorrect.fir.txt

@@ -31,9 +31,9 @@ FILE: threeReceiversCorrect.kt
                 (this@R|special/anonymous|, R|<local>/b|).R|/A.foo|.R|SubstitutionOverride<kotlin/Function1.invoke: R|kotlin/Unit|>|(R|<local>/c|)
             }
             )
-            ^ R|/with|<R|B|, R|kotlin/Unit|>(R|<local>/b|, <L> = with@fun R|B|.<anonymous>(): R|kotlin/Unit| <inline=NoInline>  {
+            R|/with|<R|B|, R|kotlin/Unit|>(R|<local>/b|, <L> = with@fun R|B|.<anonymous>(): R|kotlin/Unit| <inline=NoInline>  {
                 (this@R|special/anonymous|, this@R|special/anonymous|).R|/A.foo|.R|SubstitutionOverride<kotlin/Function1.invoke: R|kotlin/Unit|>|(R|<local>/c|)
-                ^ R|/with|<R|C|, R|kotlin/Unit|>(R|<local>/c|, <L> = with@fun R|C|.<anonymous>(): R|kotlin/Unit| <inline=NoInline>  {
+                R|/with|<R|C|, R|kotlin/Unit|>(R|<local>/c|, <L> = with@fun R|C|.<anonymous>(): R|kotlin/Unit| <inline=NoInline>  {
                     (this@R|special/anonymous|, this@R|special/anonymous|).R|/A.foo|.R|SubstitutionOverride<kotlin/Function1.invoke: R|kotlin/Unit|>|(this@R|special/anonymous|)
                 }
                 )

compiler/fir/analysis-tests/testData/resolve/inference/coercionToUnitWithEarlyReturn.fir.txt

@@ -18,7 +18,7 @@ FILE: coercionToUnitWithEarlyReturn.kt
                 }
             }
 
-            ^ R|<local>/x|?.{ $subj$.R|/A.unit|() }
+            R|<local>/x|?.{ $subj$.R|/A.unit|() }
         }
 
         R|/foo|(R|<local>/lambda|)

compiler/fir/analysis-tests/testData/resolve/smartcasts/receivers/implicitReceivers.fir.txt

@@ -51,7 +51,7 @@ FILE: implicitReceivers.kt
     }
     public final fun test_3(a: R|kotlin/Any|, b: R|kotlin/Any|, c: R|kotlin/Any|): R|kotlin/Unit| {
         R|kotlin/with|<R|kotlin/Any|, R|kotlin/Unit|>(R|<local>/a|, <L> = wa@fun R|kotlin/Any|.<anonymous>(): R|kotlin/Unit| <inline=Inline, kind=EXACTLY_ONCE>  {
-            ^ R|kotlin/with|<R|kotlin/Any|, R|kotlin/Unit|>(R|<local>/b|, <L> = wb@fun R|kotlin/Any|.<anonymous>(): R|kotlin/Unit| <inline=Inline, kind=EXACTLY_ONCE>  {
+            R|kotlin/with|<R|kotlin/Any|, R|kotlin/Unit|>(R|<local>/b|, <L> = wb@fun R|kotlin/Any|.<anonymous>(): R|kotlin/Unit| <inline=Inline, kind=EXACTLY_ONCE>  {
                 R|kotlin/with|<R|kotlin/Any|, R|kotlin/Unit|>(R|<local>/c|, <L> = wc@fun R|kotlin/Any|.<anonymous>(): R|kotlin/Unit| <inline=Inline, kind=EXACTLY_ONCE>  {
                     (this@R|special/anonymous| as R|A|)
                     this@R|special/anonymous|.R|/A.foo|()

compiler/fir/analysis-tests/testData/resolveWithStdlib/diagnostics/functionReturnTypeMismatchChecker/basic.fir.txt

@@ -48,7 +48,7 @@ FILE: basic.kt
         foo@fun <anonymous>(): R|kotlin/Unit| <inline=Unknown>  {
             ^@foo Unit
         }
-        .R|/foo|(foo@fun R|A|.<anonymous>(): R|kotlin/Unit| <inline=NoInline>  {
+        .R|/foo<Inapplicable(INAPPLICABLE): /foo>#|(foo@fun R|A|.<anonymous>(): R|kotlin/Unit| <inline=NoInline>  {
             this@R|special/anonymous|.R|/A.bar|()
             ^@foo Int(10)
         }

compiler/fir/analysis-tests/testData/resolveWithStdlib/diagnostics/functionReturnTypeMismatchChecker/basic.kt

@@ -37,7 +37,7 @@ fun errorWithLambda(): String {
         return@foo
     } foo {
         bar()
-        return@foo 10
+        return@foo <!ARGUMENT_TYPE_MISMATCH!>10<!>
     }
 
     return ""

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/ResolveUtils.kt

@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.fir.diagnostics.ConeStubDiagnostic
 import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.builder.*
+import org.jetbrains.kotlin.fir.expressions.impl.FirUnitExpression
 import org.jetbrains.kotlin.fir.references.FirErrorNamedReference
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.references.FirSuperReference
@@ -48,6 +49,7 @@ import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.builder.buildErrorTypeRef
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
+import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.name.StandardClassIds
@@ -63,6 +65,39 @@ import kotlin.contracts.contract
 fun List<FirQualifierPart>.toTypeProjections(): Array<ConeTypeProjection> =
     asReversed().flatMap { it.typeArgumentList.typeArguments.map { typeArgument -> typeArgument.toConeTypeProjection() } }.toTypedArray()
 
+fun FirAnonymousFunction.shouldReturnUnit(returnStatements: Collection<FirStatement>): Boolean =
+    isLambda && returnStatements.any { it is FirUnitExpression }
+
+fun FirAnonymousFunction.isExplicitlySuspend(session: FirSession): Boolean =
+    typeRef.coneTypeSafe<ConeKotlinType>()?.isSuspendFunctionType(session) == true
+
+fun FirAnonymousFunction.addReturnToLastStatementIfNeeded() {
+    // If this lambda's resolved, expected return type is Unit, we don't need an explicit return statement.
+    // During conversion (to backend IR), the last expression will be coerced to Unit if needed.
+    if (returnTypeRef.isUnit) return
+
+    val body = this.body ?: return
+    val lastStatement = body.statements.lastOrNull() as? FirExpression ?: return
+    if (lastStatement is FirReturnExpression) return
+
+    val returnType = (body.typeRef as? FirResolvedTypeRef) ?: return
+    if (returnType.isNothing || returnType.isUnit) return
+
+    val returnTarget = FirFunctionTarget(null, isLambda = isLambda).also { it.bind(this) }
+    val returnExpression = buildReturnExpression {
+        source = lastStatement.source?.fakeElement(KtFakeSourceElementKind.ImplicitReturn.FromLastStatement)
+        result = lastStatement
+        target = returnTarget
+    }
+    body.transformStatements(
+        object : FirTransformer<Nothing?>() {
+            override fun <E : FirElement> transformElement(element: E, data: Nothing?): E =
+                @Suppress("UNCHECKED_CAST")
+                if (element == lastStatement) returnExpression as E else element
+        }, null
+    )
+}
+
 fun FirFunction.constructFunctionalType(isSuspend: Boolean = false): ConeLookupTagBasedType {
     val receiverTypeRef = when (this) {
         is FirSimpleFunction -> receiverParameter

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/CandidateFactory.kt

@@ -167,18 +167,22 @@ class CandidateFactory private constructor(
     }
 }
 
-fun PostponedArgumentsAnalyzerContext.addSubsystemFromExpression(statement: FirStatement) {
+fun PostponedArgumentsAnalyzerContext.addSubsystemFromExpression(statement: FirStatement): Boolean {
-    when (statement) {
+    return when (statement) {
         is FirQualifiedAccessExpression,
         is FirWhenExpression,
         is FirTryExpression,
         is FirCheckNotNullCall,
-        is FirElvisExpression
+        is FirElvisExpression -> {
-        -> (statement as FirResolvable).candidate()?.let { addOtherSystem(it.system.asReadOnlyStorage()) }
+            val candidate = (statement as FirResolvable).candidate() ?: return false
+            addOtherSystem(candidate.system.asReadOnlyStorage())
+            true
+        }
 
         is FirSafeCallExpression -> addSubsystemFromExpression(statement.selector)
         is FirWrappedArgumentExpression -> addSubsystemFromExpression(statement.expression)
-        is FirBlock -> statement.returnExpressions().forEach { addSubsystemFromExpression(it) }
+        is FirBlock -> statement.returnExpressions().any { addSubsystemFromExpression(it) }
+        else -> false
     }
 }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/PostponedArgumentsAnalyzer.kt

@@ -14,12 +14,12 @@ import org.jetbrains.kotlin.fir.references.builder.buildErrorNamedReference
 import org.jetbrains.kotlin.fir.resolve.calls.*
 import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeUnresolvedReferenceError
 import org.jetbrains.kotlin.fir.resolve.inference.model.ConeLambdaArgumentConstraintPosition
+import org.jetbrains.kotlin.fir.resolve.shouldReturnUnit
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.ConeTypeVariable
 import org.jetbrains.kotlin.fir.types.builder.buildErrorTypeRef
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
-import org.jetbrains.kotlin.fir.types.coneType
 import org.jetbrains.kotlin.fir.types.isMarkedNullable
 import org.jetbrains.kotlin.resolve.calls.components.PostponedArgumentsAnalyzerContext
 import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemBuilder
@@ -170,23 +170,29 @@ class PostponedArgumentsAnalyzer(
     ) {
         val (returnArguments, inferenceSession) = results
 
-        returnArguments.forEach { c.addSubsystemFromExpression(it) }
         val checkerSink: CheckerSink = CheckerSinkImpl(candidate)
         val builder = c.getBuilder()
 
         val lastExpression = lambda.atom.body?.statements?.lastOrNull() as? FirExpression
         var hasExpressionInReturnArguments = false
-        // No constraint for return expressions of lambda if it has Unit return type.
+        val lambdaReturnType = lambda.returnType.let(substitute)
-        val lambdaReturnType = lambda.returnType.let(substitute).takeUnless { it.isUnitOrFlexibleUnit }
         returnArguments.forEach {
+            val haveSubsystem = c.addSubsystemFromExpression(it)
             if (it !is FirExpression) return@forEach
+            // If the lambda returns Unit, the last expression is not returned and should not be constrained.
+            // TODO (KT-55837) questionable moment inherited from FE1.0 (the `haveSubsystem` case):
+            //    fun <T> foo(): T
+            //    run {
+            //      if (p) return@run
+            //      foo() // T = Unit, even though there is no implicit return
+            //    }
+            //  Things get even weirder if T has an upper bound incompatible with Unit.
+            if (it == lastExpression && !haveSubsystem &&
+                (expectedReturnType?.isUnitOrFlexibleUnit == true || lambda.atom.shouldReturnUnit(returnArguments))
+            ) return@forEach
+
             hasExpressionInReturnArguments = true
-            // If it is the last expression, and the expected type is Unit, that expression will be coerced to Unit.
+            if (!builder.hasContradiction) {
-            // If the last expression is of Unit type, of course it's not coercion-to-Unit case.
-            val lastExpressionCoercedToUnit =
-                it == lastExpression && expectedReturnType?.isUnitOrFlexibleUnit == true && !it.typeRef.coneType.isUnitOrFlexibleUnit
-            // No constraint for the last expression of lambda if it will be coerced to Unit.
-            if (!lastExpressionCoercedToUnit && !builder.hasContradiction) {
                 candidate.resolveArgumentExpression(
                     builder,
                     it,
@@ -200,7 +206,7 @@ class PostponedArgumentsAnalyzer(
             }
         }
 
-        if (!hasExpressionInReturnArguments && lambdaReturnType != null) {
+        if (!hasExpressionInReturnArguments && !lambdaReturnType.isUnitOrFlexibleUnit) {
             builder.addSubtypeConstraint(
                 components.session.builtinTypes.unitType.type,
                 lambdaReturnType,

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirCallCompletionResultsWriterTransformer.kt

@@ -42,7 +42,6 @@ import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.builder.buildStarProjection
 import org.jetbrains.kotlin.fir.types.builder.buildTypeProjectionWithVariance
 import org.jetbrains.kotlin.fir.types.impl.ConeTypeParameterTypeImpl
-import org.jetbrains.kotlin.fir.types.impl.FirImplicitUnitTypeRef
 import org.jetbrains.kotlin.fir.visitors.FirDefaultTransformer
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.transformSingle
@@ -575,12 +574,10 @@ class FirCallCompletionResultsWriterTransformer(
         anonymousFunction: FirAnonymousFunction,
         data: ExpectedArgumentType?,
     ): FirStatement {
-        // This case is not common, and happens when there are anonymous function arguments that aren't mapped to any parameter in the call
+        // The case where we can't find any return expressions not common, and happens when there are anonymous function arguments
-        // So, we don't run body resolve transformation for them, thus there's no control flow info either
+        // that aren't mapped to any parameter in the call. So, we don't run body resolve transformation for them, thus there's
-        // Control flow info is necessary prerequisite because we collect return expressions in that function
+        // no control flow info either. Example: second lambda in the call like list.filter({}, {})
-        //
+        val returnStatements = dataFlowAnalyzer.returnExpressionsOfAnonymousFunctionOrNull(anonymousFunction)
-        // Example: second lambda in the call like list.filter({}, {})
-        val returnExpressionsOfAnonymousFunction = dataFlowAnalyzer.returnExpressionsOfAnonymousFunctionOrNull(anonymousFunction)
             ?: return transformImplicitTypeRefInAnonymousFunction(anonymousFunction)
 
         val expectedType = data?.getExpectedType(anonymousFunction)?.let { expectedArgumentType ->
@@ -609,7 +606,7 @@ class FirCallCompletionResultsWriterTransformer(
             }
         }
 
-        var needUpdateLambdaType = false
+        var needUpdateLambdaType = anonymousFunction.typeRef is FirImplicitTypeRef
 
         val receiverParameter = anonymousFunction.receiverParameter
         val initialReceiverType = receiverParameter?.typeRef?.coneTypeSafe<ConeKotlinType>()
@@ -619,64 +616,39 @@ class FirCallCompletionResultsWriterTransformer(
             needUpdateLambdaType = true
         }
 
-        val initialType = anonymousFunction.returnTypeRef.coneTypeSafe<ConeKotlinType>()
+        val initialReturnType = anonymousFunction.returnTypeRef.coneTypeSafe<ConeKotlinType>()
+        val expectedReturnType = initialReturnType?.let { finalSubstitutor.substituteOrSelf(it) }
+            ?: expectedType?.returnType(session) as? ConeClassLikeType
+            ?: (data as? ExpectedArgumentType.ArgumentsMap)?.lambdasReturnTypes?.get(anonymousFunction)
 
-        val finalType = if (anonymousFunction.isLambda) {
+        val newData = expectedReturnType?.toExpectedType()
-            expectedType?.returnType(session) as? ConeClassLikeType
+        val result = transformElement(anonymousFunction, newData)
-                ?: (data as? ExpectedArgumentType.ArgumentsMap)?.lambdasReturnTypes?.get(anonymousFunction)
+        for (expression in returnStatements) {
-                ?: initialType?.let(finalSubstitutor::substituteOrSelf)
+            expression.transformSingle(this, newData)
-        } else {
-            initialType?.let(finalSubstitutor::substituteOrSelf)
-                ?: expectedType?.returnType(session) as? ConeClassLikeType
-                ?: (data as? ExpectedArgumentType.ArgumentsMap)?.lambdasReturnTypes?.get(anonymousFunction)
         }
 
-        if (finalType != null) {
+        // Prefer the expected type over the inferred one - the latter is a subtype of the former in valid code,
-            if (anonymousFunction.returnTypeRef !is FirImplicitUnitTypeRef) {
+        // and there will be ARGUMENT_TYPE_MISMATCH errors on the lambda's return expressions in invalid code.
-                val resultType = anonymousFunction.returnTypeRef.withReplacedConeType(finalType)
+        val resultReturnType = expectedReturnType
-                anonymousFunction.replaceReturnTypeRef(resultType)
+            ?: session.typeContext.commonSuperTypeOrNull(returnStatements.mapNotNull { (it as? FirExpression)?.resultType?.coneType })
-            }
+            ?: session.builtinTypes.unitType.type
+
+        if (initialReturnType != resultReturnType) {
+            result.replaceReturnTypeRef(result.returnTypeRef.resolvedTypeFromPrototype(resultReturnType))
+            session.lookupTracker?.recordTypeResolveAsLookup(result.returnTypeRef, result.source, context.file.source)
             needUpdateLambdaType = true
         }
 
         if (needUpdateLambdaType) {
-            val resolvedTypeRef =
+            val isSuspend = expectedType?.isSuspendFunctionType(session) ?: result.isExplicitlySuspend(session)
-                anonymousFunction.constructFunctionalTypeRef(
+            result.replaceTypeRef(result.constructFunctionalTypeRef(isSuspend))
-                    isSuspend = expectedType?.isSuspendFunctionType(session) == true ||
+            session.lookupTracker?.recordTypeResolveAsLookup(result.typeRef, result.source, context.file.source)
-                            (expectedType == null && anonymousFunction.isSuspendFunctionType())
-                )
-            anonymousFunction.replaceTypeRef(resolvedTypeRef)
-            session.lookupTracker?.recordTypeResolveAsLookup(resolvedTypeRef, anonymousFunction.source, context.file.source)
         }
-
+        // Have to delay this until the type is written to avoid adding a return if the type is Unit.
-        val result = transformElement(anonymousFunction, null)
+        result.addReturnToLastStatementIfNeeded()
-
-        for (expression in returnExpressionsOfAnonymousFunction) {
-            expression.transform<FirElement, ExpectedArgumentType?>(this, finalType?.toExpectedType())
-        }
-
-        if (result.returnTypeRef.coneTypeSafe<ConeIntegerLiteralType>() != null) {
-            val lastExpressionType =
-                (returnExpressionsOfAnonymousFunction.lastOrNull() as? FirExpression)
-                    ?.typeRef?.coneTypeSafe<ConeKotlinType>()
-
-            val newReturnTypeRef = result.returnTypeRef.withReplacedConeType(lastExpressionType)
-            result.replaceReturnTypeRef(newReturnTypeRef)
-            val resolvedTypeRef =
-                result.constructFunctionalTypeRef(isSuspend = expectedType?.isSuspendFunctionType(session) == true)
-            result.replaceTypeRef(resolvedTypeRef)
-            session.lookupTracker?.let {
-                it.recordTypeResolveAsLookup(newReturnTypeRef, anonymousFunction.source, context.file.source)
-                it.recordTypeResolveAsLookup(resolvedTypeRef, anonymousFunction.source, context.file.source)
-            }
-        }
-
         return result
     }
 
-    private fun FirAnonymousFunction.isSuspendFunctionType() =
-        typeRef.coneTypeSafe<ConeKotlinType>()?.isSuspendFunctionType(session) == true
-
     private fun transformImplicitTypeRefInAnonymousFunction(
         anonymousFunction: FirAnonymousFunction
     ): FirStatement {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt

@@ -22,8 +22,6 @@ import org.jetbrains.kotlin.fir.declarations.utils.*
 import org.jetbrains.kotlin.fir.diagnostics.ConeSimpleDiagnostic
 import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind
 import org.jetbrains.kotlin.fir.expressions.*
-import org.jetbrains.kotlin.fir.expressions.builder.buildReturnExpression
-import org.jetbrains.kotlin.fir.expressions.builder.buildUnitExpression
 import org.jetbrains.kotlin.fir.expressions.impl.FirLazyBlock
 import org.jetbrains.kotlin.fir.references.FirResolvedErrorReference
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
@@ -36,7 +34,6 @@ import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeLocalVariableNoTypeOrIni
 import org.jetbrains.kotlin.fir.resolve.inference.FirStubTypeTransformer
 import org.jetbrains.kotlin.fir.resolve.inference.ResolvedLambdaAtom
 import org.jetbrains.kotlin.fir.resolve.inference.extractLambdaInfoFromFunctionalType
-import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.substitution.createTypeSubstitutorByTypeConstructor
 import org.jetbrains.kotlin.fir.resolve.transformers.FirCallCompletionResultsWriterTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.FirStatusResolver
@@ -49,8 +46,6 @@ import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.builder.buildErrorTypeRef
 import org.jetbrains.kotlin.fir.types.builder.buildImplicitTypeRef
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
-import org.jetbrains.kotlin.fir.types.impl.FirImplicitUnitTypeRef
-import org.jetbrains.kotlin.fir.visitors.FirDefaultTransformer
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.transformSingle
 import org.jetbrains.kotlin.name.Name
@@ -590,34 +585,6 @@ open class FirDeclarationsResolveTransformer(transformer: FirAbstractBodyResolve
         return result
     }
 
-    private fun transformAnonymousFunctionWithLambdaResolution(
-        anonymousFunction: FirAnonymousFunction, lambdaResolution: ResolutionMode.LambdaResolution
-    ): FirAnonymousFunction {
-        val expectedReturnType =
-            lambdaResolution.expectedReturnTypeRef ?: anonymousFunction.returnTypeRef.takeUnless { it is FirImplicitTypeRef }
-        val result = transformFunction(anonymousFunction, withExpectedType(expectedReturnType)) as FirAnonymousFunction
-        val body = result.body
-        if (result.returnTypeRef is FirImplicitTypeRef && body != null) {
-            // TODO: This part seems unnecessary because for lambdas in dependent context will be completed and their type
-            //  should be replaced there properly
-            val returnType =
-                dataFlowAnalyzer.returnExpressionsOfAnonymousFunction(result)
-                    .firstNotNullOfOrNull { (it as? FirExpression)?.resultType?.coneTypeSafe() }
-            val resolutionMode = if (returnType != null) {
-                withExpectedType(returnType)
-            } else {
-                withExpectedType(buildErrorTypeRef {
-                    diagnostic =
-                        ConeSimpleDiagnostic("Unresolved lambda return type", DiagnosticKind.InferenceError)
-                })
-            }
-
-            result.transformReturnTypeRef(transformer, resolutionMode)
-        }
-
-        return result
-    }
-
     override fun transformSimpleFunction(
         simpleFunction: FirSimpleFunction,
         data: ResolutionMode
@@ -785,39 +752,40 @@ open class FirDeclarationsResolveTransformer(transformer: FirAbstractBodyResolve
         return when (data) {
             is ResolutionMode.ContextDependent, is ResolutionMode.ContextDependentDelegate -> {
                 context.withAnonymousFunction(anonymousFunction, components, data) {
-                    anonymousFunction.addReturn()
+                    anonymousFunction
                 }
             }
             is ResolutionMode.LambdaResolution -> {
-                context.withAnonymousFunction(anonymousFunction, components, data) {
+                val expectedReturnTypeRef =
-                    withFullBodyResolve {
+                    data.expectedReturnTypeRef ?: anonymousFunction.returnTypeRef.takeUnless { it is FirImplicitTypeRef }
-                        transformAnonymousFunctionWithLambdaResolution(anonymousFunction, data).addReturn()
+                transformAnonymousFunctionBody(anonymousFunction, expectedReturnTypeRef, data)
-                    }
-                }
             }
-            is ResolutionMode.WithExpectedType,
+            is ResolutionMode.WithExpectedType ->
-            is ResolutionMode.ContextIndependent,
+                transformAnonymousFunctionWithExpectedType(anonymousFunction, data.expectedTypeRef, data)
-            is ResolutionMode.ReceiverResolution,
+            is ResolutionMode.WithSuggestedType ->
-            is ResolutionMode.WithSuggestedType -> {
+                transformAnonymousFunctionWithExpectedType(anonymousFunction, data.suggestedTypeRef, data)
-                val expectedTypeRef = when (data) {
+            is ResolutionMode.ContextIndependent, is ResolutionMode.ReceiverResolution ->
-                    is ResolutionMode.WithExpectedType -> {
+                transformAnonymousFunctionWithExpectedType(anonymousFunction, buildImplicitTypeRef(), data)
-                        data.expectedTypeRef
+            is ResolutionMode.WithStatus, is ResolutionMode.WithExpectedTypeFromCast ->
-                    }
-                    is ResolutionMode.WithSuggestedType -> {
-                        data.suggestedTypeRef
-                    }
-                    else -> {
-                        buildImplicitTypeRef()
-                    }
-                }
-                transformAnonymousFunctionWithExpectedType(anonymousFunction, expectedTypeRef, data)
-            }
-            is ResolutionMode.WithStatus, is ResolutionMode.WithExpectedTypeFromCast -> {
                 throw AssertionError("Should not be here in WithStatus/WithExpectedTypeFromCast mode")
-            }
         }
     }
 
+
+    private fun transformAnonymousFunctionBody(
+        anonymousFunction: FirAnonymousFunction,
+        expectedReturnTypeRef: FirTypeRef?,
+        data: ResolutionMode
+    ): FirAnonymousFunction {
+        // `transformFunction` will replace both `typeRef` and `returnTypeRef`, so make sure to keep the former.
+        val lambdaType = anonymousFunction.typeRef
+        return context.withAnonymousFunction(anonymousFunction, components, data) {
+            withFullBodyResolve {
+                transformFunction(anonymousFunction, withExpectedType(expectedReturnTypeRef)) as FirAnonymousFunction
+            }
+        }.apply { replaceTypeRef(lambdaType) }
+    }
+
     private fun transformAnonymousFunctionWithExpectedType(
         anonymousFunction: FirAnonymousFunction,
         expectedTypeRef: FirTypeRef,
@@ -829,12 +797,10 @@ open class FirDeclarationsResolveTransformer(transformer: FirAbstractBodyResolve
             )
         }
         var lambda = anonymousFunction
-        val initialReturnTypeRef = lambda.returnTypeRef
         val valueParameters = when {
             resolvedLambdaAtom != null -> obtainValueParametersFromResolvedLambdaAtom(resolvedLambdaAtom, lambda)
             else -> obtainValueParametersFromExpectedType(expectedTypeRef.coneTypeSafe(), lambda)
         }
-        val returnTypeRefFromResolvedAtom = resolvedLambdaAtom?.returnType?.let { lambda.returnTypeRef.resolvedTypeFromPrototype(it) }
         lambda = buildAnonymousFunctionCopy(lambda) {
             receiverParameter = lambda.receiverParameter?.takeIf { it.typeRef !is FirImplicitTypeRef }
                 ?: resolvedLambdaAtom?.receiver?.let { coneKotlinType ->
@@ -857,71 +823,38 @@ open class FirDeclarationsResolveTransformer(transformer: FirAbstractBodyResolve
 
             this.valueParameters.clear()
             this.valueParameters.addAll(valueParameters)
-            returnTypeRef = (lambda.returnTypeRef as? FirResolvedTypeRef)
+        }.transformValueParameters(ImplicitToErrorTypeTransformer, null)
-                ?: returnTypeRefFromResolvedAtom
-                        ?: lambda.returnTypeRef
-        }
-        lambda = lambda.transformValueParameters(ImplicitToErrorTypeTransformer, null)
-        val bodyExpectedType = returnTypeRefFromResolvedAtom ?: expectedTypeRef
-        context.withAnonymousFunction(lambda, components, data) {
-            withFullBodyResolve {
-                lambda = transformFunction(lambda, withExpectedType(bodyExpectedType)) as FirAnonymousFunction
-            }
-        }
-        // To separate function and separate commit
-        val writer = FirCallCompletionResultsWriterTransformer(
-            session,
-            ConeSubstitutor.Empty,
-            components.returnTypeCalculator,
-            session.typeApproximator,
-            dataFlowAnalyzer,
-            components.integerLiteralAndOperatorApproximationTransformer,
-            components.context
-        )
-        lambda.transformSingle(writer, expectedTypeRef.coneTypeSafe<ConeKotlinType>()?.toExpectedType())
 
-        val returnStatements = dataFlowAnalyzer.returnExpressionsOfAnonymousFunction(lambda)
+        val initialReturnTypeRef = lambda.returnTypeRef as? FirResolvedTypeRef
-        val returnExpressionsExceptLast =
+        val expectedReturnTypeRef = initialReturnTypeRef
-            if (returnStatements.size > 1)
+            ?: resolvedLambdaAtom?.returnType?.let { lambda.returnTypeRef.resolvedTypeFromPrototype(it) }
-                returnStatements - lambda.body?.statements?.lastOrNull()
+        lambda = transformAnonymousFunctionBody(lambda, expectedReturnTypeRef ?: components.noExpectedType, data)
-            else
+
-                returnStatements
+        if (initialReturnTypeRef == null) {
-        val implicitReturns = returnExpressionsExceptLast.filter {
+            lambda.replaceReturnTypeRef(lambda.computeReturnTypeRef(expectedReturnTypeRef))
-            (it as? FirExpression)?.typeRef is FirImplicitUnitTypeRef
+            session.lookupTracker?.recordTypeResolveAsLookup(lambda.returnTypeRef, lambda.source, context.file.source)
         }
 
-        val returnType = when {
+        lambda.replaceTypeRef(lambda.constructFunctionalTypeRef(resolvedLambdaAtom?.isSuspend == true))
-            initialReturnTypeRef is FirResolvedTypeRef -> {
+        session.lookupTracker?.recordTypeResolveAsLookup(lambda.typeRef, lambda.source, context.file.source)
-                initialReturnTypeRef.coneType
+        lambda.addReturnToLastStatementIfNeeded()
-            }
+        return lambda
-            implicitReturns.isNotEmpty() || (lambda.returnType?.isUnit == true && lambda.isLambda) -> {
+    }
-                // i.e., early return, e.g., l@{ ... return@l ... }
+
-                // Note that the last statement will be coerced to Unit if needed.
+    private fun FirAnonymousFunction.computeReturnTypeRef(expected: FirResolvedTypeRef?): FirResolvedTypeRef {
-                // also we don't coerce to Unit anonymous functions, only lambdas
+        // Any lambda expression assigned to `(...) -> Unit` returns Unit
-                session.builtinTypes.unitType.type
+        if (isLambda && expected?.type?.isUnit == true) return expected
-            }
+        // `lambda@ { return@lambda }` always returns Unit
-            else -> {
+        val returnStatements = dataFlowAnalyzer.returnExpressionsOfAnonymousFunction(this)
-                // Otherwise, compute the common super type of all possible return expressions
+        if (shouldReturnUnit(returnStatements)) return session.builtinTypes.unitType
-                session.typeContext.commonSuperTypeOrNull(
+        // Here is a questionable moment where we could prefer the expected type over an inferred one.
-                    returnStatements.mapNotNull { (it as? FirExpression)?.resultType?.coneType }
+        // In correct code this doesn't matter, as all return expression types should be subtypes of the expected type.
-                ) ?: session.builtinTypes.unitType.type
+        // In incorrect code, this would change diagnostics: we can get errors either on the entire lambda, or only on its
-            }
+        // return statements. The former kind of makes more sense, but the latter is more readable.
-        }
+        val inferredFromReturnStatements =
-        if (lambda.returnTypeRef !is FirImplicitUnitTypeRef) {
+            session.typeContext.commonSuperTypeOrNull(returnStatements.mapNotNull { (it as? FirExpression)?.resultType?.coneType })
-            lambda.replaceReturnTypeRef(
+        return inferredFromReturnStatements?.let { returnTypeRef.resolvedTypeFromPrototype(it) }
-                initialReturnTypeRef.resolvedTypeFromPrototype(returnType).also {
+            ?: session.builtinTypes.unitType // Empty lambda returns Unit
-                    session.lookupTracker?.recordTypeResolveAsLookup(it, lambda.source, components.file.source)
-                }
-            )
-        }
-        lambda.replaceTypeRef(
-            lambda.constructFunctionalTypeRef(
-                isSuspend = expectedTypeRef.coneTypeSafe<ConeKotlinType>()?.isSuspendFunctionType(session) == true
-            ).also {
-                session.lookupTracker?.recordTypeResolveAsLookup(it, lambda.source, components.file.source)
-            }
-        )
-        return lambda.addReturn()
     }
 
     private fun obtainValueParametersFromResolvedLambdaAtom(
@@ -982,47 +915,6 @@ open class FirDeclarationsResolveTransformer(transformer: FirAbstractBodyResolve
         }
     }
 
-    private fun FirAnonymousFunction.addReturn(): FirAnonymousFunction {
-        // If this lambda's resolved, expected return type is Unit, we don't need an explicit return statement.
-        // During conversion (to backend IR), the last expression will be coerced to Unit if needed.
-        // As per KT-41005, we should not force coercion to Unit for nullable return type, though.
-        if (returnTypeRef.isUnit && body?.typeRef?.isMarkedNullable == false) {
-            return this
-        }
-        val lastStatement = body?.statements?.lastOrNull()
-        val returnType = (body?.typeRef as? FirResolvedTypeRef) ?: return this
-        val returnNothing = returnType.isNothing || returnType.isUnit
-        if (lastStatement is FirExpression && !returnNothing) {
-            body?.transformChildren(
-                object : FirDefaultTransformer<FirExpression>() {
-                    override fun <E : FirElement> transformElement(element: E, data: FirExpression): E {
-                        if (element == lastStatement) {
-                            val returnExpression = buildReturnExpression {
-                                source = element.source?.fakeElement(KtFakeSourceElementKind.ImplicitReturn.FromLastStatement)
-                                result = lastStatement
-                                target = FirFunctionTarget(null, isLambda = this@addReturn.isLambda).also {
-                                    it.bind(this@addReturn)
-                                }
-                            }
-                            @Suppress("UNCHECKED_CAST")
-                            return (returnExpression as E)
-                        }
-                        return element
-                    }
-
-                    override fun transformReturnExpression(
-                        returnExpression: FirReturnExpression,
-                        data: FirExpression
-                    ): FirStatement {
-                        return returnExpression
-                    }
-                },
-                buildUnitExpression()
-            )
-        }
-        return this
-    }
-
     override fun transformBackingField(
         backingField: FirBackingField,
         data: ResolutionMode,

compiler/testData/diagnostics/tests/controlStructures/ifToAnyDiscriminatingUsages.fir.kt

@@ -27,7 +27,7 @@ fun testResultOfAnonFun2() =
 
 fun testReturnFromAnonFun() =
         run(fun () {
-            return <!RETURN_TYPE_MISMATCH!>if (true) 42 else println()<!>
+            return <!ARGUMENT_TYPE_MISMATCH, RETURN_TYPE_MISMATCH!>if (true) 42 else println()<!>
         })
 
 fun <!IMPLICIT_NOTHING_RETURN_TYPE!>testReturn1<!>() =

compiler/testData/diagnostics/tests/inference/coercionToUnit/afterBareReturn.fir.kt

@@ -34,7 +34,7 @@ fun test4() = <!INFERRED_TYPE_VARIABLE_INTO_EMPTY_INTERSECTION!>run<!> {
 
 fun main() {
     // all ok
-    expectUnit(<!ARGUMENT_TYPE_MISMATCH!>test1()<!>)
+    expectUnit(test1())
     expectUnit(test2())
     expectUnit(test3())
     expectUnit(test4())

compiler/testData/diagnostics/tests/inference/completion/postponedArgumentsAnalysis/basic.fir.kt

@@ -139,8 +139,8 @@ fun main() {
      * K <: (C) -> Unit -> TypeVariable(_RP1) >: C
      * K == (B) -> Unit -> TypeVariable(_RP1) == B
      */
-    val x17: (C) -> Unit = <!INITIALIZER_TYPE_MISMATCH, TYPE_MISMATCH!>selectB(id <!ARGUMENT_TYPE_MISMATCH!>{ <!DEBUG_INFO_EXPRESSION_TYPE("B")!>it<!> }<!>, id <!ARGUMENT_TYPE_MISMATCH!>{ it }<!>, id<(B) -> Unit> { x -> x })<!>
+    val x17: (C) -> Unit = <!INITIALIZER_TYPE_MISMATCH, TYPE_MISMATCH!>selectB(id <!ARGUMENT_TYPE_MISMATCH, ARGUMENT_TYPE_MISMATCH!>{ <!DEBUG_INFO_EXPRESSION_TYPE("B")!>it<!> }<!>, id <!ARGUMENT_TYPE_MISMATCH, ARGUMENT_TYPE_MISMATCH!>{ it }<!>, id<(B) -> Unit> { x -> x })<!>
-    val x18: (C) -> Unit = <!TYPE_MISMATCH!>select(id <!ARGUMENT_TYPE_MISMATCH!>{ <!DEBUG_INFO_EXPRESSION_TYPE("C")!>it<!> }<!>, <!ARGUMENT_TYPE_MISMATCH!>{ <!DEBUG_INFO_EXPRESSION_TYPE("C")!>it<!> }<!>, id<(B) -> Unit> { x -> x })<!>
+    val x18: (C) -> Unit = <!TYPE_MISMATCH!>select(id <!ARGUMENT_TYPE_MISMATCH!>{ <!DEBUG_INFO_EXPRESSION_TYPE("C")!>it<!> }<!>, <!ARGUMENT_TYPE_MISMATCH, ARGUMENT_TYPE_MISMATCH!>{ <!DEBUG_INFO_EXPRESSION_TYPE("C")!>it<!> }<!>, id<(B) -> Unit> { x -> x })<!>
 
     // Resolution of extension/non-extension functions combination
     val x19: String.() -> Unit = select(<!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Function1<kotlin.String, kotlin.Unit>")!>id { <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.String")!>this<!> }<!>, <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Function1<kotlin.String, kotlin.Unit>")!>id(fun(x: String) {})<!>)

compiler/testData/diagnostics/tests/regressions/resolveCollectionLiteralInsideLambda.fir.kt

@@ -3,5 +3,5 @@
 fun foo(l: () -> Unit) {}
 fun bar(l: () -> String) {}
 
-val a = foo { <!UNSUPPORTED!>[]<!> }
+val a = foo { <!ARGUMENT_TYPE_MISMATCH, UNSUPPORTED!>[]<!> }
-val b = bar { <!UNSUPPORTED!>[]<!> }
+val b = bar { <!ARGUMENT_TYPE_MISMATCH, UNSUPPORTED!>[]<!> }

compiler/testData/diagnostics/testsWithStdLib/inference/completion/postponedArgumentsAnalysis/suspendFunctions.fir.kt

@@ -19,7 +19,7 @@ fun main() {
     val x1: suspend (Int) -> Unit = takeSuspend(<!DEBUG_INFO_EXPRESSION_TYPE("kotlin.coroutines.SuspendFunction1<kotlin.Int, kotlin.Unit>")!>id { it }<!>, <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.coroutines.SuspendFunction1<kotlin.Int, kotlin.Unit>")!>{ x -> x }<!>)
 
     // Here, the error should be
-    val x2: (Int) -> Unit = <!INITIALIZER_TYPE_MISMATCH!>takeSuspend(<!DEBUG_INFO_EXPRESSION_TYPE("kotlin.coroutines.SuspendFunction1<kotlin.Int, kotlin.Unit>")!>id <!ARGUMENT_TYPE_MISMATCH, ARGUMENT_TYPE_MISMATCH!>{ it }<!><!>, <!ARGUMENT_TYPE_MISMATCH, ARGUMENT_TYPE_MISMATCH, DEBUG_INFO_EXPRESSION_TYPE("kotlin.coroutines.SuspendFunction1<kotlin.Int, kotlin.Unit>")!>{ x -> x }<!>)<!>
+    val x2: (Int) -> Unit = <!INITIALIZER_TYPE_MISMATCH!>takeSuspend(<!DEBUG_INFO_EXPRESSION_TYPE("kotlin.coroutines.SuspendFunction1<kotlin.Int, kotlin.Unit>")!>id <!ARGUMENT_TYPE_MISMATCH, ARGUMENT_TYPE_MISMATCH, ARGUMENT_TYPE_MISMATCH!>{ it }<!><!>, <!ARGUMENT_TYPE_MISMATCH, ARGUMENT_TYPE_MISMATCH, ARGUMENT_TYPE_MISMATCH, ARGUMENT_TYPE_MISMATCH, DEBUG_INFO_EXPRESSION_TYPE("kotlin.coroutines.SuspendFunction1<kotlin.Int, kotlin.Unit>")!>{ x -> x }<!>)<!>
-    val x3: suspend (Int) -> Unit = takeSimpleFunction(<!DEBUG_INFO_EXPRESSION_TYPE("kotlin.coroutines.SuspendFunction1<kotlin.Int, kotlin.Unit>")!>id <!ARGUMENT_TYPE_MISMATCH, ARGUMENT_TYPE_MISMATCH!>{ it }<!><!>, <!ARGUMENT_TYPE_MISMATCH, ARGUMENT_TYPE_MISMATCH, DEBUG_INFO_EXPRESSION_TYPE("kotlin.coroutines.SuspendFunction1<kotlin.Int, kotlin.Unit>")!>{ x -> x }<!>)
+    val x3: suspend (Int) -> Unit = takeSimpleFunction(<!DEBUG_INFO_EXPRESSION_TYPE("kotlin.coroutines.SuspendFunction1<kotlin.Int, kotlin.Unit>")!>id <!ARGUMENT_TYPE_MISMATCH, ARGUMENT_TYPE_MISMATCH, ARGUMENT_TYPE_MISMATCH!>{ it }<!><!>, <!ARGUMENT_TYPE_MISMATCH, ARGUMENT_TYPE_MISMATCH, ARGUMENT_TYPE_MISMATCH, ARGUMENT_TYPE_MISMATCH, DEBUG_INFO_EXPRESSION_TYPE("kotlin.coroutines.SuspendFunction1<kotlin.Int, kotlin.Unit>")!>{ x -> x }<!>)
     val x4: (Int) -> Unit = takeSimpleFunction(<!ARGUMENT_TYPE_MISMATCH!>id<suspend (Int) -> Unit> {}<!>, <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Function1<kotlin.Int, kotlin.Unit>")!>{}<!>)
 }

compiler/testData/ir/irText/declarations/contextReceivers/functionalType.fir.ir.txt

@@ -172,78 +172,77 @@ FILE fqName:<root> fileName:/functionalType.kt
           FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> ($receiver:<root>.C) returnType:kotlin.Unit
             $receiver: VALUE_PARAMETER name:$this$with type:<root>.C
             BLOCK_BODY
-              RETURN type=kotlin.Nothing from='local final fun <anonymous> (): kotlin.Unit declared in <root>.test'
+              CALL 'public final fun with <T, R> (receiver: T of kotlin.StandardKt.with, block: @[ExtensionFunctionType] kotlin.Function1<T of kotlin.StandardKt.with, R of kotlin.StandardKt.with>): R of kotlin.StandardKt.with [inline] declared in kotlin.StandardKt' type=kotlin.Unit origin=null
-                CALL 'public final fun with <T, R> (receiver: T of kotlin.StandardKt.with, block: @[ExtensionFunctionType] kotlin.Function1<T of kotlin.StandardKt.with, R of kotlin.StandardKt.with>): R of kotlin.StandardKt.with [inline] declared in kotlin.StandardKt' type=kotlin.Unit origin=null
+                <T>: <root>.R
-                  <T>: <root>.R
+                <R>: kotlin.Unit
-                  <R>: kotlin.Unit
+                receiver: CONSTRUCTOR_CALL 'public constructor <init> () [primary] declared in <root>.R' type=<root>.R origin=null
-                  receiver: CONSTRUCTOR_CALL 'public constructor <init> () [primary] declared in <root>.R' type=<root>.R origin=null
+                block: FUN_EXPR type=@[ExtensionFunctionType] kotlin.Function1<<root>.R, kotlin.Unit> origin=LAMBDA
-                  block: FUN_EXPR type=@[ExtensionFunctionType] kotlin.Function1<<root>.R, kotlin.Unit> origin=LAMBDA
+                  FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> ($receiver:<root>.R) returnType:kotlin.Unit
-                    FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> ($receiver:<root>.R) returnType:kotlin.Unit
+                    $receiver: VALUE_PARAMETER name:$this$with type:<root>.R
-                      $receiver: VALUE_PARAMETER name:$this$with type:<root>.R
+                    BLOCK_BODY
-                      BLOCK_BODY
+                      CALL 'public final fun f1 (_context_receiver_0: <root>.C, g: @[ExtensionFunctionType] kotlin.Function3<<root>.C, <root>.R, <root>.Param, kotlin.Unit>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
-                        CALL 'public final fun f1 (_context_receiver_0: <root>.C, g: @[ExtensionFunctionType] kotlin.Function3<<root>.C, <root>.R, <root>.Param, kotlin.Unit>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
+                        $receiver: GET_VAR '$this$with: <root>.R declared in <root>.test.<anonymous>.<anonymous>' type=<root>.R origin=null
-                          $receiver: GET_VAR '$this$with: <root>.R declared in <root>.test.<anonymous>.<anonymous>' type=<root>.R origin=null
+                        _context_receiver_0: GET_VAR '$this$with: <root>.C declared in <root>.test.<anonymous>' type=<root>.C origin=null
-                          _context_receiver_0: GET_VAR '$this$with: <root>.C declared in <root>.test.<anonymous>' type=<root>.C origin=null
+                        g: GET_VAR 'val lf1: @[ExtensionFunctionType] kotlin.Function3<<root>.C, <root>.R, <root>.Param, kotlin.Unit> [val] declared in <root>.test' type=@[ExtensionFunctionType] kotlin.Function3<<root>.C, <root>.R, <root>.Param, kotlin.Unit> origin=null
-                          g: GET_VAR 'val lf1: @[ExtensionFunctionType] kotlin.Function3<<root>.C, <root>.R, <root>.Param, kotlin.Unit> [val] declared in <root>.test' type=@[ExtensionFunctionType] kotlin.Function3<<root>.C, <root>.R, <root>.Param, kotlin.Unit> origin=null
+                      CALL 'public final fun f1 (_context_receiver_0: <root>.C, g: @[ExtensionFunctionType] kotlin.Function3<<root>.C, <root>.R, <root>.Param, kotlin.Unit>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
-                        CALL 'public final fun f1 (_context_receiver_0: <root>.C, g: @[ExtensionFunctionType] kotlin.Function3<<root>.C, <root>.R, <root>.Param, kotlin.Unit>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
+                        $receiver: GET_VAR '$this$with: <root>.R declared in <root>.test.<anonymous>.<anonymous>' type=<root>.R origin=null
-                          $receiver: GET_VAR '$this$with: <root>.R declared in <root>.test.<anonymous>.<anonymous>' type=<root>.R origin=null
+                        _context_receiver_0: GET_VAR '$this$with: <root>.C declared in <root>.test.<anonymous>' type=<root>.C origin=null
-                          _context_receiver_0: GET_VAR '$this$with: <root>.C declared in <root>.test.<anonymous>' type=<root>.C origin=null
+                        g: FUN_EXPR type=@[ExtensionFunctionType] kotlin.Function3<<root>.C, <root>.R, <root>.Param, kotlin.Unit> origin=LAMBDA
-                          g: FUN_EXPR type=@[ExtensionFunctionType] kotlin.Function3<<root>.C, <root>.R, <root>.Param, kotlin.Unit> origin=LAMBDA
+                          FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> ($receiver:<root>.R, _context_receiver_0:<root>.C, <unused var>:<root>.Param) returnType:kotlin.Unit
-                            FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> ($receiver:<root>.R, _context_receiver_0:<root>.C, <unused var>:<root>.Param) returnType:kotlin.Unit
+                            contextReceiverParametersCount: 1
-                              contextReceiverParametersCount: 1
+                            $receiver: VALUE_PARAMETER name:$this$f1 type:<root>.R
-                              $receiver: VALUE_PARAMETER name:$this$f1 type:<root>.R
+                            VALUE_PARAMETER name:_context_receiver_0 index:0 type:<root>.C
-                              VALUE_PARAMETER name:_context_receiver_0 index:0 type:<root>.C
+                            VALUE_PARAMETER UNDERSCORE_PARAMETER name:<unused var> index:1 type:<root>.Param
-                              VALUE_PARAMETER UNDERSCORE_PARAMETER name:<unused var> index:1 type:<root>.Param
+                            BLOCK_BODY
-                              BLOCK_BODY
+                              TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit
-                                TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit
+                                CALL 'public final fun <get-r> (): kotlin.Int declared in <root>.R' type=kotlin.Int origin=GET_PROPERTY
-                                  CALL 'public final fun <get-r> (): kotlin.Int declared in <root>.R' type=kotlin.Int origin=GET_PROPERTY
+                                  $this: GET_VAR '$this$f1: <root>.R declared in <root>.test.<anonymous>.<anonymous>.<anonymous>' type=<root>.R origin=null
-                                    $this: GET_VAR '$this$f1: <root>.R declared in <root>.test.<anonymous>.<anonymous>.<anonymous>' type=<root>.R origin=null
+                              TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit
-                                TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit
+                                CALL 'public final fun <get-c> (): kotlin.Int declared in <root>.C' type=kotlin.Int origin=GET_PROPERTY
-                                  CALL 'public final fun <get-c> (): kotlin.Int declared in <root>.C' type=kotlin.Int origin=GET_PROPERTY
+                                  $this: GET_VAR '$this$with: <root>.C declared in <root>.test.<anonymous>' type=<root>.C origin=null
-                                    $this: GET_VAR '$this$with: <root>.C declared in <root>.test.<anonymous>' type=<root>.C origin=null
+                      CALL 'public final fun f2 (_context_receiver_0: <root>.C, g: kotlin.Function2<<root>.C, <root>.Param, kotlin.Unit>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
-                        CALL 'public final fun f2 (_context_receiver_0: <root>.C, g: kotlin.Function2<<root>.C, <root>.Param, kotlin.Unit>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
+                        _context_receiver_0: GET_VAR '$this$with: <root>.C declared in <root>.test.<anonymous>' type=<root>.C origin=null
-                          _context_receiver_0: GET_VAR '$this$with: <root>.C declared in <root>.test.<anonymous>' type=<root>.C origin=null
+                        g: GET_VAR 'val lf2: kotlin.Function2<<root>.C, <root>.Param, kotlin.Unit> [val] declared in <root>.test' type=kotlin.Function2<<root>.C, <root>.Param, kotlin.Unit> origin=null
-                          g: GET_VAR 'val lf2: kotlin.Function2<<root>.C, <root>.Param, kotlin.Unit> [val] declared in <root>.test' type=kotlin.Function2<<root>.C, <root>.Param, kotlin.Unit> origin=null
+                      CALL 'public final fun f2 (_context_receiver_0: <root>.C, g: kotlin.Function2<<root>.C, <root>.Param, kotlin.Unit>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
-                        CALL 'public final fun f2 (_context_receiver_0: <root>.C, g: kotlin.Function2<<root>.C, <root>.Param, kotlin.Unit>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
+                        _context_receiver_0: GET_VAR '$this$with: <root>.C declared in <root>.test.<anonymous>' type=<root>.C origin=null
-                          _context_receiver_0: GET_VAR '$this$with: <root>.C declared in <root>.test.<anonymous>' type=<root>.C origin=null
+                        g: FUN_EXPR type=kotlin.Function2<<root>.C, <root>.Param, kotlin.Unit> origin=LAMBDA
-                          g: FUN_EXPR type=kotlin.Function2<<root>.C, <root>.Param, kotlin.Unit> origin=LAMBDA
+                          FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> (_context_receiver_0:<root>.C, <unused var>:<root>.Param) returnType:kotlin.Unit
-                            FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> (_context_receiver_0:<root>.C, <unused var>:<root>.Param) returnType:kotlin.Unit
+                            contextReceiverParametersCount: 1
-                              contextReceiverParametersCount: 1
+                            VALUE_PARAMETER name:_context_receiver_0 index:0 type:<root>.C
-                              VALUE_PARAMETER name:_context_receiver_0 index:0 type:<root>.C
+                            VALUE_PARAMETER UNDERSCORE_PARAMETER name:<unused var> index:1 type:<root>.Param
-                              VALUE_PARAMETER UNDERSCORE_PARAMETER name:<unused var> index:1 type:<root>.Param
+                            BLOCK_BODY
-                              BLOCK_BODY
+                              TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit
-                                TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit
+                                CALL 'public final fun <get-c> (): kotlin.Int declared in <root>.C' type=kotlin.Int origin=GET_PROPERTY
-                                  CALL 'public final fun <get-c> (): kotlin.Int declared in <root>.C' type=kotlin.Int origin=GET_PROPERTY
+                                  $this: GET_VAR '$this$with: <root>.C declared in <root>.test.<anonymous>' type=<root>.C origin=null
-                                    $this: GET_VAR '$this$with: <root>.C declared in <root>.test.<anonymous>' type=<root>.C origin=null
+                      CALL 'public final fun f3 (_context_receiver_0: <root>.C, g: @[ExtensionFunctionType] kotlin.Function2<<root>.C, <root>.R, kotlin.Unit>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
-                        CALL 'public final fun f3 (_context_receiver_0: <root>.C, g: @[ExtensionFunctionType] kotlin.Function2<<root>.C, <root>.R, kotlin.Unit>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
+                        $receiver: GET_VAR '$this$with: <root>.R declared in <root>.test.<anonymous>.<anonymous>' type=<root>.R origin=null
-                          $receiver: GET_VAR '$this$with: <root>.R declared in <root>.test.<anonymous>.<anonymous>' type=<root>.R origin=null
+                        _context_receiver_0: GET_VAR '$this$with: <root>.C declared in <root>.test.<anonymous>' type=<root>.C origin=null
-                          _context_receiver_0: GET_VAR '$this$with: <root>.C declared in <root>.test.<anonymous>' type=<root>.C origin=null
+                        g: GET_VAR 'val lf3: @[ExtensionFunctionType] kotlin.Function2<<root>.C, <root>.R, kotlin.Unit> [val] declared in <root>.test' type=@[ExtensionFunctionType] kotlin.Function2<<root>.C, <root>.R, kotlin.Unit> origin=null
-                          g: GET_VAR 'val lf3: @[ExtensionFunctionType] kotlin.Function2<<root>.C, <root>.R, kotlin.Unit> [val] declared in <root>.test' type=@[ExtensionFunctionType] kotlin.Function2<<root>.C, <root>.R, kotlin.Unit> origin=null
+                      CALL 'public final fun f3 (_context_receiver_0: <root>.C, g: @[ExtensionFunctionType] kotlin.Function2<<root>.C, <root>.R, kotlin.Unit>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
-                        CALL 'public final fun f3 (_context_receiver_0: <root>.C, g: @[ExtensionFunctionType] kotlin.Function2<<root>.C, <root>.R, kotlin.Unit>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
+                        $receiver: GET_VAR '$this$with: <root>.R declared in <root>.test.<anonymous>.<anonymous>' type=<root>.R origin=null
-                          $receiver: GET_VAR '$this$with: <root>.R declared in <root>.test.<anonymous>.<anonymous>' type=<root>.R origin=null
+                        _context_receiver_0: GET_VAR '$this$with: <root>.C declared in <root>.test.<anonymous>' type=<root>.C origin=null
-                          _context_receiver_0: GET_VAR '$this$with: <root>.C declared in <root>.test.<anonymous>' type=<root>.C origin=null
+                        g: FUN_EXPR type=@[ExtensionFunctionType] kotlin.Function2<<root>.C, <root>.R, kotlin.Unit> origin=LAMBDA
-                          g: FUN_EXPR type=@[ExtensionFunctionType] kotlin.Function2<<root>.C, <root>.R, kotlin.Unit> origin=LAMBDA
+                          FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> ($receiver:<root>.R, _context_receiver_0:<root>.C) returnType:kotlin.Unit
-                            FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> ($receiver:<root>.R, _context_receiver_0:<root>.C) returnType:kotlin.Unit
+                            contextReceiverParametersCount: 1
-                              contextReceiverParametersCount: 1
+                            $receiver: VALUE_PARAMETER name:$this$f3 type:<root>.R
-                              $receiver: VALUE_PARAMETER name:$this$f3 type:<root>.R
+                            VALUE_PARAMETER name:_context_receiver_0 index:0 type:<root>.C
-                              VALUE_PARAMETER name:_context_receiver_0 index:0 type:<root>.C
+                            BLOCK_BODY
-                              BLOCK_BODY
+                              TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit
-                                TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit
+                                CALL 'public final fun <get-r> (): kotlin.Int declared in <root>.R' type=kotlin.Int origin=GET_PROPERTY
-                                  CALL 'public final fun <get-r> (): kotlin.Int declared in <root>.R' type=kotlin.Int origin=GET_PROPERTY
+                                  $this: GET_VAR '$this$f3: <root>.R declared in <root>.test.<anonymous>.<anonymous>.<anonymous>' type=<root>.R origin=null
-                                    $this: GET_VAR '$this$f3: <root>.R declared in <root>.test.<anonymous>.<anonymous>.<anonymous>' type=<root>.R origin=null
+                              TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit
-                                TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit
+                                CALL 'public final fun <get-c> (): kotlin.Int declared in <root>.C' type=kotlin.Int origin=GET_PROPERTY
-                                  CALL 'public final fun <get-c> (): kotlin.Int declared in <root>.C' type=kotlin.Int origin=GET_PROPERTY
+                                  $this: GET_VAR '$this$with: <root>.C declared in <root>.test.<anonymous>' type=<root>.C origin=null
-                                    $this: GET_VAR '$this$with: <root>.C declared in <root>.test.<anonymous>' type=<root>.C origin=null
+                      CALL 'public final fun f4 (_context_receiver_0: <root>.C, g: kotlin.Function1<<root>.C, kotlin.Unit>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
-                        CALL 'public final fun f4 (_context_receiver_0: <root>.C, g: kotlin.Function1<<root>.C, kotlin.Unit>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
+                        _context_receiver_0: GET_VAR '$this$with: <root>.C declared in <root>.test.<anonymous>' type=<root>.C origin=null
-                          _context_receiver_0: GET_VAR '$this$with: <root>.C declared in <root>.test.<anonymous>' type=<root>.C origin=null
+                        g: GET_VAR 'val lf4: kotlin.Function1<<root>.C, kotlin.Unit> [val] declared in <root>.test' type=kotlin.Function1<<root>.C, kotlin.Unit> origin=null
-                          g: GET_VAR 'val lf4: kotlin.Function1<<root>.C, kotlin.Unit> [val] declared in <root>.test' type=kotlin.Function1<<root>.C, kotlin.Unit> origin=null
+                      CALL 'public final fun f4 (_context_receiver_0: <root>.C, g: kotlin.Function1<<root>.C, kotlin.Unit>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
-                        CALL 'public final fun f4 (_context_receiver_0: <root>.C, g: kotlin.Function1<<root>.C, kotlin.Unit>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
+                        _context_receiver_0: GET_VAR '$this$with: <root>.C declared in <root>.test.<anonymous>' type=<root>.C origin=null
-                          _context_receiver_0: GET_VAR '$this$with: <root>.C declared in <root>.test.<anonymous>' type=<root>.C origin=null
+                        g: FUN_EXPR type=kotlin.Function1<<root>.C, kotlin.Unit> origin=LAMBDA
-                          g: FUN_EXPR type=kotlin.Function1<<root>.C, kotlin.Unit> origin=LAMBDA
+                          FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> (_context_receiver_0:<root>.C) returnType:kotlin.Unit
-                            FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> (_context_receiver_0:<root>.C) returnType:kotlin.Unit
+                            contextReceiverParametersCount: 1
-                              contextReceiverParametersCount: 1
+                            VALUE_PARAMETER name:_context_receiver_0 index:0 type:<root>.C
-                              VALUE_PARAMETER name:_context_receiver_0 index:0 type:<root>.C
+                            BLOCK_BODY
-                              BLOCK_BODY
+                              TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit
-                                TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit
+                                CALL 'public final fun <get-c> (): kotlin.Int declared in <root>.C' type=kotlin.Int origin=GET_PROPERTY
-                                  CALL 'public final fun <get-c> (): kotlin.Int declared in <root>.C' type=kotlin.Int origin=GET_PROPERTY
+                                  $this: GET_VAR '$this$with: <root>.C declared in <root>.test.<anonymous>' type=<root>.C origin=null
-                                    $this: GET_VAR '$this$with: <root>.C declared in <root>.test.<anonymous>' type=<root>.C origin=null

compiler/testData/ir/irText/declarations/contextReceivers/functionalType.fir.kt.txt

@@ -69,7 +69,7 @@ fun test() {
   }
 
   with<C, Unit>(receiver = C(), block = local fun C.<anonymous>() {
-    return with<R, Unit>(receiver = R(), block = local fun R.<anonymous>() {
+    with<R, Unit>(receiver = R(), block = local fun R.<anonymous>() {
       $this$with.f1(_context_receiver_0 = $this$with, g = lf1)
       $this$with.f1(_context_receiver_0 = $this$with, g = local fun R.<anonymous>(_context_receiver_0: C, <unused var>: Param) {
         $this$f1.<get-r>() /*~> Unit */

compiler/testData/ir/irText/declarations/contextReceivers/lazy.fir.ir.txt

@@ -1,160 +0,0 @@
-FILE fqName:<root> fileName:/lazy.kt
-  CLASS INTERFACE name:Lazy modality:ABSTRACT visibility:public superTypes:[kotlin.Any]
-    $this: VALUE_PARAMETER INSTANCE_RECEIVER name:<this> type:<root>.Lazy<T of <root>.Lazy>
-    TYPE_PARAMETER name:T index:0 variance: superTypes:[kotlin.Any?] reified:false
-    FUN FAKE_OVERRIDE name:equals visibility:public modality:OPEN <> ($this:kotlin.Any, other:kotlin.Any?) returnType:kotlin.Boolean [fake_override,operator]
-      overridden:
-        public open fun equals (other: kotlin.Any?): kotlin.Boolean [operator] declared in kotlin.Any
-      $this: VALUE_PARAMETER name:<this> type:kotlin.Any
-      VALUE_PARAMETER name:other index:0 type:kotlin.Any?
-    FUN FAKE_OVERRIDE name:hashCode visibility:public modality:OPEN <> ($this:kotlin.Any) returnType:kotlin.Int [fake_override]
-      overridden:
-        public open fun hashCode (): kotlin.Int declared in kotlin.Any
-      $this: VALUE_PARAMETER name:<this> type:kotlin.Any
-    FUN FAKE_OVERRIDE name:toString visibility:public modality:OPEN <> ($this:kotlin.Any) returnType:kotlin.String [fake_override]
-      overridden:
-        public open fun toString (): kotlin.String declared in kotlin.Any
-      $this: VALUE_PARAMETER name:<this> type:kotlin.Any
-  FUN name:test1 visibility:public modality:FINAL <> (_context_receiver_0:<root>.Lazy<kotlin.Int>, _context_receiver_1:<root>.Lazy<kotlin.CharSequence>) returnType:kotlin.Unit
-    contextReceiverParametersCount: 2
-    VALUE_PARAMETER name:_context_receiver_0 index:0 type:<root>.Lazy<kotlin.Int>
-    VALUE_PARAMETER name:_context_receiver_1 index:1 type:<root>.Lazy<kotlin.CharSequence>
-    BLOCK_BODY
-  FUN name:test2 visibility:public modality:FINAL <T> ($receiver:<root>.Lazy<kotlin.Int>, _context_receiver_0:<root>.Lazy<T of <root>.test2>) returnType:kotlin.Unit
-    TYPE_PARAMETER name:T index:0 variance: superTypes:[kotlin.Any?] reified:false
-    contextReceiverParametersCount: 1
-    $receiver: VALUE_PARAMETER name:<this> type:<root>.Lazy<kotlin.Int>
-    VALUE_PARAMETER name:_context_receiver_0 index:0 type:<root>.Lazy<T of <root>.test2>
-    BLOCK_BODY
-  FUN name:test3 visibility:public modality:FINAL <T> ($receiver:<root>.Lazy<kotlin.Int>, _context_receiver_0:<root>.Lazy<<root>.Lazy<T of <root>.test3>>) returnType:kotlin.Unit
-    TYPE_PARAMETER name:T index:0 variance: superTypes:[kotlin.Any?] reified:false
-    contextReceiverParametersCount: 1
-    $receiver: VALUE_PARAMETER name:<this> type:<root>.Lazy<kotlin.Int>
-    VALUE_PARAMETER name:_context_receiver_0 index:0 type:<root>.Lazy<<root>.Lazy<T of <root>.test3>>
-    BLOCK_BODY
-  FUN name:f visibility:public modality:FINAL <T> (lazy1:<root>.Lazy<kotlin.Int>, lazy2:<root>.Lazy<kotlin.CharSequence>, lazyT:<root>.Lazy<T of <root>.f>, lazyLazyT:<root>.Lazy<<root>.Lazy<T of <root>.f>>) returnType:kotlin.Unit
-    TYPE_PARAMETER name:T index:0 variance: superTypes:[kotlin.Any?] reified:false
-    VALUE_PARAMETER name:lazy1 index:0 type:<root>.Lazy<kotlin.Int>
-    VALUE_PARAMETER name:lazy2 index:1 type:<root>.Lazy<kotlin.CharSequence>
-    VALUE_PARAMETER name:lazyT index:2 type:<root>.Lazy<T of <root>.f>
-    VALUE_PARAMETER name:lazyLazyT index:3 type:<root>.Lazy<<root>.Lazy<T of <root>.f>>
-    BLOCK_BODY
-      CALL 'public final fun with <T, R> (receiver: T of kotlin.StandardKt.with, block: @[ExtensionFunctionType] kotlin.Function1<T of kotlin.StandardKt.with, R of kotlin.StandardKt.with>): R of kotlin.StandardKt.with [inline] declared in kotlin.StandardKt' type=kotlin.Unit origin=null
-        <T>: <root>.Lazy<kotlin.Int>
-        <R>: kotlin.Unit
-        receiver: GET_VAR 'lazy1: <root>.Lazy<kotlin.Int> declared in <root>.f' type=<root>.Lazy<kotlin.Int> origin=null
-        block: FUN_EXPR type=@[ExtensionFunctionType] kotlin.Function1<<root>.Lazy<kotlin.Int>, kotlin.Unit> origin=LAMBDA
-          FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> ($receiver:<root>.Lazy<kotlin.Int>) returnType:kotlin.Unit
-            $receiver: VALUE_PARAMETER name:$this$with type:<root>.Lazy<kotlin.Int>
-            BLOCK_BODY
-              RETURN type=kotlin.Nothing from='local final fun <anonymous> (): kotlin.Unit declared in <root>.f'
-                CALL 'public final fun with <T, R> (receiver: T of kotlin.StandardKt.with, block: @[ExtensionFunctionType] kotlin.Function1<T of kotlin.StandardKt.with, R of kotlin.StandardKt.with>): R of kotlin.StandardKt.with [inline] declared in kotlin.StandardKt' type=kotlin.Unit origin=null
-                  <T>: <root>.Lazy<kotlin.CharSequence>
-                  <R>: kotlin.Unit
-                  receiver: GET_VAR 'lazy2: <root>.Lazy<kotlin.CharSequence> declared in <root>.f' type=<root>.Lazy<kotlin.CharSequence> origin=null
-                  block: FUN_EXPR type=@[ExtensionFunctionType] kotlin.Function1<<root>.Lazy<kotlin.CharSequence>, kotlin.Unit> origin=LAMBDA
-                    FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> ($receiver:<root>.Lazy<kotlin.CharSequence>) returnType:kotlin.Unit
-                      $receiver: VALUE_PARAMETER name:$this$with type:<root>.Lazy<kotlin.CharSequence>
-                      BLOCK_BODY
-                        CALL 'public final fun test1 (_context_receiver_0: <root>.Lazy<kotlin.Int>, _context_receiver_1: <root>.Lazy<kotlin.CharSequence>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
-                          _context_receiver_0: GET_VAR '$this$with: <root>.Lazy<kotlin.Int> declared in <root>.f.<anonymous>' type=<root>.Lazy<kotlin.Int> origin=null
-                          _context_receiver_1: GET_VAR '$this$with: <root>.Lazy<kotlin.CharSequence> declared in <root>.f.<anonymous>.<anonymous>' type=<root>.Lazy<kotlin.CharSequence> origin=null
-                        CALL 'public final fun test2 <T> (_context_receiver_0: <root>.Lazy<T of <root>.test2>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
-                          <T>: kotlin.CharSequence
-                          $receiver: GET_VAR '$this$with: <root>.Lazy<kotlin.Int> declared in <root>.f.<anonymous>' type=<root>.Lazy<kotlin.Int> origin=null
-                          _context_receiver_0: GET_VAR '$this$with: <root>.Lazy<kotlin.CharSequence> declared in <root>.f.<anonymous>.<anonymous>' type=<root>.Lazy<kotlin.CharSequence> origin=null
-      CALL 'public final fun with <T, R> (receiver: T of kotlin.StandardKt.with, block: @[ExtensionFunctionType] kotlin.Function1<T of kotlin.StandardKt.with, R of kotlin.StandardKt.with>): R of kotlin.StandardKt.with [inline] declared in kotlin.StandardKt' type=kotlin.Unit origin=null
-        <T>: <root>.Lazy<kotlin.CharSequence>
-        <R>: kotlin.Unit
-        receiver: GET_VAR 'lazy2: <root>.Lazy<kotlin.CharSequence> declared in <root>.f' type=<root>.Lazy<kotlin.CharSequence> origin=null
-        block: FUN_EXPR type=@[ExtensionFunctionType] kotlin.Function1<<root>.Lazy<kotlin.CharSequence>, kotlin.Unit> origin=LAMBDA
-          FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> ($receiver:<root>.Lazy<kotlin.CharSequence>) returnType:kotlin.Unit
-            $receiver: VALUE_PARAMETER name:$this$with type:<root>.Lazy<kotlin.CharSequence>
-            BLOCK_BODY
-              RETURN type=kotlin.Nothing from='local final fun <anonymous> (): kotlin.Unit declared in <root>.f'
-                CALL 'public final fun with <T, R> (receiver: T of kotlin.StandardKt.with, block: @[ExtensionFunctionType] kotlin.Function1<T of kotlin.StandardKt.with, R of kotlin.StandardKt.with>): R of kotlin.StandardKt.with [inline] declared in kotlin.StandardKt' type=kotlin.Unit origin=null
-                  <T>: <root>.Lazy<kotlin.Int>
-                  <R>: kotlin.Unit
-                  receiver: GET_VAR 'lazy1: <root>.Lazy<kotlin.Int> declared in <root>.f' type=<root>.Lazy<kotlin.Int> origin=null
-                  block: FUN_EXPR type=@[ExtensionFunctionType] kotlin.Function1<<root>.Lazy<kotlin.Int>, kotlin.Unit> origin=LAMBDA
-                    FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> ($receiver:<root>.Lazy<kotlin.Int>) returnType:kotlin.Unit
-                      $receiver: VALUE_PARAMETER name:$this$with type:<root>.Lazy<kotlin.Int>
-                      BLOCK_BODY
-                        CALL 'public final fun test1 (_context_receiver_0: <root>.Lazy<kotlin.Int>, _context_receiver_1: <root>.Lazy<kotlin.CharSequence>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
-                          _context_receiver_0: GET_VAR '$this$with: <root>.Lazy<kotlin.Int> declared in <root>.f.<anonymous>.<anonymous>' type=<root>.Lazy<kotlin.Int> origin=null
-                          _context_receiver_1: GET_VAR '$this$with: <root>.Lazy<kotlin.CharSequence> declared in <root>.f.<anonymous>' type=<root>.Lazy<kotlin.CharSequence> origin=null
-                        CALL 'public final fun test2 <T> (_context_receiver_0: <root>.Lazy<T of <root>.test2>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
-                          <T>: kotlin.Int
-                          $receiver: GET_VAR '$this$with: <root>.Lazy<kotlin.Int> declared in <root>.f.<anonymous>.<anonymous>' type=<root>.Lazy<kotlin.Int> origin=null
-                          _context_receiver_0: GET_VAR '$this$with: <root>.Lazy<kotlin.Int> declared in <root>.f.<anonymous>.<anonymous>' type=<root>.Lazy<kotlin.Int> origin=null
-      CALL 'public final fun with <T, R> (receiver: T of kotlin.StandardKt.with, block: @[ExtensionFunctionType] kotlin.Function1<T of kotlin.StandardKt.with, R of kotlin.StandardKt.with>): R of kotlin.StandardKt.with [inline] declared in kotlin.StandardKt' type=kotlin.Unit origin=null
-        <T>: <root>.Lazy<T of <root>.f>
-        <R>: kotlin.Unit
-        receiver: GET_VAR 'lazyT: <root>.Lazy<T of <root>.f> declared in <root>.f' type=<root>.Lazy<T of <root>.f> origin=null
-        block: FUN_EXPR type=@[ExtensionFunctionType] kotlin.Function1<<root>.Lazy<T of <root>.f>, kotlin.Unit> origin=LAMBDA
-          FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> ($receiver:<root>.Lazy<T of <root>.f>) returnType:kotlin.Unit
-            $receiver: VALUE_PARAMETER name:$this$with type:<root>.Lazy<T of <root>.f>
-            BLOCK_BODY
-              RETURN type=kotlin.Nothing from='local final fun <anonymous> (): kotlin.Unit declared in <root>.f'
-                CALL 'public final fun with <T, R> (receiver: T of kotlin.StandardKt.with, block: @[ExtensionFunctionType] kotlin.Function1<T of kotlin.StandardKt.with, R of kotlin.StandardKt.with>): R of kotlin.StandardKt.with [inline] declared in kotlin.StandardKt' type=kotlin.Unit origin=null
-                  <T>: <root>.Lazy<kotlin.Int>
-                  <R>: kotlin.Unit
-                  receiver: GET_VAR 'lazy1: <root>.Lazy<kotlin.Int> declared in <root>.f' type=<root>.Lazy<kotlin.Int> origin=null
-                  block: FUN_EXPR type=@[ExtensionFunctionType] kotlin.Function1<<root>.Lazy<kotlin.Int>, kotlin.Unit> origin=LAMBDA
-                    FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> ($receiver:<root>.Lazy<kotlin.Int>) returnType:kotlin.Unit
-                      $receiver: VALUE_PARAMETER name:$this$with type:<root>.Lazy<kotlin.Int>
-                      BLOCK_BODY
-                        CALL 'public final fun test2 <T> (_context_receiver_0: <root>.Lazy<T of <root>.test2>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
-                          <T>: kotlin.Int
-                          $receiver: GET_VAR '$this$with: <root>.Lazy<kotlin.Int> declared in <root>.f.<anonymous>.<anonymous>' type=<root>.Lazy<kotlin.Int> origin=null
-                          _context_receiver_0: GET_VAR '$this$with: <root>.Lazy<kotlin.Int> declared in <root>.f.<anonymous>.<anonymous>' type=<root>.Lazy<kotlin.Int> origin=null
-      CALL 'public final fun with <T, R> (receiver: T of kotlin.StandardKt.with, block: @[ExtensionFunctionType] kotlin.Function1<T of kotlin.StandardKt.with, R of kotlin.StandardKt.with>): R of kotlin.StandardKt.with [inline] declared in kotlin.StandardKt' type=kotlin.Unit origin=null
-        <T>: <root>.Lazy<<root>.Lazy<T of <root>.f>>
-        <R>: kotlin.Unit
-        receiver: GET_VAR 'lazyLazyT: <root>.Lazy<<root>.Lazy<T of <root>.f>> declared in <root>.f' type=<root>.Lazy<<root>.Lazy<T of <root>.f>> origin=null
-        block: FUN_EXPR type=@[ExtensionFunctionType] kotlin.Function1<<root>.Lazy<<root>.Lazy<T of <root>.f>>, kotlin.Unit> origin=LAMBDA
-          FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> ($receiver:<root>.Lazy<<root>.Lazy<T of <root>.f>>) returnType:kotlin.Unit
-            $receiver: VALUE_PARAMETER name:$this$with type:<root>.Lazy<<root>.Lazy<T of <root>.f>>
-            BLOCK_BODY
-              RETURN type=kotlin.Nothing from='local final fun <anonymous> (): kotlin.Unit declared in <root>.f'
-                CALL 'public final fun with <T, R> (receiver: T of kotlin.StandardKt.with, block: @[ExtensionFunctionType] kotlin.Function1<T of kotlin.StandardKt.with, R of kotlin.StandardKt.with>): R of kotlin.StandardKt.with [inline] declared in kotlin.StandardKt' type=kotlin.Unit origin=null
-                  <T>: <root>.Lazy<kotlin.Int>
-                  <R>: kotlin.Unit
-                  receiver: GET_VAR 'lazy1: <root>.Lazy<kotlin.Int> declared in <root>.f' type=<root>.Lazy<kotlin.Int> origin=null
-                  block: FUN_EXPR type=@[ExtensionFunctionType] kotlin.Function1<<root>.Lazy<kotlin.Int>, kotlin.Unit> origin=LAMBDA
-                    FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> ($receiver:<root>.Lazy<kotlin.Int>) returnType:kotlin.Unit
-                      $receiver: VALUE_PARAMETER name:$this$with type:<root>.Lazy<kotlin.Int>
-                      BLOCK_BODY
-                        CALL 'public final fun test2 <T> (_context_receiver_0: <root>.Lazy<T of <root>.test2>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
-                          <T>: kotlin.Int
-                          $receiver: GET_VAR '$this$with: <root>.Lazy<kotlin.Int> declared in <root>.f.<anonymous>.<anonymous>' type=<root>.Lazy<kotlin.Int> origin=null
-                          _context_receiver_0: GET_VAR '$this$with: <root>.Lazy<kotlin.Int> declared in <root>.f.<anonymous>.<anonymous>' type=<root>.Lazy<kotlin.Int> origin=null
-                        CALL 'public final fun test3 <T> (_context_receiver_0: <root>.Lazy<<root>.Lazy<T of <root>.test3>>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
-                          <T>: T of <root>.f
-                          $receiver: GET_VAR '$this$with: <root>.Lazy<kotlin.Int> declared in <root>.f.<anonymous>.<anonymous>' type=<root>.Lazy<kotlin.Int> origin=null
-                          _context_receiver_0: GET_VAR '$this$with: <root>.Lazy<<root>.Lazy<T of <root>.f>> declared in <root>.f.<anonymous>' type=<root>.Lazy<<root>.Lazy<T of <root>.f>> origin=null
-      CALL 'public final fun with <T, R> (receiver: T of kotlin.StandardKt.with, block: @[ExtensionFunctionType] kotlin.Function1<T of kotlin.StandardKt.with, R of kotlin.StandardKt.with>): R of kotlin.StandardKt.with [inline] declared in kotlin.StandardKt' type=kotlin.Unit origin=null
-        <T>: <root>.Lazy<kotlin.Int>
-        <R>: kotlin.Unit
-        receiver: GET_VAR 'lazy1: <root>.Lazy<kotlin.Int> declared in <root>.f' type=<root>.Lazy<kotlin.Int> origin=null
-        block: FUN_EXPR type=@[ExtensionFunctionType] kotlin.Function1<<root>.Lazy<kotlin.Int>, kotlin.Unit> origin=LAMBDA
-          FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> ($receiver:<root>.Lazy<kotlin.Int>) returnType:kotlin.Unit
-            $receiver: VALUE_PARAMETER name:$this$with type:<root>.Lazy<kotlin.Int>
-            BLOCK_BODY
-              RETURN type=kotlin.Nothing from='local final fun <anonymous> (): kotlin.Unit declared in <root>.f'
-                CALL 'public final fun with <T, R> (receiver: T of kotlin.StandardKt.with, block: @[ExtensionFunctionType] kotlin.Function1<T of kotlin.StandardKt.with, R of kotlin.StandardKt.with>): R of kotlin.StandardKt.with [inline] declared in kotlin.StandardKt' type=kotlin.Unit origin=null
-                  <T>: <root>.Lazy<<root>.Lazy<T of <root>.f>>
-                  <R>: kotlin.Unit
-                  receiver: GET_VAR 'lazyLazyT: <root>.Lazy<<root>.Lazy<T of <root>.f>> declared in <root>.f' type=<root>.Lazy<<root>.Lazy<T of <root>.f>> origin=null
-                  block: FUN_EXPR type=@[ExtensionFunctionType] kotlin.Function1<<root>.Lazy<<root>.Lazy<T of <root>.f>>, kotlin.Unit> origin=LAMBDA
-                    FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> ($receiver:<root>.Lazy<<root>.Lazy<T of <root>.f>>) returnType:kotlin.Unit
-                      $receiver: VALUE_PARAMETER name:$this$with type:<root>.Lazy<<root>.Lazy<T of <root>.f>>
-                      BLOCK_BODY
-                        CALL 'public final fun test2 <T> (_context_receiver_0: <root>.Lazy<T of <root>.test2>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
-                          <T>: <root>.Lazy<T of <root>.f>
-                          $receiver: GET_VAR '$this$with: <root>.Lazy<kotlin.Int> declared in <root>.f.<anonymous>' type=<root>.Lazy<kotlin.Int> origin=null
-                          _context_receiver_0: GET_VAR '$this$with: <root>.Lazy<<root>.Lazy<T of <root>.f>> declared in <root>.f.<anonymous>.<anonymous>' type=<root>.Lazy<<root>.Lazy<T of <root>.f>> origin=null
-                        CALL 'public final fun test3 <T> (_context_receiver_0: <root>.Lazy<<root>.Lazy<T of <root>.test3>>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
-                          <T>: T of <root>.f
-                          $receiver: GET_VAR '$this$with: <root>.Lazy<kotlin.Int> declared in <root>.f.<anonymous>' type=<root>.Lazy<kotlin.Int> origin=null
-                          _context_receiver_0: GET_VAR '$this$with: <root>.Lazy<<root>.Lazy<T of <root>.f>> declared in <root>.f.<anonymous>.<anonymous>' type=<root>.Lazy<<root>.Lazy<T of <root>.f>> origin=null

compiler/testData/ir/irText/declarations/contextReceivers/lazy.fir.kt.txt

@@ -1,54 +0,0 @@
-interface Lazy<T : Any?> {
-
-}
-
-fun test1(_context_receiver_0: Lazy<Int>, _context_receiver_1: Lazy<CharSequence>) {
-}
-
-fun <T : Any?> Lazy<Int>.test2(_context_receiver_0: Lazy<T>) {
-}
-
-fun <T : Any?> Lazy<Int>.test3(_context_receiver_0: Lazy<Lazy<T>>) {
-}
-
-fun <T : Any?> f(lazy1: Lazy<Int>, lazy2: Lazy<CharSequence>, lazyT: Lazy<T>, lazyLazyT: Lazy<Lazy<T>>) {
-  with<Lazy<Int>, Unit>(receiver = lazy1, block = local fun Lazy<Int>.<anonymous>() {
-    return with<Lazy<CharSequence>, Unit>(receiver = lazy2, block = local fun Lazy<CharSequence>.<anonymous>() {
-      test1(_context_receiver_0 = $this$with, _context_receiver_1 = $this$with)
-      $this$with.test2<CharSequence>(_context_receiver_0 = $this$with)
-    }
-)
-  }
-)
-  with<Lazy<CharSequence>, Unit>(receiver = lazy2, block = local fun Lazy<CharSequence>.<anonymous>() {
-    return with<Lazy<Int>, Unit>(receiver = lazy1, block = local fun Lazy<Int>.<anonymous>() {
-      test1(_context_receiver_0 = $this$with, _context_receiver_1 = $this$with)
-      $this$with.test2<Int>(_context_receiver_0 = $this$with)
-    }
-)
-  }
-)
-  with<Lazy<T>, Unit>(receiver = lazyT, block = local fun Lazy<T>.<anonymous>() {
-    return with<Lazy<Int>, Unit>(receiver = lazy1, block = local fun Lazy<Int>.<anonymous>() {
-      $this$with.test2<Int>(_context_receiver_0 = $this$with)
-    }
-)
-  }
-)
-  with<Lazy<Lazy<T>>, Unit>(receiver = lazyLazyT, block = local fun Lazy<Lazy<T>>.<anonymous>() {
-    return with<Lazy<Int>, Unit>(receiver = lazy1, block = local fun Lazy<Int>.<anonymous>() {
-      $this$with.test2<Int>(_context_receiver_0 = $this$with)
-      $this$with.test3<T>(_context_receiver_0 = $this$with)
-    }
-)
-  }
-)
-  with<Lazy<Int>, Unit>(receiver = lazy1, block = local fun Lazy<Int>.<anonymous>() {
-    return with<Lazy<Lazy<T>>, Unit>(receiver = lazyLazyT, block = local fun Lazy<Lazy<T>>.<anonymous>() {
-      $this$with.test2<Lazy<T>>(_context_receiver_0 = $this$with)
-      $this$with.test3<T>(_context_receiver_0 = $this$with)
-    }
-)
-  }
-)
-}

compiler/testData/ir/irText/declarations/contextReceivers/lazy.kt

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // !LANGUAGE: +ContextReceivers
 
 interface Lazy<T>