[FIR] Split FirAnonymousFunction to expression and declaration

compiler/fir/analysis-tests/testData/resolve/cfg/lambdas.dot

@@ -115,26 +115,27 @@ digraph lambdas_kt {
                     39 [label="Postponed enter to lambda"];
                     subgraph cluster_14 {
                         color=blue
-                        47 [label="Enter function anonymousFunction" style="filled" fillcolor=red];
+                        48 [label="Enter function anonymousFunction" style="filled" fillcolor=red];
                         subgraph cluster_15 {
                             color=blue
-                            48 [label="Enter block"];
+                            49 [label="Enter block"];
-                            49 [label="Access variable R|<local>/x|"];
+                            50 [label="Access variable R|<local>/x|"];
-                            50 [label="Function call: R|<local>/x|.R|kotlin/Int.inc|()"];
+                            51 [label="Function call: R|<local>/x|.R|kotlin/Int.inc|()"];
-                            51 [label="Exit block"];
+                            52 [label="Exit block"];
                         }
-                        52 [label="Exit function anonymousFunction" style="filled" fillcolor=red];
+                        53 [label="Exit function anonymousFunction" style="filled" fillcolor=red];
                     }
                     40 [label="Postponed exit from lambda"];
-                    41 [label="Variable declaration: lval lambda: R|() -> kotlin/Int|"];
+                    41 [label="Exit anonymous function expression"];
-                    42 [label="Exit block"];
+                    42 [label="Variable declaration: lval lambda: R|() -> kotlin/Int|"];
+                    43 [label="Exit block"];
                 }
-                43 [label="Exit when branch result"];
+                44 [label="Exit when branch result"];
-                44 [label="Exit when"];
+                45 [label="Exit when"];
             }
-            45 [label="Exit block"];
+            46 [label="Exit block"];
         }
-        46 [label="Exit function test_2" style="filled" fillcolor=red];
+        47 [label="Exit function test_2" style="filled" fillcolor=red];
     }
     29 -> {30};
     30 -> {31};
@@ -143,148 +144,149 @@ digraph lambdas_kt {
     33 -> {34};
     34 -> {35};
     35 -> {37 36};
-    36 -> {44};
+    36 -> {45};
     37 -> {38};
     38 -> {39};
-    39 -> {40 47};
+    39 -> {40 48};
-    39 -> {47} [style=dashed];
+    39 -> {48} [style=dashed];
     40 -> {41};
     41 -> {42};
     42 -> {43};
     43 -> {44};
     44 -> {45};
     45 -> {46};
-    47 -> {48};
+    46 -> {47};
     48 -> {49};
     49 -> {50};
     50 -> {51};
     51 -> {52};
+    52 -> {53};
 
     subgraph cluster_16 {
         color=red
-        53 [label="Enter function getInt" style="filled" fillcolor=red];
+        54 [label="Enter function getInt" style="filled" fillcolor=red];
         subgraph cluster_17 {
             color=blue
-            54 [label="Enter block"];
+            55 [label="Enter block"];
-            55 [label="Function call: R|<local>/block|.R|SubstitutionOverride<kotlin/Function0.invoke: R|kotlin/Unit|>|()"];
+            56 [label="Function call: R|<local>/block|.R|SubstitutionOverride<kotlin/Function0.invoke: R|kotlin/Unit|>|()"];
-            56 [label="Const: Int(1)"];
+            57 [label="Const: Int(1)"];
-            57 [label="Jump: ^getInt Int(1)"];
+            58 [label="Jump: ^getInt Int(1)"];
-            58 [label="Stub" style="filled" fillcolor=gray];
+            59 [label="Stub" style="filled" fillcolor=gray];
-            59 [label="Exit block" style="filled" fillcolor=gray];
+            60 [label="Exit block" style="filled" fillcolor=gray];
         }
-        60 [label="Exit function getInt" style="filled" fillcolor=red];
+        61 [label="Exit function getInt" style="filled" fillcolor=red];
     }
-    53 -> {54};
     54 -> {55};
     55 -> {56};
     56 -> {57};
-    57 -> {60};
+    57 -> {58};
-    57 -> {58} [style=dotted];
+    58 -> {61};
     58 -> {59} [style=dotted];
     59 -> {60} [style=dotted];
+    60 -> {61} [style=dotted];
 
     subgraph cluster_18 {
         color=red
-        61 [label="Enter function test_3" style="filled" fillcolor=red];
+        62 [label="Enter function test_3" style="filled" fillcolor=red];
         subgraph cluster_19 {
             color=blue
-            62 [label="Enter block"];
+            63 [label="Enter block"];
-            63 [label="Postponed enter to lambda"];
+            64 [label="Postponed enter to lambda"];
             subgraph cluster_20 {
                 color=blue
-                70 [label="Enter function anonymousFunction" style="filled" fillcolor=red];
+                71 [label="Enter function anonymousFunction" style="filled" fillcolor=red];
                 subgraph cluster_21 {
                     color=blue
-                    71 [label="Enter block"];
+                    72 [label="Enter block"];
-                    72 [label="Const: Int(1)"];
+                    73 [label="Const: Int(1)"];
-                    73 [label="Jump: ^test_3 Int(1)"];
+                    74 [label="Jump: ^test_3 Int(1)"];
-                    74 [label="Stub" style="filled" fillcolor=gray];
+                    75 [label="Stub" style="filled" fillcolor=gray];
-                    75 [label="Exit block" style="filled" fillcolor=gray];
+                    76 [label="Exit block" style="filled" fillcolor=gray];
                 }
-                76 [label="Exit function anonymousFunction" style="filled" fillcolor=red];
+                77 [label="Exit function anonymousFunction" style="filled" fillcolor=red];
             }
-            64 [label="Postponed exit from lambda"];
+            65 [label="Postponed exit from lambda"];
-            65 [label="Function call: R|/getInt|(...)"];
+            66 [label="Function call: R|/getInt|(...)"];
-            66 [label="Jump: ^test_3 R|/getInt|(<L> = getInt@fun <anonymous>(): R|kotlin/Unit| <inline=Inline, kind=UNKNOWN>  {
+            67 [label="Jump: ^test_3 R|/getInt|(<L> = getInt@fun <anonymous>(): R|kotlin/Unit| <inline=Inline, kind=UNKNOWN>  {
     ^test_3 Int(1)
 }
 )"];
-            67 [label="Stub" style="filled" fillcolor=gray];
+            68 [label="Stub" style="filled" fillcolor=gray];
-            68 [label="Exit block" style="filled" fillcolor=gray];
+            69 [label="Exit block" style="filled" fillcolor=gray];
         }
-        69 [label="Exit function test_3" style="filled" fillcolor=red];
+        70 [label="Exit function test_3" style="filled" fillcolor=red];
     }
-    61 -> {62};
     62 -> {63};
-    63 -> {70};
+    63 -> {64};
-    63 -> {64} [color=red];
+    64 -> {71};
-    63 -> {70} [style=dashed];
+    64 -> {65} [color=red];
-    64 -> {65};
+    64 -> {71} [style=dashed];
     65 -> {66};
-    66 -> {69};
+    66 -> {67};
-    66 -> {67} [style=dotted];
+    67 -> {70};
     67 -> {68} [style=dotted];
     68 -> {69} [style=dotted];
-    70 -> {76 71};
+    69 -> {70} [style=dotted];
-    71 -> {72};
+    71 -> {77 72};
     72 -> {73};
-    73 -> {69};
+    73 -> {74};
-    73 -> {74} [style=dotted];
+    74 -> {70};
     74 -> {75} [style=dotted];
     75 -> {76} [style=dotted];
-    76 -> {64} [color=green];
+    76 -> {77} [style=dotted];
-    76 -> {70} [color=green style=dashed];
+    77 -> {65} [color=green];
+    77 -> {71} [color=green style=dashed];
 
     subgraph cluster_22 {
         color=red
-        77 [label="Enter function test_4" style="filled" fillcolor=red];
+        78 [label="Enter function test_4" style="filled" fillcolor=red];
         subgraph cluster_23 {
             color=blue
-            78 [label="Enter block"];
+            79 [label="Enter block"];
-            79 [label="Postponed enter to lambda"];
+            80 [label="Postponed enter to lambda"];
             subgraph cluster_24 {
                 color=blue
-                86 [label="Enter function anonymousFunction" style="filled" fillcolor=red];
+                87 [label="Enter function anonymousFunction" style="filled" fillcolor=red];
                 subgraph cluster_25 {
                     color=blue
-                    87 [label="Enter block"];
+                    88 [label="Enter block"];
-                    88 [label="Const: Int(1)"];
+                    89 [label="Const: Int(1)"];
-                    89 [label="Jump: ^test_4 Int(1)"];
+                    90 [label="Jump: ^test_4 Int(1)"];
-                    90 [label="Stub" style="filled" fillcolor=gray];
+                    91 [label="Stub" style="filled" fillcolor=gray];
-                    91 [label="Exit block" style="filled" fillcolor=gray];
+                    92 [label="Exit block" style="filled" fillcolor=gray];
                 }
-                92 [label="Exit function anonymousFunction" style="filled" fillcolor=red];
+                93 [label="Exit function anonymousFunction" style="filled" fillcolor=red];
             }
-            80 [label="Postponed exit from lambda"];
+            81 [label="Postponed exit from lambda"];
-            81 [label="Function call: R|/getInt|(...)"];
+            82 [label="Function call: R|/getInt|(...)"];
-            82 [label="Jump: ^test_4 R|/getInt|(block = getInt@fun <anonymous>(): R|kotlin/Unit| <inline=Inline, kind=UNKNOWN>  {
+            83 [label="Jump: ^test_4 R|/getInt|(block = getInt@fun <anonymous>(): R|kotlin/Unit| <inline=Inline, kind=UNKNOWN>  {
     ^test_4 Int(1)
 }
 )"];
-            83 [label="Stub" style="filled" fillcolor=gray];
+            84 [label="Stub" style="filled" fillcolor=gray];
-            84 [label="Exit block" style="filled" fillcolor=gray];
+            85 [label="Exit block" style="filled" fillcolor=gray];
         }
-        85 [label="Exit function test_4" style="filled" fillcolor=red];
+        86 [label="Exit function test_4" style="filled" fillcolor=red];
     }
-    77 -> {78};
     78 -> {79};
-    79 -> {86};
+    79 -> {80};
-    79 -> {80} [color=red];
+    80 -> {87};
-    79 -> {86} [style=dashed];
+    80 -> {81} [color=red];
-    80 -> {81};
+    80 -> {87} [style=dashed];
     81 -> {82};
-    82 -> {85};
+    82 -> {83};
-    82 -> {83} [style=dotted];
+    83 -> {86};
     83 -> {84} [style=dotted];
     84 -> {85} [style=dotted];
-    86 -> {92 87};
+    85 -> {86} [style=dotted];
-    87 -> {88};
+    87 -> {93 88};
     88 -> {89};
-    89 -> {85};
+    89 -> {90};
-    89 -> {90} [style=dotted];
+    90 -> {86};
     90 -> {91} [style=dotted];
     91 -> {92} [style=dotted];
-    92 -> {80} [color=green];
+    92 -> {93} [style=dotted];
-    92 -> {86} [color=green style=dashed];
+    93 -> {81} [color=green];
+    93 -> {87} [color=green style=dashed];
 
 }

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/cfa/CfaUtils.kt

@@ -87,7 +87,7 @@ class LocalPropertyAndCapturedWriteCollector private constructor() : ControlFlow
     }
 
     override fun visitPostponedLambdaExitNode(node: PostponedLambdaExitNode) {
-        lambdaOrLocalFunctionStack.remove(node.fir)
+        lambdaOrLocalFunctionStack.remove(node.fir.anonymousFunction)
     }
 
     override fun visitLocalFunctionDeclarationNode(node: LocalFunctionDeclarationNode, data: Nothing?) {

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/expression/FirNotNullAssertionChecker.kt

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
 import org.jetbrains.kotlin.fir.analysis.diagnostics.reportOn
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
+import org.jetbrains.kotlin.fir.expressions.FirAnonymousFunctionExpression
 import org.jetbrains.kotlin.fir.expressions.FirCallableReferenceAccess
 import org.jetbrains.kotlin.fir.expressions.FirCheckNotNullCall
 import org.jetbrains.kotlin.fir.expressions.FirResolvedQualifier
@@ -19,7 +20,7 @@ import org.jetbrains.kotlin.fir.types.*
 object FirNotNullAssertionChecker : FirCheckNotNullCallChecker() {
     override fun check(expression: FirCheckNotNullCall, context: CheckerContext, reporter: DiagnosticReporter) {
         val argument = expression.argumentList.arguments.singleOrNull() ?: return
-        if (argument is FirAnonymousFunction && argument.isLambda) {
+        if (argument is FirAnonymousFunctionExpression && argument.anonymousFunction.isLambda) {
             reporter.reportOn(expression.source, FirErrors.NOT_NULL_ASSERTION_ON_LAMBDA_EXPRESSION, context)
             return
         }

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/expression/FirUninitializedEnumChecker.kt

@@ -15,10 +15,7 @@ import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
 import org.jetbrains.kotlin.fir.analysis.diagnostics.reportOn
 import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
+import org.jetbrains.kotlin.fir.expressions.*
-import org.jetbrains.kotlin.fir.expressions.FirLambdaArgumentExpression
-import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression
-import org.jetbrains.kotlin.fir.expressions.FirResolvedQualifier
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 
 object FirUninitializedEnumChecker : FirQualifiedAccessExpressionChecker() {
@@ -222,6 +219,6 @@ object FirUninitializedEnumChecker : FirQualifiedAccessExpressionChecker() {
                 (delegateCall.calleeReference as? FirResolvedNamedReference)?.resolvedSymbol?.fir as? FirSimpleFunction ?: return null
             if (callee.symbol.callableId.asSingleFqName().asString() != "kotlin.lazy") return null
             val lazyCallArgument = delegateCall.argumentList.arguments.singleOrNull() as? FirLambdaArgumentExpression ?: return null
-            return lazyCallArgument.expression as? FirAnonymousFunction
+            return (lazyCallArgument.expression as? FirAnonymousFunctionExpression)?.anonymousFunction
         }
 }

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/collectors/AbstractDiagnosticCollectorVisitor.kt

@@ -99,6 +99,10 @@ abstract class AbstractDiagnosticCollectorVisitor(
         }
     }
 
+    override fun visitAnonymousFunctionExpression(anonymousFunctionExpression: FirAnonymousFunctionExpression, data: Nothing?) {
+        visitAnonymousFunction(anonymousFunctionExpression.anonymousFunction, data)
+    }
+
     override fun visitAnonymousFunction(anonymousFunction: FirAnonymousFunction, data: Nothing?) {
         withSuppressedDiagnostics(anonymousFunction) {
             val labelName = anonymousFunction.label?.name?.let { Name.identifier(it) }

compiler/fir/dump/src/org/jetbrains/kotlin/fir/dump/HtmlFirDump.kt

@@ -1513,6 +1513,7 @@ class HtmlFirDump internal constructor(private var linkResolver: FirLinkResolver
                 is FirContinueExpression -> generate(expression)
                 is FirBreakExpression -> generate(expression)
                 is FirAnonymousObjectExpression -> generate(expression.anonymousObject)
+                is FirAnonymousFunctionExpression -> generate(expression.anonymousFunction, isStatement = false)
                 is FirUnitExpression -> generate(expression)
                 is FirStringConcatenationCall -> generate(expression)
                 is FirAnonymousFunction -> generate(expression, isStatement = false)

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrImplicitCastInserter.kt

@@ -46,7 +46,9 @@ class Fir2IrImplicitCastInserter(
 
     override fun visitAnonymousObjectExpression(anonymousObjectExpression: FirAnonymousObjectExpression, data: IrElement): IrElement = data
 
-    override fun visitAnonymousFunction(anonymousFunction: FirAnonymousFunction, data: IrElement): IrElement = data
+    override fun visitAnonymousFunctionExpression(anonymousFunctionExpression: FirAnonymousFunctionExpression, data: IrElement): IrElement {
+        return data
+    }
 
     override fun visitBinaryLogicExpression(binaryLogicExpression: FirBinaryLogicExpression, data: IrElement): IrElement = data
 

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrVisitor.kt

@@ -225,6 +225,10 @@ class Fir2IrVisitor(
         }
     }
 
+    override fun visitAnonymousFunctionExpression(anonymousFunctionExpression: FirAnonymousFunctionExpression, data: Any?): IrElement {
+        return visitAnonymousFunction(anonymousFunctionExpression.anonymousFunction, data)
+    }
+
     override fun visitAnonymousFunction(anonymousFunction: FirAnonymousFunction, data: Any?): IrElement {
         return anonymousFunction.convertWithOffsets { startOffset, endOffset ->
             val irFunction = declarationStorage.createIrFunction(

compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/DeclarationsConverter.kt

@@ -95,7 +95,7 @@ class DeclarationsConverter(
                 PACKAGE_DIRECTIVE -> context.packageFqName = convertPackageName(it)
                 IMPORT_LIST -> importList += convertImportDirectives(it)
                 CLASS -> firDeclarationList += convertClass(it)
-                FUN -> firDeclarationList += convertFunctionDeclaration(it)
+                FUN -> firDeclarationList += convertFunctionDeclaration(it) as FirDeclaration
                 PROPERTY -> firDeclarationList += convertPropertyDeclaration(it)
                 TYPEALIAS -> firDeclarationList += convertTypeAlias(it)
                 OBJECT_DECLARATION -> firDeclarationList += convertClass(it)
@@ -125,7 +125,7 @@ class DeclarationsConverter(
         val firStatements = block.forEachChildrenReturnList<FirStatement> { node, container ->
             when (node.tokenType) {
                 CLASS, OBJECT_DECLARATION -> container += convertClass(node) as FirStatement
-                FUN -> container += convertFunctionDeclaration(node) as FirStatement
+                FUN -> container += convertFunctionDeclaration(node)
                 PROPERTY -> container += convertPropertyDeclaration(node) as FirStatement
                 DESTRUCTURING_DECLARATION -> container += convertDestructingDeclaration(node).toFirDestructingDeclaration(baseModuleData)
                 TYPEALIAS -> container += convertTypeAlias(node) as FirStatement
@@ -735,7 +735,7 @@ class DeclarationsConverter(
             when (node.tokenType) {
                 ENUM_ENTRY -> container += convertEnumEntry(node, classWrapper)
                 CLASS -> container += convertClass(node)
-                FUN -> container += convertFunctionDeclaration(node, classWrapper)
+                FUN -> container += convertFunctionDeclaration(node, classWrapper) as FirDeclaration
                 PROPERTY -> container += convertPropertyDeclaration(node, classWrapper)
                 TYPEALIAS -> container += convertTypeAlias(node)
                 OBJECT_DECLARATION -> container += convertClass(node)
@@ -1338,7 +1338,7 @@ class DeclarationsConverter(
     /**
      * @see org.jetbrains.kotlin.parsing.KotlinParsing.parseFunction
      */
-    fun convertFunctionDeclaration(functionDeclaration: LighterASTNode, classWrapper: ClassWrapper? = null): FirDeclaration {
+    fun convertFunctionDeclaration(functionDeclaration: LighterASTNode, classWrapper: ClassWrapper? = null): FirStatement {
         var modifiers = Modifier()
         var identifier: String? = null
         val firTypeParameters = mutableListOf<FirTypeParameter>()
@@ -1378,11 +1378,12 @@ class DeclarationsConverter(
         val parentNode = functionDeclaration.getParent()
         val isLocal = !(parentNode?.tokenType == KT_FILE || parentNode?.tokenType == CLASS_BODY)
         val target: FirFunctionTarget
+        val functionSource = functionDeclaration.toFirSourceElement()
         val functionBuilder = if (identifier == null && isLocal) {
             val labelName = functionDeclaration.getLabelName() ?: context.calleeNamesForLambda.lastOrNull()?.identifier
             target = FirFunctionTarget(labelName = labelName, isLambda = false)
             FirAnonymousFunctionBuilder().apply {
-                source = functionDeclaration.toFirSourceElement()
+                source = functionSource
                 receiverTypeRef = receiverType
                 symbol = FirAnonymousFunctionSymbol()
                 isLambda = false
@@ -1392,7 +1393,7 @@ class DeclarationsConverter(
             val labelName = runIf(!functionName.isSpecial) { functionName.identifier }
             target = FirFunctionTarget(labelName, isLambda = false)
             FirSimpleFunctionBuilder().apply {
-                source = functionDeclaration.toFirSourceElement()
+                source = functionSource
                 receiverTypeRef = receiverType
                 name = functionName
                 status = FirDeclarationStatusImpl(
@@ -1415,7 +1416,7 @@ class DeclarationsConverter(
             }
         }
 
-        return functionBuilder.apply {
+        val function = functionBuilder.apply {
             moduleData = baseModuleData
             origin = FirDeclarationOrigin.Source
             returnTypeRef = returnType!!
@@ -1451,6 +1452,14 @@ class DeclarationsConverter(
                 fillDanglingConstraintsTo(firTypeParameters, typeConstraints, it)
             }
         }
+        return if (function is FirAnonymousFunction) {
+            buildAnonymousFunctionExpression {
+                source = functionSource
+                anonymousFunction = function
+            }
+        } else {
+            function
+        }
     }
 
     /**

compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/ExpressionsConverter.kt

@@ -133,7 +133,7 @@ class ExpressionsConverter(
 
         val expressionSource = lambdaExpression.toFirSourceElement()
         val target: FirFunctionTarget
-        return buildAnonymousFunction {
+        val anonymousFunction = buildAnonymousFunction {
             source = expressionSource
             moduleData = baseModuleData
             origin = FirDeclarationOrigin.Source
@@ -199,6 +199,10 @@ class ExpressionsConverter(
         }.also { 
             target.bind(it)
         }
+        return buildAnonymousFunctionExpression {
+            source = expressionSource
+            this.anonymousFunction = anonymousFunction
+        }
     }
 
     /**

compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt

@@ -1048,8 +1048,9 @@ open class RawFirBuilder(
             }
 
             val target = FirFunctionTarget(labelName, isLambda = false)
-            return functionBuilder.apply {
+            val functionSource = function.toFirSourceElement()
-                source = function.toFirSourceElement()
+            val firFunction = functionBuilder.apply {
+                source = functionSource
                 moduleData = baseModuleData
                 origin = FirDeclarationOrigin.Source
                 returnTypeRef = returnType
@@ -1085,6 +1086,14 @@ open class RawFirBuilder(
                     function.fillDanglingConstraintsTo(it)
                 }
             }
+            return if (firFunction is FirAnonymousFunction) {
+                buildAnonymousFunctionExpression {
+                    source = functionSource
+                    anonymousFunction = firFunction
+                }
+            } else {
+                firFunction
+            }
         }
 
         private fun KtDeclarationWithBody.obtainContractDescription(): FirContractDescription? {
@@ -1114,7 +1123,7 @@ open class RawFirBuilder(
             }
 
             val target: FirFunctionTarget
-            return buildAnonymousFunction {
+            val anonymousFunction = buildAnonymousFunction {
                 source = literalSource
                 moduleData = baseModuleData
                 origin = FirDeclarationOrigin.Source
@@ -1190,6 +1199,10 @@ open class RawFirBuilder(
             }.also {
                 bindFunctionTarget(target, it)
             }
+            return buildAnonymousFunctionExpression {
+                source = expression.toFirPsiSourceElement()
+                this.anonymousFunction = anonymousFunction
+            }
         }
 
         private fun KtSecondaryConstructor.toFirConstructor(

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

@@ -107,7 +107,7 @@ fun Candidate.resolveArgumentExpression(
             else
                 preprocessCallableReference(argument, expectedType, context)
         // TODO:!
-        is FirAnonymousFunction -> preprocessLambdaArgument(csBuilder, argument, expectedType, expectedTypeRef, context, sink)
+        is FirAnonymousFunctionExpression -> preprocessLambdaArgument(csBuilder, argument, expectedType, expectedTypeRef, context, sink)
         // TODO:!
         //TODO: Collection literal
         is FirWrappedArgumentExpression -> resolveArgumentExpression(
@@ -508,7 +508,7 @@ fun FirExpression.isFunctional(
     expectedFunctionType: ConeKotlinType?,
 ): Boolean {
     when (unwrapArgument()) {
-        is FirAnonymousFunction, is FirCallableReferenceAccess -> return true
+        is FirAnonymousFunctionExpression, is FirCallableReferenceAccess -> return true
         else -> {
             // Either a functional type or a subtype of a class that has a contributed `invoke`.
             val coneType = typeRef.coneTypeSafe<ConeKotlinType>() ?: return false

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt

@@ -272,14 +272,19 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         return FirControlFlowGraphReferenceImpl(graph)
     }
 
-    fun visitPostponedAnonymousFunction(anonymousFunction: FirAnonymousFunction) {
+    fun visitPostponedAnonymousFunction(anonymousFunctionExpression: FirAnonymousFunctionExpression) {
+        val anonymousFunction = anonymousFunctionExpression.anonymousFunction
         getOrCreateLocalVariableAssignmentAnalyzer(anonymousFunction)?.visitPostponedAnonymousFunction(anonymousFunction)
-        val (enterNode, exitNode) = graphBuilder.visitPostponedAnonymousFunction(anonymousFunction)
+        val (enterNode, exitNode) = graphBuilder.visitPostponedAnonymousFunction(anonymousFunctionExpression)
         enterNode.mergeIncomingFlow()
         exitNode.mergeIncomingFlow()
         enterNode.flow = enterNode.flow.fork()
     }
 
+    fun exitAnonymousFunctionExpression(anonymousFunctionExpression: FirAnonymousFunctionExpression) {
+        graphBuilder.exitAnonymousFunctionExpression(anonymousFunctionExpression).mergeIncomingFlow()
+    }
+
     // ----------------------------------- Classes -----------------------------------
 
     fun enterClass() {
@@ -890,7 +895,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
     }
 
     fun enterFunctionCall(functionCall: FirFunctionCall) {
-        val lambdaArgs = functionCall.arguments.filterIsInstance<FirAnonymousFunction>()
+        val lambdaArgs = functionCall.arguments.mapNotNull { (it as? FirAnonymousFunctionExpression)?.anonymousFunction }
         if (lambdaArgs.size > 1) {
             getOrCreateLocalVariableAssignmentAnalyzer(lambdaArgs.first())?.enterFunctionCallWithMultipleLambdaArgs(lambdaArgs)
         }
@@ -898,7 +903,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
 
     @OptIn(PrivateForInline::class)
     fun exitFunctionCall(functionCall: FirFunctionCall, callCompleted: Boolean) {
-        val lambdaArgs = functionCall.arguments.filterIsInstance<FirAnonymousFunction>()
+        val lambdaArgs = functionCall.arguments.mapNotNull { (it as? FirAnonymousFunctionExpression)?.anonymousFunction }
         if (lambdaArgs.size > 1) {
             getOrCreateLocalVariableAssignmentAnalyzer(lambdaArgs.first())?.enterFunctionCallWithMultipleLambdaArgs(lambdaArgs)
         }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/CFGNode.kt

@@ -206,7 +206,7 @@ class PostponedLambdaEnterNode(owner: ControlFlowGraph, override val fir: FirAno
         return visitor.visitPostponedLambdaEnterNode(this, data)
     }
 }
-class PostponedLambdaExitNode(owner: ControlFlowGraph, override val fir: FirAnonymousFunction, level: Int, id: Int) : CFGNode<FirAnonymousFunction>(owner, level, id) {
+class PostponedLambdaExitNode(owner: ControlFlowGraph, override val fir: FirAnonymousFunctionExpression, level: Int, id: Int) : CFGNode<FirAnonymousFunctionExpression>(owner, level, id) {
     override fun <R, D> accept(visitor: ControlFlowGraphVisitor<R, D>, data: D): R {
         return visitor.visitPostponedLambdaExitNode(this, data)
     }
@@ -217,6 +217,12 @@ class UnionFunctionCallArgumentsNode(owner: ControlFlowGraph, override val fir:
     }
 }
 
+class AnonymousFunctionExpressionExitNode(owner: ControlFlowGraph, override val fir: FirAnonymousFunctionExpression, level: Int, id: Int) : CFGNode<FirAnonymousFunctionExpression>(owner, level, id) {
+    override fun <R, D> accept(visitor: ControlFlowGraphVisitor<R, D>, data: D): R {
+        return visitor.visitAnonymousFunctionExpressionExitNode(this, data)
+    }
+}
+
 // ----------------------------------- Classes -----------------------------------
 
 class ClassEnterNode(owner: ControlFlowGraph, override val fir: FirClass<*>, level: Int, id: Int) : CFGNode<FirClass<*>>(owner, level, id), EnterNodeMarker {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/CFGNodeRenderer.kt

@@ -102,6 +102,8 @@ fun CFGNode<*>.render(): String =
                 is PostponedLambdaEnterNode -> "Postponed enter to lambda"
                 is PostponedLambdaExitNode -> "Postponed exit from lambda"
 
+                is AnonymousFunctionExpressionExitNode -> "Exit anonymous function expression"
+
                 is UnionFunctionCallArgumentsNode -> "Call arguments union"
 
                 is ClassEnterNode -> "Enter class ${owner.name}"

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphBuilder.kt

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSymbolOwner
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.*
+import org.jetbrains.kotlin.fir.expressions.builder.buildAnonymousFunctionExpression
 import org.jetbrains.kotlin.fir.references.FirControlFlowGraphReference
 import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.dfa.*
@@ -228,9 +229,10 @@ class ControlFlowGraphBuilder {
 
     // ----------------------------------- Anonymous function -----------------------------------
 
-    fun visitPostponedAnonymousFunction(anonymousFunction: FirAnonymousFunction): Pair<PostponedLambdaEnterNode, PostponedLambdaExitNode> {
+    fun visitPostponedAnonymousFunction(anonymousFunctionExpression: FirAnonymousFunctionExpression): Pair<PostponedLambdaEnterNode, PostponedLambdaExitNode> {
+        val anonymousFunction = anonymousFunctionExpression.anonymousFunction
         val enterNode = createPostponedLambdaEnterNode(anonymousFunction)
-        val exitNode = createPostponedLambdaExitNode(anonymousFunction)
+        val exitNode = createPostponedLambdaExitNode(anonymousFunctionExpression)
         val symbol = anonymousFunction.symbol
         postponedLambdas += symbol
         entersToPostponedAnonymousFunctions[symbol] = enterNode
@@ -256,7 +258,11 @@ class ControlFlowGraphBuilder {
 
         if (previousNodeIsNew) {
             assert(symbol !in exitsFromPostponedAnonymousFunctions)
-            val exitFromLambda = createPostponedLambdaExitNode(anonymousFunction).also {
+            val lambdaExpression = buildAnonymousFunctionExpression {
+                source = anonymousFunction.source
+                this.anonymousFunction = anonymousFunction
+            }
+            val exitFromLambda = createPostponedLambdaExitNode(lambdaExpression).also {
                 exitsFromPostponedAnonymousFunctions[symbol] = it
             }
             addEdge(previousNode, exitFromLambda)
@@ -352,6 +358,12 @@ class ControlFlowGraphBuilder {
         }
     }
 
+    fun exitAnonymousFunctionExpression(anonymousFunctionExpression: FirAnonymousFunctionExpression): AnonymousFunctionExpressionExitNode {
+        return createAnonymousFunctionExpressionExitNode(anonymousFunctionExpression).also {
+            addNewSimpleNode(it)
+        }
+    }
+
     // ----------------------------------- Classes -----------------------------------
 
     fun enterClass() {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphNodeBuilder.kt

@@ -215,12 +215,15 @@ fun ControlFlowGraphBuilder.createExitSafeCallNode(fir: FirSafeCallExpression):
 fun ControlFlowGraphBuilder.createEnterSafeCallNode(fir: FirSafeCallExpression): EnterSafeCallNode =
     EnterSafeCallNode(currentGraph, fir, levelCounter, createId())
 
-fun ControlFlowGraphBuilder.createPostponedLambdaExitNode(fir: FirAnonymousFunction): PostponedLambdaExitNode =
+fun ControlFlowGraphBuilder.createPostponedLambdaExitNode(fir: FirAnonymousFunctionExpression): PostponedLambdaExitNode =
     PostponedLambdaExitNode(currentGraph, fir, levelCounter, createId())
 
 fun ControlFlowGraphBuilder.createPostponedLambdaEnterNode(fir: FirAnonymousFunction): PostponedLambdaEnterNode =
     PostponedLambdaEnterNode(currentGraph, fir, levelCounter, createId())
 
+fun ControlFlowGraphBuilder.createAnonymousFunctionExpressionExitNode(fir: FirAnonymousFunctionExpression): AnonymousFunctionExpressionExitNode =
+    AnonymousFunctionExpressionExitNode(currentGraph, fir, levelCounter, createId())
+
 fun ControlFlowGraphBuilder.createAnonymousObjectExitNode(fir: FirAnonymousObject): AnonymousObjectExitNode =
     AnonymousObjectExitNode(currentGraph, fir, levelCounter, createId())
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphVisitor.kt

@@ -46,6 +46,10 @@ abstract class ControlFlowGraphVisitor<out R, in D> {
         return visitNode(node, data)
     }
 
+    open fun visitAnonymousFunctionExpressionExitNode(node: AnonymousFunctionExpressionExitNode, data: D): R {
+        return visitNode(node, data)
+    }
+
     // ----------------------------------- Classes -----------------------------------
 
     open fun visitAnonymousObjectExitNode(node: AnonymousObjectExitNode, data: D): R {

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

@@ -62,7 +62,7 @@ object ConeConstraintSystemUtilContext : ConstraintSystemUtilContext {
         require(declaration is PostponedResolvedAtom)
         return when (declaration) {
             is LambdaWithTypeVariableAsExpectedTypeAtom -> {
-                val atom = declaration.atom
+                val atom = declaration.atom.anonymousFunction
                 return if (atom.isLambda) { // lambda - must return null in case of absent parameters
                     if (atom.valueParameters.isNotEmpty())
                         atom.collectDeclaredValueParameterTypes()
@@ -83,12 +83,12 @@ object ConeConstraintSystemUtilContext : ConstraintSystemUtilContext {
 
     override fun PostponedAtomWithRevisableExpectedType.isAnonymousFunction(): Boolean {
         require(this is PostponedResolvedAtom)
-        return this is LambdaWithTypeVariableAsExpectedTypeAtom && !this.atom.isLambda
+        return this is LambdaWithTypeVariableAsExpectedTypeAtom && !this.atom.anonymousFunction.isLambda
     }
 
     override fun PostponedAtomWithRevisableExpectedType.isFunctionExpressionWithReceiver(): Boolean {
         require(this is PostponedResolvedAtom)
-        return this is LambdaWithTypeVariableAsExpectedTypeAtom && !this.atom.isLambda && this.atom.receiverTypeRef?.coneType != null
+        return this is LambdaWithTypeVariableAsExpectedTypeAtom && !this.atom.anonymousFunction.isLambda && this.atom.anonymousFunction.receiverTypeRef?.coneType != null
     }
 
     override fun createTypeVariableForLambdaReturnType(): TypeVariableMarker {

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

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.fir.resolve.inference
 
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
+import org.jetbrains.kotlin.fir.expressions.FirAnonymousFunctionExpression
 import org.jetbrains.kotlin.fir.expressions.FirCallableReferenceAccess
 import org.jetbrains.kotlin.fir.resolve.calls.*
 import org.jetbrains.kotlin.fir.resolve.createFunctionalType
@@ -15,12 +16,11 @@ import org.jetbrains.kotlin.fir.typeContext
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemBuilder
 import org.jetbrains.kotlin.resolve.calls.inference.addSubtypeConstraintIfCompatible
-import org.jetbrains.kotlin.types.AbstractTypeChecker
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 fun Candidate.preprocessLambdaArgument(
     csBuilder: ConstraintSystemBuilder,
-    argument: FirAnonymousFunction,
+    argument: FirAnonymousFunctionExpression,
     expectedType: ConeKotlinType?,
     expectedTypeRef: FirTypeRef?,
     context: ResolutionContext,
@@ -32,15 +32,17 @@ fun Candidate.preprocessLambdaArgument(
         return LambdaWithTypeVariableAsExpectedTypeAtom(argument, expectedType, expectedTypeRef, this)
     }
 
+    val anonymousFunction = argument.anonymousFunction
+
     val resolvedArgument =
         extractLambdaInfoFromFunctionalType(
             expectedType,
             expectedTypeRef,
-            argument,
+            anonymousFunction,
             returnTypeVariable,
             context.bodyResolveComponents,
             this
-        ) ?: extraLambdaInfo(expectedType, argument, csBuilder, context.session, this)
+        ) ?: extraLambdaInfo(expectedType, anonymousFunction, csBuilder, context.session, this)
 
     if (expectedType != null) {
         // TODO: add SAM conversion processing

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

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.fir.resolve.inference
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
+import org.jetbrains.kotlin.fir.expressions.FirAnonymousFunctionExpression
 import org.jetbrains.kotlin.fir.expressions.FirCallableReferenceAccess
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.references.FirNamedReference
@@ -78,7 +79,7 @@ class ResolvedLambdaAtom(
 }
 
 class LambdaWithTypeVariableAsExpectedTypeAtom(
-    override val atom: FirAnonymousFunction,
+    override val atom: FirAnonymousFunctionExpression,
     private val initialExpectedTypeType: ConeKotlinType,
     val expectedTypeRef: FirTypeRef,
     val candidateOfOuterCall: Candidate,

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

@@ -190,7 +190,7 @@ class FirCallCompletionResultsWriterTransformer(
                         val functionIsInline =
                             (symbol as? FirNamedFunctionSymbol)?.fir?.isInline == true || symbol.isArrayConstructorWithLambda
                         for ((argument, parameter) in newArgumentList.mapping) {
-                            val lambda = argument.unwrapArgument() as? FirAnonymousFunction ?: continue
+                            val lambda = (argument.unwrapArgument() as? FirAnonymousFunctionExpression)?.anonymousFunction ?: continue
                             val inlineStatus = when {
                                 parameter.isCrossinline && functionIsInline -> InlineStatus.CrossInline
                                 parameter.isNoinline -> InlineStatus.NoInline
@@ -412,7 +412,11 @@ class FirCallCompletionResultsWriterTransformer(
             } else {
                 valueParameter.returnTypeRef.substitute(this)
             }
-            argument.unwrapArgument() to expectedType
+            val unwrappedArgument: FirElement = when (val unwrappedArgument = argument.unwrapArgument()) {
+                is FirAnonymousFunctionExpression -> unwrappedArgument.anonymousFunction
+                else -> unwrappedArgument
+            }
+            unwrappedArgument to expectedType
         }?.toMap()
 
 
@@ -495,6 +499,13 @@ class FirCallCompletionResultsWriterTransformer(
             }
     }
 
+    override fun transformAnonymousFunctionExpression(
+        anonymousFunctionExpression: FirAnonymousFunctionExpression,
+        data: ExpectedArgumentType?
+    ): FirStatement {
+        return anonymousFunctionExpression.transformAnonymousFunction(this, data)
+    }
+
     override fun transformAnonymousFunction(
         anonymousFunction: FirAnonymousFunction,
         data: ExpectedArgumentType?,
@@ -770,7 +781,7 @@ class FirCallCompletionResultsWriterTransformer(
 
 sealed class ExpectedArgumentType {
     class ArgumentsMap(
-        val map: Map<FirExpression, ConeKotlinType>,
+        val map: Map<FirElement, ConeKotlinType>,
         val lambdasReturnTypes: Map<FirAnonymousFunction, ConeKotlinType>
     ) : ExpectedArgumentType()
 
@@ -778,7 +789,7 @@ sealed class ExpectedArgumentType {
     object NoApproximation : ExpectedArgumentType()
 }
 
-private fun ExpectedArgumentType.getExpectedType(argument: FirExpression): ConeKotlinType? = when (this) {
+private fun ExpectedArgumentType.getExpectedType(argument: FirElement): ConeKotlinType? = when (this) {
     is ExpectedArgumentType.ArgumentsMap -> map[argument]
     is ExpectedArgumentType.ExpectedType -> type
     ExpectedArgumentType.NoApproximation -> null

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

@@ -22,11 +22,19 @@ object InvocationKindTransformer : FirTransformer<Any?>() {
             return element
         }
 
+        override fun transformAnonymousFunctionExpression(
+            anonymousFunctionExpression: FirAnonymousFunctionExpression,
+            data: Pair<Map<FirExpression, EventOccurrencesRange>, EventOccurrencesRange?>
+        ): FirStatement {
+            val kind = data.second ?: data.first[anonymousFunctionExpression]
+            return anonymousFunctionExpression.transformAnonymousFunction(this, emptyMap<FirExpression, EventOccurrencesRange>() to kind)
+        }
+
         override fun transformAnonymousFunction(
             anonymousFunction: FirAnonymousFunction,
             data: Pair<Map<FirExpression, EventOccurrencesRange>, EventOccurrencesRange?>
         ): FirStatement {
-            val kind = data.second ?: data.first[anonymousFunction]
+            val kind = data.second
             if (kind != null) {
                 anonymousFunction.replaceInvocationKind(kind)
             }

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

@@ -310,6 +310,13 @@ open class FirBodyResolveTransformer(
         return declarationsTransformer.transformAnonymousFunction(anonymousFunction, data)
     }
 
+    override fun transformAnonymousFunctionExpression(
+        anonymousFunctionExpression: FirAnonymousFunctionExpression,
+        data: ResolutionMode
+    ): FirStatement {
+        return expressionsTransformer.transformAnonymousFunctionExpression(anonymousFunctionExpression, data)
+    }
+
     override fun transformValueParameter(valueParameter: FirValueParameter, data: ResolutionMode): FirStatement {
         return declarationsTransformer.transformValueParameter(valueParameter, data)
     }

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

@@ -639,7 +639,6 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
         return when (data) {
             is ResolutionMode.ContextDependent, is ResolutionMode.ContextDependentDelegate -> {
                 context.withAnonymousFunction(anonymousFunction, components, data) {
-                    dataFlowAnalyzer.visitPostponedAnonymousFunction(anonymousFunction)
                     anonymousFunction.addReturn()
                 }
             }

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

@@ -1078,6 +1078,22 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
         return anonymousObjectExpression
     }
 
+    override fun transformAnonymousFunctionExpression(
+        anonymousFunctionExpression: FirAnonymousFunctionExpression,
+        data: ResolutionMode
+    ): FirStatement {
+        anonymousFunctionExpression.transformAnonymousFunction(transformer, data)
+        when (data) {
+            is ResolutionMode.ContextDependent, is ResolutionMode.ContextDependentDelegate -> {
+                dataFlowAnalyzer.visitPostponedAnonymousFunction(anonymousFunctionExpression)
+            }
+            else -> {
+                dataFlowAnalyzer.exitAnonymousFunctionExpression(anonymousFunctionExpression)
+            }
+        }
+        return anonymousFunctionExpression
+    }
+
     // ------------------------------------------------------------------------------------------------
 
     internal fun <T> storeTypeFromCallee(access: T) where T : FirQualifiedAccess, T : FirExpression {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/contracts/FirContractResolveTransformer.kt

@@ -18,10 +18,7 @@ import org.jetbrains.kotlin.fir.declarations.builder.buildAnonymousFunction
 import org.jetbrains.kotlin.fir.declarations.synthetic.FirSyntheticProperty
 import org.jetbrains.kotlin.fir.errorTypeFromPrototype
 import org.jetbrains.kotlin.fir.expressions.*
-import org.jetbrains.kotlin.fir.expressions.builder.buildArgumentList
+import org.jetbrains.kotlin.fir.expressions.builder.*
-import org.jetbrains.kotlin.fir.expressions.builder.buildBlock
-import org.jetbrains.kotlin.fir.expressions.builder.buildFunctionCall
-import org.jetbrains.kotlin.fir.expressions.builder.buildLambdaArgumentExpression
 import org.jetbrains.kotlin.fir.moduleData
 import org.jetbrains.kotlin.fir.references.builder.buildSimpleNamedReference
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
@@ -146,7 +143,8 @@ open class FirContractResolveTransformer(
             if (resolvedId != FirContractsDslNames.CONTRACT) return transformOwnerWithUnresolvedContract(owner)
             if (contractCall.arguments.size != 1) return transformOwnerOfErrorContract(owner)
             val argument = contractCall.argument as? FirLambdaArgumentExpression ?: return transformOwnerOfErrorContract(owner)
-            val lambdaBody = (argument.expression as FirAnonymousFunction).body ?: return transformOwnerOfErrorContract(owner)
+            val lambdaBody = (argument.expression as FirAnonymousFunctionExpression).anonymousFunction.body
+                ?: return transformOwnerOfErrorContract(owner)
 
             val resolvedContractDescription = buildResolvedContractDescription {
                 val effectExtractor = ConeEffectExtractor(session, owner, valueParameters)
@@ -185,7 +183,9 @@ open class FirContractResolveTransformer(
             }
 
             val lambdaArgument = buildLambdaArgumentExpression {
-                expression = effectsBlock
+                expression = buildAnonymousFunctionExpression {
+                    anonymousFunction = effectsBlock
+                }
             }
 
             val contractCall = buildFunctionCall {

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/FirAnonymousFunction.kt

@@ -12,7 +12,6 @@ import org.jetbrains.kotlin.fir.FirModuleData
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.expressions.FirBlock
-import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.references.FirControlFlowGraphReference
 import org.jetbrains.kotlin.fir.symbols.impl.FirAnonymousFunctionSymbol
 import org.jetbrains.kotlin.fir.types.FirTypeRef
@@ -23,7 +22,7 @@ import org.jetbrains.kotlin.fir.visitors.*
  * DO NOT MODIFY IT MANUALLY
  */
 
-abstract class FirAnonymousFunction : FirFunction<FirAnonymousFunction>, FirExpression(), FirTypeParametersOwner {
+abstract class FirAnonymousFunction : FirFunction<FirAnonymousFunction>(), FirTypeParametersOwner {
     abstract override val source: FirSourceElement?
     abstract override val moduleData: FirModuleData
     abstract override val resolvePhase: FirResolvePhase
@@ -35,13 +34,13 @@ abstract class FirAnonymousFunction : FirFunction<FirAnonymousFunction>, FirExpr
     abstract override val controlFlowGraphReference: FirControlFlowGraphReference?
     abstract override val valueParameters: List<FirValueParameter>
     abstract override val body: FirBlock?
-    abstract override val typeRef: FirTypeRef
     abstract override val symbol: FirAnonymousFunctionSymbol
     abstract val label: FirLabel?
     abstract val invocationKind: EventOccurrencesRange?
     abstract val inlineStatus: InlineStatus
     abstract val isLambda: Boolean
     abstract override val typeParameters: List<FirTypeParameter>
+    abstract val typeRef: FirTypeRef
 
     override fun <R, D> accept(visitor: FirVisitor<R, D>, data: D): R = visitor.visitAnonymousFunction(this, data)
 
@@ -61,12 +60,12 @@ abstract class FirAnonymousFunction : FirFunction<FirAnonymousFunction>, FirExpr
 
     abstract override fun replaceBody(newBody: FirBlock?)
 
-    abstract override fun replaceTypeRef(newTypeRef: FirTypeRef)
-
     abstract fun replaceInvocationKind(newInvocationKind: EventOccurrencesRange?)
 
     abstract fun replaceInlineStatus(newInlineStatus: InlineStatus)
 
+    abstract fun replaceTypeRef(newTypeRef: FirTypeRef)
+
     abstract override fun <D> transformAnnotations(transformer: FirTransformer<D>, data: D): FirAnonymousFunction
 
     abstract override fun <D> transformReturnTypeRef(transformer: FirTransformer<D>, data: D): FirAnonymousFunction

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/FirConstructor.kt

@@ -7,7 +7,6 @@ package org.jetbrains.kotlin.fir.declarations
 
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirModuleData
-import org.jetbrains.kotlin.fir.FirPureAbstractElement
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.expressions.FirBlock
@@ -24,7 +23,7 @@ import org.jetbrains.kotlin.fir.visitors.*
  * DO NOT MODIFY IT MANUALLY
  */
 
-abstract class FirConstructor : FirPureAbstractElement(), FirFunction<FirConstructor>, FirCallableMemberDeclaration<FirConstructor>, FirTypeParameterRefsOwner {
+abstract class FirConstructor : FirFunction<FirConstructor>(), FirCallableMemberDeclaration<FirConstructor>, FirTypeParameterRefsOwner {
     abstract override val source: FirSourceElement?
     abstract override val moduleData: FirModuleData
     abstract override val resolvePhase: FirResolvePhase

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/FirErrorFunction.kt

@@ -7,7 +7,6 @@ package org.jetbrains.kotlin.fir.declarations
 
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirModuleData
-import org.jetbrains.kotlin.fir.FirPureAbstractElement
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.diagnostics.ConeDiagnostic
 import org.jetbrains.kotlin.fir.diagnostics.FirDiagnosticHolder
@@ -23,7 +22,7 @@ import org.jetbrains.kotlin.fir.visitors.*
  * DO NOT MODIFY IT MANUALLY
  */
 
-abstract class FirErrorFunction : FirPureAbstractElement(), FirFunction<FirErrorFunction>, FirDiagnosticHolder, FirTypeParametersOwner {
+abstract class FirErrorFunction : FirFunction<FirErrorFunction>(), FirDiagnosticHolder, FirTypeParametersOwner {
     abstract override val source: FirSourceElement?
     abstract override val moduleData: FirModuleData
     abstract override val resolvePhase: FirResolvePhase

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/FirFunction.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.fir.declarations
 
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirModuleData
+import org.jetbrains.kotlin.fir.FirPureAbstractElement
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.FirTargetElement
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
@@ -22,20 +23,20 @@ import org.jetbrains.kotlin.fir.visitors.*
  * DO NOT MODIFY IT MANUALLY
  */
 
-sealed interface FirFunction<F : FirFunction<F>> : FirCallableDeclaration<F>, FirTargetElement, FirTypeParameterRefsOwner, FirControlFlowGraphOwner, FirStatement {
+sealed class FirFunction<F : FirFunction<F>> : FirPureAbstractElement(), FirCallableDeclaration<F>, FirTargetElement, FirTypeParameterRefsOwner, FirControlFlowGraphOwner, FirStatement {
-    override val source: FirSourceElement?
+    abstract override val source: FirSourceElement?
-    override val moduleData: FirModuleData
+    abstract override val moduleData: FirModuleData
-    override val resolvePhase: FirResolvePhase
+    abstract override val resolvePhase: FirResolvePhase
-    override val origin: FirDeclarationOrigin
+    abstract override val origin: FirDeclarationOrigin
-    override val attributes: FirDeclarationAttributes
+    abstract override val attributes: FirDeclarationAttributes
-    override val annotations: List<FirAnnotationCall>
+    abstract override val annotations: List<FirAnnotationCall>
-    override val returnTypeRef: FirTypeRef
+    abstract override val returnTypeRef: FirTypeRef
-    override val receiverTypeRef: FirTypeRef?
+    abstract override val receiverTypeRef: FirTypeRef?
-    override val typeParameters: List<FirTypeParameterRef>
+    abstract override val typeParameters: List<FirTypeParameterRef>
-    override val controlFlowGraphReference: FirControlFlowGraphReference?
+    abstract override val controlFlowGraphReference: FirControlFlowGraphReference?
-    override val symbol: FirFunctionSymbol<F>
+    abstract override val symbol: FirFunctionSymbol<F>
-    val valueParameters: List<FirValueParameter>
+    abstract val valueParameters: List<FirValueParameter>
-    val body: FirBlock?
+    abstract val body: FirBlock?
 
     override fun <R, D> accept(visitor: FirVisitor<R, D>, data: D): R = visitor.visitFunction(this, data)
 
@@ -43,27 +44,27 @@ sealed interface FirFunction<F : FirFunction<F>> : FirCallableDeclaration<F>, Fi
     override fun <E: FirElement, D> transform(transformer: FirTransformer<D>, data: D): E = 
         transformer.transformFunction(this, data) as E
 
-    override fun replaceResolvePhase(newResolvePhase: FirResolvePhase)
+    abstract override fun replaceResolvePhase(newResolvePhase: FirResolvePhase)
 
-    override fun replaceReturnTypeRef(newReturnTypeRef: FirTypeRef)
+    abstract override fun replaceReturnTypeRef(newReturnTypeRef: FirTypeRef)
 
-    override fun replaceReceiverTypeRef(newReceiverTypeRef: FirTypeRef?)
+    abstract override fun replaceReceiverTypeRef(newReceiverTypeRef: FirTypeRef?)
 
-    override fun replaceControlFlowGraphReference(newControlFlowGraphReference: FirControlFlowGraphReference?)
+    abstract override fun replaceControlFlowGraphReference(newControlFlowGraphReference: FirControlFlowGraphReference?)
 
-    fun replaceValueParameters(newValueParameters: List<FirValueParameter>)
+    abstract fun replaceValueParameters(newValueParameters: List<FirValueParameter>)
 
-    fun replaceBody(newBody: FirBlock?)
+    abstract fun replaceBody(newBody: FirBlock?)
 
-    override fun <D> transformAnnotations(transformer: FirTransformer<D>, data: D): FirFunction<F>
+    abstract override fun <D> transformAnnotations(transformer: FirTransformer<D>, data: D): FirFunction<F>
 
-    override fun <D> transformReturnTypeRef(transformer: FirTransformer<D>, data: D): FirFunction<F>
+    abstract override fun <D> transformReturnTypeRef(transformer: FirTransformer<D>, data: D): FirFunction<F>
 
-    override fun <D> transformReceiverTypeRef(transformer: FirTransformer<D>, data: D): FirFunction<F>
+    abstract override fun <D> transformReceiverTypeRef(transformer: FirTransformer<D>, data: D): FirFunction<F>
 
-    override fun <D> transformTypeParameters(transformer: FirTransformer<D>, data: D): FirFunction<F>
+    abstract override fun <D> transformTypeParameters(transformer: FirTransformer<D>, data: D): FirFunction<F>
 
-    fun <D> transformValueParameters(transformer: FirTransformer<D>, data: D): FirFunction<F>
+    abstract fun <D> transformValueParameters(transformer: FirTransformer<D>, data: D): FirFunction<F>
 
-    fun <D> transformBody(transformer: FirTransformer<D>, data: D): FirFunction<F>
+    abstract fun <D> transformBody(transformer: FirTransformer<D>, data: D): FirFunction<F>
 }

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/FirPropertyAccessor.kt

@@ -7,7 +7,6 @@ package org.jetbrains.kotlin.fir.declarations
 
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirModuleData
-import org.jetbrains.kotlin.fir.FirPureAbstractElement
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.contracts.FirContractDescription
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
@@ -24,7 +23,7 @@ import org.jetbrains.kotlin.fir.visitors.*
  * DO NOT MODIFY IT MANUALLY
  */
 
-abstract class FirPropertyAccessor : FirPureAbstractElement(), FirFunction<FirPropertyAccessor>, FirCallableMemberDeclaration<FirPropertyAccessor>, FirContractDescriptionOwner, FirTypeParametersOwner {
+abstract class FirPropertyAccessor : FirFunction<FirPropertyAccessor>(), FirCallableMemberDeclaration<FirPropertyAccessor>, FirContractDescriptionOwner, FirTypeParametersOwner {
     abstract override val source: FirSourceElement?
     abstract override val moduleData: FirModuleData
     abstract override val resolvePhase: FirResolvePhase

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/FirSimpleFunction.kt

@@ -7,7 +7,6 @@ package org.jetbrains.kotlin.fir.declarations
 
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirModuleData
-import org.jetbrains.kotlin.fir.FirPureAbstractElement
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.contracts.FirContractDescription
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
@@ -25,7 +24,7 @@ import org.jetbrains.kotlin.fir.visitors.*
  * DO NOT MODIFY IT MANUALLY
  */
 
-abstract class FirSimpleFunction : FirPureAbstractElement(), FirFunction<FirSimpleFunction>, FirCallableMemberDeclaration<FirSimpleFunction>, FirContractDescriptionOwner, FirTypeParametersOwner {
+abstract class FirSimpleFunction : FirFunction<FirSimpleFunction>(), FirCallableMemberDeclaration<FirSimpleFunction>, FirContractDescriptionOwner, FirTypeParametersOwner {
     abstract override val source: FirSourceElement?
     abstract override val moduleData: FirModuleData
     abstract override val resolvePhase: FirResolvePhase

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/builder/FirAnonymousFunctionBuilder.kt

@@ -23,7 +23,6 @@ import org.jetbrains.kotlin.fir.declarations.builder.FirFunctionBuilder
 import org.jetbrains.kotlin.fir.declarations.impl.FirAnonymousFunctionImpl
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.expressions.FirBlock
-import org.jetbrains.kotlin.fir.expressions.builder.FirExpressionBuilder
 import org.jetbrains.kotlin.fir.references.FirControlFlowGraphReference
 import org.jetbrains.kotlin.fir.symbols.impl.FirAnonymousFunctionSymbol
 import org.jetbrains.kotlin.fir.types.FirTypeRef
@@ -36,7 +35,7 @@ import org.jetbrains.kotlin.fir.visitors.*
  */
 
 @FirBuilderDsl
-class FirAnonymousFunctionBuilder : FirFunctionBuilder, FirAnnotationContainerBuilder, FirExpressionBuilder {
+class FirAnonymousFunctionBuilder : FirFunctionBuilder, FirAnnotationContainerBuilder {
     override var source: FirSourceElement? = null
     override lateinit var moduleData: FirModuleData
     override lateinit var origin: FirDeclarationOrigin
@@ -47,13 +46,13 @@ class FirAnonymousFunctionBuilder : FirFunctionBuilder, FirAnnotationContainerBu
     var controlFlowGraphReference: FirControlFlowGraphReference? = null
     override val valueParameters: MutableList<FirValueParameter> = mutableListOf()
     override var body: FirBlock? = null
-    override var typeRef: FirTypeRef = FirImplicitTypeRefImpl(null)
     lateinit var symbol: FirAnonymousFunctionSymbol
     var label: FirLabel? = null
     var invocationKind: EventOccurrencesRange? = null
     var inlineStatus: InlineStatus = InlineStatus.Unknown
     var isLambda: Boolean by kotlin.properties.Delegates.notNull<Boolean>()
     val typeParameters: MutableList<FirTypeParameter> = mutableListOf()
+    var typeRef: FirTypeRef = FirImplicitTypeRefImpl(null)
 
     override fun build(): FirAnonymousFunction {
         return FirAnonymousFunctionImpl(
@@ -67,13 +66,13 @@ class FirAnonymousFunctionBuilder : FirFunctionBuilder, FirAnnotationContainerBu
             controlFlowGraphReference,
             valueParameters,
             body,
-            typeRef,
             symbol,
             label,
             invocationKind,
             inlineStatus,
             isLambda,
             typeParameters,
+            typeRef,
         )
     }
 

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/impl/FirAnonymousFunctionImpl.kt

@@ -39,13 +39,13 @@ internal class FirAnonymousFunctionImpl(
     override var controlFlowGraphReference: FirControlFlowGraphReference?,
     override val valueParameters: MutableList<FirValueParameter>,
     override var body: FirBlock?,
-    override var typeRef: FirTypeRef,
     override val symbol: FirAnonymousFunctionSymbol,
     override var label: FirLabel?,
     override var invocationKind: EventOccurrencesRange?,
     override var inlineStatus: InlineStatus,
     override val isLambda: Boolean,
     override val typeParameters: MutableList<FirTypeParameter>,
+    override var typeRef: FirTypeRef,
 ) : FirAnonymousFunction() {
     override var resolvePhase: FirResolvePhase = FirResolvePhase.DECLARATIONS
 
@@ -60,9 +60,9 @@ internal class FirAnonymousFunctionImpl(
         controlFlowGraphReference?.accept(visitor, data)
         valueParameters.forEach { it.accept(visitor, data) }
         body?.accept(visitor, data)
-        typeRef.accept(visitor, data)
         label?.accept(visitor, data)
         typeParameters.forEach { it.accept(visitor, data) }
+        typeRef.accept(visitor, data)
     }
 
     override fun <D> transformChildren(transformer: FirTransformer<D>, data: D): FirAnonymousFunctionImpl {
@@ -72,9 +72,9 @@ internal class FirAnonymousFunctionImpl(
         controlFlowGraphReference = controlFlowGraphReference?.transform(transformer, data)
         transformValueParameters(transformer, data)
         transformBody(transformer, data)
-        typeRef = typeRef.transform(transformer, data)
         label = label?.transform(transformer, data)
         transformTypeParameters(transformer, data)
+        typeRef = typeRef.transform(transformer, data)
         return this
     }
 
@@ -133,10 +133,6 @@ internal class FirAnonymousFunctionImpl(
         body = newBody
     }
 
-    override fun replaceTypeRef(newTypeRef: FirTypeRef) {
-        typeRef = newTypeRef
-    }
-
     override fun replaceInvocationKind(newInvocationKind: EventOccurrencesRange?) {
         invocationKind = newInvocationKind
     }
@@ -144,4 +140,8 @@ internal class FirAnonymousFunctionImpl(
     override fun replaceInlineStatus(newInlineStatus: InlineStatus) {
         inlineStatus = newInlineStatus
     }
+
+    override fun replaceTypeRef(newTypeRef: FirTypeRef) {
+        typeRef = newTypeRef
+    }
 }

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/expressions/FirAnonymousFunctionExpression.kt

@@ -0,0 +1,36 @@
+/*
+ * Copyright 2010-2021 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.fir.expressions
+
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirSourceElement
+import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
+import org.jetbrains.kotlin.fir.types.FirTypeRef
+import org.jetbrains.kotlin.fir.visitors.*
+
+/*
+ * This file was generated automatically
+ * DO NOT MODIFY IT MANUALLY
+ */
+
+abstract class FirAnonymousFunctionExpression : FirExpression() {
+    abstract override val source: FirSourceElement?
+    abstract override val typeRef: FirTypeRef
+    abstract override val annotations: List<FirAnnotationCall>
+    abstract val anonymousFunction: FirAnonymousFunction
+
+    override fun <R, D> accept(visitor: FirVisitor<R, D>, data: D): R = visitor.visitAnonymousFunctionExpression(this, data)
+
+    @Suppress("UNCHECKED_CAST")
+    override fun <E: FirElement, D> transform(transformer: FirTransformer<D>, data: D): E = 
+        transformer.transformAnonymousFunctionExpression(this, data) as E
+
+    abstract override fun replaceTypeRef(newTypeRef: FirTypeRef)
+
+    abstract override fun <D> transformAnnotations(transformer: FirTransformer<D>, data: D): FirAnonymousFunctionExpression
+
+    abstract fun <D> transformAnonymousFunction(transformer: FirTransformer<D>, data: D): FirAnonymousFunctionExpression
+}

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/visitors/FirTransformer.kt

@@ -45,6 +45,7 @@ import org.jetbrains.kotlin.fir.declarations.FirPropertyAccessor
 import org.jetbrains.kotlin.fir.declarations.FirConstructor
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
+import org.jetbrains.kotlin.fir.expressions.FirAnonymousFunctionExpression
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousObject
 import org.jetbrains.kotlin.fir.expressions.FirAnonymousObjectExpression
 import org.jetbrains.kotlin.fir.diagnostics.FirDiagnosticHolder
@@ -299,6 +300,10 @@ abstract class FirTransformer<in D> : FirVisitor<FirElement, D>() {
         return transformElement(anonymousFunction, data)
     }
 
+    open fun transformAnonymousFunctionExpression(anonymousFunctionExpression: FirAnonymousFunctionExpression, data: D): FirStatement {
+        return transformElement(anonymousFunctionExpression, data)
+    }
+
     open fun transformAnonymousObject(anonymousObject: FirAnonymousObject, data: D): FirStatement {
         return transformElement(anonymousObject, data)
     }
@@ -811,6 +816,10 @@ abstract class FirTransformer<in D> : FirVisitor<FirElement, D>() {
         return transformAnonymousFunction(anonymousFunction, data)
     }
 
+    final override fun visitAnonymousFunctionExpression(anonymousFunctionExpression: FirAnonymousFunctionExpression, data: D): FirStatement {
+        return transformAnonymousFunctionExpression(anonymousFunctionExpression, data)
+    }
+
     final override fun visitAnonymousObject(anonymousObject: FirAnonymousObject, data: D): FirStatement {
         return transformAnonymousObject(anonymousObject, data)
     }

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/visitors/FirVisitor.kt

@@ -45,6 +45,7 @@ import org.jetbrains.kotlin.fir.declarations.FirPropertyAccessor
 import org.jetbrains.kotlin.fir.declarations.FirConstructor
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
+import org.jetbrains.kotlin.fir.expressions.FirAnonymousFunctionExpression
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousObject
 import org.jetbrains.kotlin.fir.expressions.FirAnonymousObjectExpression
 import org.jetbrains.kotlin.fir.diagnostics.FirDiagnosticHolder
@@ -220,6 +221,8 @@ abstract class FirVisitor<out R, in D> {
 
     open fun visitAnonymousFunction(anonymousFunction: FirAnonymousFunction, data: D): R  = visitElement(anonymousFunction, data)
 
+    open fun visitAnonymousFunctionExpression(anonymousFunctionExpression: FirAnonymousFunctionExpression, data: D): R  = visitElement(anonymousFunctionExpression, data)
+
     open fun visitAnonymousObject(anonymousObject: FirAnonymousObject, data: D): R  = visitElement(anonymousObject, data)
 
     open fun visitAnonymousObjectExpression(anonymousObjectExpression: FirAnonymousObjectExpression, data: D): R  = visitElement(anonymousObjectExpression, data)

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/visitors/FirVisitorVoid.kt

@@ -45,6 +45,7 @@ import org.jetbrains.kotlin.fir.declarations.FirPropertyAccessor
 import org.jetbrains.kotlin.fir.declarations.FirConstructor
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
+import org.jetbrains.kotlin.fir.expressions.FirAnonymousFunctionExpression
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousObject
 import org.jetbrains.kotlin.fir.expressions.FirAnonymousObjectExpression
 import org.jetbrains.kotlin.fir.diagnostics.FirDiagnosticHolder
@@ -298,6 +299,10 @@ abstract class FirVisitorVoid : FirVisitor<Unit, Nothing?>() {
         visitElement(anonymousFunction)
     }
 
+    open fun visitAnonymousFunctionExpression(anonymousFunctionExpression: FirAnonymousFunctionExpression) {
+        visitElement(anonymousFunctionExpression)
+    }
+
     open fun visitAnonymousObject(anonymousObject: FirAnonymousObject) {
         visitElement(anonymousObject)
     }
@@ -810,6 +815,10 @@ abstract class FirVisitorVoid : FirVisitor<Unit, Nothing?>() {
         visitAnonymousFunction(anonymousFunction)
     }
 
+    final override fun visitAnonymousFunctionExpression(anonymousFunctionExpression: FirAnonymousFunctionExpression, data: Nothing?) {
+        visitAnonymousFunctionExpression(anonymousFunctionExpression)
+    }
+
     final override fun visitAnonymousObject(anonymousObject: FirAnonymousObject, data: Nothing?) {
         visitAnonymousObject(anonymousObject)
     }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirRenderer.kt

@@ -580,6 +580,10 @@ class FirRenderer(builder: StringBuilder, private val mode: RenderMode = RenderM
         propertyAccessor.body?.renderBody()
     }
 
+    override fun visitAnonymousFunctionExpression(anonymousFunctionExpression: FirAnonymousFunctionExpression) {
+        visitAnonymousFunction(anonymousFunctionExpression.anonymousFunction)
+    }
+
     override fun visitAnonymousFunction(anonymousFunction: FirAnonymousFunction) {
         anonymousFunction.renderDeclarationData()
         anonymousFunction.annotations.renderAnnotations()

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/builder/FirAnonymousFunctionExpressionBuilder.kt

@@ -0,0 +1,52 @@
+/*
+ * Copyright 2010-2021 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.fir.expressions.builder
+
+import kotlin.contracts.*
+import org.jetbrains.kotlin.fir.FirSourceElement
+import org.jetbrains.kotlin.fir.builder.FirAnnotationContainerBuilder
+import org.jetbrains.kotlin.fir.builder.FirBuilderDsl
+import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
+import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
+import org.jetbrains.kotlin.fir.expressions.FirAnonymousFunctionExpression
+import org.jetbrains.kotlin.fir.expressions.builder.FirExpressionBuilder
+import org.jetbrains.kotlin.fir.expressions.impl.FirAnonymousFunctionExpressionImpl
+import org.jetbrains.kotlin.fir.types.FirTypeRef
+import org.jetbrains.kotlin.fir.types.impl.FirImplicitTypeRefImpl
+import org.jetbrains.kotlin.fir.visitors.*
+
+@FirBuilderDsl
+class FirAnonymousFunctionExpressionBuilder : FirAnnotationContainerBuilder, FirExpressionBuilder {
+    override var source: FirSourceElement? = null
+    lateinit var anonymousFunction: FirAnonymousFunction
+
+    override fun build(): FirAnonymousFunctionExpression {
+        return FirAnonymousFunctionExpressionImpl(
+            source,
+            anonymousFunction
+        )
+    }
+
+
+    @Deprecated("Modification of 'typeRef' has no impact for FirAnonymousFunctionExpressionBuilder", level = DeprecationLevel.HIDDEN)
+    override var typeRef: FirTypeRef
+        get() = throw IllegalStateException()
+        set(_) {
+            throw IllegalStateException()
+        }
+
+    @Deprecated("Modification of 'annotations' has no impact for FirAnonymousFunctionExpressionBuilder", level = DeprecationLevel.HIDDEN)
+    override val annotations: MutableList<FirAnnotationCall>
+        get() = throw IllegalStateException()
+}
+
+@OptIn(ExperimentalContracts::class)
+inline fun buildAnonymousFunctionExpression(init: FirAnonymousFunctionExpressionBuilder.() -> Unit = {}): FirAnonymousFunctionExpression {
+    contract {
+        callsInPlace(init, InvocationKind.EXACTLY_ONCE)
+    }
+    return FirAnonymousFunctionExpressionBuilder().apply(init).build()
+}

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/impl/FirAnonymousFunctionExpressionImpl.kt

@@ -0,0 +1,46 @@
+/*
+ * Copyright 2010-2021 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.fir.expressions.impl
+
+import org.jetbrains.kotlin.fir.FirSourceElement
+import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
+import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
+import org.jetbrains.kotlin.fir.expressions.FirAnonymousFunctionExpression
+import org.jetbrains.kotlin.fir.types.FirTypeRef
+import org.jetbrains.kotlin.fir.types.impl.FirImplicitTypeRefImpl
+import org.jetbrains.kotlin.fir.visitors.*
+
+internal class FirAnonymousFunctionExpressionImpl(
+    override val source: FirSourceElement?,
+    override var anonymousFunction: FirAnonymousFunction
+) : FirAnonymousFunctionExpression() {
+    override val typeRef: FirTypeRef
+        get() = anonymousFunction.typeRef
+
+    override val annotations: List<FirAnnotationCall>
+        get() = anonymousFunction.annotations
+
+    override fun <R, D> acceptChildren(visitor: FirVisitor<R, D>, data: D) {
+        anonymousFunction.accept(visitor, data)
+    }
+
+    override fun <D> transformChildren(transformer: FirTransformer<D>, data: D): FirAnonymousFunctionExpressionImpl {
+        return transformAnonymousFunction(transformer, data)
+    }
+
+    override fun <D> transformAnnotations(transformer: FirTransformer<D>, data: D): FirAnonymousFunctionExpressionImpl {
+        return this
+    }
+
+    override fun <D> transformAnonymousFunction(transformer: FirTransformer<D>, data: D): FirAnonymousFunctionExpressionImpl {
+        anonymousFunction = anonymousFunction.transform(transformer, data)
+        return this
+    }
+
+    override fun replaceTypeRef(newTypeRef: FirTypeRef) {
+        anonymousFunction.replaceTypeRef(newTypeRef)
+    }
+}

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/impl/FirArgumentListImpl.kt

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.fir.expressions.impl
 import org.jetbrains.kotlin.fir.FirPureAbstractElement
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
+import org.jetbrains.kotlin.fir.expressions.FirAnonymousFunctionExpression
 import org.jetbrains.kotlin.fir.expressions.FirArgumentList
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
@@ -28,12 +29,12 @@ internal class FirArgumentListImpl(
 
     override fun <D> transformArguments(transformer: FirTransformer<D>, data: D): FirArgumentListImpl {
         // Transform all normal arguments first and then lambda to make CFG correct. See KT-46825
-        val postponedFunctionArgs = mutableListOf<Pair<Int, FirAnonymousFunction>>()
+        val postponedFunctionArgs = mutableListOf<Pair<Int, FirAnonymousFunctionExpression>>()
         val iterator = arguments.listIterator()
         while (iterator.hasNext()) {
             val index = iterator.nextIndex()
             val next = iterator.next() as FirPureAbstractElement
-            if (next is FirAnonymousFunction) {
+            if (next is FirAnonymousFunctionExpression) {
                 postponedFunctionArgs += (index to next)
                 continue
             }

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/FirTreeBuilder.kt

@@ -58,7 +58,8 @@ object FirTreeBuilder : AbstractFirTreeBuilder() {
     val constructor = element("Constructor", Declaration, function, callableMemberDeclaration, typeParameterRefsOwner)
     val file = element("File", Declaration, annotatedDeclaration)
 
-    val anonymousFunction = element("AnonymousFunction", Declaration, function, expression, typeParametersOwner)
+    val anonymousFunction = element("AnonymousFunction", Declaration, function, typeParametersOwner)
+    val anonymousFunctionExpression = element("AnonymousFunctionExpression", Expression, expression)
 
     val anonymousObject = element("AnonymousObject", Declaration, klass, controlFlowGraphOwner)
     val anonymousObjectExpression = element("AnonymousObjectExpression", Expression, expression)

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/ImplementationConfigurator.kt

@@ -320,6 +320,8 @@ object ImplementationConfigurator : AbstractFirTreeImplementationConfigurator()
             default("resolvePhase", "FirResolvePhase.DECLARATIONS")
         }
 
+        noImpl(anonymousFunctionExpression)
+
         impl(propertyAccessor) {
             default("receiverTypeRef") {
                 value = "null"

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/NodeConfigurator.kt

@@ -294,6 +294,11 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
             }
             +booleanField("isLambda")
             +typeParameters
+            +field(typeRef, withReplace = true)
+        }
+
+        anonymousFunctionExpression.configure {
+            +field(anonymousFunction).withTransform()
         }
 
         typeParameter.configure {

compiler/visualizer/render-fir/src/org/jetbrains/kotlin/compiler/visualizer/FirVisualizer.kt

@@ -264,7 +264,7 @@ class FirVisualizer(private val firFile: FirFile) : BaseRenderer() {
             expression.children.filter { it.node.elementType != KtNodeTypes.REFERENCE_EXPRESSION }.forEach { psi ->
                 when (psi) {
                     is KtLambdaArgument -> {
-                        val firLambda = psi.firstOfType<FirLambdaArgumentExpression>()?.expression as? FirAnonymousFunction
+                        val firLambda = (psi.firstOfType<FirLambdaArgumentExpression>()?.expression as? FirAnonymousFunctionExpression)?.anonymousFunction
                         firLambda?.receiverTypeRef?.let {
                             lastCallWithLambda = psi.getLambdaExpression()?.firstOfType<FirLabel>()?.name
                             implicitReceivers += it.coneType