[FIR] Do not dereference delegate in lazy mode

Tenth step for ^KT-52615

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/FirLazyBodiesCalculator.kt

@@ -150,14 +150,7 @@ internal object FirLazyBodiesCalculator {
             val newDelegate = newProperty.delegate as? FirWrappedDelegateExpression
             check(newDelegate != null) { "Invalid replacement delegate" }
             delegate.replaceExpression(newDelegate.expression)
-
+            delegate.replaceDelegateProvider(newDelegate.delegateProvider)
-            val delegateProviderCall = delegate.delegateProvider as? FirFunctionCall
-            val delegateProviderExplicitReceiver = delegateProviderCall?.explicitReceiver
-            if (delegateProviderExplicitReceiver is FirLazyExpression) {
-                val newDelegateProviderExplicitReceiver = (newDelegate.delegateProvider as? FirFunctionCall)?.explicitReceiver
-                check(newDelegateProviderExplicitReceiver != null) { "Invalid replacement expression" }
-                delegateProviderCall.replaceExplicitReceiver(newDelegateProviderExplicitReceiver)
-            }
         }
     }
 

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/declaration/FirDelegatedPropertyChecker.kt

@@ -31,16 +31,16 @@ object FirDelegatedPropertyChecker : FirPropertyChecker() {
     override fun check(declaration: FirProperty, context: CheckerContext, reporter: DiagnosticReporter) {
         val delegate = declaration.delegate ?: return
         val delegateType = delegate.typeRef.coneType
+        val source = delegate.source;
 
         // TODO: Also suppress delegate issue if type inference failed. For example, in
         //  compiler/testData/diagnostics/tests/delegatedProperty/inference/differentDelegatedExpressions.fir.kt, no delegate issues are
         //  reported due to the inference issue.
         if (delegateType is ConeErrorType) {
-            val delegateSource = delegate.source
             // Implicit recursion type is not reported since the type ref does not have a real source.
-            if (delegateSource != null && (delegateType.diagnostic as? ConeSimpleDiagnostic)?.kind == DiagnosticKind.RecursionInImplicitTypes) {
+            if (source != null && (delegateType.diagnostic as? ConeSimpleDiagnostic)?.kind == DiagnosticKind.RecursionInImplicitTypes) {
                 // skip reporting other issues in this case
-                reporter.reportOn(delegateSource, FirErrors.RECURSION_IN_IMPLICIT_TYPES, context)
+                reporter.reportOn(source, FirErrors.RECURSION_IN_IMPLICIT_TYPES, context)
             }
             return
         }
@@ -74,7 +74,7 @@ object FirDelegatedPropertyChecker : FirPropertyChecker() {
 
                 fun reportInapplicableDiagnostics(candidates: Collection<FirBasedSymbol<*>>) {
                     reporter.reportOn(
-                        reference.source,
+                        source,
                         FirErrors.DELEGATE_SPECIAL_FUNCTION_NONE_APPLICABLE,
                         expectedFunctionSignature,
                         candidates,
@@ -85,7 +85,7 @@ object FirDelegatedPropertyChecker : FirPropertyChecker() {
                 var errorReported = true
                 when (diagnostic) {
                     is ConeUnresolvedNameError -> reporter.reportOn(
-                        reference.source,
+                        source,
                         FirErrors.DELEGATE_SPECIAL_FUNCTION_MISSING,
                         expectedFunctionSignature,
                         delegateType,
@@ -97,7 +97,7 @@ object FirDelegatedPropertyChecker : FirPropertyChecker() {
                         if (diagnostic.applicability.isSuccess) {
                             // Match is successful but there are too many matches! So we report DELEGATE_SPECIAL_FUNCTION_AMBIGUITY.
                             reporter.reportOn(
-                                reference.source,
+                                source,
                                 FirErrors.DELEGATE_SPECIAL_FUNCTION_AMBIGUITY,
                                 expectedFunctionSignature,
                                 diagnostic.candidates.map { it.symbol },
@@ -109,7 +109,7 @@ object FirDelegatedPropertyChecker : FirPropertyChecker() {
                     }
 
                     is ConeInapplicableWrongReceiver -> reporter.reportOn(
-                        reference.source,
+                        source,
                         FirErrors.DELEGATE_SPECIAL_FUNCTION_MISSING,
                         expectedFunctionSignature,
                         delegateType,
@@ -131,7 +131,7 @@ object FirDelegatedPropertyChecker : FirPropertyChecker() {
                 val propertyType = declaration.returnTypeRef.coneType
                 if (!AbstractTypeChecker.isSubtypeOf(context.session.typeContext, returnType, propertyType)) {
                     reporter.reportOn(
-                        delegate.source,
+                        source,
                         FirErrors.DELEGATE_SPECIAL_FUNCTION_RETURN_TYPE_MISMATCH,
                         "getValue",
                         propertyType,

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

@@ -7,8 +7,10 @@ package org.jetbrains.kotlin.fir.analysis.checkers.expression
 
 import org.jetbrains.kotlin.KtFakeSourceElementKind
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
+import org.jetbrains.kotlin.fakeElement
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.analysis.checkers.isLhsOfAssignment
+import org.jetbrains.kotlin.fir.declarations.FirProperty
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.impl.FirNoReceiverExpression
 import org.jetbrains.kotlin.fir.references.toResolvedBaseSymbol
@@ -37,7 +39,13 @@ object FirOptInUsageAccessChecker : FirBasicExpressionChecker() {
 
                 val experimentalities = resolvedSymbol.loadExperimentalities(context, fromSetter = false, dispatchReceiverType) +
                         loadExperimentalitiesFromTypeArguments(context, expression.typeArguments)
-                reportNotAcceptedExperimentalities(experimentalities, expression, context, reporter)
+                val source = if (expression.source?.kind == KtFakeSourceElementKind.DelegatedPropertyAccessor) {
+                    val property = context.containingDeclarations.lastOrNull { it is FirProperty } as? FirProperty ?: return
+                    property.delegate?.source?.fakeElement(KtFakeSourceElementKind.DelegatedPropertyAccessor) ?: return
+                } else {
+                    expression.source
+                }
+                reportNotAcceptedExperimentalities(experimentalities, expression, context, reporter, source)
             }
         }
     }

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

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.analysis.checkers.expression
 
+import org.jetbrains.kotlin.KtSourceElement
 import org.jetbrains.kotlin.config.AnalysisFlags
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.diagnostics.reportOn
@@ -262,7 +263,8 @@ object FirOptInUsageBaseChecker {
         experimentalities: Collection<Experimentality>,
         element: FirElement,
         context: CheckerContext,
-        reporter: DiagnosticReporter
+        reporter: DiagnosticReporter,
+        source: KtSourceElement? = element.source,
     ) {
         for ((annotationClassId, severity, message, _, fromSupertype) in experimentalities) {
             if (!isExperimentalityAcceptableInContext(annotationClassId, context, fromSupertype)) {
@@ -273,7 +275,7 @@ object FirOptInUsageBaseChecker {
                 val fqName = annotationClassId.asSingleFqName()
                 val reportedMessage = message?.takeIf { it.isNotBlank() }
                     ?: OptInNames.buildDefaultDiagnosticMessage(OptInNames.buildMessagePrefix(verb), fqName.asString())
-                reporter.reportOn(element.source, diagnostic, fqName, reportedMessage, context)
+                reporter.reportOn(source, diagnostic, fqName, reportedMessage, context)
             }
         }
     }

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

@@ -11,9 +11,11 @@ import org.jetbrains.kotlin.KtNodeTypes
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
+import org.jetbrains.kotlin.fakeElement
 import org.jetbrains.kotlin.fir.analysis.diagnostics.toFirDiagnostics
 import org.jetbrains.kotlin.fir.declarations.FirErrorFunction
 import org.jetbrains.kotlin.fir.declarations.FirErrorImport
+import org.jetbrains.kotlin.fir.declarations.FirProperty
 import org.jetbrains.kotlin.fir.diagnostics.ConeAmbiguousSuper
 import org.jetbrains.kotlin.fir.diagnostics.ConeDiagnostic
 import org.jetbrains.kotlin.fir.diagnostics.ConeSimpleDiagnostic
@@ -59,7 +61,7 @@ class ErrorNodeDiagnosticCollectorComponent(
     }
 
     private fun processErrorReference(reference: FirNamedReference, diagnostic: ConeDiagnostic, context: CheckerContext) {
-        val source = reference.source ?: return
+        var source = reference.source ?: return
         val qualifiedAccessOrAnnotationCall = context.qualifiedAccessOrAssignmentsOrAnnotationCalls.lastOrNull()?.takeIf {
             // Use the source of the enclosing FirQualifiedAccess if it is exactly the call to the erroneous callee.
             it.calleeReference == reference
@@ -79,6 +81,11 @@ class ErrorNodeDiagnosticCollectorComponent(
             ) return
         }
 
+        if (source.kind == KtFakeSourceElementKind.DelegatedPropertyAccessor) {
+            val property = context.containingDeclarations.lastOrNull { it is FirProperty } as? FirProperty ?: return
+            source = property.delegate?.source?.fakeElement(KtFakeSourceElementKind.DelegatedPropertyAccessor) ?: return
+        }
+
         reportFirDiagnostic(diagnostic, source, context, qualifiedAccessOrAnnotationCall?.source)
     }
 

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

@@ -1803,18 +1803,18 @@ open class RawFirBuilder(
                         }
 
                         if (hasDelegate()) {
-                            fun extractDelegateExpression() = this@toFirProperty.delegate?.expression?.let { expression ->
+                            fun extractDelegateExpression() = buildOrLazyExpression(this@toFirProperty.toFirSourceElement(KtFakeSourceElementKind.WrappedDelegate)) {
-                                buildOrLazyExpression(expression.toFirSourceElement()) {
+                                this@toFirProperty.delegate?.expression?.let { expression ->
                                     expression.toFirExpression("Should have delegate")
+                                } ?: buildErrorExpression {
+                                    diagnostic = ConeSimpleDiagnostic("Should have delegate", DiagnosticKind.ExpressionExpected)
                                 }
-                            } ?: buildErrorExpression {
-                                diagnostic = ConeSimpleDiagnostic("Should have delegate", DiagnosticKind.ExpressionExpected)
                             }
 
                             val delegateBuilder = FirWrappedDelegateExpressionBuilder().apply {
                                 val delegateExpression = extractDelegateExpression()
                                 source = delegateExpression.source?.fakeElement(KtFakeSourceElementKind.WrappedDelegate)
-                                expression = extractDelegateExpression()
+                                expression = delegateExpression
                             }
 
                             generateAccessorsByDelegate(

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

@@ -34,5 +34,7 @@ abstract class FirWrappedDelegateExpression : FirWrappedExpression() {
 
     abstract override fun replaceExpression(newExpression: FirExpression)
 
+    abstract fun replaceDelegateProvider(newDelegateProvider: FirExpression)
+
     abstract override fun <D> transformAnnotations(transformer: FirTransformer<D>, data: D): FirWrappedDelegateExpression
 }

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/expressions/impl/FirWrappedDelegateExpressionImpl.kt

@@ -56,4 +56,8 @@ internal class FirWrappedDelegateExpressionImpl(
     override fun replaceExpression(newExpression: FirExpression) {
         expression = newExpression
     }
+
+    override fun replaceDelegateProvider(newDelegateProvider: FirExpression) {
+        delegateProvider = newDelegateProvider
+    }
 }

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

@@ -628,7 +628,7 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
         }
 
         wrappedDelegateExpression.configure {
-            +field("delegateProvider", expression)
+            +field("delegateProvider", expression).withReplace()
         }
 
         namedReference.configure {