[FIR] Introduce FirAssignExpressionAltererExtension

This extension can be used to override resolution of assign statements
  with custom statement

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/diagnostics/KtFirDataClassConverters.kt

@@ -363,6 +363,15 @@ internal val KT_DIAGNOSTIC_CONVERTER = KtDiagnosticConverterBuilder.buildConvert
             token,
         )
     }
+    add(FirErrors.AMBIGUOUS_ALTERED_ASSIGN) { firDiagnostic ->
+        AmbiguousAlteredAssignImpl(
+            firDiagnostic.a.map { string ->
+                string
+            },
+            firDiagnostic as KtPsiDiagnostic,
+            token,
+        )
+    }
     add(FirErrors.SUPER_IS_NOT_AN_EXPRESSION) { firDiagnostic ->
         SuperIsNotAnExpressionImpl(
             firDiagnostic as KtPsiDiagnostic,

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/diagnostics/KtFirDiagnostics.kt

@@ -291,6 +291,11 @@ sealed class KtFirDiagnostic<PSI : PsiElement> : KtDiagnosticWithPsi<PSI> {
         abstract val classSymbol: KtClassLikeSymbol
     }
 
+    abstract class AmbiguousAlteredAssign : KtFirDiagnostic<PsiElement>() {
+        override val diagnosticClass get() = AmbiguousAlteredAssign::class
+        abstract val altererNames: List<String?>
+    }
+
     abstract class SuperIsNotAnExpression : KtFirDiagnostic<PsiElement>() {
         override val diagnosticClass get() = SuperIsNotAnExpression::class
     }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/diagnostics/KtFirDiagnosticsImpl.kt

@@ -336,6 +336,12 @@ internal class ResolutionToClassifierImpl(
     override val token: KtLifetimeToken,
 ) : KtFirDiagnostic.ResolutionToClassifier(), KtAbstractFirDiagnostic<PsiElement>
 
+internal class AmbiguousAlteredAssignImpl(
+    override val altererNames: List<String?>,
+    override val firDiagnostic: KtPsiDiagnostic,
+    override val token: KtLifetimeToken,
+) : KtFirDiagnostic.AmbiguousAlteredAssign(), KtAbstractFirDiagnostic<PsiElement>
+
 internal class SuperIsNotAnExpressionImpl(
     override val firDiagnostic: KtPsiDiagnostic,
     override val token: KtLifetimeToken,

compiler/fir/checkers/checkers-component-generator/src/org/jetbrains/kotlin/fir/checkers/generator/diagnostics/FirDiagnosticsList.kt

@@ -142,6 +142,9 @@ object DIAGNOSTICS_LIST : DiagnosticList("FirErrors") {
         val RESOLUTION_TO_CLASSIFIER by error<PsiElement> {
             parameter<FirRegularClassSymbol>("classSymbol")
         }
+        val AMBIGUOUS_ALTERED_ASSIGN by error<PsiElement> {
+            parameter<List<String?>>("altererNames")
+        }
     }
 
     val SUPER by object : DiagnosticGroup("Super") {

compiler/fir/checkers/gen/org/jetbrains/kotlin/fir/analysis/diagnostics/FirErrors.kt

@@ -160,6 +160,7 @@ object FirErrors {
     val NO_RECEIVER_ALLOWED by error0<PsiElement>()
     val FUNCTION_EXPECTED by error2<PsiElement, String, ConeKotlinType>(SourceElementPositioningStrategies.REFERENCED_NAME_BY_QUALIFIED)
     val RESOLUTION_TO_CLASSIFIER by error1<PsiElement, FirRegularClassSymbol>()
+    val AMBIGUOUS_ALTERED_ASSIGN by error1<PsiElement, List<String?>>()
 
     // Super
     val SUPER_IS_NOT_AN_EXPRESSION by error0<PsiElement>(SourceElementPositioningStrategies.REFERENCED_NAME_BY_QUALIFIED)

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/FirErrorsDefaultMessages.kt

@@ -56,6 +56,7 @@ import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.ACTUAL_TYPE_ALIAS
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.ACTUAL_TYPE_ALIAS_WITH_USE_SITE_VARIANCE
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.ACTUAL_WITHOUT_EXPECT
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.AMBIGUOUS_ACTUALS
+import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.AMBIGUOUS_ALTERED_ASSIGN
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.AMBIGUOUS_ANONYMOUS_TYPE_INFERRED
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.AMBIGUOUS_CALL_WITH_IMPLICIT_CONTEXT_RECEIVER
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.AMBIGUOUS_EXPECTS
@@ -674,6 +675,12 @@ object FirErrorsDefaultMessages : BaseDiagnosticRendererFactory() {
             "Constructor of inner class {0} can be called only with receiver of containing class",
             SYMBOL
         )
+        map.put(
+            AMBIGUOUS_ALTERED_ASSIGN,
+            "Multiple extensions tried to alter this assignement at the same time. Extensions: {0}",
+            COLLECTION(NULLABLE_STRING)
+        )
+
         map.put(ILLEGAL_SELECTOR, "The expression cannot be a selector (occur after a dot)")
         map.put(NO_RECEIVER_ALLOWED, "No receiver can be passed to this function or property")
 

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/coneDiagnosticToFirDiagnostic.kt

@@ -133,6 +133,7 @@ private fun ConeDiagnostic.toKtDiagnostic(
     is ConeUnderscoreUsageWithoutBackticks -> FirErrors.UNDERSCORE_USAGE_WITHOUT_BACKTICKS.createOn(this.source)
     is ConeAmbiguousSuper -> FirErrors.AMBIGUOUS_SUPER.createOn(source, this.candidateTypes)
     is ConeUnresolvedParentInImport -> null // reported in FirUnresolvedImportChecker
+    is ConeAmbiguousAlteredAssign -> FirErrors.AMBIGUOUS_ALTERED_ASSIGN.createOn(source, this.altererNames)
     else -> throw IllegalArgumentException("Unsupported diagnostic type: ${this.javaClass}")
 }
 

compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/extensions/FirExtensionRegistrar.kt

@@ -29,6 +29,7 @@ abstract class FirExtensionRegistrar : FirExtensionRegistrarAdapter() {
             FirExpressionResolutionExtension::class,
             FirExtensionSessionComponent::class,
             FirSamConversionTransformerExtension::class,
+            FirAssignExpressionAltererExtension::class,
         )
     }
 
@@ -77,6 +78,11 @@ abstract class FirExtensionRegistrar : FirExtensionRegistrarAdapter() {
             registerExtension(FirSamConversionTransformerExtension::class, this)
         }
 
+        @JvmName("plusAssignExpressionAltererExtension")
+        operator fun (FirAssignExpressionAltererExtension.Factory).unaryPlus() {
+            registerExtension(FirAssignExpressionAltererExtension::class, this)
+        }
+
         // ------------------ reference methods ------------------
 
         @JvmName("plusStatusTransformerExtension")
@@ -119,6 +125,11 @@ abstract class FirExtensionRegistrar : FirExtensionRegistrarAdapter() {
             FirSamConversionTransformerExtension.Factory { this.invoke(it) }.unaryPlus()
         }
 
+        @JvmName("plusAssignExpressionAltererExtension")
+        operator fun ((FirSession) -> FirAssignExpressionAltererExtension).unaryPlus() {
+            FirAssignExpressionAltererExtension.Factory { this.invoke(it) }.unaryPlus()
+        }
+
         // ------------------ utilities ------------------
 
         @JvmName("bindLeft")

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirAssignExpressionAltererExtension.kt

@@ -0,0 +1,36 @@
+/*
+ * Copyright 2010-2022 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.extensions
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.expressions.FirStatement
+import org.jetbrains.kotlin.fir.expressions.FirVariableAssignment
+import kotlin.reflect.KClass
+
+abstract class FirAssignExpressionAltererExtension(session: FirSession) : FirExtension(session) {
+    companion object {
+        val NAME = FirExtensionPointName("AssignAlterer")
+    }
+
+    final override val name: FirExtensionPointName
+        get() = NAME
+
+    final override val extensionType: KClass<out FirExtension> = FirAssignExpressionAltererExtension::class
+
+    /**
+     * At this point [variableAssignment] contains resolved and completed lhs and calleeReference(lvalue)
+     *   and unresolved rValue expression
+     *
+     * It's allowed to transform [variableAssignment] into any kind of statement. This state should be unresolved
+     *   (modulo usages of already resolved parts, like lValue). Later this statement will be resolved by compiler
+     *   itself using regular resolution algorithms
+     */
+    abstract fun transformVariableAssignment(variableAssignment: FirVariableAssignment): FirStatement?
+
+    fun interface Factory : FirExtension.Factory<FirAssignExpressionAltererExtension>
+}
+
+val FirExtensionService.assignAltererExtensions: List<FirAssignExpressionAltererExtension> by FirExtensionService.registeredExtensions()

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

@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.builder.*
 import org.jetbrains.kotlin.fir.expressions.impl.FirResolvedArgumentList
 import org.jetbrains.kotlin.fir.expressions.impl.toAnnotationArgumentMapping
+import org.jetbrains.kotlin.fir.extensions.assignAltererExtensions
 import org.jetbrains.kotlin.fir.extensions.expressionResolutionExtensions
 import org.jetbrains.kotlin.fir.extensions.extensionService
 import org.jetbrains.kotlin.fir.references.*
@@ -60,6 +61,7 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
     var containingSafeCallExpression: FirSafeCallExpression? = null
 
     private val expressionResolutionExtensions = session.extensionService.expressionResolutionExtensions.takeIf { it.isNotEmpty() }
+    private val assignAltererExtensions = session.extensionService.assignAltererExtensions.takeIf { it.isNotEmpty() }
 
     init {
         @Suppress("LeakingThis")
@@ -845,6 +847,30 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
         val resolvedAssignment = callResolver.resolveVariableAccessAndSelectCandidate(variableAssignment, isUsedAsReceiver = false)
         val result = if (resolvedAssignment is FirVariableAssignment) {
             val completeAssignment = callCompleter.completeCall(resolvedAssignment, noExpectedType).result // TODO: check
+            val resolvedReference = completeAssignment.calleeReference
+            if (assignAltererExtensions != null && resolvedReference is FirResolvedNamedReference) {
+                val alteredAssignments = assignAltererExtensions.mapNotNull { alterer ->
+                    alterer.transformVariableAssignment(variableAssignment)?.let { it to alterer }
+                }
+                when (alteredAssignments.size) {
+                    0 -> {}
+                    1 -> {
+                        val transformedAssignment = alteredAssignments.first().first
+                        return transformedAssignment.transform(transformer, ResolutionMode.ContextIndependent)
+                    }
+
+                    else -> {
+                        val altererNames = alteredAssignments.map { it.second::class.qualifiedName }
+                        val errorReference = buildErrorNamedReference {
+                            source = resolvedReference.source
+                            candidateSymbol = resolvedReference.resolvedSymbol
+                            diagnostic = ConeAmbiguousAlteredAssign(altererNames)
+                        }
+                        completeAssignment.replaceCalleeReference(errorReference)
+                    }
+                }
+            }
+
             completeAssignment.transformRValue(
                 transformer,
                 withExpectedType(variableAssignment.lValueTypeRef, expectedTypeMismatchIsReportedInChecker = true),

compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/diagnostics/ConeDiagnostics.kt

@@ -16,7 +16,6 @@ import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.FirTypeRef
-import org.jetbrains.kotlin.fir.types.renderForDebugging
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.calls.tower.CandidateApplicability
@@ -233,3 +232,8 @@ private fun describeSymbol(symbol: FirBasedSymbol<*>): String {
         else -> "$symbol"
     }
 }
+
+class ConeAmbiguousAlteredAssign(val altererNames: List<String?>) : ConeDiagnostic {
+    override val reason: String
+        get() = "Assign altered by multiple extensions"
+}