From da7d48a0e13f82c84265d289592a3ccf8fecbd90 Mon Sep 17 00:00:00 2001
From: Roman Golyshev <roman.golyshev@jetbrains.com>
Date: Sun, 5 Jul 2026 19:43:03 +0200
Subject: [PATCH] KT-86909 [kotlin] Cleanup: pass `KaSession` as context
 parameter instead of extension receiver

GitOrigin-RevId: c2871483af3df855532aa6b7902f556e760330c2
---
 .../references/impl/KaBaseKDocReference.kt    |  2 +-
 .../idea/references/impl/KaBaseReference.kt   | 22 +++++++++++--------
 2 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/plugins/kotlin/base/analysis-api-platform/src/org/jetbrains/kotlin/idea/references/impl/KaBaseKDocReference.kt b/plugins/kotlin/base/analysis-api-platform/src/org/jetbrains/kotlin/idea/references/impl/KaBaseKDocReference.kt
index 52f687a4591b..73a1c4595ce7 100644
--- a/plugins/kotlin/base/analysis-api-platform/src/org/jetbrains/kotlin/idea/references/impl/KaBaseKDocReference.kt
+++ b/plugins/kotlin/base/analysis-api-platform/src/org/jetbrains/kotlin/idea/references/impl/KaBaseKDocReference.kt
@@ -29,7 +29,7 @@ internal class KaBaseKDocReference(element: KDocName) : KDocReference(element),
         referenceTargetSymbols.flatMap { symbol ->
             when (symbol) {
                 is KaSyntheticJavaPropertySymbol -> listOfNotNull(symbol.javaGetterSymbol.psi, symbol.javaSetterSymbol?.psi)
-                is KaSymbol -> getPsiDeclarations(symbol)
+                is KaSymbol -> symbol.getPsiDeclarations()
             }
         }
     }
diff --git a/plugins/kotlin/base/analysis-api-platform/src/org/jetbrains/kotlin/idea/references/impl/KaBaseReference.kt b/plugins/kotlin/base/analysis-api-platform/src/org/jetbrains/kotlin/idea/references/impl/KaBaseReference.kt
index b8649a78df0a..49fc3b4001e6 100644
--- a/plugins/kotlin/base/analysis-api-platform/src/org/jetbrains/kotlin/idea/references/impl/KaBaseReference.kt
+++ b/plugins/kotlin/base/analysis-api-platform/src/org/jetbrains/kotlin/idea/references/impl/KaBaseReference.kt
@@ -10,6 +10,8 @@ import org.jetbrains.kotlin.analysis.api.symbols.KaBackingFieldSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KaCallableSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KaSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KaSymbolOrigin
+import org.jetbrains.kotlin.analysis.api.symbols.containingSymbol
+import org.jetbrains.kotlin.analysis.api.symbols.intersectionOverriddenSymbols
 import org.jetbrains.kotlin.idea.references.KtReference
 import org.jetbrains.kotlin.idea.references.isConstructorOf
 import org.jetbrains.kotlin.idea.references.mainReference
@@ -29,7 +31,7 @@ internal sealed interface KaBaseReference : KtReference, KaResolvableReferenceBr
     fun getResolvedToPsi(analysisSession: KaSession, referenceTargetSymbols: Collection<KaSymbol>): Collection<PsiElement> =
         with(analysisSession) {
             referenceTargetSymbols.flatMap { symbol ->
-                getPsiDeclarations(symbol)
+                symbol.getPsiDeclarations()
             }
         }
 
@@ -68,22 +70,24 @@ internal sealed interface KaBaseReference : KtReference, KaResolvableReferenceBr
     }
 }
 
-internal fun KaSession.getPsiDeclarations(symbol: KaSymbol): Collection<PsiElement> {
+context(_: KaSession)
+internal fun KaSymbol.getPsiDeclarations(): Collection<PsiElement> {
     val intersectionOverriddenSymbolsOrSingle = when {
-        symbol.origin == KaSymbolOrigin.INTERSECTION_OVERRIDE && symbol is KaCallableSymbol -> symbol.intersectionOverriddenSymbols
-        else -> listOf(symbol)
+        this.origin == KaSymbolOrigin.INTERSECTION_OVERRIDE && this is KaCallableSymbol -> this.intersectionOverriddenSymbols
+        else -> listOf(this)
     }
 
-    return intersectionOverriddenSymbolsOrSingle.mapNotNull { unwrapPsi(it) }
+    return intersectionOverriddenSymbolsOrSingle.mapNotNull { it.unwrapPsi() }
 }
 
-private fun KaSession.unwrapPsi(symbol: KaSymbol): PsiElement? = when {
-    symbol is KaBackingFieldSymbol -> unwrapPsi(symbol.owningProperty)
+context(_: KaSession)
+private fun KaSymbol.unwrapPsi(): PsiElement? = when {
+    this is KaBackingFieldSymbol -> this.owningProperty.unwrapPsi()
 
     // For declarations generated by compiler plugins, `symbol.psi` is undefined and most probably is `null`.
     // To make the navigation work reasonably for the plugins which generate declarations inside existing
     // user-defined declarations, we use `containingSymbol` to navigate to the containing declaration.
-    symbol.origin == KaSymbolOrigin.PLUGIN && symbol.psi == null -> symbol.containingSymbol?.let { unwrapPsi(it) }
+    this.origin == KaSymbolOrigin.PLUGIN && this.psi == null -> this.containingSymbol?.unwrapPsi()
 
-    else -> symbol.psi
+    else -> this.psi
 }