FIR (WIP): Initial support for proper lambda analysis

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

@@ -6,12 +6,19 @@
 package org.jetbrains.kotlin.fir
 
 import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
+import org.jetbrains.kotlin.fir.declarations.FirValueParameter
+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.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
 import org.jetbrains.kotlin.fir.expressions.impl.FirFunctionCallImpl
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.FirTypeProjection
 import org.jetbrains.kotlin.fir.types.FirTypeRef
+import org.jetbrains.kotlin.fir.types.impl.FirResolvedTypeRefImpl
 
 fun FirFunctionCall.copy(
     annotations: List<FirAnnotationCall> = this.annotations,
@@ -36,3 +43,31 @@ fun FirFunctionCall.copy(
     }
 }
 
+fun FirAnonymousFunction.copy(
+    receiverTypeRef: FirTypeRef? = this.receiverTypeRef,
+    psi: PsiElement? = this.psi,
+    session: FirSession = this.session,
+    returnTypeRef: FirTypeRef = this.returnTypeRef,
+    valueParameters: List<FirValueParameter> = this.valueParameters,
+    body: FirBlock? = this.body,
+    annotations: List<FirAnnotationCall> = this.annotations,
+    typeRef: FirTypeRef = this.typeRef,
+    label: FirLabel? = this.label
+): FirAnonymousFunction {
+    return FirAnonymousFunctionImpl(session, psi, returnTypeRef, receiverTypeRef).apply {
+        this.valueParameters.addAll(valueParameters)
+        this.body = body
+        this.annotations.addAll(annotations)
+        this.typeRef = typeRef
+        this.label = label
+    }
+}
+
+
+fun FirTypeRef.resolvedTypeFromPrototype(
+    type: ConeKotlinType
+): FirResolvedTypeRef {
+    return FirResolvedTypeRefImpl(
+        session, psi, type, false, annotations
+    )
+}

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

@@ -7,15 +7,13 @@ package org.jetbrains.kotlin.fir.resolve.calls
 
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
 import org.jetbrains.kotlin.fir.declarations.FirValueParameter
-import org.jetbrains.kotlin.fir.expressions.FirCallableReferenceAccess
+import org.jetbrains.kotlin.fir.expressions.*
-import org.jetbrains.kotlin.fir.expressions.FirExpression
-import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
-import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression
 import org.jetbrains.kotlin.fir.resolve.withNullability
 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.resolve.calls.inference.model.SimpleConstraintSystemConstraintPosition
+import org.jetbrains.kotlin.resolve.calls.model.PostponedResolvedAtomMarker
 
 
 fun resolveArgumentExpression(
@@ -31,16 +29,42 @@ fun resolveArgumentExpression(
     expectedType: ConeKotlinType,
     sink: CheckerSink,
     isReceiver: Boolean,
+    acceptLambdaAtoms: (PostponedResolvedAtomMarker) -> Unit,
     typeProvider: (FirExpression) -> FirTypeRef?
 ) {
     return when (argument) {
-        is FirQualifiedAccessExpression, is FirFunctionCall -> checkPlainExpressionArgument(csBuilder, argument, expectedType, sink, isReceiver, typeProvider)
+        is FirQualifiedAccessExpression, is FirFunctionCall -> checkPlainExpressionArgument(
+            csBuilder,
+            argument,
+            expectedType,
+            sink,
+            isReceiver,
+            typeProvider
+        )
         // TODO:!
-        is FirAnonymousFunction -> Unit
+        is FirAnonymousFunction -> preprocessLambdaArgument(csBuilder, argument, expectedType, acceptLambdaAtoms)
         // TODO:!
         is FirCallableReferenceAccess -> Unit
         // TODO:!
         //TODO: Collection literal
+        is FirLambdaArgumentExpression -> resolveArgumentExpression(
+            csBuilder,
+            argument.expression,
+            expectedType,
+            sink,
+            isReceiver,
+            acceptLambdaAtoms,
+            typeProvider
+        )
+        is FirNamedArgumentExpression -> resolveArgumentExpression(
+            csBuilder,
+            argument.expression,
+            expectedType,
+            sink,
+            isReceiver,
+            acceptLambdaAtoms,
+            typeProvider
+        )
         else -> checkPlainExpressionArgument(csBuilder, argument, expectedType, sink, isReceiver, typeProvider)
     }
 }
@@ -79,7 +103,15 @@ internal fun Candidate.resolveArgument(
 ) {
 
     val expectedType = prepareExpectedType(argument, parameter)
-    resolveArgumentExpression(this.system.getBuilder(), argument, expectedType, sink, isReceiver, typeProvider)
+    resolveArgumentExpression(
+        this.system.getBuilder(),
+        argument,
+        expectedType,
+        sink,
+        isReceiver,
+        { this.postponedAtoms += it },
+        typeProvider
+    )
 }
 
 private fun Candidate.prepareExpectedType(argument: FirExpression, parameter: FirValueParameter): ConeKotlinType {

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

@@ -5,6 +5,8 @@
 
 package org.jetbrains.kotlin.fir.resolve.calls
 
+import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.resolve.calls.inference.components.KotlinConstraintSystemCompleter
@@ -12,9 +14,12 @@ import org.jetbrains.kotlin.resolve.calls.inference.components.TypeVariableDirec
 import org.jetbrains.kotlin.resolve.calls.inference.components.VariableFixationFinder
 import org.jetbrains.kotlin.resolve.calls.inference.model.VariableWithConstraints
 import org.jetbrains.kotlin.resolve.calls.model.PostponedResolvedAtom
+import org.jetbrains.kotlin.resolve.calls.model.PostponedResolvedAtomMarker
 import org.jetbrains.kotlin.types.model.KotlinTypeMarker
 import org.jetbrains.kotlin.types.model.isIntegerLiteralTypeConstructor
 import org.jetbrains.kotlin.types.model.typeConstructor
+import org.jetbrains.kotlin.utils.addIfNotNull
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 
 fun Candidate.computeCompletionMode(
@@ -68,19 +73,21 @@ class ConstraintSystemCompleter(val components: InferenceComponents) {
     fun complete(
         c: KotlinConstraintSystemCompleter.Context,
         completionMode: KotlinConstraintSystemCompleter.ConstraintSystemCompletionMode,
-        candidateReturnType: ConeKotlinType
+        topLevelAtoms: List<FirExpression>,
+        candidateReturnType: ConeKotlinType,
+        analyze: (PostponedResolvedAtomMarker) -> Unit
     ) {
 
         while (true) {
-//            if (analyzePostponeArgumentIfPossible(c, topLevelAtoms, analyze)) continue
+            if (analyzePostponeArgumentIfPossible(c, topLevelAtoms, analyze)) continue
 
 
 //            val allTypeVariables = getOrderedAllTypeVariables(c, collectVariablesFromContext, topLevelAtoms)
             val allTypeVariables = c.notFixedTypeVariables.keys.toList()
-//            val postponedKtPrimitives = getOrderedNotAnalyzedPostponedArguments(topLevelAtoms)
+            val postponedKtPrimitives = getOrderedNotAnalyzedPostponedArguments(topLevelAtoms)
             val variableForFixation =
                 variableFixationFinder.findFirstVariableForFixation(
-                    c, allTypeVariables, emptyList(), completionMode, candidateReturnType
+                    c, allTypeVariables, postponedKtPrimitives, completionMode, candidateReturnType
                 ) ?: break
 
 //            if (shouldForceCallableReferenceOrLambdaResolution(completionMode, variableForFixation)) {
@@ -104,7 +111,7 @@ class ConstraintSystemCompleter(val components: InferenceComponents) {
 
         if (completionMode == KotlinConstraintSystemCompleter.ConstraintSystemCompletionMode.FULL) {
             // force resolution for all not-analyzed argument's
-//            getOrderedNotAnalyzedPostponedArguments(topLevelAtoms).forEach(analyze)
+            getOrderedNotAnalyzedPostponedArguments(topLevelAtoms).forEach(analyze)
 //
 //            if (c.notFixedTypeVariables.isNotEmpty() && c.postponedTypeVariables.isEmpty()) {
 //                runCompletion(c, completionMode, topLevelAtoms, topLevelType, analyze)
@@ -131,4 +138,49 @@ class ConstraintSystemCompleter(val components: InferenceComponents) {
         c.fixVariable(variableWithConstraints.typeVariable, resultType)
     }
 
+    private fun analyzePostponeArgumentIfPossible(
+        c: KotlinConstraintSystemCompleter.Context,
+        topLevelAtoms: List<FirExpression>,
+        analyze: (PostponedResolvedAtomMarker) -> Unit
+    ): Boolean {
+        for (argument in getOrderedNotAnalyzedPostponedArguments(topLevelAtoms)) {
+            if (canWeAnalyzeIt(c, argument)) {
+                analyze(argument)
+                return true
+            }
+        }
+        return false
+    }
+
+    private fun getOrderedNotAnalyzedPostponedArguments(topLevelAtoms: List<FirExpression>): List<PostponedResolvedAtomMarker> {
+        fun FirExpression.process(to: MutableList<PostponedResolvedAtomMarker>) {
+            when (this) {
+                is FirFunctionCall -> {
+                    val candidate = (this.calleeReference as? FirNamedReferenceWithCandidate)?.candidate
+                    candidate?.postponedAtoms?.forEach {
+                        to.addIfNotNull(it.safeAs<PostponedResolvedAtomMarker>()?.takeUnless { it.analyzed })
+                    }
+                    this.arguments.forEach { it.process(to) }
+                }
+                // TOOD: WTF?
+            }
+//            if (analyzed) {
+//                subResolvedAtoms.forEach { it.process(to) }
+//            }
+        }
+
+        val notAnalyzedArguments = arrayListOf<PostponedResolvedAtomMarker>()
+        for (primitive in topLevelAtoms) {
+            primitive.process(notAnalyzedArguments)
+        }
+
+        return notAnalyzedArguments
+    }
+
+    private fun canWeAnalyzeIt(c: KotlinConstraintSystemCompleter.Context, argument: PostponedResolvedAtomMarker): Boolean {
+        if (argument.analyzed) return false
+
+        return argument.inputTypes.all { c.containsOnlyFixedOrPostponedVariables(it) }
+    }
+
 }

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

@@ -0,0 +1,167 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. 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.resolve.calls
+
+import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
+import org.jetbrains.kotlin.fir.types.*
+import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemBuilder
+import org.jetbrains.kotlin.resolve.calls.model.PostponedResolvedAtomMarker
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
+
+fun preprocessLambdaArgument(
+    csBuilder: ConstraintSystemBuilder,
+    argument: FirAnonymousFunction,
+    expectedType: ConeKotlinType?,
+    acceptLambdaAtoms: (PostponedResolvedAtomMarker) -> Unit,
+    forceResolution: Boolean = false
+) {
+    if (expectedType != null && !forceResolution && csBuilder.isTypeVariable(expectedType)) {
+        //return LambdaWithTypeVariableAsExpectedTypeAtom(argument, expectedType)
+    }
+
+    val resolvedArgument = extractLambdaInfoFromFunctionalType(expectedType, argument) ?: extraLambdaInfo(expectedType, argument, csBuilder)
+
+    if (expectedType != null) {
+//        val lambdaType = createFunctionType(
+//            csBuilder.builtIns, Annotations.EMPTY, resolvedArgument.receiver,
+//            resolvedArgument.parameters, null, resolvedArgument.returnType, resolvedArgument.isSuspend
+//        )
+//        csBuilder.addSubtypeConstraint(lambdaType, expectedType, ArgumentConstraintPosition(argument))
+    }
+
+    acceptLambdaAtoms(resolvedArgument)
+}
+
+
+private val ConeKotlinType.isBuiltinFunctionalType: Boolean
+    get () {
+        val type = this
+        return when (type) {
+            is ConeClassType -> type.lookupTag.classId.asString().startsWith("kotlin/Function")
+            else -> false
+        }
+    }
+
+
+private val ConeKotlinType.isSuspendFunctionType: Boolean
+    get() {
+        val type = this
+        return when (type) {
+            is ConeClassType -> {
+                val classId = type.lookupTag.classId
+                classId.packageFqName.asString() == "kotlin" && classId.relativeClassName.asString().startsWith("SuspendFunction")
+            }
+            else -> false
+        }
+    }
+
+private val ConeKotlinType.receiverType: ConeKotlinType? get() = null
+private val ConeKotlinType.returnType: ConeKotlinType
+    get() {
+        require(this is ConeClassType)
+        val projection = typeArguments.last()
+        require(projection is ConeTypedProjection)
+        return projection.type
+    }
+
+private val ConeKotlinType.valueParameterTypes: List<ConeKotlinType?>
+    get() {
+        require(this is ConeClassType)
+        return typeArguments.dropLast(1).map {
+            (it as? ConeTypedProjection)?.type
+        }
+    }
+
+private val FirAnonymousFunction.returnType get() = returnTypeRef.coneTypeSafe<ConeKotlinType>()
+private val FirAnonymousFunction.receiverType get() = receiverTypeRef?.coneTypeSafe<ConeKotlinType>()
+
+
+private fun extraLambdaInfo(
+    expectedType: ConeKotlinType?,
+    argument: FirAnonymousFunction,
+    csBuilder: ConstraintSystemBuilder
+): ResolvedLambdaAtom {
+//    val builtIns = csBuilder.builtIns
+    val isSuspend = expectedType?.isSuspendFunctionType ?: false
+
+    val isFunctionSupertype =
+        expectedType != null && expectedType.isBuiltinFunctionalType//isNotNullOrNullableFunctionSupertype(expectedType)
+    val argumentAsFunctionExpression = argument//.safeAs<FunctionExpression>()
+
+    val typeVariable = TypeVariableForLambdaReturnType(argument, "_L")
+
+    val receiverType = null//argumentAsFunctionExpression?.receiverType
+    val returnType =
+        argumentAsFunctionExpression?.returnType
+            ?: expectedType?.typeArguments?.singleOrNull()?.safeAs<ConeTypedProjection>()?.type?.takeIf { isFunctionSupertype }
+            ?: typeVariable.defaultType
+
+    val parameters = argument.valueParameters?.map { it.returnTypeRef.coneTypeSafe<ConeKotlinType>() ?: TODO("!nothing") } ?: emptyList()
+
+    val newTypeVariableUsed = returnType == typeVariable.defaultType
+    if (newTypeVariableUsed) csBuilder.registerVariable(typeVariable)
+
+    return ResolvedLambdaAtom(argument, isSuspend, receiverType, parameters, returnType, typeVariable.takeIf { newTypeVariableUsed })
+}
+
+private fun extractLambdaInfoFromFunctionalType(expectedType: ConeKotlinType?, argument: FirAnonymousFunction): ResolvedLambdaAtom? {
+    if (expectedType == null || !expectedType.isBuiltinFunctionalType) return null
+    val parameters = extractLambdaParameters(expectedType, argument)
+
+    val argumentAsFunctionExpression = argument//.safeAs<FunctionExpression>()
+    val receiverType = argumentAsFunctionExpression?.receiverType ?: expectedType.receiverType
+    val returnType = argumentAsFunctionExpression?.returnType ?: expectedType.returnType
+
+    return ResolvedLambdaAtom(
+        argument,
+        expectedType.isSuspendFunctionType,
+        receiverType,
+        parameters,
+        returnType,
+        typeVariableForLambdaReturnType = null
+    )
+}
+
+private fun extractLambdaParameters(expectedType: ConeKotlinType, argument: FirAnonymousFunction): List<ConeKotlinType> {
+    val parameters = argument.valueParameters
+    val expectedParameters = expectedType.valueParameterTypes
+    if (parameters.isEmpty()) {
+        return expectedParameters.map { it?.type ?: TODO("any?") }
+    }
+
+    return parameters.mapIndexed { index, parameter ->
+        parameter.returnTypeRef.coneTypeSafe() ?: expectedParameters.getOrNull(index) ?: TODO("any?")//expectedType.builtIns.nullableAnyType
+    }
+}
+
+class TypeVariableForLambdaReturnType(val argument: FirAnonymousFunction, name: String) : ConeTypeVariable(name)
+
+
+class ResolvedLambdaAtom(
+    val atom: FirAnonymousFunction,
+    val isSuspend: Boolean,
+    val receiver: ConeKotlinType?,
+    val parameters: List<ConeKotlinType>,
+    val returnType: ConeKotlinType,
+    val typeVariableForLambdaReturnType: TypeVariableForLambdaReturnType?
+) : PostponedResolvedAtomMarker {
+
+    override var analyzed: Boolean = false
+
+//    lateinit var resultArguments: List<KotlinCallArgument>
+//        private set
+
+//    fun setAnalyzedResults(
+//        resultArguments: List<KotlinCallArgument>,
+//        subResolvedAtoms: List<ResolvedAtom>
+//    ) {
+//        this.resultArguments = resultArguments
+//        setAnalyzedResults(subResolvedAtoms)
+//    }
+
+    override val inputTypes: Collection<ConeKotlinType> get() = receiver?.let { parameters + it } ?: parameters
+    override val outputType: ConeKotlinType get() = returnType
+}

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

@@ -0,0 +1,153 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. 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.resolve.calls
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
+import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.resolve.constructType
+import org.jetbrains.kotlin.fir.service
+import org.jetbrains.kotlin.fir.symbols.StandardClassIds.Unit
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.FirTypeRef
+import org.jetbrains.kotlin.resolve.calls.components.InferenceSession
+import org.jetbrains.kotlin.resolve.calls.components.PostponedArgumentsAnalyzer
+import org.jetbrains.kotlin.resolve.calls.inference.model.SimpleConstraintSystemConstraintPosition
+import org.jetbrains.kotlin.resolve.calls.model.PostponedResolvedAtomMarker
+import org.jetbrains.kotlin.types.model.*
+import org.jetbrains.kotlin.fir.symbols.invoke
+
+interface LambdaAnalyzer {
+    fun analyzeAndGetLambdaReturnArguments(
+        lambdaArgument: FirAnonymousFunction,
+        isSuspend: Boolean,
+        receiverType: ConeKotlinType?,
+        parameters: List<ConeKotlinType>,
+        expectedReturnType: ConeKotlinType?, // null means, that return type is not proper i.e. it depends on some type variables
+        stubsForPostponedVariables: Map<TypeVariableMarker, StubTypeMarker>
+    ): Pair<List<FirExpression>, InferenceSession>
+}
+
+
+class PostponedArgumentsAnalyzer(
+    val lambdaAnalyzer: LambdaAnalyzer,
+    val typeProvider: (FirExpression) -> FirTypeRef?,
+    val session: FirSession
+) {
+
+    fun analyze(
+        c: PostponedArgumentsAnalyzer.Context,
+//        resolutionCallbacks: KotlinResolutionCallbacks,
+        argument: PostponedResolvedAtomMarker
+        //diagnosticsHolder: KotlinDiagnosticsHolder
+    ) {
+        when (argument) {
+            is ResolvedLambdaAtom ->
+                analyzeLambda(c, argument/*, diagnosticsHolder*/)
+
+//            is LambdaWithTypeVariableAsExpectedTypeAtom ->
+//                analyzeLambda(
+//                    c, resolutionCallbacks, argument.transformToResolvedLambda(c.getBuilder()), diagnosticsHolder
+//                )
+
+//            is ResolvedCallableReferenceAtom ->
+//                callableReferenceResolver.processCallableReferenceArgument(c.getBuilder(), argument, diagnosticsHolder)
+//
+//            is ResolvedCollectionLiteralAtom -> TODO("Not supported")
+
+            else -> error("Unexpected resolved primitive: ${argument.javaClass.canonicalName}")
+        }
+    }
+
+
+    private fun analyzeLambda(
+        c: PostponedArgumentsAnalyzer.Context,
+        lambda: ResolvedLambdaAtom//,
+        //diagnosticHolder: KotlinDiagnosticsHolder
+    ) {
+        val unitType = Unit(session.service()).constructType(emptyArray(), false)
+        val stubsForPostponedVariables = c.bindingStubsForPostponedVariables()
+        val currentSubstitutor = c.buildCurrentSubstitutor(stubsForPostponedVariables.mapKeys { it.key.freshTypeConstructor(c) })
+
+        fun substitute(type: ConeKotlinType) = currentSubstitutor.safeSubstitute(c, type) as ConeKotlinType
+
+        val receiver = lambda.receiver?.let(::substitute)
+        val parameters = lambda.parameters.map(::substitute)
+        val rawReturnType = lambda.returnType
+
+        val expectedTypeForReturnArguments = when {
+            c.canBeProper(rawReturnType) -> substitute(rawReturnType)
+
+            // For Unit-coercion
+            c.hasUpperOrEqualUnitConstraint(rawReturnType) -> unitType
+
+            else -> null
+        }
+
+        val (returnArguments, inferenceSession) = lambdaAnalyzer.analyzeAndGetLambdaReturnArguments(
+            lambda.atom,
+            lambda.isSuspend,
+            receiver,
+            parameters,
+            expectedTypeForReturnArguments,
+            stubsForPostponedVariables
+        )
+
+        returnArguments.forEach { c.addSubsystemFromExpression(it) }
+
+        val checkerSink: CheckerSink = CheckerSinkImpl()
+
+        val subResolvedKtPrimitives = returnArguments.map {
+            var atom: PostponedResolvedAtomMarker? = null
+            resolveArgumentExpression(
+                c.getBuilder(),
+                it,
+                lambda.returnType.let(::substitute),
+                checkerSink,
+                false,
+                { atom = it },
+                typeProvider
+            )
+//            resolveKtPrimitive(
+//                c.getBuilder(), it, lambda.returnType.let(::substitute), diagnosticHolder, isReceiver = false
+//            )
+        }
+
+        if (returnArguments.isEmpty()) {
+//            val unitType =
+            val lambdaReturnType = lambda.returnType.let(::substitute)
+            c.getBuilder().addSubtypeConstraint(
+                lambdaReturnType,
+                unitType, /*LambdaArgumentConstraintPosition(lambda)*/
+                SimpleConstraintSystemConstraintPosition
+            )
+            c.getBuilder().addSubtypeConstraint(
+                unitType,
+                lambdaReturnType, /*LambdaArgumentConstraintPosition(lambda)*/
+                SimpleConstraintSystemConstraintPosition
+            )
+        }
+
+        lambda.analyzed = true
+        //lambda.setAnalyzedResults(returnArguments, subResolvedKtPrimitives)
+
+//        if (inferenceSession != null) {
+//            val storageSnapshot = c.getBuilder().currentStorage()
+//
+//            val postponedVariables = inferenceSession.inferPostponedVariables(lambda, storageSnapshot)
+//
+//            for ((constructor, resultType) in postponedVariables) {
+//                val variableWithConstraints = storageSnapshot.notFixedTypeVariables[constructor] ?: continue
+//                val variable = variableWithConstraints.typeVariable
+//
+//                c.getBuilder().unmarkPostponedVariable(variable)
+//                c.getBuilder().addEqualityConstraint(variable.defaultType(c), resultType, CoroutinePosition())
+//            }
+//        }
+    }
+}
+
+

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

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.fir.resolve.transformers
 import com.google.common.collect.LinkedHashMultimap
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.declarations.impl.FirValueParameterImpl
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.references.FirErrorNamedReference
 import org.jetbrains.kotlin.fir.references.FirSimpleNamedReference
@@ -29,6 +30,7 @@ import org.jetbrains.kotlin.ir.expressions.IrConstKind
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.calls.NewCommonSuperTypeCalculator
+import org.jetbrains.kotlin.resolve.calls.components.InferenceSession
 import org.jetbrains.kotlin.resolve.calls.inference.buildAbstractResultingSubstitutor
 import org.jetbrains.kotlin.resolve.calls.inference.components.KotlinConstraintSystemCompleter
 import org.jetbrains.kotlin.resolve.calls.results.TypeSpecificityComparator
@@ -249,6 +251,40 @@ open class FirBodyResolveTransformer(val session: FirSession, val implicitTypeOn
         return transformCallee(variableAssignment).compose()
     }
 
+    override fun transformAnonymousFunction(anonymousFunction: FirAnonymousFunction, data: Any?): CompositeTransformResult<FirDeclaration> {
+        if (data == null) return anonymousFunction.compose()
+        if (data is LambdaResolution) return transformAnonymousFunction(anonymousFunction, data).compose()
+        return super.transformAnonymousFunction(anonymousFunction, data)
+    }
+
+    fun transformAnonymousFunction(anonymousFunction: FirAnonymousFunction, lambdaResolution: LambdaResolution): FirAnonymousFunction {
+        val receiverTypeRef = anonymousFunction.receiverTypeRef
+        fun transform(): FirAnonymousFunction {
+            return withScopeCleanup(scopes) {
+                scopes.addIfNotNull(receiverTypeRef?.coneTypeSafe()?.scope(session, ScopeSession()))
+                val result =
+                    super.transformAnonymousFunction(
+                        anonymousFunction,
+                        lambdaResolution.expectedReturnTypeRef ?: anonymousFunction.returnTypeRef
+                    ).single as FirAnonymousFunction
+                val body = result.body
+                if (result.returnTypeRef is FirImplicitTypeRef && body != null) {
+                    result.transformReturnTypeRef(this, body.resultType)
+                    result
+                } else {
+                    result
+                }
+            }
+        }
+
+        val label = anonymousFunction.label
+        return if (label != null && receiverTypeRef != null) {
+            withLabel(Name.identifier(label.name), receiverTypeRef.coneTypeUnsafe()) { transform() }
+        } else {
+            transform()
+        }
+    }
+
     private fun resolveCallAndSelectCandidate(functionCall: FirFunctionCall, expectedTypeRef: FirTypeRef?): FirFunctionCall {
 
         val functionCall = functionCall.transformChildren(this, null) as FirFunctionCall
@@ -314,6 +350,8 @@ open class FirBodyResolveTransformer(val session: FirSession, val implicitTypeOn
         return resultExpression
     }
 
+    data class LambdaResolution(val expectedReturnTypeRef: FirResolvedTypeRef?)
+
     private fun completeTypeInference(functionCall: FirFunctionCall, expectedTypeRef: FirTypeRef?): FirFunctionCall {
         val typeRef = typeFromCallee(functionCall)
         if (typeRef.type is ConeKotlinErrorType) {
@@ -327,7 +365,61 @@ open class FirBodyResolveTransformer(val session: FirSession, val implicitTypeOn
 
         val completionMode = candidate.computeCompletionMode(inferenceComponents, expectedTypeRef, initialType)
         val completer = ConstraintSystemCompleter(inferenceComponents)
-        completer.complete(candidate.system.asConstraintSystemCompleterContext(), completionMode, initialType)
+        val replacements = mutableMapOf<FirExpression, FirExpression>()
+
+        val analyzer = PostponedArgumentsAnalyzer(object : LambdaAnalyzer {
+            override fun analyzeAndGetLambdaReturnArguments(
+                lambdaArgument: FirAnonymousFunction,
+                isSuspend: Boolean,
+                receiverType: ConeKotlinType?,
+                parameters: List<ConeKotlinType>,
+                expectedReturnType: ConeKotlinType?,
+                stubsForPostponedVariables: Map<TypeVariableMarker, StubTypeMarker>
+            ): Pair<List<FirExpression>, InferenceSession> {
+
+                val itParam = when {
+                    lambdaArgument.valueParameters.isEmpty() && parameters.size == 1 ->
+                        FirValueParameterImpl(
+                            session,
+                            null,
+                            Name.identifier("it"),
+                            FirResolvedTypeRefImpl(session, null, parameters.single(), false, emptyList()),
+                            null,
+                            false,
+                            false,
+                            false
+                        )
+                    else -> null
+                }
+
+                val newLambdaExpression = lambdaArgument.copy(
+                    receiverTypeRef = receiverType?.let { lambdaArgument.receiverTypeRef!!.resolvedTypeFromPrototype(it) },
+                    valueParameters = lambdaArgument.valueParameters.mapIndexed { index, parameter ->
+                        parameter.transformReturnTypeRef(StoreType, parameter.returnTypeRef.resolvedTypeFromPrototype(parameters[index]))
+                        parameter
+                    } + listOfNotNull(itParam)
+                )
+
+
+                val expectedReturnTypeRef = expectedReturnType?.let { newLambdaExpression.returnTypeRef.resolvedTypeFromPrototype(it) }
+                replacements[lambdaArgument] =
+                    newLambdaExpression.transformSingle(this@FirBodyResolveTransformer, LambdaResolution(expectedReturnTypeRef))
+
+
+                return listOfNotNull(newLambdaExpression.body?.statements?.lastOrNull() as? FirExpression) to InferenceSession.default
+            }
+
+        }, { it.resultType }, session)
+
+        completer.complete(candidate.system.asConstraintSystemCompleterContext(), completionMode, listOf(functionCall), initialType) {
+            analyzer.analyze(
+                candidate.system.asPostponedArgumentsAnalyzerContext(),
+                it
+//                diagnosticsHolder
+            )
+        }
+
+        functionCall.transformChildren(ReplaceInArguments, replacements.toMap())
 
 
         if (completionMode == KotlinConstraintSystemCompleter.ConstraintSystemCompletionMode.FULL) {
@@ -712,4 +804,41 @@ private object StoreNameReference : FirTransformer<FirNamedReference>() {
     ): CompositeTransformResult<FirNamedReference> {
         return data.compose()
     }
+}
+
+private object StoreType : FirTransformer<FirResolvedTypeRef>() {
+    override fun <E : FirElement> transformElement(element: E, data: FirResolvedTypeRef): CompositeTransformResult<E> {
+        return element.compose()
+    }
+
+    override fun transformTypeRef(typeRef: FirTypeRef, data: FirResolvedTypeRef): CompositeTransformResult<FirTypeRef> {
+        return data.compose()
+    }
+}
+
+private object ReplaceInArguments : FirTransformer<Map<FirElement, FirElement>>() {
+    override fun <E : FirElement> transformElement(element: E, data: Map<FirElement, FirElement>): CompositeTransformResult<E> {
+        return ((data[element] ?: element) as E).compose()
+    }
+
+    override fun transformFunctionCall(
+        functionCall: FirFunctionCall,
+        data: Map<FirElement, FirElement>
+    ): CompositeTransformResult<FirStatement> {
+        return (functionCall.transformChildren(this, data) as FirStatement).compose()
+    }
+
+    override fun transformNamedArgumentExpression(
+        namedArgumentExpression: FirNamedArgumentExpression,
+        data: Map<FirElement, FirElement>
+    ): CompositeTransformResult<FirStatement> {
+        return (namedArgumentExpression.transformChildren(this, data) as FirStatement).compose()
+    }
+
+    override fun transformLambdaArgumentExpression(
+        lambdaArgumentExpression: FirLambdaArgumentExpression,
+        data: Map<FirElement, FirElement>
+    ): CompositeTransformResult<FirStatement> {
+        return (lambdaArgumentExpression.transformChildren(this, data) as FirStatement).compose()
+    }
 }

compiler/fir/resolve/testData/resolve/arguments/lambda.txt

@@ -6,122 +6,71 @@ FILE: lambda.kt
     public final fun baz(f: R|kotlin/Function0<kotlin/Unit>|, other: R|kotlin/Boolean| = Boolean(true)): R|kotlin/Unit| {
     }
     public final fun test(): R|kotlin/Unit| {
-        R|/foo|(<L> = foo@fun <implicit>.<anonymous>(): <implicit> {
+        R|/foo|(<L> = foo@fun <anonymous>(): R|kotlin/Unit| {
-            ^  {
+            Unit
-                Unit
-            }
-
         }
         )
-        R|/foo|(<L> = foo@fun <implicit>.<anonymous>(): <implicit> {
+        R|/foo|(<L> = foo@fun <anonymous>(): R|kotlin/Unit| {
-            ^  {
+            Unit
-                Unit
-            }
-
         }
         )
-        R|/foo|(foo@fun <implicit>.<anonymous>(): <implicit> {
+        R|/foo|(foo@fun <anonymous>(): R|kotlin/Unit| {
-            ^  {
+            Unit
-                Unit
-            }
-
         }
         )
         <Inapplicable(PARAMETER_MAPPING_ERROR): [/foo]>#(Int(1), <L> = foo@fun <implicit>.<anonymous>(): <implicit> {
-            ^  {
+            Unit
-                Unit
-            }
-
         }
         )
         <Inapplicable(PARAMETER_MAPPING_ERROR): [/foo]>#(f = foo@fun <implicit>.<anonymous>(): <implicit> {
-            ^  {
+            Unit
-                Unit
-            }
-
         }
         , <L> = foo@fun <implicit>.<anonymous>(): <implicit> {
-            ^  {
+            Unit
-                Unit
-            }
-
         }
         )
-        R|/bar|(Int(1), <L> = bar@fun <implicit>.<anonymous>(): <implicit> {
+        R|/bar|(Int(1), <L> = bar@fun <anonymous>(): R|kotlin/Unit| {
-            ^  {
+            Unit
-                Unit
-            }
-
         }
         )
-        R|/bar|(x = Int(1), <L> = bar@fun <implicit>.<anonymous>(): <implicit> {
+        R|/bar|(x = Int(1), <L> = bar@fun <anonymous>(): R|kotlin/Unit| {
-            ^  {
+            Unit
-                Unit
-            }
-
         }
         )
-        R|/bar|(Int(1), bar@fun <implicit>.<anonymous>(): <implicit> {
+        R|/bar|(Int(1), bar@fun <anonymous>(): R|kotlin/Unit| {
-            ^  {
+            Unit
-                Unit
-            }
-
         }
         )
-        R|/bar|(x = Int(1), f = bar@fun <implicit>.<anonymous>(): <implicit> {
+        R|/bar|(x = Int(1), f = bar@fun <anonymous>(): R|kotlin/Unit| {
-            ^  {
+            Unit
-                Unit
-            }
-
         }
         )
         <Inapplicable(PARAMETER_MAPPING_ERROR): [/bar]>#(<L> = bar@fun <implicit>.<anonymous>(): <implicit> {
-            ^  {
+            Unit
-                Unit
-            }
-
         }
         )
         <Inapplicable(PARAMETER_MAPPING_ERROR): [/bar]>#(bar@fun <implicit>.<anonymous>(): <implicit> {
-            ^  {
+            Unit
-                Unit
-            }
-
         }
         )
-        R|/baz|(other = Boolean(false), f = baz@fun <implicit>.<anonymous>(): <implicit> {
+        R|/baz|(other = Boolean(false), f = baz@fun <anonymous>(): R|kotlin/Unit| {
-            ^  {
+            Unit
-                Unit
-            }
-
         }
         )
-        R|/baz|(baz@fun <implicit>.<anonymous>(): <implicit> {
+        R|/baz|(baz@fun <anonymous>(): R|kotlin/Unit| {
-            ^  {
+            Unit
-                Unit
-            }
-
         }
         , Boolean(false))
         <Inapplicable(PARAMETER_MAPPING_ERROR): [/baz]>#(<L> = baz@fun <implicit>.<anonymous>(): <implicit> {
-            ^  {
+            Unit
-                Unit
-            }
-
         }
         )
         <Inapplicable(PARAMETER_MAPPING_ERROR): [/baz]>#(<L> = baz@fun <implicit>.<anonymous>(): <implicit> {
-            ^  {
+            Unit
-                Unit
-            }
-
         }
         )
         <Inapplicable(PARAMETER_MAPPING_ERROR): [/baz]>#(other = Boolean(false), <L> = baz@fun <implicit>.<anonymous>(): <implicit> {
-            ^  {
+            Unit
-                Unit
-            }
-
         }
         )
     }

compiler/fir/resolve/testData/resolve/expresssions/lambda.kt

@@ -0,0 +1,12 @@
+fun foo(block: () -> Unit) {}
+fun bar(block: () -> String) {}
+fun itIs(block: (String) -> String) {}
+fun multipleArgs(block: (String, String) -> String) {}
+
+fun main() {
+    foo { "This is test" }
+    bar { "This is also test" }
+    itIs { "this is $it test" }
+    multipleArgs { a, b -> "This is test of $a, $b" }
+}
+

compiler/fir/resolve/testData/resolve/expresssions/lambda.txt

@@ -0,0 +1,27 @@
+FILE: lambda.kt
+    public final fun foo(block: R|kotlin/Function0<kotlin/Unit>|): R|kotlin/Unit| {
+    }
+    public final fun bar(block: R|kotlin/Function0<kotlin/String>|): R|kotlin/Unit| {
+    }
+    public final fun itIs(block: R|kotlin/Function1<kotlin/String, kotlin/String>|): R|kotlin/Unit| {
+    }
+    public final fun multipleArgs(block: R|kotlin/Function2<kotlin/String, kotlin/String, kotlin/String>|): R|kotlin/Unit| {
+    }
+    public final fun main(): R|kotlin/Unit| {
+        R|/foo|(<L> = foo@fun <anonymous>(): R|kotlin/Unit| {
+            String(This is test)
+        }
+        )
+        R|/bar|(<L> = bar@fun <anonymous>(): R|kotlin/String| {
+            String(This is also test)
+        }
+        )
+        R|/itIs|(<L> = itIs@fun <anonymous>(it: R|kotlin/String|): R|kotlin/String| {
+            <strcat>(String(this is ), R|<local>/it|, String( test))
+        }
+        )
+        R|/multipleArgs|(<L> = multipleArgs@fun <anonymous>(a: R|kotlin/String|, b: R|kotlin/String|): R|kotlin/String| {
+            <strcat>(String(This is test of ), R|<local>/a|, String(, ), R|<local>/b|)
+        }
+        )
+    }

compiler/fir/resolve/testData/resolve/stdlib/functionX.txt

@@ -1,17 +1,11 @@
 FILE: functionX.kt
     public final val x: R|kotlin/jvm/functions/Function0<kotlin/Int>| = fun R|kotlin/jvm/functions/Function0<kotlin/Int>|.<anonymous>(): R|kotlin/jvm/functions/Function0<kotlin/Int>| {
-        ^  {
+        Int(42)
-            Int(42)
-        }
-
     }
 
         public get(): R|kotlin/jvm/functions/Function0<kotlin/Int>|
     public final val y: R|kotlin/Function1<kotlin/String, kotlin/String>| = fun R|kotlin/Function1<kotlin/String, kotlin/String>|.<anonymous>(): R|kotlin/Function1<kotlin/String, kotlin/String>| {
-        ^  {
+        <Unresolved name: it>#
-            <Unresolved name: it>#
-        }
-
     }
 
         public get(): R|kotlin/Function1<kotlin/String, kotlin/String>|

compiler/fir/resolve/testData/resolve/stdlib/mapList.kt

@@ -0,0 +1,5 @@
+fun main() {
+    val x = List(10) {
+        "number = $it"
+    }
+}

compiler/fir/resolve/testData/resolve/stdlib/mapList.txt

@@ -0,0 +1,7 @@
+FILE: mapList.kt
+    public final fun main(): R|kotlin/Unit| {
+        lval x: R|kotlin/collections/List<kotlin/String>| = R|kotlin/collections/List|<R|kotlin/String|>(Int(10), <L> = List@fun <anonymous>(it: R|kotlin/Int|): R|kotlin/String| {
+            <strcat>(String(number = ), R|<local>/it|)
+        }
+        )
+    }

compiler/fir/resolve/tests/org/jetbrains/kotlin/fir/FirResolveTestCaseGenerated.java

@@ -212,6 +212,11 @@ public class FirResolveTestCaseGenerated extends AbstractFirResolveTestCase {
             runTest("compiler/fir/resolve/testData/resolve/expresssions/dispatchReceiver.kt");
         }
 
+        @TestMetadata("lambda.kt")
+        public void testLambda() throws Exception {
+            runTest("compiler/fir/resolve/testData/resolve/expresssions/lambda.kt");
+        }
+
         @TestMetadata("localImplicitBodies.kt")
         public void testLocalImplicitBodies() throws Exception {
             runTest("compiler/fir/resolve/testData/resolve/expresssions/localImplicitBodies.kt");

compiler/fir/resolve/tests/org/jetbrains/kotlin/fir/FirResolveTestCaseWithStdlibGenerated.java

@@ -44,6 +44,11 @@ public class FirResolveTestCaseWithStdlibGenerated extends AbstractFirResolveTes
         runTest("compiler/fir/resolve/testData/resolve/stdlib/helloWorld.kt");
     }
 
+    @TestMetadata("mapList.kt")
+    public void testMapList() throws Exception {
+        runTest("compiler/fir/resolve/testData/resolve/stdlib/mapList.kt");
+    }
+
     @TestMetadata("reflectionClass.kt")
     public void testReflectionClass() throws Exception {
         runTest("compiler/fir/resolve/testData/resolve/stdlib/reflectionClass.kt");