From ba421aa320a3992cd7f1e09ebd4b61cea0f4c621 Mon Sep 17 00:00:00 2001
From: Dmitrii Gridin <dmitry.gridin@jetbrains.com>
Date: Tue, 12 Mar 2024 17:08:40 +0100
Subject: [PATCH] [Analysis API FIR] KtFirPsiTypeProvider: support type
 parameter position in `asKtType`

^KT-66530
---
 .../api/fir/components/KtFirPsiTypeProvider.kt       | 12 +++++++++---
 .../methodTypeParameterWithTypeParameterUseSite.txt  |  2 +-
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirPsiTypeProvider.kt b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirPsiTypeProvider.kt
index 8827baa12f3..3b3d8167302 100644
--- a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirPsiTypeProvider.kt
+++ b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirPsiTypeProvider.kt
@@ -58,8 +58,10 @@ import org.jetbrains.kotlin.platform.jvm.JvmPlatform
 import org.jetbrains.kotlin.psi
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.psiUtil.parents
+import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
 import org.jetbrains.kotlin.types.model.SimpleTypeMarker
 import org.jetbrains.kotlin.types.updateArgumentModeFromAnnotations
+import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 import java.text.StringCharacterIterator
 
 internal class KtFirPsiTypeProvider(
@@ -144,7 +146,7 @@ internal class KtFirPsiTypeProvider(
         val javaTypeParameterStack = MutableJavaTypeParameterStack()
 
         var psiClass = PsiTreeUtil.getContextOfType(useSitePosition, PsiClass::class.java, false)
-        while (psiClass != null && psiClass.name == null) {
+        while (psiClass != null && psiClass.name == null || psiClass is PsiTypeParameter) {
             psiClass = PsiTreeUtil.getContextOfType(psiClass, PsiClass::class.java, true)
         }
         if (psiClass != null) {
@@ -161,9 +163,13 @@ internal class KtFirPsiTypeProvider(
                     ),
                     javaClass
                 )
+
                 if (containingClassSymbol != null) {
-                    val member =
-                        PsiTreeUtil.getContextOfType(useSitePosition, PsiTypeParameterListOwner::class.java, false, PsiClass::class.java)
+                    val member = useSitePosition.parentsWithSelf
+                        .filterNot { it is PsiTypeParameter }
+                        .takeWhile { it !is PsiClass }
+                        .firstIsInstanceOrNull<PsiTypeParameterListOwner>()
+
                     if (member != null) {
                         val memberSymbol = containingClassSymbol.declarationSymbols.find { it.findPsi() == member } as? FirCallableSymbol<*>
                         if (memberSymbol != null) {
diff --git a/analysis/analysis-api/testData/components/psiTypeProvider/psiType/asKtType/methodTypeParameterWithTypeParameterUseSite.txt b/analysis/analysis-api/testData/components/psiTypeProvider/psiType/asKtType/methodTypeParameterWithTypeParameterUseSite.txt
index 13ef5fe42c5..d0d2edf5c71 100644
--- a/analysis/analysis-api/testData/components/psiTypeProvider/psiType/asKtType/methodTypeParameterWithTypeParameterUseSite.txt
+++ b/analysis/analysis-api/testData/components/psiTypeProvider/psiType/asKtType/methodTypeParameterWithTypeParameterUseSite.txt
@@ -1,2 +1,2 @@
 PsiType: PsiType:TYPE
-KtType: TYPE
+KtType: TYPE?