Skip to content

Commit 3209dcd

Browse files
Jammy-LouieJammy-Louie
committed
Simplify buildBlock method and improve handling of comments
1 parent e6f509e commit 3209dcd

File tree

1 file changed

+60
-18
lines changed

1 file changed

+60
-18
lines changed

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

+60-18
Original file line numberDiff line numberDiff line change
@@ -69,20 +69,42 @@ private static class NeedBracesVisitor extends JavaIsoVisitor<ExecutionContext>
6969
* We can use that to our advantage by saying if you aren't a block (e.g. a single {@link Statement}, etc.),
7070
* then we're going to make this into a block. That's how we'll get the code bodies surrounded in braces.
7171
*/
72-
private <T extends Statement> J.Block buildBlock(Statement owner, T element) {
73-
J j = getCursor().getParentTreeCursor().getValue();
72+
private <T extends Statement> J.Block buildBlock(T element) {
73+
J rootElement = null;
7474
Space end = Space.EMPTY;
75-
if (j instanceof J.Block) {
76-
J.Block block = (J.Block) j;
75+
76+
Cursor currentCursor = getCursor();
77+
while (
78+
currentCursor != null &&
79+
currentCursor.getParent() != null &&
80+
!(currentCursor.getParent().getValue() instanceof J.Block)
81+
) {
82+
currentCursor = currentCursor.getParent();
83+
}
84+
85+
if (currentCursor != null && currentCursor.getValue() instanceof JRightPadded) {
86+
JRightPadded<J> paddedIf = currentCursor.getValue();
87+
rootElement = paddedIf.getElement();
88+
}
89+
90+
if (rootElement instanceof Statement) {
91+
Cursor blockParentCursor = currentCursor.getParent();
92+
J.Block block = blockParentCursor.getValue();
7793
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);
94+
int currentIndex = statements.indexOf(rootElement);
95+
boolean last = currentIndex == statements.size() - 1;
96+
Space trailingComment = last ? block.getEnd() : statements.get(currentIndex + 1).getPrefix();
97+
98+
if (!trailingComment.isEmpty() && !trailingComment.getWhitespace().contains("\n")) {
99+
end = trailingComment;
100+
if (last) {
101+
blockParentCursor.putMessage("removeEndComments", true);
102+
} else {
103+
blockParentCursor.<List<Integer>>computeMessageIfAbsent("replaced", k -> new ArrayList<>()).add(currentIndex);
104+
}
84105
}
85106
}
107+
86108
return new J.Block(
87109
Tree.randomId(),
88110
Space.EMPTY,
@@ -111,6 +133,11 @@ private <T extends Statement> J.Block buildBlock(Statement owner, T element) {
111133
@Override
112134
public J.Block visitBlock(J.Block block, ExecutionContext ctx) {
113135
J.Block bl = super.visitBlock(block, ctx);
136+
Boolean removeEndComments = getCursor().pollMessage("removeEndComments");
137+
if (removeEndComments != null) {
138+
bl = bl.withEnd(bl.getEnd().withComments(Collections.emptyList()));
139+
bl = maybeAutoFormat(block, bl, ctx);
140+
}
114141
List<Integer> indexes = getCursor().pollMessage("replaced");
115142
if (indexes != null) {
116143
for (int index : indexes) {
@@ -140,7 +167,15 @@ public J.If visitIf(J.If iff, ExecutionContext ctx) {
140167
J.If elem = super.visitIf(iff, ctx);
141168
boolean hasAllowableBodyType = elem.getThenPart() instanceof J.Block;
142169
if (!needBracesStyle.getAllowSingleLineStatement() && !hasAllowableBodyType) {
143-
J.Block b = buildBlock(elem, elem.getThenPart());
170+
J.Block b;
171+
if (elem.getElsePart() != null && !elem.getElsePart().getPrefix().getComments().isEmpty()) {
172+
Space end = elem.getElsePart().getPrefix();
173+
elem = elem.withElsePart(elem.getElsePart().withPrefix(Space.EMPTY));
174+
b = buildBlock(elem.getThenPart()).withEnd(end);
175+
} else {
176+
b = buildBlock(elem.getThenPart());
177+
}
178+
144179
elem = maybeAutoFormat(elem, elem.withThenPart(b), ctx);
145180
}
146181
return elem;
@@ -155,7 +190,14 @@ public J.If.Else visitElse(J.If.Else else_, ExecutionContext ctx) {
155190
J.If.Else elem = super.visitElse(else_, ctx);
156191
boolean hasAllowableBodyType = elem.getBody() instanceof J.Block || elem.getBody() instanceof J.If;
157192
if (!needBracesStyle.getAllowSingleLineStatement() && !hasAllowableBodyType) {
158-
J.Block b = buildBlock(getCursor().getParentTreeCursor().getValue(), elem.getBody());
193+
Space prefix = elem.getPrefix();
194+
Statement body = elem.getBody();
195+
196+
if (!prefix.getComments().isEmpty() && prefix.getWhitespace().contains("\n")) {
197+
body = body.withPrefix(prefix);
198+
}
199+
200+
J.Block b = buildBlock(body);
159201
elem = maybeAutoFormat(elem, elem.withBody(b), ctx);
160202
}
161203
return elem;
@@ -168,10 +210,10 @@ public J.WhileLoop visitWhileLoop(J.WhileLoop whileLoop, ExecutionContext ctx) {
168210
elem.getBody() instanceof J.Block || elem.getBody() instanceof J.Empty :
169211
elem.getBody() instanceof J.Block;
170212
if (!needBracesStyle.getAllowEmptyLoopBody() && elem.getBody() instanceof J.Empty) {
171-
J.Block b = buildBlock(elem, elem.getBody());
213+
J.Block b = buildBlock(elem.getBody());
172214
elem = maybeAutoFormat(elem, elem.withBody(b), ctx);
173215
} else if (!needBracesStyle.getAllowSingleLineStatement() && !hasAllowableBodyType) {
174-
J.Block b = buildBlock(elem, elem.getBody());
216+
J.Block b = buildBlock(elem.getBody());
175217
elem = maybeAutoFormat(elem, elem.withBody(b), ctx);
176218
}
177219
return elem;
@@ -184,10 +226,10 @@ public J.DoWhileLoop visitDoWhileLoop(J.DoWhileLoop doWhileLoop, ExecutionContex
184226
elem.getBody() instanceof J.Block || elem.getBody() instanceof J.Empty :
185227
elem.getBody() instanceof J.Block;
186228
if (!needBracesStyle.getAllowEmptyLoopBody() && elem.getBody() instanceof J.Empty) {
187-
J.Block b = buildBlock(elem, elem.getBody());
229+
J.Block b = buildBlock(elem.getBody());
188230
elem = maybeAutoFormat(elem, elem.withBody(b), ctx);
189231
} else if (!needBracesStyle.getAllowSingleLineStatement() && !hasAllowableBodyType) {
190-
J.Block b = buildBlock(elem, elem.getBody());
232+
J.Block b = buildBlock(elem.getBody());
191233
elem = maybeAutoFormat(elem, elem.withBody(b), ctx);
192234
}
193235
return elem;
@@ -200,10 +242,10 @@ public J.ForLoop visitForLoop(J.ForLoop forLoop, ExecutionContext ctx) {
200242
elem.getBody() instanceof J.Block || elem.getBody() instanceof J.Empty :
201243
elem.getBody() instanceof J.Block;
202244
if (!needBracesStyle.getAllowEmptyLoopBody() && elem.getBody() instanceof J.Empty) {
203-
J.Block b = buildBlock(elem, elem.getBody());
245+
J.Block b = buildBlock(elem.getBody());
204246
elem = maybeAutoFormat(elem, elem.withBody(b), ctx);
205247
} else if (!needBracesStyle.getAllowSingleLineStatement() && !hasAllowableBodyType) {
206-
J.Block b = buildBlock(elem, elem.getBody());
248+
J.Block b = buildBlock(elem.getBody());
207249
elem = maybeAutoFormat(elem, elem.withBody(b), ctx);
208250
}
209251
return elem;

0 commit comments

Comments
 (0)