Preserve nullable values in UseMapOf prose maps (#1204)

Author: martinfrancois

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

@@ -29,6 +29,7 @@
 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.TypeUtils;
 
@@ -78,7 +79,9 @@ public J visitNewClass(J.NewClass newClass, ExecutionContext ctx) {
                                         new StringJoiner(", ", "Map.ofEntries(", ")") :
                                         new StringJoiner(", ", "Map.of(", ")");
                                 for (J.MethodInvocation put : puts) {
-                                    args.addAll(put.getArguments());
+                                    for (Expression arg : put.getArguments()) {
+                                        args.add(arg.unwrap());
+                                    }
                                     if (useEntries) {
                                         inner.add("Map.entry(#{any()}, #{any()})");
                                     } else {
@@ -317,7 +320,7 @@ private String matchingTargetName(J.VariableDeclarations decl) {
                     /**
                      * If {@code stmt} is {@code targetName.put(k, v)} matching {@link #MAP_PUT},
                      * returns [key, value] as a list; otherwise {@code null}. Returns {@code null}
-                     * if either argument is the {@code null} literal, since {@code Map.of(..)} and
+                     * unless both arguments are proven non-null, since {@code Map.of(..)} and
                      * {@code Map.entry(..)} reject nulls.
                      */
                     private List<Expression> matchPutCallOn(Statement stmt, String targetName) {
@@ -338,13 +341,35 @@ private List<Expression> matchPutCallOn(Statement stmt, String targetName) {
                             return null;
                         }
                         for (Expression arg : mi.getArguments()) {
-                            if (J.Literal.isLiteralValue( arg, null )) {
+                            if (!isProvablyNonNull(arg)) {
                                 return null;
                             }
                         }
                         return mi.getArguments();
                     }
 
+                    private boolean isProvablyNonNull(Expression expression) {
+                        expression = expression.unwrap();
+                        if (expression instanceof J.Literal) {
+                            return !J.Literal.isLiteralValue(expression, null);
+                        }
+                        if (expression instanceof J.NewClass ||
+                                expression instanceof J.NewArray ||
+                                expression instanceof J.Lambda ||
+                                expression instanceof J.MemberReference) {
+                            return true;
+                        }
+                        JavaType type = expression.getType();
+                        return type == JavaType.Primitive.Boolean ||
+                               type == JavaType.Primitive.Byte ||
+                               type == JavaType.Primitive.Char ||
+                               type == JavaType.Primitive.Double ||
+                               type == JavaType.Primitive.Float ||
+                               type == JavaType.Primitive.Int ||
+                               type == JavaType.Primitive.Long ||
+                               type == JavaType.Primitive.Short;
+                    }
+
                     private boolean expressionReferences(Expression expr, String name) {
                         return new JavaIsoVisitor<AtomicBoolean>() {
                             @Override

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

@@ -574,4 +574,119 @@ void m() {
           )
         );
     }
+
+    @Test
+    void proseMapsWithProvablyNonNullValuesAreCollapsed() {
+        rewriteRun(
+          //language=java
+          java(
+            """
+              import java.util.HashMap;
+              import java.util.Map;
+
+              class PrimitiveValues {
+                  Map<String, Integer> values(int first, int second) {
+                      Map<String, Integer> values = new HashMap<>();
+                      values.put("first", first);
+                      values.put("second", second);
+                      return values;
+                  }
+              }
+              """,
+            """
+              import java.util.HashMap;
+              import java.util.Map;
+
+              class PrimitiveValues {
+                  Map<String, Integer> values(int first, int second) {
+                      Map<String, Integer> values = new HashMap<>(Map.of(
+                              "first", first,
+                              "second", second));
+                      return values;
+                  }
+              }
+              """
+          ),
+          //language=java
+          java(
+            """
+              import java.util.HashMap;
+              import java.util.Map;
+
+              class ParenthesizedValues {
+                  Map<String, Integer> values(int first, int second) {
+                      Map<String, Integer> values = new HashMap<>();
+                      values.put(("first"), first);
+                      values.put(("second"), second);
+                      return values;
+                  }
+              }
+              """,
+            """
+              import java.util.HashMap;
+              import java.util.Map;
+
+              class ParenthesizedValues {
+                  Map<String, Integer> values(int first, int second) {
+                      Map<String, Integer> values = new HashMap<>(Map.of(
+                              "first", first,
+                              "second", second));
+                      return values;
+                  }
+              }
+              """
+          ),
+          //language=java
+          java(
+            """
+              import java.util.HashMap;
+              import java.util.Map;
+
+              class NewClassValues {
+                  Map<String, Object> values() {
+                      Map<String, Object> values = new HashMap<>();
+                      values.put("first", new Object());
+                      values.put("second", new Object());
+                      return values;
+                  }
+              }
+              """,
+            """
+              import java.util.HashMap;
+              import java.util.Map;
+
+              class NewClassValues {
+                  Map<String, Object> values() {
+                      Map<String, Object> values = new HashMap<>(Map.of(
+                              "first", new Object(),
+                              "second", new Object()));
+                      return values;
+                  }
+              }
+              """
+          )
+        );
+    }
+
+    @Test
+    void proseMapWithUnknownNullableValueIsLeftAlone() {
+        rewriteRun(
+          //language=java
+          java(
+            """
+              import java.util.HashMap;
+              import java.util.Map;
+
+              class Test {
+                  Map<String, Object> prompt(Object card, Object attempt) {
+                      Map<String, Object> values = new HashMap<>();
+                      values.put("card", card);
+                      values.put("attempt", attempt);
+                      return values;
+                  }
+              }
+              """
+          )
+        );
+    }
 }