Set constructor type on enum values.

openrewrite · Dec 12, 2023 · 2c1b918 · 2c1b918
1 parent c992df9
commit 2c1b918
Show file tree

Hide file tree

Showing 4 changed files with 55 additions and 4 deletions.
diff --git a/src/main/java/org/openrewrite/kotlin/internal/KotlinTreeParserVisitor.java b/src/main/java/org/openrewrite/kotlin/internal/KotlinTreeParserVisitor.java
@@ -494,14 +494,13 @@ public J visitEnumEntry(KtEnumEntry enumEntry, ExecutionContext data) {
 
         Set<PsiElement> consumedSpaces = preConsumedInfix(enumEntry);
 
-        // TODO: in java the EnumValue has a type of JavaType.Variable with a null fieldType.
         J.Identifier name = createIdentifier(requireNonNull(enumEntry.getNameIdentifier()), type(enumEntry), consumedSpaces);
         J.NewClass initializer = null;
 
+        JavaType.Method mt = methodDeclarationType(enumEntry);
         if (enumEntry.getInitializerList() != null) {
-            // TODO: this creates an empty init with no type attribution: enum class Foo { BAR() }
-            //   Add constructor method type.
             initializer = (J.NewClass) enumEntry.getInitializerList().accept(this, data);
+            initializer = initializer.withMethodType(mt);
         }
 
         if (enumEntry.getBody() != null) {
@@ -522,7 +521,7 @@ public J visitEnumEntry(KtEnumEntry enumEntry, ExecutionContext data) {
                         null,
                         args,
                         body,
-                        null
+                        mt
                 );
             }
         }

diff --git a/src/main/kotlin/org/openrewrite/kotlin/KotlinTypeMapping.kt b/src/main/kotlin/org/openrewrite/kotlin/KotlinTypeMapping.kt
@@ -28,6 +28,7 @@ import org.jetbrains.kotlin.fir.analysis.checkers.modality
 import org.jetbrains.kotlin.fir.analysis.checkers.toRegularClassSymbol
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.impl.FirOuterClassTypeParameterRef
+import org.jetbrains.kotlin.fir.declarations.impl.FirPrimaryConstructor
 import org.jetbrains.kotlin.fir.declarations.utils.isLocal
 import org.jetbrains.kotlin.fir.declarations.utils.isStatic
 import org.jetbrains.kotlin.fir.declarations.utils.modality
@@ -482,6 +483,23 @@ class KotlinTypeMapping(
         return clazz
     }
 
+    @OptIn(SymbolInternals::class)
+    fun methodDeclarationType(enumEntry: FirEnumEntry): Method? {
+        val type = when (val fir = enumEntry.symbol.getContainingClassSymbol(firSession)?.fir) {
+            is FirClass -> {
+                when (val primary = fir.declarations.firstOrNull { it is FirPrimaryConstructor }) {
+                    is FirPrimaryConstructor -> type(primary as FirFunction)
+                    else -> null
+                }
+            }
+            else -> null
+        }
+        return when (type) {
+            is Method -> type
+            else -> null
+        }
+    }
+
     fun methodDeclarationType(function: FirFunction, parent: Any?): Method {
         val signature = signatureBuilder.methodSignature(function, parent)
         val existing = typeCache.get<Method>(signature)

diff --git a/src/main/kotlin/org/openrewrite/kotlin/internal/PsiElementAssociations.kt b/src/main/kotlin/org/openrewrite/kotlin/internal/PsiElementAssociations.kt
@@ -159,6 +159,7 @@ class PsiElementAssociations(val typeMapping: KotlinTypeMapping, val file: FirFi
 
     fun methodDeclarationType(psi: PsiElement): JavaType.Method? {
         return when (val fir = primary(psi)) {
+            is FirEnumEntry -> typeMapping.methodDeclarationType(fir)
             is FirFunction -> typeMapping.methodDeclarationType(fir, null)
             is FirAnonymousFunctionExpression -> typeMapping.methodDeclarationType(fir.anonymousFunction, null)
             else -> {

diff --git a/src/test/java/org/openrewrite/kotlin/KotlinTypeMappingTest.java b/src/test/java/org/openrewrite/kotlin/KotlinTypeMappingTest.java
@@ -1462,5 +1462,38 @@ public J.Identifier visitIdentifier(J.Identifier identifier, Integer integer) {
               )
             );
         }
+
+        @Issue("https://github.com/openrewrite/rewrite-kotlin/issues/533")
+        @Test
+        void enumConstructorType() {
+            rewriteRun(
+              kotlin(
+                """
+                  enum class Code {
+                      YES ,
+                  }
+                  enum class Test ( val arg: Code , ) {
+                      FOO ( Code.YES , ) {
+                          // Body is required to reproduce issue
+                      }
+                  }
+                  """, spec -> spec.afterRecipe(cu -> {
+                    AtomicBoolean found = new AtomicBoolean(false);
+                    new KotlinIsoVisitor<Integer>() {
+                        @Override
+                        public J.EnumValue visitEnumValue(J.EnumValue _enum, Integer integer) {
+                            if ("FOO".equals(_enum.getName().getSimpleName())) {
+                                assertThat(_enum.getInitializer().getConstructorType().toString())
+                                  .isEqualTo("Test{name=<constructor>,return=Test,parameters=[Code]}");
+                                found.set(true);
+                            }
+                            return super.visitEnumValue(_enum, integer);
+                        }
+                    }.visit(cu, 0);
+                    assertThat(found.get()).isTrue();
+                })
+              )
+            );
+        }
     }
 }