[FIR] Introduce FirAssignExpressionAltererExtension

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

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),