From 7a3e90f9075fe1162fbf85b9060d201fa5bf3499 Mon Sep 17 00:00:00 2001
From: Jaebaek Seo <jaebaek@google.com>
Date: Fri, 8 Jul 2022 19:00:49 +0000
Subject: [PATCH] FIR find usage: separate FirSimpleNameReference objects for
 read/write

FirSimpleNameReference class handles the references to the Java
synthetic property. It causes an issue in `processReference(..)` method
of `MethodTextOccurrenceProcessor` when the property is used by
read/write.  Since read/write has two results of the reference resolution
i.e., getter for read and setter for write, `ref.solve()` in the method
returns `null` because it uses `resolve()` method in
PsiPolyVariantReferenceBase that returns null when it has multiple
results. This commit separates FirSimpleNameReference objects for
read/write to fix the issue.
---
 .../fir/references/FirSimpleNameReference.kt    | 10 ++++------
 .../references/KotlinFirReferenceContributor.kt | 17 ++++++++++++++++-
 2 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/FirSimpleNameReference.kt b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/FirSimpleNameReference.kt
index 278de6c9cc1..ef0bed2ca76 100644
--- a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/FirSimpleNameReference.kt
+++ b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/FirSimpleNameReference.kt
@@ -16,7 +16,8 @@ import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.resolve.references.ReferenceAccess
 
 internal class KtFirSimpleNameReference(
-    expression: KtSimpleNameExpression
+    expression: KtSimpleNameExpression,
+    val isRead: Boolean,
 ) : KtSimpleNameReference(expression), KtFirReference {
 
     private val isAnnotationCall: Boolean
@@ -51,11 +52,8 @@ internal class KtFirSimpleNameReference(
         referenceTargetSymbols.flatMap { symbol ->
             when (symbol) {
                 is KtFirSyntheticJavaPropertySymbol ->
-                    when (expression.readWriteAccess(useResolveForReadWrite = true)) {
-                        ReferenceAccess.READ -> listOfNotNull(symbol.javaGetterSymbol.psi)
-                        ReferenceAccess.WRITE -> listOfNotNull(symbol.javaSetterSymbol?.psi)
-                        ReferenceAccess.READ_WRITE -> listOfNotNull(symbol.javaGetterSymbol.psi, symbol.javaSetterSymbol?.psi)
-                    }
+                    if (isRead) listOfNotNull(symbol.javaGetterSymbol.psi)
+                    else listOfNotNull(symbol.javaSetterSymbol?.psi)
                 else -> listOfNotNull(symbol.psi)
             }
         }
diff --git a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/KotlinFirReferenceContributor.kt b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/KotlinFirReferenceContributor.kt
index 10645d520c3..3609d06628f 100644
--- a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/KotlinFirReferenceContributor.kt
+++ b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/KotlinFirReferenceContributor.kt
@@ -5,10 +5,14 @@
 
 package org.jetbrains.kotlin.idea.references
 
+import com.intellij.psi.PsiReference
+import org.jetbrains.kotlin.psi.KtNameReferenceExpression
+import org.jetbrains.kotlin.psi.KtSimpleNameExpression
+import org.jetbrains.kotlin.resolve.references.ReferenceAccess
+
 class KotlinFirReferenceContributor : KotlinReferenceProviderContributor {
     override fun registerReferenceProviders(registrar: KotlinPsiReferenceRegistrar) {
         with(registrar) {
-            registerProvider(factory = ::KtFirSimpleNameReference)
             registerProvider(factory = ::KtFirForLoopInReference)
             registerProvider(factory = ::KtFirInvokeFunctionReference)
             registerProvider(factory = ::KtFirPropertyDelegationMethodsReference)
@@ -16,6 +20,17 @@ class KotlinFirReferenceContributor : KotlinReferenceProviderContributor {
             registerProvider(factory = ::KtFirArrayAccessReference)
             registerProvider(factory = ::KtFirConstructorDelegationReference)
             registerProvider(factory = ::KtFirCollectionLiteralReference)
+
+            registerMultiProvider<KtSimpleNameExpression> { nameReferenceExpression ->
+                when (nameReferenceExpression.readWriteAccess(useResolveForReadWrite = true)) {
+                    ReferenceAccess.READ -> arrayOf(KtFirSimpleNameReference(nameReferenceExpression, isRead = true))
+                    ReferenceAccess.WRITE -> arrayOf(KtFirSimpleNameReference(nameReferenceExpression, isRead = false))
+                    ReferenceAccess.READ_WRITE -> arrayOf(
+                        KtFirSimpleNameReference(nameReferenceExpression, isRead = true),
+                        KtFirSimpleNameReference(nameReferenceExpression, isRead = false),
+                    )
+                }
+            }
         }
     }
 }