From 137a78a00eaf9d7cbd342b465edb1ad3a91189ab Mon Sep 17 00:00:00 2001
From: Roman Golyshev <roman.golyshev@jetbrains.com>
Date: Tue, 30 Jul 2024 15:30:10 +0200
Subject: [PATCH] KTIJ-26752 [kotlin] Introduce `UsedReference` and
 `UsedSymbol` classes to simplify `UsedReferencesCollector`

GitOrigin-RevId: c8d332abe0f717c68d7937bcd88e026f14f6f9d4
---
 .../imports/UsedReferencesCollector.kt        | 66 ++++++++++++++-----
 1 file changed, 48 insertions(+), 18 deletions(-)

diff --git a/plugins/kotlin/code-insight/kotlin.code-insight.k2/src/org/jetbrains/kotlin/idea/k2/codeinsight/imports/UsedReferencesCollector.kt b/plugins/kotlin/code-insight/kotlin.code-insight.k2/src/org/jetbrains/kotlin/idea/k2/codeinsight/imports/UsedReferencesCollector.kt
index 80468f1225db..9753f13b9f42 100644
--- a/plugins/kotlin/code-insight/kotlin.code-insight.k2/src/org/jetbrains/kotlin/idea/k2/codeinsight/imports/UsedReferencesCollector.kt
+++ b/plugins/kotlin/code-insight/kotlin.code-insight.k2/src/org/jetbrains/kotlin/idea/k2/codeinsight/imports/UsedReferencesCollector.kt
@@ -64,25 +64,15 @@ internal class UsedReferencesCollector(private val file: KtFile) {
         if (element is KtLabelReferenceExpression) return
 
         val references = element.references
+            .filterIsInstance<KtReference>()
+            .mapNotNull { UsedReference.run { createFrom(it) } }
+
         if (references.isEmpty()) return
 
         for (reference in references) {
             ProgressIndicatorProvider.checkCanceled()
 
-            if (reference !is KtReference) continue
-            if (isEmptyInvokeReference(reference)) continue
-
-            val resolvedSymbols = reference.resolveToSymbols()
-            val isResolved = resolvedSymbols.isNotEmpty()
-
-            ProgressIndicatorProvider.checkCanceled()
-
-            val dispatchReceiver = resolveDispatchReceiver(element)
-            if (dispatchReceiver != null && isDispatchedCall(element, dispatchReceiver, file)) {
-                continue
-            }
-
-            ProgressIndicatorProvider.checkCanceled()
+            val isResolved = reference.run { isResolved() }
 
             val names = reference.resolvesByNames
             if (!isResolved) {
@@ -90,12 +80,12 @@ internal class UsedReferencesCollector(private val file: KtFile) {
                 continue
             }
 
-            val importableSymbols = resolvedSymbols.mapNotNull { toImportableSymbol(it, reference, file) }
+            val symbols = reference.run { resolveToImportableSymbols() }
 
-            for (symbol in importableSymbols) {
-                if (!canBeResolvedViaImport(reference, symbol)) continue
+            for (symbol in symbols) {
+                if (!symbol.run { isResolvedWithImport() }) continue
 
-                val importableName = computeImportableName(symbol, dispatchReceiver as? KaImplicitReceiverValue) ?: continue
+                val importableName = symbol.run { computeImportableFqName() } ?: continue
 
                 ProgressIndicatorProvider.checkCanceled()
 
@@ -106,6 +96,46 @@ internal class UsedReferencesCollector(private val file: KtFile) {
     }
 }
 
+internal class UsedReference private constructor(val reference: KtReference) {
+    val resolvesByNames: Collection<Name>
+        get() = reference.resolvesByNames
+
+    fun KaSession.isResolved(): Boolean {
+        val resolvedSymbols = reference.resolveToSymbols()
+        val isResolved = resolvedSymbols.isNotEmpty()
+
+        return isResolved
+    }
+
+    fun KaSession.resolveToImportableSymbols(): Collection<UsedSymbol> {
+        return reference.resolveToSymbols().mapNotNull { toImportableSymbol(it, reference) }.map { UsedSymbol(reference, it) }
+    }
+
+    companion object {
+        fun KaSession.createFrom(reference: KtReference): UsedReference? {
+            if (isEmptyInvokeReference(reference)) return null
+
+            return UsedReference(reference)
+        }
+    }
+}
+
+internal class UsedSymbol(val reference: KtReference, val symbol: KaSymbol) {
+    val name: Name?
+        get() = symbol.name
+
+    fun KaSession.computeImportableFqName(): FqName? {
+        return computeImportableName(symbol, resolveDispatchReceiver(reference.element) as? KaImplicitReceiverValue?)
+    }
+
+    fun KaSession.isResolvedWithImport(): Boolean {
+        val dispatchReceiver = resolveDispatchReceiver(reference.element)
+        if (dispatchReceiver != null && isDispatchedCall(reference.element, dispatchReceiver)) return false
+
+        return canBeResolvedViaImport(reference, symbol)
+    }
+}
+
 /**
  * In K2, every call in the form of `foo()` has `KtInvokeFunctionReference` on it.
  *