Fixed types on TypeParameter names.

Updated fir() to return the correctly associated fir in PsiElementAssociations
openrewrite · Nov 22, 2023 · 7ec070e · 7ec070e
1 parent 075e40a
commit 7ec070e
Show file tree

Hide file tree

Showing 3 changed files with 94 additions and 9 deletions.
diff --git a/src/main/java/org/openrewrite/kotlin/internal/KotlinTreeParserVisitor.java b/src/main/java/org/openrewrite/kotlin/internal/KotlinTreeParserVisitor.java
@@ -1303,6 +1303,8 @@ public J visitTypeArgumentList(KtTypeArgumentList typeArgumentList, ExecutionCon
     public J visitTypeConstraint(KtTypeConstraint constraint, ExecutionContext data) {
         List<J.Annotation> annotations = new ArrayList<>();
         J.Identifier typeParamName = (J.Identifier) requireNonNull(constraint.getSubjectTypeParameterName()).accept(this, data);
+        PsiElement ref = PsiTreeUtil.getChildOfType(constraint, KtTypeReference.class);
+        typeParamName = typeParamName.withType(psiElementAssociations.type(ref, owner(constraint)));
         TypeTree typeTree = (TypeTree) requireNonNull(constraint.getBoundTypeReference()).accept(this, data);
 
         return new J.TypeParameter(

diff --git a/src/main/kotlin/org/openrewrite/kotlin/internal/PsiElementAssociations.kt b/src/main/kotlin/org/openrewrite/kotlin/internal/PsiElementAssociations.kt
@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirPackageDirective
 import org.jetbrains.kotlin.fir.declarations.*
+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
@@ -55,8 +56,11 @@ class PsiElementAssociations(val typeMapping: KotlinTypeMapping, val file: FirFi
                 depth++
                 element.acceptChildren(this, data)
                 if (element is FirResolvedTypeRef) {
-                    // not sure why this isn't taken care of by `FirResolvedTypeRefImpl#acceptChildren()`
-                    element.delegatedTypeRef?.accept(this, data)
+                    // Do not visit FirUserTypeRef, since it's not mappable to a type.
+                    if (element.delegatedTypeRef != null && element.delegatedTypeRef !is FirUserTypeRef) {
+                        // not sure why this isn't taken care of by `FirResolvedTypeRefImpl#acceptChildren()`
+                        element.delegatedTypeRef?.accept(this, data)
+                    }
                 }
                 depth--
             }
@@ -69,12 +73,7 @@ class PsiElementAssociations(val typeMapping: KotlinTypeMapping, val file: FirFi
     }
 
     fun primary(psiElement: PsiElement) =
-        fir(psiElement) { filterFirElement(it) }
-
-    private fun filterFirElement(firElement : FirElement) : Boolean {
-        return firElement.source is KtRealPsiSourceElement &&
-                firElement !is FirUserTypeRef
-    }
+        fir(psiElement) { it.source is KtRealPsiSourceElement }
 
     fun methodDeclarationType(psi: PsiElement): JavaType.Method? {
         return when (val fir = primary(psi)) {
@@ -173,7 +172,43 @@ class PsiElementAssociations(val typeMapping: KotlinTypeMapping, val file: FirFi
         val directFirInfos = allFirInfos.filter { filter.invoke(it.fir) }
         return if (directFirInfos.isNotEmpty())
             // It might be more reliable to have explicit mappings in case something changes.
-            directFirInfos[0].fir
+            return when {
+                directFirInfos.size == 1 -> directFirInfos[0].fir
+                else -> {
+                    return when (p) {
+                        is KtConstantExpression -> {
+                            directFirInfos.firstOrNull { it.fir is FirConstExpression<*> }?.fir
+                        }
+                        is KtImportDirective -> {
+                            directFirInfos.firstOrNull { it.fir is FirImport && it.fir !is FirErrorImport }?.fir
+                        }
+                        is KtNamedFunction -> {
+                            val found = directFirInfos.firstOrNull { it.fir is FirFunction }?.fir
+//                            if (found == null) {
+//                                // Review how to expose unmatched types without causing an error.
+//                            }
+                            found
+                        }
+                        is KtNameReferenceExpression, is KtTypeReference -> {
+                            val found = directFirInfos.firstOrNull { it.fir is FirResolvedTypeRef || it.fir is FirResolvedNamedReference }?.fir
+//                            if (found == null) {
+//                                // Review how to expose unmatched types without causing an error.
+//                            }
+                            found
+                        }
+                        is KtPropertyAccessor -> {
+                            val found = directFirInfos.firstOrNull { it.fir is FirDefaultPropertySetter }?.fir
+//                            if (found == null) {
+//                                // Review how to expose unmatched types without causing an error.
+//                            }
+                            found
+                        }
+                        else -> {
+                            directFirInfos[0].fir
+                        }
+                    }
+                }
+            }
         else if (allFirInfos.isNotEmpty()) {
             return when {
                 allFirInfos.size == 1 -> allFirInfos[0].fir

diff --git a/src/test/java/org/openrewrite/kotlin/KotlinTypeMappingTest.java b/src/test/java/org/openrewrite/kotlin/KotlinTypeMappingTest.java
@@ -40,6 +40,7 @@
 import static java.util.Collections.singletonList;
 import static java.util.Objects.requireNonNull;
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.in;
 import static org.openrewrite.ExecutionContext.REQUIRE_PRINT_EQUALS_INPUT;
 import static org.openrewrite.java.tree.JavaType.GenericTypeVariable.Variance.*;
 import static org.openrewrite.kotlin.Assertions.kotlin;
@@ -863,6 +864,53 @@ class A : RemoteStub()
             );
         }
 
+        @Issue("https://github.com/openrewrite/rewrite-kotlin/issues/461")
+        @Test
+        void multipleBounds() {
+            rewriteRun(
+              kotlin(
+                """
+                  interface A
+                  interface B
+                  interface C
+                  interface D
+                  
+                  class KotlinTypeGoat<T, S>  where   S : A, T : D, S : B, T : C
+                  """, spec -> spec.afterRecipe(cu -> {
+                      AtomicBoolean found = new AtomicBoolean(false);
+                      new KotlinIsoVisitor<AtomicBoolean>() {
+                          @Override
+                          public K.ClassDeclaration visitClassDeclaration(K.ClassDeclaration classDeclaration, AtomicBoolean atomicBoolean) {
+                              if ("KotlinTypeGoat".equals(classDeclaration.getClassDeclaration().getSimpleName())) {
+                                  List<J.TypeParameter> constraints = classDeclaration.getTypeConstraints().getConstraints();
+                                  for (int i = 0; i < constraints.size(); i++) {
+                                      J.TypeParameter p = constraints.get(i);
+                                      switch (i) {
+                                          case 0:
+                                              assertThat(p.getName().getType().toString()).isEqualTo("A");
+                                              break;
+                                          case 1:
+                                              assertThat(p.getName().getType().toString()).isEqualTo("D");
+                                              break;
+                                          case 2:
+                                              assertThat(p.getName().getType().toString()).isEqualTo("B");
+                                              break;
+                                          case 3:
+                                              assertThat(p.getName().getType().toString()).isEqualTo("C");
+                                              found.set(true);
+                                              break;
+                                      }
+                                  }
+                              }
+                              return super.visitClassDeclaration(classDeclaration, atomicBoolean);
+                          }
+                    }.visit(cu, found);
+                    assertThat(found.get()).isTrue();
+                })
+              )
+            );
+        }
+
         @Test
         void nullJavaClassifierType() {
             rewriteRun(