KT-57468 Kotlin assignment plugin: operation name cannot be found

The problem results in broken import quick fix and import optimizer on
the IDE side [1].

`AssignResolutionAltererExtension` introduced a possibility to override
 resolution of assignment statements. The inconsistency though is
 that `KtSimpleNameReference.getResolvesByNames` doesn't return a name
 for the overridden `=`. Kotlin as a language doesn't support this [2].

This commit eliminates the drawback above:
1. It fixes the name `assign` the `=` can be resolved to [3].
   This eliminates the need to search for the name, bypassing the
   plugins.
2. `KtSimpleNameReference.getResolvesByNames` returns `assign` among
   other names in case it deals with binary `=` and assignment is
   resolved.
3. `KtCompilerPluginsProvider` was extended to check plugins' presence.
   K1 implementation added.

----------------------------------------------------------------
[1]: https://youtrack.jetbrains.com/issue/KTIJ-24390
[2]: OperatorConventions.getNameForOperationSymbol
     https://kotlinlang.org/docs/operator-overloading.html#augmented-assignments
[3]: OperatorConventions#ASSIGN_METHOD + AssignmentPluginNames

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/KtFe10CompilerPluginsProvider.kt

@@ -0,0 +1,28 @@
+/*
+ * Copyright 2010-2023 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.analysis.api.descriptors
+
+import org.jetbrains.kotlin.analysis.project.structure.KtCompilerPluginsProvider
+import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
+import org.jetbrains.kotlin.extensions.ProjectExtensionDescriptor
+import org.jetbrains.kotlin.extensions.internal.InternalNonStableExtensionPoints
+import org.jetbrains.kotlin.resolve.extensions.AssignResolutionAltererExtension
+
+@Suppress("unused")
+class KtFe10CompilerPluginsProvider : KtCompilerPluginsProvider() {
+    override fun <T : Any> getRegisteredExtensions(module: KtSourceModule, extensionType: ProjectExtensionDescriptor<T>): List<T> {
+        return extensionType.getInstances(module.project)
+    }
+
+    @OptIn(InternalNonStableExtensionPoints::class)
+    override fun isPluginOfTypeRegistered(module: KtSourceModule, pluginType: CompilerPluginType): Boolean {
+        val extension = when (pluginType) {
+            CompilerPluginType.ASSIGNMENT -> AssignResolutionAltererExtension
+            else -> return false
+        }
+        return extension.getInstances(module.project).isNotEmpty()
+    }
+}

analysis/analysis-api-fir/build.gradle.kts

@@ -15,6 +15,7 @@ dependencies {
     api(project(":compiler:fir:checkers:checkers.js"))
     api(project(":compiler:fir:checkers:checkers.native"))
     api(project(":compiler:fir:java"))
+    api(project(":compiler:fir:entrypoint"))
     api(project(":analysis:low-level-api-fir"))
     api(project(":analysis:analysis-api"))
     api(project(":analysis:analysis-api-impl-base"))

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/FirSimpleNameReference.kt

@@ -75,4 +75,4 @@ internal class KtFirSimpleNameReference(
         // TODO: Implement.
         return null
     }
-}
+}

analysis/kt-references/build.gradle.kts

@@ -7,6 +7,8 @@ dependencies {
     implementation(project(":compiler:psi"))
     implementation(project(":analysis:light-classes-base"))
     implementation(intellijCore())
+    implementation(project(":analysis:analysis-api-providers"))
+    implementation(project(":analysis:project-structure"))
 
     compileOnly(commonDependency("com.google.guava:guava"))
 }

analysis/kt-references/kt-references-fe10/src/org/jetbrains/kotlin/references/fe10/KtFe10SimpleNameReference.kt

@@ -7,14 +7,14 @@ package org.jetbrains.kotlin.references.fe10
 
 import com.intellij.psi.PsiElement
 import com.intellij.util.SmartList
-import org.jetbrains.kotlin.references.fe10.base.KtFe10Reference
-import org.jetbrains.kotlin.references.fe10.base.KtFe10ReferenceResolutionHelper
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.idea.references.KtSimpleNameReference
 import org.jetbrains.kotlin.idea.references.readWriteAccess
 import org.jetbrains.kotlin.load.java.descriptors.JavaPropertyDescriptor
 import org.jetbrains.kotlin.plugin.references.SimpleNameReferenceExtension
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.references.fe10.base.KtFe10Reference
+import org.jetbrains.kotlin.references.fe10.base.KtFe10ReferenceResolutionHelper
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.ImportedFromObjectCallableDescriptor
 import org.jetbrains.kotlin.resolve.descriptorUtil.getImportableDescriptor
@@ -98,4 +98,4 @@ class KtFe10SimpleNameReference(expression: KtSimpleNameExpression) : KtSimpleNa
         }
         return null
     }
-}
+}

analysis/kt-references/src/org/jetbrains/kotlin/idea/references/KtSimpleNameReference.kt

@@ -5,8 +5,12 @@
 
 package org.jetbrains.kotlin.idea.references
 
+import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.analysis.project.structure.KtCompilerPluginsProvider
+import org.jetbrains.kotlin.analysis.project.structure.KtCompilerPluginsProvider.CompilerPluginType
+import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
@@ -14,6 +18,9 @@ import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getParentOfTypeAndBranch
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
 import org.jetbrains.kotlin.types.expressions.OperatorConventions
+import org.jetbrains.kotlin.types.expressions.OperatorConventions.ASSIGN_METHOD
+import org.jetbrains.kotlin.utils.addToStdlib.runIf
+import org.jetbrains.kotlin.analysis.project.structure.getKtModule
 
 abstract class KtSimpleNameReference(expression: KtSimpleNameExpression) : KtSimpleReference<KtSimpleNameExpression>(expression) {
     // Extension point used by deprecated android extensions.
@@ -65,7 +72,12 @@ abstract class KtSimpleNameReference(expression: KtSimpleNameExpression) : KtSim
                 if (tokenType != null) {
                     val name = OperatorConventions.getNameForOperationSymbol(
                         tokenType, element.parent is KtUnaryExpression, element.parent is KtBinaryExpression
-                    ) ?: return emptyList()
+                    )
+                        ?: (expression.parent as? KtBinaryExpression)?.let {
+                            runIf(it.operationToken == KtTokens.EQ && isAssignmentResolved(element.project, it)) { ASSIGN_METHOD }
+                        }
+                        ?: return emptyList()
+
                     val counterpart = OperatorConventions.ASSIGNMENT_OPERATION_COUNTERPARTS[tokenType]
                     return if (counterpart != null) {
                         val counterpartName = OperatorConventions.getNameForOperationSymbol(counterpart, false, true)!!
@@ -80,4 +92,13 @@ abstract class KtSimpleNameReference(expression: KtSimpleNameExpression) : KtSim
         }
 
     abstract fun getImportAlias(): KtImportAlias?
-}
+
+    private fun isAssignmentResolved(project: Project, binaryExpression: KtBinaryExpression): Boolean {
+        val sourceModule = binaryExpression.getKtModule(element.project) as? KtSourceModule ?: return false
+        val reference = binaryExpression.operationReference.reference ?: return false
+        val pluginPresenceService = project.getService(KtCompilerPluginsProvider::class.java)
+            ?: error("KtAssignResolutionPresenceService is not available as a service")
+        return pluginPresenceService.isPluginOfTypeRegistered(sourceModule, CompilerPluginType.ASSIGNMENT)
+                && (reference.resolve() as? KtNamedFunction)?.nameAsName == ASSIGN_METHOD
+    }
+}

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/services/NoOpKtCompilerPluginsProvider.kt

@@ -12,4 +12,6 @@ import org.jetbrains.kotlin.extensions.ProjectExtensionDescriptor
 internal object NoOpKtCompilerPluginsProvider : KtCompilerPluginsProvider() {
     override fun <T : Any> getRegisteredExtensions(module: KtSourceModule, extensionType: ProjectExtensionDescriptor<T>): List<T> =
         emptyList()
+
+    override fun isPluginOfTypeRegistered(module: KtSourceModule, pluginType: CompilerPluginType): Boolean = false
 }

analysis/project-structure/src/org/jetbrains/kotlin/analysis/project/structure/KtCompilerPluginsProvider.kt

@@ -11,10 +11,20 @@ import org.jetbrains.kotlin.extensions.ProjectExtensionDescriptor
  * A service which can return extensions which are registered for some module
  */
 public abstract class KtCompilerPluginsProvider {
+
+    enum class CompilerPluginType {
+        ASSIGNMENT
+    }
+
     /**
      * Returns a list of extensions of a base [extensionType] which are registered for [module]
      *
      * These extensions are used in addition to those provided by the extension descriptor's [ProjectExtensionDescriptor.getInstances].
      */
     public abstract fun <T : Any> getRegisteredExtensions(module: KtSourceModule, extensionType: ProjectExtensionDescriptor<T>): List<T>
-}
+
+    /**
+     * Returns `true` if at least one plugin with requested `pluginType` is registered, `false` otherwise
+     */
+    public abstract fun isPluginOfTypeRegistered(module: KtSourceModule, pluginType: CompilerPluginType): Boolean
+}