Map types in PsiElementAssociations

The types themselves are not FIR elements and thus not visited by the visitor used to populate the associations map. However, these can be visited separately, so that they later can be retrieved using `type()`.
openrewrite · Nov 30, 2023 · 66b546a · 66b546a
1 parent 1f94c7e
commit 66b546a
Showing 1 changed file with 30 additions and 2 deletions.
diff --git a/src/main/kotlin/org/openrewrite/kotlin/internal/PsiElementAssociations.kt b/src/main/kotlin/org/openrewrite/kotlin/internal/PsiElementAssociations.kt
@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.fir.declarations.impl.FirDefaultPropertySetter
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.impl.FirElseIfTrueCondition
 import org.jetbrains.kotlin.fir.expressions.impl.FirSingleExpressionBlock
+import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.references.FirErrorNamedReference
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.references.resolved
@@ -32,8 +33,7 @@ import org.jetbrains.kotlin.fir.symbols.impl.FirConstructorSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
-import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
-import org.jetbrains.kotlin.fir.types.FirUserTypeRef
+import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.visitors.FirDefaultVisitor
 import org.jetbrains.kotlin.psi
 import org.jetbrains.kotlin.psi.*
@@ -43,6 +43,7 @@ import org.openrewrite.kotlin.KotlinTypeMapping
 class PsiElementAssociations(val typeMapping: KotlinTypeMapping, val file: FirFile) {
 
     private val elementMap: MutableMap<PsiElement, MutableList<FirInfo>> = HashMap()
+    private val typeMap: MutableMap<PsiElement, ConeTypeProjection> = HashMap()
 
     fun initialize() {
         var depth = 0
@@ -64,10 +65,37 @@ class PsiElementAssociations(val typeMapping: KotlinTypeMapping, val file: FirFi
                 }
                 depth--
             }
+
+            override fun visitResolvedTypeRef(
+                resolvedTypeRef: FirResolvedTypeRef,
+                data: MutableMap<PsiElement, MutableList<FirInfo>>
+            ) {
+                super.visitResolvedTypeRef(resolvedTypeRef, data)
+                if (resolvedTypeRef.type is ConeClassLikeType && resolvedTypeRef.type.typeArguments.isNotEmpty() && resolvedTypeRef.psi is KtTypeReference) {
+                    visitType(resolvedTypeRef.type, resolvedTypeRef.psi as KtTypeReference)
+                }
+            }
+
+            private fun visitType(firType: ConeTypeProjection, psiType: KtTypeReference) {
+                if (firType is ConeClassLikeType) {
+                    val psiTypeArguments = psiType.typeElement!!.typeArgumentsAsTypes
+                    for ((index, typeArgument) in firType.typeArguments.withIndex()) {
+                        visitType(typeArgument, psiTypeArguments[index])
+                        typeMap[psiTypeArguments[index]] = typeArgument
+                    }
+                }
+            }
         }.visitFile(file, elementMap)
     }
 
     fun type(psiElement: PsiElement, owner: FirElement?): JavaType? {
+        if (typeMap.isNotEmpty() && psiElement is KtNameReferenceExpression) {
+            // TODO can / should we make this more generic?
+            val type = typeMap[psiElement.parent.parent]
+            if (type != null) {
+                return typeMapping.type(type, owner)
+            }
+        }
         val fir = primary(psiElement)
         return if (fir != null) typeMapping.type(fir, owner) else null
     }