Add explicit type arguments to generic methods in UseVarForGenericMethodInvocations (#869)

* Added tests for issue #868

* Fix tests

* Minimize implementation

---------

Co-authored-by: Tim te Beek <[email protected]>

Author: jagodevreede

src/main/java/org/openrewrite/java/migrate/lang/var/UseVarForGenericMethodInvocations.java

@@ -23,11 +23,10 @@
 import org.openrewrite.internal.ListUtils;
 import org.openrewrite.java.JavaIsoVisitor;
 import org.openrewrite.java.search.UsesJavaVersion;
-import org.openrewrite.java.tree.Expression;
-import org.openrewrite.java.tree.J;
-import org.openrewrite.java.tree.JavaType;
-import org.openrewrite.java.tree.TypeTree;
+import org.openrewrite.java.tree.*;
+import org.openrewrite.marker.Markers;
 
+import java.util.ArrayList;
 import java.util.List;
 
 import static java.util.Objects.requireNonNull;
@@ -97,6 +96,7 @@ public J.VariableDeclarations visitVariableDeclarations(J.VariableDeclarations v
         /**
          * Makes nested generic types explicit by replacing diamond operators in constructor calls
          * with explicit type parameters based on the variable declaration type.
+         * Also adds explicit type parameters to the method invocation itself when needed.
          */
         private J.MethodInvocation makeNestedGenericsExplicit(J.MethodInvocation mi, J.VariableDeclarations vd) {
             // Extract type parameters from the variable declaration
@@ -109,6 +109,16 @@ private J.MethodInvocation makeNestedGenericsExplicit(J.MethodInvocation mi, J.V
                 return mi;
             }
 
+            // Add explicit type parameters when the method is generic and the return type's type parameter matches a type parameter from the declaring class
+            if (mi.getTypeParameters() == null && mi.getMethodType() != null && containsGenericTypeVariable(mi.getMethodType().getReturnType())) {
+                // Create JRightPadded list from leftTypeParams
+                List<JRightPadded<Expression>> typeParamsList = new ArrayList<>();
+                for (Expression typeParam : leftTypeParams) {
+                    typeParamsList.add(JRightPadded.build(typeParam));
+                }
+                mi = mi.withTypeParameters(JContainer.build(Space.EMPTY, typeParamsList, Markers.EMPTY));
+            }
+
             // Visit arguments and replace diamond operators with explicit type parameters
             return mi.withArguments(ListUtils.map(mi.getArguments(), arg -> {
                 if (arg instanceof J.NewClass) {
@@ -124,6 +134,22 @@ private J.MethodInvocation makeNestedGenericsExplicit(J.MethodInvocation mi, J.V
             }));
         }
 
+        private boolean containsGenericTypeVariable(JavaType type) {
+            if (type instanceof JavaType.GenericTypeVariable) {
+                return true;
+            }
+
+            if (type instanceof JavaType.Parameterized) {
+                for (JavaType typeParam : ((JavaType.Parameterized) type).getTypeParameters()) {
+                    if (containsGenericTypeVariable(typeParam)) {
+                        return true;
+                    }
+                }
+            }
+
+            return false;
+        }
+
         private static boolean hasTypeParams(@Nullable TypeTree clazz) {
             if (clazz instanceof J.ParameterizedType) {
                 List<Expression> typeParameters = ((J.ParameterizedType) clazz).getTypeParameters();

src/test/java/org/openrewrite/java/migrate/lang/var/UseVarForGenericMethodInvocationsTest.java

@@ -295,5 +295,79 @@ void dostuff() {
               )
             );
         }
+
+        @Issue("https://github.com/openrewrite/rewrite-migrate-java/issues/868")
+        @Test
+        void applyVarAndAddExplicitTypeArgumentForNestedGenericFactoryMethod() {
+            //language=java
+            rewriteRun(
+              version(
+                java(
+                  """
+                    class A<T> {
+                        void getX(final Root<T> root) {
+                            Path<T> x = root.get("x");
+                        }
+                    }
+                    interface Root<X>{
+                        <Y> Path<Y> get(String var1);
+                    }
+                    interface Path<X> { }
+                    """,
+                  """
+                    class A<T> {
+                        void getX(final Root<T> root) {
+                            var x = root.<T>get("x");
+                        }
+                    }
+                    interface Root<X>{
+                        <Y> Path<Y> get(String var1);
+                    }
+                    interface Path<X> { }
+                    """
+                ),
+                10
+              )
+            );
+        }
+
+        @Issue("https://github.com/openrewrite/rewrite-migrate-java/issues/868")
+        @Test
+        void genericsCollectorsRegression() {
+            //language=java
+            rewriteRun(
+              version(
+                java(
+                  """
+                    import java.time.LocalDate;
+                    import java.util.ArrayList;
+                    import java.util.Set;
+                    import java.util.stream.Collectors;
+
+                    public class A {
+                        void getX() {
+                            Set<String> strDates = new ArrayList<LocalDate>().stream()
+                                    .map(date -> date.toString()).collect(Collectors.toSet());
+                        }
+                    }
+                    """,
+                  """
+                    import java.time.LocalDate;
+                    import java.util.ArrayList;
+                    import java.util.Set;
+                    import java.util.stream.Collectors;
+
+                    public class A {
+                        void getX() {
+                            var strDates = new ArrayList<LocalDate>().stream()
+                                    .map(date -> date.toString()).collect(Collectors.toSet());
+                        }
+                    }
+                    """
+                ),
+                10
+              )
+            );
+        }
     }
 }