[FIR] Implement delegate inference

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

@@ -0,0 +1,120 @@
+/*
+ * Copyright 2010-2020 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.resolve.inference
+
+import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.FirResolvable
+import org.jetbrains.kotlin.fir.expressions.FirStatement
+import org.jetbrains.kotlin.fir.resolve.BodyResolveComponents
+import org.jetbrains.kotlin.fir.resolve.calls.Candidate
+import org.jetbrains.kotlin.fir.resolve.calls.FirNamedReferenceWithCandidate
+import org.jetbrains.kotlin.fir.resolve.calls.candidate
+import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
+import org.jetbrains.kotlin.fir.resolve.transformers.FirCallCompletionResultsWriterTransformer
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.coneTypeUnsafe
+import org.jetbrains.kotlin.fir.visitors.transformSingle
+import org.jetbrains.kotlin.resolve.calls.inference.NewConstraintSystem
+import org.jetbrains.kotlin.resolve.calls.inference.buildAbstractResultingSubstitutor
+import org.jetbrains.kotlin.resolve.calls.inference.components.KotlinConstraintSystemCompleter
+import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintStorage
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
+
+abstract class AbstractManyCandidatesInferenceSession(
+    protected val components: BodyResolveComponents,
+    initialCall: FirExpression,
+    private val postponedArgumentsAnalyzer: PostponedArgumentsAnalyzer,
+) : FirInferenceSession() {
+    private val errorCalls: MutableList<FirResolvable> = mutableListOf()
+    private val partiallyResolvedCalls: MutableList<FirResolvable> = mutableListOf()
+    private val completedCalls: MutableSet<FirResolvable> = mutableSetOf()
+
+    init {
+        val initialCandidate = (initialCall as? FirResolvable)
+            ?.calleeReference
+            ?.safeAs<FirNamedReferenceWithCandidate>()
+            ?.candidate
+        if (initialCandidate != null) {
+            partiallyResolvedCalls += initialCall as FirResolvable
+        }
+    }
+
+    private val unitType: ConeKotlinType = components.session.builtinTypes.unitType.coneTypeUnsafe()
+
+    override val currentConstraintSystem: ConstraintStorage
+        get() = partiallyResolvedCalls.lastOrNull()
+            ?.calleeReference
+            ?.safeAs<FirNamedReferenceWithCandidate>()
+            ?.candidate
+            ?.system
+            ?.currentStorage()
+            ?: ConstraintStorage.Empty
+
+    private lateinit var resultingConstraintSystem: NewConstraintSystem
+
+    override fun shouldRunCompletion(candidate: Candidate): Boolean {
+        return false
+    }
+
+    override fun <T> addCompetedCall(call: T) where T : FirResolvable, T : FirStatement {
+        // do nothing
+    }
+
+    final override fun <T> addPartiallyResolvedCall(call: T) where T : FirResolvable, T : FirStatement {
+        partiallyResolvedCalls += call
+    }
+
+    final override fun <T> addErrorCall(call: T) where T : FirResolvable, T : FirStatement {
+        errorCalls += call
+    }
+
+    final override fun <T> callCompleted(call: T): Boolean where T : FirResolvable, T : FirStatement {
+        return !completedCalls.add(call)
+    }
+
+    protected open fun prepareForCompletion(
+        commonSystem: NewConstraintSystem,
+        partiallyResolvedCalls: List<FirResolvable>
+    ) {
+        // do nothing
+    }
+
+    fun completeCandidates(): List<FirResolvable> {
+        fun runCompletion(constraintSystem: NewConstraintSystem, atoms: List<FirStatement>) {
+            components.callCompleter.completer.complete(
+                constraintSystem.asConstraintSystemCompleterContext(),
+                KotlinConstraintSystemCompleter.ConstraintSystemCompletionMode.FULL,
+                atoms,
+                unitType
+            ) {
+                postponedArgumentsAnalyzer.analyze(
+                    constraintSystem.asPostponedArgumentsAnalyzerContext(),
+                    it,
+                    (atoms.first() as FirResolvable).candidate
+                )
+            }
+        }
+
+        @Suppress("UNCHECKED_CAST")
+        val resolvedCalls = partiallyResolvedCalls as List<FirResolvable>
+        val commonSystem = components.inferenceComponents.createConstraintSystem().apply {
+            addOtherSystem(currentConstraintSystem)
+        }
+        prepareForCompletion(commonSystem, resolvedCalls)
+        @Suppress("UNCHECKED_CAST")
+        runCompletion(commonSystem, resolvedCalls as List<FirStatement>)
+        resultingConstraintSystem = commonSystem
+        return resolvedCalls
+    }
+
+    protected val FirResolvable.candidate: Candidate
+        get() = candidate()!!
+
+    fun createFinalSubstitutor(): ConeSubstitutor {
+        return resultingConstraintSystem.asReadOnlyStorage()
+            .buildAbstractResultingSubstitutor(components.inferenceComponents.ctx) as ConeSubstitutor
+    }
+}

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

@@ -118,6 +118,19 @@ class FirCallCompleter(
         }
     }
 
+    fun createCompletionResultsWriter(
+        substitutor: ConeSubstitutor,
+        mode: FirCallCompletionResultsWriterTransformer.Mode = FirCallCompletionResultsWriterTransformer.Mode.Normal
+    ): FirCallCompletionResultsWriterTransformer {
+        return FirCallCompletionResultsWriterTransformer(
+            session, substitutor, returnTypeCalculator,
+            inferenceComponents.approximator,
+            integerOperatorsTypeUpdater,
+            integerLiteralTypeApproximator,
+            mode
+        )
+    }
+
     fun createPostponedArgumentsAnalyzer(): PostponedArgumentsAnalyzer {
         return PostponedArgumentsAnalyzer(
             LambdaAnalyzerImpl(), inferenceComponents,

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

@@ -0,0 +1,92 @@
+/*
+ * Copyright 2010-2020 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.resolve.inference
+
+import org.jetbrains.kotlin.fir.declarations.FirAnonymousObject
+import org.jetbrains.kotlin.fir.declarations.FirProperty
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
+import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.FirResolvable
+import org.jetbrains.kotlin.fir.expressions.FirStatement
+import org.jetbrains.kotlin.fir.references.FirNamedReference
+import org.jetbrains.kotlin.fir.resolve.BodyResolveComponents
+import org.jetbrains.kotlin.fir.resolve.calls.Candidate
+import org.jetbrains.kotlin.fir.resolve.defaultType
+import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.ConeTypeVariableTypeConstructor
+import org.jetbrains.kotlin.fir.types.coneTypeSafe
+import org.jetbrains.kotlin.fir.types.coneTypeUnsafe
+import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemBuilder
+import org.jetbrains.kotlin.resolve.calls.inference.NewConstraintSystem
+import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintStorage
+import org.jetbrains.kotlin.resolve.calls.inference.model.SimpleConstraintSystemConstraintPosition
+import org.jetbrains.kotlin.util.OperatorNameConventions
+
+class FirDelegatedPropertyInferenceSession(
+    val property: FirProperty,
+    initialCall: FirExpression,
+    components: BodyResolveComponents,
+    postponedArgumentsAnalyzer: PostponedArgumentsAnalyzer,
+) : AbstractManyCandidatesInferenceSession(components, initialCall, postponedArgumentsAnalyzer) {
+    val expectedType: ConeKotlinType? by lazy { property.returnTypeRef.coneTypeSafe() }
+
+    override fun inferPostponedVariables(
+        lambda: ResolvedLambdaAtom,
+        initialStorage: ConstraintStorage
+    ): Map<ConeTypeVariableTypeConstructor, ConeKotlinType> = emptyMap()
+
+    override fun <T> shouldCompleteResolvedSubAtomsOf(call: T): Boolean where T : FirResolvable, T : FirStatement = true
+
+    override fun prepareForCompletion(commonSystem: NewConstraintSystem, partiallyResolvedCalls: List<FirResolvable>) {
+        val csBuilder = commonSystem.getBuilder()
+        for (call in partiallyResolvedCalls) {
+            val candidate = call.candidate
+            when ((call.calleeReference as FirNamedReference).name) {
+                OperatorNameConventions.GET_VALUE -> candidate.addConstraintsForGetValueMethod(csBuilder)
+                OperatorNameConventions.SET_VALUE -> candidate.addConstraintsForSetValueMethod(csBuilder)
+            }
+        }
+    }
+
+    private fun Candidate.addConstraintsForGetValueMethod(commonSystem: ConstraintSystemBuilder) {
+        if (expectedType != null) {
+            val accessor = symbol.fir as? FirSimpleFunction ?: return
+            val unsubstitutedReturnType = accessor.returnTypeRef.coneTypeSafe<ConeKotlinType>() ?: return
+
+            val substitutedReturnType = substitutor.substituteOrSelf(unsubstitutedReturnType)
+            commonSystem.addSubtypeConstraint(substitutedReturnType, expectedType!!, SimpleConstraintSystemConstraintPosition)
+        }
+
+        addConstraintForThis(commonSystem)
+    }
+
+
+    private fun Candidate.addConstraintsForSetValueMethod(commonSystem: ConstraintSystemBuilder) {
+        if (expectedType != null) {
+            val accessor = symbol.fir as? FirSimpleFunction ?: return
+            val unsubstitutedParameterType = accessor.valueParameters.getOrNull(2)?.returnTypeRef?.coneTypeSafe<ConeKotlinType>() ?: return
+
+            val substitutedReturnType = substitutor.substituteOrSelf(unsubstitutedParameterType)
+            commonSystem.addSubtypeConstraint(substitutedReturnType, expectedType!!, SimpleConstraintSystemConstraintPosition)
+        }
+
+        addConstraintForThis(commonSystem)
+    }
+
+    private fun Candidate.addConstraintForThis(commonSystem: ConstraintSystemBuilder) {
+        val typeOfThis: ConeKotlinType = property.receiverTypeRef?.coneTypeUnsafe()
+            ?: when (val container = components.container) {
+                is FirRegularClass -> container.defaultType()
+                is FirAnonymousObject -> container.defaultType()
+                else -> null
+            } ?: components.session.builtinTypes.nullableNothingType.coneTypeUnsafe()
+        val valueParameterForThis = (symbol as? FirFunctionSymbol<*>)?.fir?.valueParameters?.firstOrNull() ?: return
+        val substitutedType = substitutor.substituteOrSelf(valueParameterForThis.returnTypeRef.coneTypeUnsafe<ConeKotlinType>())
+        commonSystem.addSubtypeConstraint(typeOfThis, substitutedType, SimpleConstraintSystemConstraintPosition)
+    }
+}

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

@@ -19,7 +19,6 @@ import org.jetbrains.kotlin.fir.resolve.constructFunctionalTypeRef
 import org.jetbrains.kotlin.fir.resolve.inference.isBuiltinFunctionalType
 import org.jetbrains.kotlin.fir.resolve.inference.returnType
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
-import org.jetbrains.kotlin.fir.resolve.substitution.substituteOrNull
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.resultType
 import org.jetbrains.kotlin.fir.scopes.impl.FirIntegerOperatorCall
 import org.jetbrains.kotlin.fir.scopes.impl.withReplacedConeType
@@ -28,9 +27,7 @@ import org.jetbrains.kotlin.fir.types.builder.buildErrorTypeRef
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.builder.buildTypeProjectionWithVariance
 import org.jetbrains.kotlin.fir.types.impl.ConeTypeParameterTypeImpl
-import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
+import org.jetbrains.kotlin.fir.visitors.*
-import org.jetbrains.kotlin.fir.visitors.compose
-import org.jetbrains.kotlin.fir.visitors.transformSingle
 import org.jetbrains.kotlin.types.AbstractTypeApproximator
 import org.jetbrains.kotlin.types.TypeApproximatorConfiguration
 import org.jetbrains.kotlin.types.Variance
@@ -43,8 +40,15 @@ class FirCallCompletionResultsWriterTransformer(
     private val typeApproximator: AbstractTypeApproximator,
     private val integerOperatorsTypeUpdater: IntegerOperatorsTypeUpdater,
     private val integerApproximator: IntegerLiteralTypeApproximationTransformer,
+    private val mode: Mode = Mode.Normal
 ) : FirAbstractTreeTransformer<ExpectedArgumentType?>(phase = FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE) {
 
+    private val declarationWriter by lazy { FirDeclarationCompletionResultsWriter(finalSubstitutor) }
+
+    enum class Mode {
+        Normal, DelegatedPropertyCompletion
+    }
+
     override fun transformQualifiedAccessExpression(
         qualifiedAccessExpression: FirQualifiedAccessExpression,
         data: ExpectedArgumentType?,
@@ -87,7 +91,7 @@ class FirCallCompletionResultsWriterTransformer(
     private fun FirResolvedTypeRef.substituteTypeRef(
         candidate: Candidate,
     ): FirResolvedTypeRef {
-        val initialType = candidate.substitutor.substituteOrNull(type)
+        val initialType = candidate.substitutor.substituteOrSelf(type)
         val finalType = finalSubstitutor.substituteOrNull(initialType)?.let { substitutedType ->
             typeApproximator.approximateToSuperType(
                 substitutedType, TypeApproximatorConfiguration.FinalApproximationAfterResolutionAndInference,
@@ -208,9 +212,34 @@ class FirCallCompletionResultsWriterTransformer(
         result.replaceTypeRef(resultType)
         result.replaceTypeArguments(typeArguments)
 
+        if (mode == Mode.DelegatedPropertyCompletion) {
+            calleeReference.candidateSymbol.fir.transformSingle(declarationWriter, null)
+            val typeUpdater = TypeUpdaterForDelegateArguments()
+            result.transformArguments(typeUpdater, null)
+            result.transformExplicitReceiver(typeUpdater, null)
+        }
+
         return result.compose()
     }
 
+    private inner class TypeUpdaterForDelegateArguments : FirTransformer<Nothing?>() {
+        override fun <E : FirElement> transformElement(element: E, data: Nothing?): CompositeTransformResult<E> {
+            return element.compose()
+        }
+
+        override fun transformQualifiedAccessExpression(
+            qualifiedAccessExpression: FirQualifiedAccessExpression,
+            data: Nothing?
+        ): CompositeTransformResult<FirStatement> {
+            val originalType = qualifiedAccessExpression.typeRef.coneTypeUnsafe<ConeKotlinType>()
+            val substitutedReceiverType = finalSubstitutor.substituteOrNull(originalType) ?: return qualifiedAccessExpression.compose()
+            qualifiedAccessExpression.replaceTypeRef(
+                qualifiedAccessExpression.typeRef.resolvedTypeFromPrototype(substitutedReceiverType)
+            )
+            return qualifiedAccessExpression.compose()
+        }
+    }
+
     private fun FirTypeRef.substitute(candidate: Candidate): ConeKotlinType =
         coneTypeUnsafe<ConeKotlinType>()
             .let { candidate.substitutor.substituteOrSelf(it) }
@@ -383,3 +412,47 @@ private fun ConeKotlinType.approximateIfPossible(expectedType: ConeKotlinType?)
 } else {
     this
 }
+
+class FirDeclarationCompletionResultsWriter(private val finalSubstitutor: ConeSubstitutor) : FirDefaultTransformer<Nothing?>() {
+    override fun <E : FirElement> transformElement(element: E, data: Nothing?): CompositeTransformResult<E> {
+        return element.compose()
+    }
+
+    override fun transformSimpleFunction(simpleFunction: FirSimpleFunction, data: Nothing?): CompositeTransformResult<FirDeclaration> {
+        simpleFunction.transformReturnTypeRef(this, data)
+        simpleFunction.transformValueParameters(this, data)
+        simpleFunction.transformReceiverTypeRef(this, data)
+        return simpleFunction.compose()
+    }
+
+    override fun transformProperty(property: FirProperty, data: Nothing?): CompositeTransformResult<FirDeclaration> {
+        property.transformGetter(this, data)
+        property.transformSetter(this, data)
+        property.transformReturnTypeRef(this, data)
+        property.transformReceiverTypeRef(this, data)
+        return property.compose()
+    }
+
+    override fun transformPropertyAccessor(
+        propertyAccessor: FirPropertyAccessor,
+        data: Nothing?
+    ): CompositeTransformResult<FirStatement> {
+        propertyAccessor.transformReturnTypeRef(this, data)
+        propertyAccessor.transformValueParameters(this, data)
+        return propertyAccessor.compose()
+    }
+
+    override fun transformValueParameter(
+        valueParameter: FirValueParameter,
+        data: Nothing?
+    ): CompositeTransformResult<FirStatement> {
+        valueParameter.transformReturnTypeRef(this, data)
+        return valueParameter.compose()
+    }
+
+    override fun transformTypeRef(typeRef: FirTypeRef, data: Nothing?): CompositeTransformResult<FirTypeRef> {
+        return finalSubstitutor.substituteOrNull(typeRef.coneTypeUnsafe())?.let {
+            typeRef.resolvedTypeFromPrototype(it)
+        }?.compose() ?: typeRef.compose()
+    }
+}

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

@@ -150,7 +150,7 @@ open class FirBodyResolveTransformer(
         wrappedDelegateExpression: FirWrappedDelegateExpression,
         data: ResolutionMode
     ): CompositeTransformResult<FirStatement> {
-        return expressionsTransformer.transformWrappedDelegateExpression(wrappedDelegateExpression, data)
+        return declarationsTransformer.transformWrappedDelegateExpression(wrappedDelegateExpression, data)
     }
 
     override fun <T> transformConstExpression(constExpression: FirConstExpression<T>, data: ResolutionMode): CompositeTransformResult<FirStatement> {

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

@@ -14,16 +14,13 @@ import org.jetbrains.kotlin.fir.declarations.impl.FirDefaultPropertyAccessor
 import org.jetbrains.kotlin.fir.declarations.synthetic.FirSyntheticProperty
 import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind
 import org.jetbrains.kotlin.fir.diagnostics.FirSimpleDiagnostic
-import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.*
-import org.jetbrains.kotlin.fir.expressions.FirReturnExpression
-import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.expressions.builder.buildReturnExpression
 import org.jetbrains.kotlin.fir.expressions.builder.buildUnitExpression
+import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.*
-import org.jetbrains.kotlin.fir.resolve.calls.ImplicitDispatchReceiverKind
+import org.jetbrains.kotlin.fir.resolve.calls.*
-import org.jetbrains.kotlin.fir.resolve.calls.ImplicitDispatchReceiverValue
+import org.jetbrains.kotlin.fir.resolve.inference.FirDelegatedPropertyInferenceSession
-import org.jetbrains.kotlin.fir.resolve.calls.ImplicitExtensionReceiverValue
-import org.jetbrains.kotlin.fir.resolve.calls.ImplicitReceiverValue
 import org.jetbrains.kotlin.fir.resolve.inference.extractLambdaInfoFromFunctionalType
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.transformers.*
@@ -104,21 +101,25 @@ class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransformer)
             dataFlowAnalyzer.enterProperty(property)
             withFullBodyResolve {
                 withScopeCleanup(localScopes) {
-
                     components.withContainer(property) {
-                        withScopeCleanup(localScopes) {
+                        if (property.delegate != null) {
                             localScopes.addIfNotNull(primaryConstructorParametersScope)
-                            property.transformChildrenWithoutAccessors(returnTypeRef)
+                            transformPropertyWithDelegate(property)
-                            property.transformInitializer(integerLiteralTypeApproximator, null)
+                        } else {
-                        }
+                            withScopeCleanup(localScopes) {
-                        if (property.initializer != null) {
+                                localScopes.addIfNotNull(primaryConstructorParametersScope)
-                            storeVariableReturnType(property)
+                                property.transformChildrenWithoutAccessors(returnTypeRef)
-                        }
+                                property.transformInitializer(integerLiteralTypeApproximator, null)
-                        withScopeCleanup(localScopes) {
+                            }
-                            localScopes.add(FirLocalScope().apply {
+                            if (property.initializer != null) {
-                                storeBackingField(property)
+                                storeVariableReturnType(property)
-                            })
+                            }
-                            property.transformAccessors()
+                            withScopeCleanup(localScopes) {
+                                localScopes.add(FirLocalScope().apply {
+                                    storeBackingField(property)
+                                })
+                                property.transformAccessors()
+                            }
                         }
                     }
                     property.replaceResolvePhase(transformerPhase)
@@ -130,17 +131,68 @@ class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransformer)
         }
     }
 
+    private fun transformPropertyWithDelegate(property: FirProperty) {
+        property.transformDelegate(transformer, ResolutionMode.ContextDependent)
+
+        val delegateExpression = property.delegate!!
+
+        val inferenceSession = FirDelegatedPropertyInferenceSession(
+            property,
+            delegateExpression,
+            components,
+            callCompleter.createPostponedArgumentsAnalyzer()
+        )
+
+        components.inferenceComponents.withInferenceSession(inferenceSession) {
+            property.transformAccessors()
+            val completedCalls = inferenceSession.completeCandidates()
+            val finalSubstitutor = inferenceSession.createFinalSubstitutor()
+            val callCompletionResultsWriter = components.callCompleter.createCompletionResultsWriter(
+                finalSubstitutor,
+                mode = FirCallCompletionResultsWriterTransformer.Mode.DelegatedPropertyCompletion
+            )
+            completedCalls.forEach {
+                it.transformSingle(callCompletionResultsWriter, null)
+            }
+            val declarationCompletionResultsWriter = FirDeclarationCompletionResultsWriter(finalSubstitutor)
+            property.transformSingle(declarationCompletionResultsWriter, null)
+        }
+        property.transformOtherChildren(transformer, ResolutionMode.ContextIndependent)
+    }
+
+    override fun transformWrappedDelegateExpression(
+        wrappedDelegateExpression: FirWrappedDelegateExpression,
+        data: ResolutionMode,
+    ): CompositeTransformResult<FirStatement> {
+        val delegateProvider = wrappedDelegateExpression.delegateProvider.transformSingle(transformer, ResolutionMode.ContextDependent)
+        when (val calleeReference = (delegateProvider as FirResolvable).calleeReference) {
+            is FirResolvedNamedReference -> return delegateProvider.compose()
+            is FirNamedReferenceWithCandidate -> {
+                val candidate = calleeReference.candidate
+                if (!candidate.system.hasContradiction) {
+                    return delegateProvider.compose()
+                }
+            }
+        }
+
+        return wrappedDelegateExpression.expression.transform(transformer, ResolutionMode.ContextDependent)
+    }
+
     private fun transformLocalVariable(variable: FirProperty): CompositeTransformResult<FirProperty> {
         assert(variable.isLocal)
-        val resolutionMode = withExpectedType(variable.returnTypeRef)
+        if (variable.delegate != null) {
-        variable.transformInitializer(transformer, resolutionMode)
+            transformPropertyWithDelegate(variable)
-            .transformDelegate(transformer, resolutionMode)
+        } else {
-            .transformOtherChildren(transformer, resolutionMode)
+            val resolutionMode = withExpectedType(variable.returnTypeRef)
-            .transformInitializer(integerLiteralTypeApproximator, null)
+            variable.transformInitializer(transformer, resolutionMode)
-        if (variable.initializer != null) {
+                .transformDelegate(transformer, resolutionMode)
-            storeVariableReturnType(variable)
+                .transformOtherChildren(transformer, resolutionMode)
+                .transformInitializer(integerLiteralTypeApproximator, null)
+            if (variable.initializer != null) {
+                storeVariableReturnType(variable)
+            }
+            variable.transformAccessors()
         }
-        variable.transformAccessors()
         localScopes.lastOrNull()?.storeVariable(variable)
         variable.replaceResolvePhase(transformerPhase)
         dataFlowAnalyzer.exitLocalVariableDeclaration(variable)
@@ -191,7 +243,7 @@ class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransformer)
             }
         } else {
             transformFunctionWithGivenSignature(accessor, expectedReturnTypeRef)
-        } as CompositeTransformResult<FirStatement>
+        }
     }
 
     private fun FirDeclaration.resolveStatus(status: FirDeclarationStatus, containingClass: FirClass<*>? = null): FirDeclarationStatus {
@@ -263,6 +315,7 @@ class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransformer)
             }
             containingClass = oldContainingClass
             primaryConstructorParametersScope = oldConstructorScope
+            @Suppress("UNCHECKED_CAST")
             result as CompositeTransformResult<FirStatement>
         }
     }
@@ -278,6 +331,7 @@ class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransformer)
             val result = withLabelAndReceiverType(null, anonymousObject, type) {
                 transformDeclaration(anonymousObject, data)
             }
+            @Suppress("UNCHECKED_CAST")
             result as CompositeTransformResult<FirStatement>
         }
     }
@@ -344,7 +398,7 @@ class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransformer)
 
     private fun <F : FirFunction<F>> transformFunctionWithGivenSignature(
         function: F,
-        returnTypeRef: FirTypeRef
+        returnTypeRef: FirTypeRef,
     ): CompositeTransformResult<F> {
         if (function is FirSimpleFunction) {
             localScopes.lastOrNull()?.storeFunction(function)
@@ -372,6 +426,7 @@ class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransformer)
         return withScopeCleanup(localScopes) {
             localScopes += FirLocalScope()
             dataFlowAnalyzer.enterFunction(function)
+            @Suppress("UNCHECKED_CAST")
             transformDeclaration(function, data).also {
                 val result = it.single as FirFunction<*>
                 dataFlowAnalyzer.exitFunction(result)?.let { controlFlowGraph ->
@@ -383,6 +438,7 @@ class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransformer)
 
     override fun transformConstructor(constructor: FirConstructor, data: ResolutionMode): CompositeTransformResult<FirDeclaration> {
         if (implicitTypeOnly) return constructor.compose()
+        @Suppress("UNCHECKED_CAST")
         return transformFunction(constructor, data) as CompositeTransformResult<FirDeclaration>
     }
 

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

@@ -37,7 +37,10 @@ import org.jetbrains.kotlin.fir.symbols.invoke
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.builder.*
 import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
-import org.jetbrains.kotlin.fir.visitors.*
+import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
+import org.jetbrains.kotlin.fir.visitors.FirDefaultTransformer
+import org.jetbrains.kotlin.fir.visitors.compose
+import org.jetbrains.kotlin.fir.visitors.transformSingle
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.types.Variance
@@ -478,20 +481,6 @@ class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransformer) :
         return transformedGetClassCall.compose()
     }
 
-    override fun transformWrappedDelegateExpression(
-        wrappedDelegateExpression: FirWrappedDelegateExpression,
-        data: ResolutionMode,
-    ): CompositeTransformResult<FirStatement> {
-        transformExpression(wrappedDelegateExpression, data)
-        with(wrappedDelegateExpression) {
-            val delegateProviderTypeRef = delegateProvider.typeRef
-            val useDelegateProvider = delegateProviderTypeRef is FirResolvedTypeRef &&
-                    delegateProviderTypeRef !is FirErrorTypeRef &&
-                    delegateProviderTypeRef.type !is ConeKotlinErrorType
-            return if (useDelegateProvider) delegateProvider.compose() else expression.compose()
-        }
-    }
-
     override fun <T> transformConstExpression(
         constExpression: FirConstExpression<T>,
         data: ResolutionMode,

compiler/fir/resolve/testData/resolve/delegates/delegateInference.kt

@@ -0,0 +1,27 @@
+import kotlin.reflect.KProperty
+
+class FreezableVar<T>(private var value: T)  {
+    operator fun getValue(thisRef: Any, property: KProperty<*>): T  = value
+
+    operator fun setValue(thisRef: Any, property: KProperty<*>, value: T) {
+        this.value = value
+    }
+}
+
+class LocalFreezableVar<T>(private var value: T)  {
+    operator fun getValue(thisRef: Nothing?, property: KProperty<*>): T  = value
+
+    operator fun setValue(thisRef: Nothing?, property: KProperty<*>, value: T) {
+        this.value = value
+    }
+}
+
+class Test {
+    var x: Boolean by FreezableVar(true)
+    var y by FreezableVar("")
+}
+
+fun test() {
+    var x: Boolean by LocalFreezableVar(true)
+    var y by LocalFreezableVar("")
+}

compiler/fir/resolve/testData/resolve/delegates/delegateInference.txt

@@ -0,0 +1,63 @@
+FILE: delegateInference.kt
+    public final class FreezableVar<T> : R|kotlin/Any| {
+        public constructor<T>(value: R|T|): R|FreezableVar<T>| {
+            super<R|kotlin/Any|>()
+        }
+
+        private final var value: R|T| = R|<local>/value|
+            private get(): R|T|
+            private set(value: R|T|): R|kotlin/Unit|
+
+        public final operator fun getValue(thisRef: R|kotlin/Any|, property: R|kotlin/reflect/KProperty<*>|): R|T| {
+            ^getValue this@R|/FreezableVar|.R|/FreezableVar.value|
+        }
+
+        public final operator fun setValue(thisRef: R|kotlin/Any|, property: R|kotlin/reflect/KProperty<*>|, value: R|T|): R|kotlin/Unit| {
+            this@R|/FreezableVar|.R|/FreezableVar.value| = R|<local>/value|
+        }
+
+    }
+    public final class LocalFreezableVar<T> : R|kotlin/Any| {
+        public constructor<T>(value: R|T|): R|LocalFreezableVar<T>| {
+            super<R|kotlin/Any|>()
+        }
+
+        private final var value: R|T| = R|<local>/value|
+            private get(): R|T|
+            private set(value: R|T|): R|kotlin/Unit|
+
+        public final operator fun getValue(thisRef: R|kotlin/Nothing?|, property: R|kotlin/reflect/KProperty<*>|): R|T| {
+            ^getValue this@R|/LocalFreezableVar|.R|/LocalFreezableVar.value|
+        }
+
+        public final operator fun setValue(thisRef: R|kotlin/Nothing?|, property: R|kotlin/reflect/KProperty<*>|, value: R|T|): R|kotlin/Unit| {
+            this@R|/LocalFreezableVar|.R|/LocalFreezableVar.value| = R|<local>/value|
+        }
+
+    }
+    public final class Test : R|kotlin/Any| {
+        public constructor(): R|Test| {
+            super<R|kotlin/Any|>()
+        }
+
+        public final var x: R|kotlin/Boolean|by R|/FreezableVar.FreezableVar|<R|kotlin/Boolean|>(Boolean(true))
+            public get(): R|kotlin/Boolean| {
+                ^ D|/Test.x|.R|FakeOverride</FreezableVar.getValue: R|kotlin/Boolean|>|(this@R|/Test|, ::R|/Test.x|)
+            }
+            public set(<set-?>: R|kotlin/Boolean|): R|kotlin/Unit| {
+                D|/Test.x|.R|FakeOverride</FreezableVar.setValue: R|kotlin/Unit|>|(this@R|/Test|, ::R|/Test.x|, R|<local>/x|)
+            }
+
+        public final var y: R|kotlin/String|by R|/FreezableVar.FreezableVar|<R|kotlin/String|>(String())
+            public get(): R|kotlin/String| {
+                ^ D|/Test.y|.R|FakeOverride</FreezableVar.getValue: R|kotlin/String|>|(this@R|/Test|, ::R|/Test.y|)
+            }
+            public set(<set-?>: R|kotlin/String|): R|kotlin/Unit| {
+                D|/Test.y|.R|FakeOverride</FreezableVar.setValue: R|kotlin/Unit|>|(this@R|/Test|, ::R|/Test.y|, R|<local>/y|)
+            }
+
+    }
+    public final fun test(): R|kotlin/Unit| {
+        lvar x: R|kotlin/Boolean|by R|/LocalFreezableVar.LocalFreezableVar|<R|kotlin/Boolean|>(Boolean(true))
+        lvar y: R|kotlin/String|by R|/LocalFreezableVar.LocalFreezableVar|<R|kotlin/String|>(String())
+    }

compiler/fir/resolve/testData/resolve/delegates/delegateWithLambda.kt

@@ -0,0 +1,16 @@
+import kotlin.reflect.KProperty
+
+class LazyDelegate<T>(val value: T) {
+    operator fun getValue(thisRef: Any?, property: KProperty<*>): T = value
+}
+
+fun <T> lazy(block: () -> T): LazyDelegate<T> = LazyDelegate(block())
+
+fun getAny(): Any? = null
+
+class Test {
+    val x by lazy {
+        val y = getAny() as? String ?: ""
+        y
+    }
+}

compiler/fir/resolve/testData/resolve/delegates/delegateWithLambda.txt

@@ -0,0 +1,43 @@
+FILE: delegateWithLambda.kt
+    public final class LazyDelegate<T> : R|kotlin/Any| {
+        public constructor<T>(value: R|T|): R|LazyDelegate<T>| {
+            super<R|kotlin/Any|>()
+        }
+
+        public final val value: R|T| = R|<local>/value|
+            public get(): R|T|
+
+        public final operator fun getValue(thisRef: R|kotlin/Any?|, property: R|kotlin/reflect/KProperty<*>|): R|T| {
+            ^getValue this@R|/LazyDelegate|.R|/LazyDelegate.value|
+        }
+
+    }
+    public final fun <T> lazy(block: R|() -> T|): R|LazyDelegate<T>| {
+        ^lazy R|/LazyDelegate.LazyDelegate|<R|T|>(R|<local>/block|.R|FakeOverride<kotlin/Function0.invoke: R|T|>|())
+    }
+    public final fun getAny(): R|kotlin/Any?| {
+        ^getAny Null(null)
+    }
+    public final class Test : R|kotlin/Any| {
+        public constructor(): R|Test| {
+            super<R|kotlin/Any|>()
+        }
+
+        public final val x: R|kotlin/String|by R|/lazy|<R|kotlin/String|>(<L> = lazy@fun <anonymous>(): R|kotlin/String| {
+            lval y: R|kotlin/String| = when (lval <elvis>: R|kotlin/String?| = (R|/getAny|() as? R|kotlin/String|)) {
+                ==($subj$, Null(null)) ->  {
+                    String()
+                }
+                else ->  {
+                    R|<local>/<elvis>|
+                }
+            }
+
+            ^ R|<local>/y|
+        }
+        )
+            public get(): R|kotlin/String| {
+                ^ D|/Test.x|.R|FakeOverride</LazyDelegate.getValue: R|kotlin/String|>|(this@R|/Test|, ::R|/Test.x|)
+            }
+
+    }

compiler/fir/resolve/testData/resolve/delegates/provideDelegate.kt

@@ -0,0 +1,19 @@
+import kotlin.reflect.KProperty
+
+class Delegate<T>(var value: T) {
+    operator fun getValue(thisRef: Any?, property: KProperty<*>): T = value
+
+    operator fun setValue(thisRef: Any?, property: KProperty<*>, newValue: T) {
+        value = newValue
+    }
+}
+
+class DelegateProvider<T>(val value: T) {
+    operator fun provideDelegate(thisRef: Any?, property: KProperty<*>): Delegate<T> = Delegate(value)
+}
+
+fun <T> delegate(value: T): DelegateProvider<T> = DelegateProvider(value)
+
+class A {
+    val x by delegate(1)
+}

compiler/fir/resolve/testData/resolve/delegates/provideDelegate.txt

@@ -0,0 +1,46 @@
+FILE: provideDelegate.kt
+    public final class Delegate<T> : R|kotlin/Any| {
+        public constructor<T>(value: R|T|): R|Delegate<T>| {
+            super<R|kotlin/Any|>()
+        }
+
+        public final var value: R|T| = R|<local>/value|
+            public get(): R|T|
+            public set(value: R|T|): R|kotlin/Unit|
+
+        public final operator fun getValue(thisRef: R|kotlin/Any?|, property: R|kotlin/reflect/KProperty<*>|): R|T| {
+            ^getValue this@R|/Delegate|.R|/Delegate.value|
+        }
+
+        public final operator fun setValue(thisRef: R|kotlin/Any?|, property: R|kotlin/reflect/KProperty<*>|, newValue: R|T|): R|kotlin/Unit| {
+            this@R|/Delegate|.R|/Delegate.value| = R|<local>/newValue|
+        }
+
+    }
+    public final class DelegateProvider<T> : R|kotlin/Any| {
+        public constructor<T>(value: R|T|): R|DelegateProvider<T>| {
+            super<R|kotlin/Any|>()
+        }
+
+        public final val value: R|T| = R|<local>/value|
+            public get(): R|T|
+
+        public final operator fun provideDelegate(thisRef: R|kotlin/Any?|, property: R|kotlin/reflect/KProperty<*>|): R|Delegate<T>| {
+            ^provideDelegate R|/Delegate.Delegate|<R|T|>(this@R|/DelegateProvider|.R|/DelegateProvider.value|)
+        }
+
+    }
+    public final fun <T> delegate(value: R|T|): R|DelegateProvider<T>| {
+        ^delegate R|/DelegateProvider.DelegateProvider|<R|T|>(R|<local>/value|)
+    }
+    public final class A : R|kotlin/Any| {
+        public constructor(): R|A| {
+            super<R|kotlin/Any|>()
+        }
+
+        public final val x: R|kotlin/Int|by R|/delegate|<R|kotlin/Int|>(Int(1)).R|FakeOverride</DelegateProvider.provideDelegate: R|Delegate<kotlin/Int>|>|(this@R|/A|, ::R|/A.x|)
+            public get(): R|kotlin/Int| {
+                ^ D|/A.x|.R|FakeOverride</Delegate.getValue: R|kotlin/Int|>|(this@R|/A|, ::R|/A.x|)
+            }
+
+    }

compiler/fir/resolve/testData/resolveWithStdlib/delegateTypeMismatch.txt

@@ -33,12 +33,12 @@ FILE: delegateTypeMismatch.kt
             )
         }
 
-        public final var classifierNamePolicy: R|ClassifierNamePolicy|by this@R|/A|.R|/A.property|<R|ClassifierNamePolicy.SOURCE_CODE_QUALIFIED|>(Q|ClassifierNamePolicy.SOURCE_CODE_QUALIFIED|)
+        public final var classifierNamePolicy: R|ClassifierNamePolicy|by this@R|/A|.R|/A.property|<R|ClassifierNamePolicy|>(Q|ClassifierNamePolicy.SOURCE_CODE_QUALIFIED|)
             public get(): R|ClassifierNamePolicy| {
-                ^ D|/A.classifierNamePolicy|.R|FakeOverride<kotlin/properties/ReadWriteProperty.getValue: R|ClassifierNamePolicy.SOURCE_CODE_QUALIFIED|>|(this@R|/A|, ::R|/A.classifierNamePolicy|)
+                ^ D|/A.classifierNamePolicy|.R|FakeOverride<kotlin/properties/ReadWriteProperty.getValue: R|ClassifierNamePolicy|>|(this@R|/A|, ::R|/A.classifierNamePolicy|)
             }
             public set(<set-?>: R|ClassifierNamePolicy|): R|kotlin/Unit| {
-                D|/A.classifierNamePolicy|.<Inapplicable(INAPPLICABLE): [kotlin/properties/ReadWriteProperty.setValue]>#(this@R|/A|, ::R|/A.classifierNamePolicy|, R|<local>/classifierNamePolicy|)
+                D|/A.classifierNamePolicy|.R|FakeOverride<kotlin/properties/ReadWriteProperty.setValue: R|kotlin/Unit|>|(this@R|/A|, ::R|/A.classifierNamePolicy|, R|<local>/classifierNamePolicy|)
             }
 
         public final var typeNormalizer: <ERROR TYPE REF: Ambiguity: getValue, [kotlin/getValue, kotlin/collections/getValue, kotlin/collections/getValue]>by <Inapplicable(INAPPLICABLE): [/A.property]>#<R|(KotlinType) -> KotlinType|>(property@fun <anonymous>(): R|ERROR CLASS: Unresolved name: it| {

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

@@ -572,6 +572,34 @@ public class FirDiagnosticsTestGenerated extends AbstractFirDiagnosticsTest {
         }
     }
 
+    @TestMetadata("compiler/fir/resolve/testData/resolve/delegates")
+    @TestDataPath("$PROJECT_ROOT")
+    @RunWith(JUnit3RunnerWithInners.class)
+    public static class Delegates extends AbstractFirDiagnosticsTest {
+        private void runTest(String testDataFilePath) throws Exception {
+            KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
+        }
+
+        public void testAllFilesPresentInDelegates() throws Exception {
+            KotlinTestUtils.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/fir/resolve/testData/resolve/delegates"), Pattern.compile("^([^.]+)\\.kt$"), null, true);
+        }
+
+        @TestMetadata("delegateInference.kt")
+        public void testDelegateInference() throws Exception {
+            runTest("compiler/fir/resolve/testData/resolve/delegates/delegateInference.kt");
+        }
+
+        @TestMetadata("delegateWithLambda.kt")
+        public void testDelegateWithLambda() throws Exception {
+            runTest("compiler/fir/resolve/testData/resolve/delegates/delegateWithLambda.kt");
+        }
+
+        @TestMetadata("provideDelegate.kt")
+        public void testProvideDelegate() throws Exception {
+            runTest("compiler/fir/resolve/testData/resolve/delegates/provideDelegate.kt");
+        }
+    }
+
     @TestMetadata("compiler/fir/resolve/testData/resolve/diagnostics")
     @TestDataPath("$PROJECT_ROOT")
     @RunWith(JUnit3RunnerWithInners.class)

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

@@ -572,6 +572,34 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
         }
     }
 
+    @TestMetadata("compiler/fir/resolve/testData/resolve/delegates")
+    @TestDataPath("$PROJECT_ROOT")
+    @RunWith(JUnit3RunnerWithInners.class)
+    public static class Delegates extends AbstractFirDiagnosticsWithLightTreeTest {
+        private void runTest(String testDataFilePath) throws Exception {
+            KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
+        }
+
+        public void testAllFilesPresentInDelegates() throws Exception {
+            KotlinTestUtils.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/fir/resolve/testData/resolve/delegates"), Pattern.compile("^([^.]+)\\.kt$"), null, true);
+        }
+
+        @TestMetadata("delegateInference.kt")
+        public void testDelegateInference() throws Exception {
+            runTest("compiler/fir/resolve/testData/resolve/delegates/delegateInference.kt");
+        }
+
+        @TestMetadata("delegateWithLambda.kt")
+        public void testDelegateWithLambda() throws Exception {
+            runTest("compiler/fir/resolve/testData/resolve/delegates/delegateWithLambda.kt");
+        }
+
+        @TestMetadata("provideDelegate.kt")
+        public void testProvideDelegate() throws Exception {
+            runTest("compiler/fir/resolve/testData/resolve/delegates/provideDelegate.kt");
+        }
+    }
+
     @TestMetadata("compiler/fir/resolve/testData/resolve/diagnostics")
     @TestDataPath("$PROJECT_ROOT")
     @RunWith(JUnit3RunnerWithInners.class)

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

@@ -83,6 +83,11 @@ public class FirDiagnosticsWithStdlibTestGenerated extends AbstractFirDiagnostic
         runTest("compiler/fir/resolve/testData/resolveWithStdlib/concurrentMapOfAliases.kt");
     }
 
+    @TestMetadata("delegateTypeMismatch.kt")
+    public void testDelegateTypeMismatch() throws Exception {
+        runTest("compiler/fir/resolve/testData/resolveWithStdlib/delegateTypeMismatch.kt");
+    }
+
     @TestMetadata("delegationByMap.kt")
     public void testDelegationByMap() throws Exception {
         runTest("compiler/fir/resolve/testData/resolveWithStdlib/delegationByMap.kt");
@@ -685,11 +690,6 @@ public class FirDiagnosticsWithStdlibTestGenerated extends AbstractFirDiagnostic
             runTest("compiler/fir/resolve/testData/resolveWithStdlib/problems/cloneArray.kt");
         }
 
-        @TestMetadata("delegateTypeMismatch.kt")
-        public void testDelegateTypeMismatch() throws Exception {
-            runTest("compiler/fir/resolve/testData/resolveWithStdlib/problems/delegateTypeMismatch.kt");
-        }
-
         @TestMetadata("invokePriority.kt")
         public void testInvokePriority() throws Exception {
             runTest("compiler/fir/resolve/testData/resolveWithStdlib/problems/invokePriority.kt");

compiler/testData/diagnostics/testsWithStdLib/PropertyDelegateWithPrivateSet.fir.kt

@@ -8,5 +8,5 @@ class My {
     val another: String = delegate
 
     var delegateWithBackingField: String by kotlin.properties.Delegates.notNull()
-        private set(arg) { field = arg }
+        private set(arg) { <!UNRESOLVED_REFERENCE!>field<!> = arg }
 }