diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/FirCallCompleter.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/FirCallCompleter.kt
index bc11d4f2fcc..61536c281ba 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/FirCallCompleter.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/FirCallCompleter.kt
@@ -122,7 +122,7 @@ class FirCallCompleter(
             replacements[lambdaArgument] =
                 newLambdaExpression.transformSingle(transformer, FirBodyResolveTransformer.LambdaResolution(expectedReturnTypeRef))
 
-            return listOfNotNull(newLambdaExpression.body?.statements?.lastOrNull() as? FirExpression) to InferenceSession.default
+            return (newLambdaExpression.body?.returnExpressions() ?: emptyList()) to InferenceSession.default
         }
     }
 }
diff --git a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/Utils.kt b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/Utils.kt
index b7e62020ba2..2eece7ad17d 100644
--- a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/Utils.kt
+++ b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/Utils.kt
@@ -6,6 +6,8 @@
 package org.jetbrains.kotlin.fir
 
 import org.jetbrains.kotlin.analyzer.ModuleInfo
+import org.jetbrains.kotlin.fir.expressions.FirBlock
+import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 
 fun <T : FirElement, D> MutableList<T>.transformInplace(transformer: FirTransformer<D>, data: D) {
@@ -34,4 +36,7 @@ fun <T : FirElement, D> T.transformSingle(transformer: FirTransformer<D>, data:
     return (this as FirPureAbstractElement).transform<T, D>(transformer, data).single
 }
 
-fun ModuleInfo.dependenciesWithoutSelf(): Sequence<ModuleInfo> = dependencies().asSequence().filter { it != this }
\ No newline at end of file
+fun ModuleInfo.dependenciesWithoutSelf(): Sequence<ModuleInfo> = dependencies().asSequence().filter { it != this }
+
+// TODO: rewrite
+fun FirBlock.returnExpressions(): List<FirExpression> = listOfNotNull(statements.lastOrNull() as? FirExpression)
\ No newline at end of file