Widen concrete HashMap declared type when UseMapOf rewrites it to Map.of (#1148)

Author: timtebeek

src/main/java/org/openrewrite/java/migrate/util/UseMapOf.java

@@ -16,6 +16,8 @@
 package org.openrewrite.java.migrate.util;
 
 import lombok.Getter;
+import org.jspecify.annotations.Nullable;
+import org.openrewrite.Cursor;
 import org.openrewrite.ExecutionContext;
 import org.openrewrite.Preconditions;
 import org.openrewrite.Recipe;
@@ -29,7 +31,9 @@
 import org.openrewrite.java.search.UsesMethod;
 import org.openrewrite.java.tree.Expression;
 import org.openrewrite.java.tree.J;
+import org.openrewrite.java.tree.JavaType;
 import org.openrewrite.java.tree.Statement;
+import org.openrewrite.java.tree.TypeTree;
 import org.openrewrite.java.tree.TypeUtils;
 
 import java.util.ArrayList;
@@ -194,6 +198,94 @@ private J reattachPairPrefixes(J applied, List<J.MethodInvocation> puts, boolean
                         return nc.withArguments(Collections.singletonList(mapCall.withArguments(withPrefixes)));
                     }
 
+                    @Override
+                    public J visitVariableDeclarations(J.VariableDeclarations multiVariable, ExecutionContext ctx) {
+                        // The anonymous-class form replaces `new HashMap<>() {{ put(..); }}` with an immutable
+                        // `Map.of(..)`/`Map.ofEntries(..)`, which is not assignable to a concrete `HashMap`
+                        // declared type. When such an initializer will be rewritten, widen the declared type to
+                        // `Map<..>` and swap it onto the cursor before visiting the initializer, so the generated
+                        // `Map.of(..)` is type-attributed against the `Map` LHS rather than `HashMap` (#1148).
+                        J.VariableDeclarations widened = widenRewritableHashMapDeclaration(multiVariable);
+                        if (widened != null) {
+                            setCursor(new Cursor(getCursor().getParent(), widened));
+                            return super.visitVariableDeclarations(widened, ctx);
+                        }
+                        return super.visitVariableDeclarations(multiVariable, ctx);
+                    }
+
+                    /**
+                     * If {@code decl} is a single-variable, concrete {@code HashMap} declaration whose anonymous
+                     * {@code new HashMap<>() {{ put(..); }}} initializer would be rewritten to an immutable
+                     * {@code Map.of(..)}, returns a copy with the declared type widened to {@code java.util.Map}
+                     * (type expression plus the variable's own type attribution); {@code null} otherwise.
+                     */
+                    private J.@Nullable VariableDeclarations widenRewritableHashMapDeclaration(J.VariableDeclarations decl) {
+                        if (decl.getVariables().size() != 1 ||
+                                !TypeUtils.isOfClassType(decl.getType(), "java.util.HashMap")) {
+                            return null;
+                        }
+                        Expression initializer = decl.getVariables().get(0).getInitializer();
+                        if (!(initializer instanceof J.NewClass) || !rewritesToImmutableMapOf((J.NewClass) initializer)) {
+                            return null;
+                        }
+
+                        JavaType.FullyQualified mapType = JavaType.ShallowClass.build("java.util.Map");
+                        TypeTree typeExpression = decl.getTypeExpression();
+                        TypeTree widenedExpression;
+                        JavaType widenedType;
+                        if (typeExpression instanceof J.ParameterizedType) {
+                            J.ParameterizedType pt = (J.ParameterizedType) typeExpression;
+                            if (!(pt.getClazz() instanceof J.Identifier)) {
+                                return null;
+                            }
+                            widenedType = pt.getType() instanceof JavaType.Parameterized ?
+                                    ((JavaType.Parameterized) pt.getType()).withType(mapType) : mapType;
+                            J.Identifier clazz = ((J.Identifier) pt.getClazz()).withSimpleName("Map").withType(mapType);
+                            widenedExpression = pt.withClazz(clazz).withType(widenedType);
+                        } else if (typeExpression instanceof J.Identifier) {
+                            widenedType = mapType;
+                            widenedExpression = ((J.Identifier) typeExpression).withSimpleName("Map").withType(mapType);
+                        } else {
+                            return null;
+                        }
+
+                        J.VariableDeclarations.NamedVariable nv = decl.getVariables().get(0);
+                        JavaType.Variable variableType = nv.getVariableType();
+                        if (variableType != null) {
+                            variableType = variableType.withType(widenedType);
+                            nv = nv.withVariableType(variableType).withName(nv.getName().withType(widenedType).withFieldType(variableType));
+                        }
+                        return decl.withTypeExpression(widenedExpression).withVariables(Collections.singletonList(nv));
+                    }
+
+                    /**
+                     * Mirrors the anonymous-class guard in {@link #visitNewClass}: returns {@code true} when
+                     * {@code nc} is a {@code new HashMap<>() {{ put(k, v); ... }}} whose body holds only
+                     * non-null {@code put(..)} calls, and so will be replaced with an immutable {@code Map.of(..)}.
+                     */
+                    private boolean rewritesToImmutableMapOf(J.NewClass nc) {
+                        J.Block body = nc.getBody();
+                        if (!NEW_HASH_MAP.matches(nc) || body == null || body.getStatements().size() != 1 ||
+                                !TypeUtils.isOfClassType(nc.getClazz() != null ? nc.getClazz().getType() : null, "java.util.HashMap")) {
+                            return false;
+                        }
+                        Statement statement = body.getStatements().get(0);
+                        if (!(statement instanceof J.Block)) {
+                            return false;
+                        }
+                        for (Statement stat : ((J.Block) statement).getStatements()) {
+                            if (!(stat instanceof J.MethodInvocation) || !MAP_PUT.matches((Expression) stat)) {
+                                return false;
+                            }
+                            for (Expression arg : ((J.MethodInvocation) stat).getArguments()) {
+                                if (J.Literal.isLiteralValue(arg, null)) {
+                                    return false;
+                                }
+                            }
+                        }
+                        return true;
+                    }
+
                     @Override
                     public J visitBlock(J.Block block, ExecutionContext ctx) {
                         Map<UUID, List<J.MethodInvocation>> rewrites = new HashMap<>();

src/test/java/org/openrewrite/java/migrate/util/UseMapOfTest.java

@@ -205,6 +205,66 @@ class Test {
         );
     }
 
+    @Issue("https://github.com/openrewrite/rewrite-migrate-java/issues/1148")
+    @Test
+    void widenConcreteHashMapFieldTypeToMap() {
+        //language=java
+        rewriteRun(
+          spec -> spec.allSources(s -> s.markers(javaVersion(25))),
+          java(
+            """
+              import java.util.HashMap;
+
+              class Main {
+                  private static final HashMap<String, String> VALUES = new HashMap<String, String>() {{
+                      put("key", "value");
+                  }};
+              }
+              """,
+            """
+              import java.util.Map;
+
+              class Main {
+                  private static final Map<String, String> VALUES = Map.of("key", "value");
+              }
+              """
+          )
+        );
+    }
+
+    @Issue("https://github.com/openrewrite/rewrite-migrate-java/issues/1148")
+    @Test
+    void widenConcreteHashMapLocalVariableTypeToMap() {
+        //language=java
+        rewriteRun(
+          java(
+            """
+              import java.util.HashMap;
+
+              class Main {
+                  void m() {
+                      HashMap<String, Integer> ages = new HashMap<>() {{
+                          put("Bob", 42);
+                          put("alice", 30);
+                      }};
+                  }
+              }
+              """,
+            """
+              import java.util.Map;
+
+              class Main {
+                  void m() {
+                      Map<String, Integer> ages = Map.of(
+                          "Bob", 42,
+                          "alice", 30);
+                  }
+              }
+              """
+          )
+        );
+    }
+
     @Issue("https://github.com/openrewrite/rewrite-migrate-java/issues/1112")
     @Test
     void doNotChangeLinkedHashMap() {