Skip to content

Commit 60fb46d

Browse files
Jammy-LouietimtebeekJenson3210
authored
Fix Comment handling in NeedsBraces (#518)
* Initial Tests * Simplify buildBlock method and improve handling of comments * apply code suggestion * Update src/main/java/org/openrewrite/staticanalysis/NeedBraces.java Co-authored-by: Jente Sondervorst <[email protected]> * Do not move comments on do while loop --------- Co-authored-by: Tim te Beek <[email protected]> Co-authored-by: Jente Sondervorst <[email protected]>
1 parent 7358214 commit 60fb46d

File tree

2 files changed

+267
-41
lines changed

2 files changed

+267
-41
lines changed

‎src/main/java/org/openrewrite/staticanalysis/NeedBraces.java

+64-19
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,10 @@
2121
import org.openrewrite.java.JavaIsoVisitor;
2222
import org.openrewrite.java.style.Checkstyle;
2323
import org.openrewrite.java.style.NeedBracesStyle;
24-
import org.openrewrite.java.tree.*;
24+
import org.openrewrite.java.tree.J;
25+
import org.openrewrite.java.tree.JRightPadded;
26+
import org.openrewrite.java.tree.Space;
27+
import org.openrewrite.java.tree.Statement;
2528
import org.openrewrite.kotlin.tree.K;
2629
import org.openrewrite.marker.Markers;
2730

@@ -69,20 +72,43 @@ private static class NeedBracesVisitor extends JavaIsoVisitor<ExecutionContext>
6972
* We can use that to our advantage by saying if you aren't a block (e.g. a single {@link Statement}, etc.),
7073
* then we're going to make this into a block. That's how we'll get the code bodies surrounded in braces.
7174
*/
72-
private <T extends Statement> J.Block buildBlock(Statement owner, T element) {
73-
J j = getCursor().getParentTreeCursor().getValue();
75+
private <T extends Statement> J.Block buildBlock(T element) {
76+
J rootElement = null;
7477
Space end = Space.EMPTY;
75-
if (j instanceof J.Block) {
76-
J.Block block = (J.Block) j;
78+
79+
Cursor currentCursor = getCursor();
80+
while (
81+
currentCursor != null &&
82+
currentCursor.getParent() != null &&
83+
!(currentCursor.getParent().getValue() instanceof J.Block)
84+
) {
85+
currentCursor = currentCursor.getParent();
86+
}
87+
88+
if (currentCursor != null && currentCursor.getValue() instanceof JRightPadded) {
89+
JRightPadded<J> paddedIf = currentCursor.getValue();
90+
rootElement = paddedIf.getElement();
91+
}
92+
93+
// Move comments
94+
if (rootElement instanceof Statement && !(rootElement instanceof J.DoWhileLoop)) {
95+
Cursor blockParentCursor = currentCursor.getParent();
96+
J.Block block = blockParentCursor.getValue();
7797
List<Statement> statements = block.getStatements();
78-
int i = statements.indexOf(owner);
79-
boolean last = i == statements.size() - 1;
80-
Space trailingSpace = last ? block.getEnd() : statements.get(i + 1).getPrefix();
81-
if (!trailingSpace.getComments().isEmpty() && trailingSpace.getWhitespace().indexOf('\n') == -1) {
82-
end = trailingSpace;
83-
getCursor().getParentTreeCursor().<List<Integer>>computeMessageIfAbsent("replaced", k -> new ArrayList<>()).add(i);
98+
int currentIndex = statements.indexOf(rootElement);
99+
boolean last = currentIndex == statements.size() - 1;
100+
Space trailingComment = last ? block.getEnd() : statements.get(currentIndex + 1).getPrefix();
101+
102+
if (!trailingComment.isEmpty() && !trailingComment.getWhitespace().contains("\n")) {
103+
end = trailingComment;
104+
if (last) {
105+
blockParentCursor.putMessage("removeEndComments", true);
106+
} else {
107+
blockParentCursor.<List<Integer>>computeMessageIfAbsent("replaced", k -> new ArrayList<>()).add(currentIndex);
108+
}
84109
}
85110
}
111+
86112
return new J.Block(
87113
Tree.randomId(),
88114
Space.EMPTY,
@@ -111,6 +137,10 @@ private <T extends Statement> J.Block buildBlock(Statement owner, T element) {
111137
@Override
112138
public J.Block visitBlock(J.Block block, ExecutionContext ctx) {
113139
J.Block bl = super.visitBlock(block, ctx);
140+
if (Boolean.TRUE.equals(getCursor().pollMessage("removeEndComments"))) {
141+
bl = bl.withEnd(bl.getEnd().withComments(Collections.emptyList()));
142+
bl = maybeAutoFormat(block, bl, ctx);
143+
}
114144
List<Integer> indexes = getCursor().pollMessage("replaced");
115145
if (indexes != null) {
116146
for (int index : indexes) {
@@ -140,7 +170,15 @@ public J.If visitIf(J.If iff, ExecutionContext ctx) {
140170
J.If elem = super.visitIf(iff, ctx);
141171
boolean hasAllowableBodyType = elem.getThenPart() instanceof J.Block;
142172
if (!needBracesStyle.getAllowSingleLineStatement() && !hasAllowableBodyType) {
143-
J.Block b = buildBlock(elem, elem.getThenPart());
173+
J.Block b;
174+
if (elem.getElsePart() != null && !elem.getElsePart().getPrefix().getComments().isEmpty()) {
175+
Space end = elem.getElsePart().getPrefix();
176+
elem = elem.withElsePart(elem.getElsePart().withPrefix(Space.EMPTY));
177+
b = buildBlock(elem.getThenPart()).withEnd(end);
178+
} else {
179+
b = buildBlock(elem.getThenPart());
180+
}
181+
144182
elem = maybeAutoFormat(elem, elem.withThenPart(b), ctx);
145183
}
146184
return elem;
@@ -155,7 +193,14 @@ public J.If.Else visitElse(J.If.Else else_, ExecutionContext ctx) {
155193
J.If.Else elem = super.visitElse(else_, ctx);
156194
boolean hasAllowableBodyType = elem.getBody() instanceof J.Block || elem.getBody() instanceof J.If;
157195
if (!needBracesStyle.getAllowSingleLineStatement() && !hasAllowableBodyType) {
158-
J.Block b = buildBlock(getCursor().getParentTreeCursor().getValue(), elem.getBody());
196+
Space prefix = elem.getPrefix();
197+
Statement body = elem.getBody();
198+
199+
if (!prefix.getComments().isEmpty() && prefix.getWhitespace().contains("\n")) {
200+
body = body.withPrefix(prefix);
201+
}
202+
203+
J.Block b = buildBlock(body);
159204
elem = maybeAutoFormat(elem, elem.withBody(b), ctx);
160205
}
161206
return elem;
@@ -168,10 +213,10 @@ public J.WhileLoop visitWhileLoop(J.WhileLoop whileLoop, ExecutionContext ctx) {
168213
elem.getBody() instanceof J.Block || elem.getBody() instanceof J.Empty :
169214
elem.getBody() instanceof J.Block;
170215
if (!needBracesStyle.getAllowEmptyLoopBody() && elem.getBody() instanceof J.Empty) {
171-
J.Block b = buildBlock(elem, elem.getBody());
216+
J.Block b = buildBlock(elem.getBody());
172217
elem = maybeAutoFormat(elem, elem.withBody(b), ctx);
173218
} else if (!needBracesStyle.getAllowSingleLineStatement() && !hasAllowableBodyType) {
174-
J.Block b = buildBlock(elem, elem.getBody());
219+
J.Block b = buildBlock(elem.getBody());
175220
elem = maybeAutoFormat(elem, elem.withBody(b), ctx);
176221
}
177222
return elem;
@@ -184,10 +229,10 @@ public J.DoWhileLoop visitDoWhileLoop(J.DoWhileLoop doWhileLoop, ExecutionContex
184229
elem.getBody() instanceof J.Block || elem.getBody() instanceof J.Empty :
185230
elem.getBody() instanceof J.Block;
186231
if (!needBracesStyle.getAllowEmptyLoopBody() && elem.getBody() instanceof J.Empty) {
187-
J.Block b = buildBlock(elem, elem.getBody());
232+
J.Block b = buildBlock(elem.getBody());
188233
elem = maybeAutoFormat(elem, elem.withBody(b), ctx);
189234
} else if (!needBracesStyle.getAllowSingleLineStatement() && !hasAllowableBodyType) {
190-
J.Block b = buildBlock(elem, elem.getBody());
235+
J.Block b = buildBlock(elem.getBody());
191236
elem = maybeAutoFormat(elem, elem.withBody(b), ctx);
192237
}
193238
return elem;
@@ -200,10 +245,10 @@ public J.ForLoop visitForLoop(J.ForLoop forLoop, ExecutionContext ctx) {
200245
elem.getBody() instanceof J.Block || elem.getBody() instanceof J.Empty :
201246
elem.getBody() instanceof J.Block;
202247
if (!needBracesStyle.getAllowEmptyLoopBody() && elem.getBody() instanceof J.Empty) {
203-
J.Block b = buildBlock(elem, elem.getBody());
248+
J.Block b = buildBlock(elem.getBody());
204249
elem = maybeAutoFormat(elem, elem.withBody(b), ctx);
205250
} else if (!needBracesStyle.getAllowSingleLineStatement() && !hasAllowableBodyType) {
206-
J.Block b = buildBlock(elem, elem.getBody());
251+
J.Block b = buildBlock(elem.getBody());
207252
elem = maybeAutoFormat(elem, elem.withBody(b), ctx);
208253
}
209254
return elem;

0 commit comments

Comments
 (0)