Skip to content

Commit 3a497b9

Browse files
timtebeektimtebeek
committed
Extend prose comment + one-per-line formatting to UseListOf and UseSetOf
1 parent 99973be commit 3a497b9

4 files changed

Lines changed: 160 additions & 28 deletions

File tree

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

Lines changed: 42 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import org.openrewrite.java.tree.Statement;
3333

3434
import java.util.ArrayList;
35+
import java.util.Collections;
3536
import java.util.HashMap;
3637
import java.util.HashSet;
3738
import java.util.List;
@@ -71,20 +72,25 @@ public J visitNewClass(J.NewClass newClass, ExecutionContext ctx) {
7172

7273
// Prose-pattern: see if visitBlock (above us on the cursor) decided this
7374
// initializer should be wrapped with `new ArrayList<>(List.of(..))`.
74-
Map<UUID, List<Expression>> rewrites = getCursor().getNearestMessage(PROSE_REWRITES_KEY);
75+
Map<UUID, List<J.MethodInvocation>> rewrites = getCursor().getNearestMessage(PROSE_REWRITES_KEY);
7576
if (rewrites != null) {
76-
List<Expression> proseArgs = rewrites.get(n.getId());
77-
if (proseArgs != null) {
77+
List<J.MethodInvocation> adds = rewrites.get(n.getId());
78+
if (adds != null) {
79+
List<Expression> args = new ArrayList<>();
7880
StringJoiner joiner = new StringJoiner(", ", "new ArrayList<>(List.of(", "))");
79-
for (int k = 0; k < proseArgs.size(); k++) {
81+
for (J.MethodInvocation add : adds) {
82+
args.add(add.getArguments().get(0));
8083
joiner.add("#{any()}");
8184
}
8285
maybeAddImport("java.util.List");
83-
return JavaTemplate.builder(joiner.toString())
86+
J applied = JavaTemplate.builder(joiner.toString())
8487
.contextSensitive()
8588
.imports("java.util.ArrayList", "java.util.List")
8689
.build()
87-
.apply(updateCursor(n), n.getCoordinates().replace(), proseArgs.toArray());
90+
.apply(updateCursor(n), n.getCoordinates().replace(), args.toArray());
91+
// Reattach each add's prefix so the elements land one-per-line and any
92+
// leading comments survive, then autoformat to nest the indentation.
93+
return autoFormat(reattachElementPrefixes(applied, adds), ctx);
8894
}
8995
}
9096

@@ -120,14 +126,36 @@ public J visitNewClass(J.NewClass newClass, ExecutionContext ctx) {
120126
return n;
121127
}
122128

129+
/**
130+
* Re-applies the absorbed add statements' prefixes to the generated
131+
* {@code new ArrayList<>(List.of(..))} so each element keeps its own line and any
132+
* leading comments. {@code adds} holds one invocation per element, in order.
133+
*/
134+
private J reattachElementPrefixes(J applied, List<J.MethodInvocation> adds) {
135+
if (!(applied instanceof J.NewClass)) {
136+
return applied;
137+
}
138+
J.NewClass nc = (J.NewClass) applied;
139+
if (nc.getArguments().size() != 1 || !(nc.getArguments().get(0) instanceof J.MethodInvocation)) {
140+
return applied;
141+
}
142+
J.MethodInvocation listCall = (J.MethodInvocation) nc.getArguments().get(0);
143+
List<Expression> listArgs = listCall.getArguments();
144+
List<Expression> withPrefixes = new ArrayList<>(listArgs.size());
145+
for (int i = 0; i < listArgs.size(); i++) {
146+
withPrefixes.add(listArgs.get(i).withPrefix(adds.get(i).getPrefix()));
147+
}
148+
return nc.withArguments(Collections.singletonList(listCall.withArguments(withPrefixes)));
149+
}
150+
123151
@Override
124152
public J visitBlock(J.Block block, ExecutionContext ctx) {
125153
// Pre-pass: scan the ORIGINAL block to identify which initializers to
126154
// rewrite and which `add(..)` statements to absorb. UUIDs are stable
127155
// through super.visitBlock unless a child visitor rebuilds the node,
128156
// and nothing else in this recipe touches the targeted initializers
129157
// before visitNewClass fires.
130-
Map<UUID, List<Expression>> rewrites = new HashMap<>();
158+
Map<UUID, List<J.MethodInvocation>> rewrites = new HashMap<>();
131159
Set<UUID> absorbedAddIds = new HashSet<>();
132160
identifyProseRewrites(block, rewrites, absorbedAddIds);
133161

@@ -150,12 +178,12 @@ public J visitBlock(J.Block block, ExecutionContext ctx) {
150178
* ...
151179
* </pre>
152180
* For each such sequence with at least two adds, record
153-
* (initializer UUID, [args]) in {@code rewrites} and the absorbed add
154-
* statement UUIDs in {@code absorbedAddIds}.
181+
* (initializer UUID, the absorbed {@code add(..)} invocations in order) in
182+
* {@code rewrites} and the absorbed add statement UUIDs in {@code absorbedAddIds}.
155183
*/
156184
private void identifyProseRewrites(
157185
J.Block block,
158-
Map<UUID, List<Expression>> rewrites,
186+
Map<UUID, List<J.MethodInvocation>> rewrites,
159187
Set<UUID> absorbedAddIds) {
160188
List<Statement> stmts = block.getStatements();
161189
int i = 0;
@@ -174,7 +202,7 @@ private void identifyProseRewrites(
174202
J.NewClass initializer = (J.NewClass) decl.getVariables().get(0).getInitializer();
175203
// (matchingTargetName already verified the initializer is a J.NewClass)
176204

177-
List<Expression> args = new ArrayList<>();
205+
List<J.MethodInvocation> adds = new ArrayList<>();
178206
List<UUID> absorbedHere = new ArrayList<>();
179207
int j = i + 1;
180208
while (j < stmts.size()) {
@@ -183,12 +211,12 @@ private void identifyProseRewrites(
183211
if (arg == null || expressionReferences(arg, targetName)) {
184212
break;
185213
}
186-
args.add(arg);
214+
adds.add((J.MethodInvocation) next);
187215
absorbedHere.add(next.getId());
188216
j++;
189217
}
190-
if (args.size() >= 2 && initializer != null) {
191-
rewrites.put(initializer.getId(), args);
218+
if (adds.size() >= 2 && initializer != null) {
219+
rewrites.put(initializer.getId(), adds);
192220
absorbedAddIds.addAll(absorbedHere);
193221
i = j;
194222
} else {

‎src/main/java/org/openrewrite/java/migrate/util/UseSetOf.java‎

Lines changed: 40 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import org.openrewrite.java.tree.Statement;
3333

3434
import java.util.ArrayList;
35+
import java.util.Collections;
3536
import java.util.HashMap;
3637
import java.util.HashSet;
3738
import java.util.List;
@@ -71,20 +72,25 @@ public J visitNewClass(J.NewClass newClass, ExecutionContext ctx) {
7172

7273
// Prose-pattern: see if visitBlock (above us on the cursor) decided this
7374
// initializer should be wrapped with `new HashSet<>(Set.of(..))`.
74-
Map<UUID, List<Expression>> rewrites = getCursor().getNearestMessage(PROSE_REWRITES_KEY);
75+
Map<UUID, List<J.MethodInvocation>> rewrites = getCursor().getNearestMessage(PROSE_REWRITES_KEY);
7576
if (rewrites != null) {
76-
List<Expression> proseArgs = rewrites.get(n.getId());
77-
if (proseArgs != null) {
77+
List<J.MethodInvocation> adds = rewrites.get(n.getId());
78+
if (adds != null) {
79+
List<Expression> args = new ArrayList<>();
7880
StringJoiner joiner = new StringJoiner(", ", "new HashSet<>(Set.of(", "))");
79-
for (int k = 0; k < proseArgs.size(); k++) {
81+
for (J.MethodInvocation add : adds) {
82+
args.add(add.getArguments().get(0));
8083
joiner.add("#{any()}");
8184
}
8285
maybeAddImport("java.util.Set");
83-
return JavaTemplate.builder(joiner.toString())
86+
J applied = JavaTemplate.builder(joiner.toString())
8487
.contextSensitive()
8588
.imports("java.util.HashSet", "java.util.Set")
8689
.build()
87-
.apply(updateCursor(n), n.getCoordinates().replace(), proseArgs.toArray());
90+
.apply(updateCursor(n), n.getCoordinates().replace(), args.toArray());
91+
// Reattach each add's prefix so the elements land one-per-line and any
92+
// leading comments survive, then autoformat to nest the indentation.
93+
return autoFormat(reattachElementPrefixes(applied, adds), ctx);
8894
}
8995
}
9096

@@ -120,9 +126,31 @@ public J visitNewClass(J.NewClass newClass, ExecutionContext ctx) {
120126
return n;
121127
}
122128

129+
/**
130+
* Re-applies the absorbed add statements' prefixes to the generated
131+
* {@code new HashSet<>(Set.of(..))} so each element keeps its own line and any
132+
* leading comments. {@code adds} holds one invocation per element, in order.
133+
*/
134+
private J reattachElementPrefixes(J applied, List<J.MethodInvocation> adds) {
135+
if (!(applied instanceof J.NewClass)) {
136+
return applied;
137+
}
138+
J.NewClass nc = (J.NewClass) applied;
139+
if (nc.getArguments().size() != 1 || !(nc.getArguments().get(0) instanceof J.MethodInvocation)) {
140+
return applied;
141+
}
142+
J.MethodInvocation setCall = (J.MethodInvocation) nc.getArguments().get(0);
143+
List<Expression> setArgs = setCall.getArguments();
144+
List<Expression> withPrefixes = new ArrayList<>(setArgs.size());
145+
for (int i = 0; i < setArgs.size(); i++) {
146+
withPrefixes.add(setArgs.get(i).withPrefix(adds.get(i).getPrefix()));
147+
}
148+
return nc.withArguments(Collections.singletonList(setCall.withArguments(withPrefixes)));
149+
}
150+
123151
@Override
124152
public J visitBlock(J.Block block, ExecutionContext ctx) {
125-
Map<UUID, List<Expression>> rewrites = new HashMap<>();
153+
Map<UUID, List<J.MethodInvocation>> rewrites = new HashMap<>();
126154
Set<UUID> absorbedAddIds = new HashSet<>();
127155
identifyProseRewrites(block, rewrites, absorbedAddIds);
128156

@@ -138,7 +166,7 @@ public J visitBlock(J.Block block, ExecutionContext ctx) {
138166

139167
private void identifyProseRewrites(
140168
J.Block block,
141-
Map<UUID, List<Expression>> rewrites,
169+
Map<UUID, List<J.MethodInvocation>> rewrites,
142170
Set<UUID> absorbedAddIds) {
143171
List<Statement> stmts = block.getStatements();
144172
int i = 0;
@@ -156,7 +184,7 @@ private void identifyProseRewrites(
156184
}
157185
J.NewClass initializer = (J.NewClass) decl.getVariables().get(0).getInitializer();
158186

159-
List<Expression> args = new ArrayList<>();
187+
List<J.MethodInvocation> adds = new ArrayList<>();
160188
List<UUID> absorbedHere = new ArrayList<>();
161189
int j = i + 1;
162190
while (j < stmts.size()) {
@@ -165,12 +193,12 @@ private void identifyProseRewrites(
165193
if (arg == null || expressionReferences(arg, targetName)) {
166194
break;
167195
}
168-
args.add(arg);
196+
adds.add((J.MethodInvocation) next);
169197
absorbedHere.add(next.getId());
170198
j++;
171199
}
172-
if (args.size() >= 2 && initializer != null) {
173-
rewrites.put(initializer.getId(), args);
200+
if (adds.size() >= 2 && initializer != null) {
201+
rewrites.put(initializer.getId(), adds);
174202
absorbedAddIds.addAll(absorbedHere);
175203
i = j;
176204
} else {

‎src/test/java/org/openrewrite/java/migrate/util/UseListOfTest.java‎

Lines changed: 39 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -227,6 +227,41 @@ void foo() {
227227
);
228228
}
229229

230+
@Test
231+
void prosePreservesCommentsOnAddStatements() {
232+
//language=java
233+
rewriteRun(
234+
java(
235+
"""
236+
import java.util.ArrayList;
237+
import java.util.List;
238+
239+
class Test {
240+
void m() {
241+
List<String> names = new ArrayList<>();
242+
// the boss
243+
names.add("Bob");
244+
names.add("alice");
245+
}
246+
}
247+
""",
248+
"""
249+
import java.util.ArrayList;
250+
import java.util.List;
251+
252+
class Test {
253+
void m() {
254+
List<String> names = new ArrayList<>(List.of(
255+
// the boss
256+
"Bob",
257+
"alice"));
258+
}
259+
}
260+
"""
261+
)
262+
);
263+
}
264+
230265
@Test
231266
void proseAddChainCollapsedIntoArrayListConstructor() {
232267
// The mutable-ArrayList output preserves the ability to .add() / .sort() / etc. afterwards.
@@ -253,7 +288,10 @@ void m() {
253288
254289
class Test {
255290
void m() {
256-
List<String> names = new ArrayList<>(List.of("Bob", "alice", "Charlie"));
291+
List<String> names = new ArrayList<>(List.of(
292+
"Bob",
293+
"alice",
294+
"Charlie"));
257295
names.sort(null);
258296
}
259297
}

‎src/test/java/org/openrewrite/java/migrate/util/UseSetOfTest.java‎

Lines changed: 39 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -227,6 +227,41 @@ void foo() {
227227
);
228228
}
229229

230+
@Test
231+
void prosePreservesCommentsOnAddStatements() {
232+
//language=java
233+
rewriteRun(
234+
java(
235+
"""
236+
import java.util.HashSet;
237+
import java.util.Set;
238+
239+
class Test {
240+
void m() {
241+
Set<String> tags = new HashSet<>();
242+
// primary
243+
tags.add("alpha");
244+
tags.add("beta");
245+
}
246+
}
247+
""",
248+
"""
249+
import java.util.HashSet;
250+
import java.util.Set;
251+
252+
class Test {
253+
void m() {
254+
Set<String> tags = new HashSet<>(Set.of(
255+
// primary
256+
"alpha",
257+
"beta"));
258+
}
259+
}
260+
"""
261+
)
262+
);
263+
}
264+
230265
@Test
231266
void proseAddChainCollapsedIntoHashSetConstructor() {
232267
//language=java
@@ -251,7 +286,10 @@ void m() {
251286
252287
class Test {
253288
void m() {
254-
Set<String> tags = new HashSet<>(Set.of("alpha", "beta", "gamma"));
289+
Set<String> tags = new HashSet<>(Set.of(
290+
"alpha",
291+
"beta",
292+
"gamma"));
255293
}
256294
}
257295
"""

0 commit comments

Comments
 (0)