feat: support type parameters in ImplementInterface visitor (#2194)

Co-authored-by: Thomas Zub <t.zub@lvm.de>

Author: thomaszub

rewrite-java-tck/src/main/kotlin/org/openrewrite/java/ImplementInterfaceTest.kt

@@ -20,6 +20,10 @@ import org.openrewrite.ExecutionContext
 import org.openrewrite.Recipe
 import org.openrewrite.TreeVisitor
 import org.openrewrite.java.tree.J
+import org.openrewrite.java.tree.JavaType.ShallowClass
+import org.openrewrite.java.tree.Space
+import org.openrewrite.marker.Markers
+import java.util.*
 
 interface ImplementInterfaceTest : JavaRecipeTest {
     override val recipe: Recipe?
@@ -34,6 +38,41 @@ interface ImplementInterfaceTest : JavaRecipeTest {
             }
         }
 
+    val recipeTyped: Recipe
+        get() = object : TestRecipe() {
+            override fun getVisitor(): TreeVisitor<*, ExecutionContext> {
+                return object : JavaVisitor<ExecutionContext>() {
+                    override fun visitClassDeclaration(classDecl: J.ClassDeclaration, ctx: ExecutionContext): J {
+                        doAfterVisit(
+                            ImplementInterface(
+                                classDecl,
+                                "b.B",
+                                listOf(
+                                    J.Identifier(
+                                        UUID.randomUUID(),
+                                        Space.EMPTY,
+                                        Markers.EMPTY,
+                                        "String",
+                                        ShallowClass.build("java.lang.String"),
+                                        null
+                                    ),
+                                    J.Identifier(
+                                        UUID.randomUUID(),
+                                        Space.build(" ", emptyList()),
+                                        Markers.EMPTY,
+                                        "LocalDate",
+                                        ShallowClass.build("java.time.LocalDate"),
+                                        null
+                                    )
+                                )
+                            )
+                        )
+                        return classDecl
+                    }
+                }
+            }
+        }
+
     companion object {
         private const val b = "package b;\npublic interface B {}"
         private const val c = "package c;\npublic interface C {}"
@@ -73,4 +112,26 @@ interface ImplementInterfaceTest : JavaRecipeTest {
             }
         """
     )
+
+    @Test
+    fun addAnImplementedInterfaceWithTypeParameters(jp: JavaParser) = assertChanged(
+        jp,
+        recipe = recipeTyped,
+        dependsOn = arrayOf(b, c),
+        before = """
+            import c.C;
+            
+            class A implements C {
+            }
+        """,
+        after = """
+            import b.B;
+            import c.C;
+            
+            import java.time.LocalDate;
+            
+            class A implements C, B<String, LocalDate> {
+            }
+        """
+    )
 }

rewrite-java/src/main/java/org/openrewrite/java/ImplementInterface.java

@@ -16,23 +16,38 @@
 package org.openrewrite.java;
 
 import org.openrewrite.internal.ListUtils;
+import org.openrewrite.internal.lang.Nullable;
 import org.openrewrite.java.tree.*;
 import org.openrewrite.marker.Markers;
 
+import java.util.List;
+import java.util.Objects;
+import java.util.stream.Collectors;
+
 import static org.openrewrite.Tree.randomId;
 import static org.openrewrite.java.tree.Space.format;
 
 public class ImplementInterface<P> extends JavaIsoVisitor<P> {
     private final J.ClassDeclaration scope;
     private final JavaType.FullyQualified interfaceType;
+    private final @Nullable List<Expression> typeParameters;
 
-    public ImplementInterface(J.ClassDeclaration scope, JavaType.FullyQualified interfaceType) {
+    public ImplementInterface(J.ClassDeclaration scope, JavaType.FullyQualified interfaceType, @Nullable List<Expression> typeParameters) {
         this.scope = scope;
         this.interfaceType = interfaceType;
+        this.typeParameters = typeParameters;
+    }
+
+    public ImplementInterface(J.ClassDeclaration scope, String interfaze, @Nullable List<Expression> typeParameters) {
+        this(scope, JavaType.ShallowClass.build(interfaze), typeParameters);
+    }
+
+    public ImplementInterface(J.ClassDeclaration scope, JavaType.FullyQualified interfaceType) {
+        this(scope, interfaceType, null);
     }
 
     public ImplementInterface(J.ClassDeclaration scope, String interfaze) {
-        this(scope, JavaType.ShallowClass.build(interfaze));
+        this(scope, interfaze, null);
     }
 
     @Override
@@ -41,7 +56,7 @@ public J.ClassDeclaration visitClassDeclaration(J.ClassDeclaration classDecl, P
         if (c.isScope(scope) && (c.getImplements() == null || c.getImplements().stream()
                 .noneMatch(f -> TypeUtils.isAssignableTo(f.getType(), interfaceType)))) {
 
-            if(!classDecl.getSimpleName().equals(interfaceType.getClassName())) {
+            if (!classDecl.getSimpleName().equals(interfaceType.getClassName())) {
                 maybeAddImport(interfaceType);
             }
 
@@ -50,7 +65,29 @@ public J.ClassDeclaration visitClassDeclaration(J.ClassDeclaration classDecl, P
                     .withType(interfaceType)
                     .withPrefix(format(" "));
 
-            c = c.withImplements(ListUtils.concat(c.getImplements(), impl));
+            if (typeParameters != null && !typeParameters.isEmpty()) {
+                typeParameters.stream()
+                        .map(Expression::getType)
+                        .map(t -> (t instanceof JavaType.FullyQualified) ? (JavaType.FullyQualified) t : null)
+                        .filter(Objects::nonNull)
+                        .forEach(t -> maybeAddImport(t.getFullyQualifiedName()));
+
+                List<JRightPadded<Expression>> elements = typeParameters.stream()
+                        .map(t -> new JRightPadded<>(t, Space.EMPTY, Markers.EMPTY))
+                        .collect(Collectors.toList());
+
+                J.ParameterizedType typedImpl = new J.ParameterizedType(
+                        randomId(),
+                        Space.EMPTY,
+                        Markers.EMPTY,
+                        impl,
+                        JContainer.build(Space.EMPTY, elements, Markers.EMPTY)
+                );
+
+                c = c.withImplements(ListUtils.concat(c.getImplements(), typedImpl));
+            } else {
+                c = c.withImplements(ListUtils.concat(c.getImplements(), impl));
+            }
 
             JContainer<TypeTree> anImplements = c.getPadding().getImplements();
             assert anImplements != null;