Extend Use{List,Set,Map}Of to recognise prose-statement chains (#1145)

* Extend Use{List,Set,Map}Of to recognise prose-statement chains

UseListOf, UseSetOf, and UseMapOf previously matched only the
anonymous-class idiom:

    List<String> l = new ArrayList<>() {{ add("a"); add("b"); }};

Extends each recipe to also recognise the more common prose-statement
shape, where a no-arg constructor declaration is followed by a chain
of add(..) or put(..) calls on the declared variable:

    List<String> l = new ArrayList<>();
    l.add("a");
    l.add("b");

These are collapsed into the constructor with the immutable factory:

    List<String> l = new ArrayList<>(List.of("a", "b"));

Output is always wrapped in the mutable ArrayList/HashSet/HashMap so
downstream mutations (.sort(), .add(), etc.) remain valid.

Recognised shapes per recipe:

- UseListOf: target.add(arg) chains; output uses List.of(..).
- UseSetOf:  target.add(arg) chains; output uses Set.of(..).
- UseMapOf:  target.put(k, v) chains; output uses Map.of(k, v, ..)
             up to 10 pairs, Map.ofEntries(Map.entry(k, v), ..) beyond.

Bail conditions (all three): raw LHS, non-no-arg constructor (e.g.
capacity hint), intervening non-recognised statement, an argument
that references the target variable (depends on step-by-step
mutation), null literal argument. Threshold: at least 2 add/put
statements; a single call is not worth the rewrite churn.

Implementation: a visitBlock pre-pass scans for prose patterns and
records (initializer UUID -> [args]) on a cursor message; the
existing visitNewClass override consults the map and applies the
JavaTemplate at the correct cursor depth. After super.visitBlock,
the post-pass filters absorbed add/put statements from the block.

Regenerate recipes.csv to reflect the new descriptions.

* Simplify Use{List,Set,Map}Of prose helpers with ListUtils.filter and reduce

* Preserve comments and one-per-line formatting in UseMapOf prose chains

* Extend prose comment + one-per-line formatting to UseListOf and UseSetOf

---------

Co-authored-by: Tim te Beek <tim@moderne.io>

Author: steve-aom-elliott

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

@@ -20,6 +20,8 @@
 import org.openrewrite.Preconditions;
 import org.openrewrite.Recipe;
 import org.openrewrite.TreeVisitor;
+import org.openrewrite.internal.ListUtils;
+import org.openrewrite.java.JavaIsoVisitor;
 import org.openrewrite.java.JavaTemplate;
 import org.openrewrite.java.JavaVisitor;
 import org.openrewrite.java.MethodMatcher;
@@ -30,19 +32,32 @@
 import org.openrewrite.java.tree.Statement;
 
 import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
+import java.util.Set;
 import java.util.StringJoiner;
+import java.util.UUID;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 public class UseListOf extends Recipe {
     private static final MethodMatcher NEW_ARRAY_LIST = new MethodMatcher("java.util.ArrayList <constructor>()", true);
     private static final MethodMatcher LIST_ADD = new MethodMatcher("java.util.List add(..)", true);
 
+    private static final String PROSE_REWRITES_KEY = "use-list-of.prose-rewrites";
+
     @Getter
     final String displayName = "Prefer `List.of(..)`";
 
     @Getter
-    final String description = "Prefer `List.of(..)` instead of using `java.util.List#add(..)` in anonymous ArrayList initializers in Java 10 or higher. " +
-            "This recipe will not modify code where the List is later mutated since `List.of` returns an immutable list.";
+    final String description = "Prefer `List.of(..)` in Java 10 or higher. Two input shapes are recognised:\n\n" +
+            "- Anonymous-class initialization (`new ArrayList<>() {{ add(\"a\"); add(\"b\"); }}`), " +
+            "which is replaced wholesale with `List.of(\"a\", \"b\")` (immutable result, matching the " +
+            "anonymous-class idiom's typical intent).\n" +
+            "- A `new ArrayList<>()` declaration followed by a chain of `target.add(..)` statements, " +
+            "which is collapsed to `new ArrayList<>(List.of(..))` (preserving the mutable `ArrayList`).";
 
     @Override
     public TreeVisitor<?, ExecutionContext> getVisitor() {
@@ -54,6 +69,32 @@ public TreeVisitor<?, ExecutionContext> getVisitor() {
                     @Override
                     public J visitNewClass(J.NewClass newClass, ExecutionContext ctx) {
                         J.NewClass n = (J.NewClass) super.visitNewClass(newClass, ctx);
+
+                        // Prose-pattern: see if visitBlock (above us on the cursor) decided this
+                        // initializer should be wrapped with `new ArrayList<>(List.of(..))`.
+                        Map<UUID, List<J.MethodInvocation>> rewrites = getCursor().getNearestMessage(PROSE_REWRITES_KEY);
+                        if (rewrites != null) {
+                            List<J.MethodInvocation> adds = rewrites.get(n.getId());
+                            if (adds != null) {
+                                List<Expression> args = new ArrayList<>();
+                                StringJoiner joiner = new StringJoiner(", ", "new ArrayList<>(List.of(", "))");
+                                for (J.MethodInvocation add : adds) {
+                                    args.add(add.getArguments().get(0));
+                                    joiner.add("#{any()}");
+                                }
+                                maybeAddImport("java.util.List");
+                                J applied = JavaTemplate.builder(joiner.toString())
+                                        .contextSensitive()
+                                        .imports("java.util.ArrayList", "java.util.List")
+                                        .build()
+                                        .apply(updateCursor(n), n.getCoordinates().replace(), args.toArray());
+                                // Reattach each add's prefix so the elements land one-per-line and any
+                                // leading comments survive, then autoformat to nest the indentation.
+                                return autoFormat(reattachElementPrefixes(applied, adds), ctx);
+                            }
+                        }
+
+                        // Anonymous-class form (original UseListOf logic, unchanged).
                         J.Block body = n.getBody();
                         if (NEW_ARRAY_LIST.matches(n) && body != null && body.getStatements().size() == 1) {
                             Statement statement = body.getStatements().get(0);
@@ -65,7 +106,6 @@ public J visitNewClass(J.NewClass newClass, ExecutionContext ctx) {
                                         return n;
                                     }
                                     J.MethodInvocation add = (J.MethodInvocation) stat;
-                                    // List.add() takes only one argument
                                     if (add.getArguments().size() != 1) {
                                         return n;
                                     }
@@ -85,7 +125,175 @@ public J visitNewClass(J.NewClass newClass, ExecutionContext ctx) {
 
                         return n;
                     }
+
+                    /**
+                     * Re-applies the absorbed add statements' prefixes to the generated
+                     * {@code new ArrayList<>(List.of(..))} so each element keeps its own line and any
+                     * leading comments. {@code adds} holds one invocation per element, in order.
+                     */
+                    private J reattachElementPrefixes(J applied, List<J.MethodInvocation> adds) {
+                        if (!(applied instanceof J.NewClass)) {
+                            return applied;
+                        }
+                        J.NewClass nc = (J.NewClass) applied;
+                        if (nc.getArguments().size() != 1 || !(nc.getArguments().get(0) instanceof J.MethodInvocation)) {
+                            return applied;
+                        }
+                        J.MethodInvocation listCall = (J.MethodInvocation) nc.getArguments().get(0);
+                        List<Expression> listArgs = listCall.getArguments();
+                        List<Expression> withPrefixes = new ArrayList<>(listArgs.size());
+                        for (int i = 0; i < listArgs.size(); i++) {
+                            withPrefixes.add(listArgs.get(i).withPrefix(adds.get(i).getPrefix()));
+                        }
+                        return nc.withArguments(Collections.singletonList(listCall.withArguments(withPrefixes)));
+                    }
+
+                    @Override
+                    public J visitBlock(J.Block block, ExecutionContext ctx) {
+                        // Pre-pass: scan the ORIGINAL block to identify which initializers to
+                        // rewrite and which `add(..)` statements to absorb. UUIDs are stable
+                        // through super.visitBlock unless a child visitor rebuilds the node,
+                        // and nothing else in this recipe touches the targeted initializers
+                        // before visitNewClass fires.
+                        Map<UUID, List<J.MethodInvocation>> rewrites = new HashMap<>();
+                        Set<UUID> absorbedAddIds = new HashSet<>();
+                        identifyProseRewrites(block, rewrites, absorbedAddIds);
+
+                        if (!rewrites.isEmpty()) {
+                            getCursor().putMessage(PROSE_REWRITES_KEY, rewrites);
+                        }
+
+                        J.Block b = (J.Block) super.visitBlock(block, ctx);
+
+                        // Post-pass: drop the now-absorbed `add(..)` statements from the block.
+                        return b.withStatements(ListUtils.filter(b.getStatements(), s -> !absorbedAddIds.contains(s.getId())));
+                    }
+
+                    /**
+                     * Walk the block's statements looking for:
+                     * <pre>
+                     *     List&lt;T&gt; name = new ArrayList&lt;&gt;();
+                     *     name.add(x1);
+                     *     name.add(x2);
+                     *     ...
+                     * </pre>
+                     * For each such sequence with at least two adds, record
+                     * (initializer UUID, the absorbed {@code add(..)} invocations in order) in
+                     * {@code rewrites} and the absorbed add statement UUIDs in {@code absorbedAddIds}.
+                     */
+                    private void identifyProseRewrites(
+                            J.Block block,
+                            Map<UUID, List<J.MethodInvocation>> rewrites,
+                            Set<UUID> absorbedAddIds) {
+                        List<Statement> stmts = block.getStatements();
+                        int i = 0;
+                        while (i < stmts.size()) {
+                            Statement stmt = stmts.get(i);
+                            if (!(stmt instanceof J.VariableDeclarations)) {
+                                i++;
+                                continue;
+                            }
+                            J.VariableDeclarations decl = (J.VariableDeclarations) stmt;
+                            String targetName = matchingTargetName(decl);
+                            if (targetName == null) {
+                                i++;
+                                continue;
+                            }
+                            J.NewClass initializer = (J.NewClass) decl.getVariables().get(0).getInitializer();
+                            // (matchingTargetName already verified the initializer is a J.NewClass)
+
+                            List<J.MethodInvocation> adds = new ArrayList<>();
+                            List<UUID> absorbedHere = new ArrayList<>();
+                            int j = i + 1;
+                            while (j < stmts.size()) {
+                                Statement next = stmts.get(j);
+                                Expression arg = matchAddCallOn(next, targetName);
+                                if (arg == null || expressionReferences(arg, targetName)) {
+                                    break;
+                                }
+                                adds.add((J.MethodInvocation) next);
+                                absorbedHere.add(next.getId());
+                                j++;
+                            }
+                            if (adds.size() >= 2 && initializer != null) {
+                                rewrites.put(initializer.getId(), adds);
+                                absorbedAddIds.addAll(absorbedHere);
+                                i = j;
+                            } else {
+                                i++;
+                            }
+                        }
+                    }
+
+                    /**
+                     * Returns the variable name if {@code decl} is a single-variable, parameterized
+                     * {@code List<T>} declaration whose initializer is a no-arg {@code new ArrayList<>()}
+                     * with no anonymous-class body. Returns {@code null} otherwise.
+                     */
+                    private String matchingTargetName(J.VariableDeclarations decl) {
+                        if (decl.getVariables().size() != 1) {
+                            return null;
+                        }
+                        // Require parameterized LHS; for raw `List` we'd be guessing at a type argument.
+                        if (!(decl.getTypeExpression() instanceof J.ParameterizedType)) {
+                            return null;
+                        }
+                        J.VariableDeclarations.NamedVariable nv = decl.getVariables().get(0);
+                        if (!(nv.getInitializer() instanceof J.NewClass)) {
+                            return null;
+                        }
+                        J.NewClass nc = (J.NewClass) nv.getInitializer();
+                        if (!NEW_ARRAY_LIST.matches(nc)) {
+                            return null;
+                        }
+                        // A body would put us in the anonymous-class case handled by visitNewClass directly.
+                        if (nc.getBody() != null) {
+                            return null;
+                        }
+                        return nv.getSimpleName();
+                    }
+
+                    /**
+                     * If {@code stmt} is {@code targetName.add(arg)} matching {@link #LIST_ADD},
+                     * returns the single argument expression; otherwise {@code null}. Also returns
+                     * {@code null} when the argument is the {@code null} literal, since
+                     * {@code List.of(..)} rejects nulls.
+                     */
+                    private Expression matchAddCallOn(Statement stmt, String targetName) {
+                        if (!(stmt instanceof J.MethodInvocation)) {
+                            return null;
+                        }
+                        J.MethodInvocation mi = (J.MethodInvocation) stmt;
+                        if (!LIST_ADD.matches(mi)) {
+                            return null;
+                        }
+                        if (mi.getArguments().size() != 1) {
+                            return null;
+                        }
+                        if (!(mi.getSelect() instanceof J.Identifier)) {
+                            return null;
+                        }
+                        if (!targetName.equals(((J.Identifier) mi.getSelect()).getSimpleName())) {
+                            return null;
+                        }
+                        Expression arg = mi.getArguments().get(0);
+                        if (arg instanceof J.Literal && ((J.Literal) arg).getValue() == null) {
+                            return null;
+                        }
+                        return arg;
+                    }
+
+                    private boolean expressionReferences(Expression expr, String name) {
+                        return new JavaIsoVisitor<AtomicBoolean>() {
+                            @Override
+                            public J.Identifier visitIdentifier(J.Identifier id, AtomicBoolean f) {
+                                if (name.equals(id.getSimpleName())) {
+                                    f.set(true);
+                                }
+                                return id;
+                            }
+                        }.reduce(expr, new AtomicBoolean(false)).get();
+                    }
                 });
     }
-
 }