AddBraces { public static TranslationUnit transform(TranslationUnit tu) { return new StandardVisitor() { @Override public void visitIfStmt(IfStmt ifStmt) { super.visitIfStmt(ifStmt); if (!isBlock(ifStmt.getThenStmt())) { ifStmt.setThenStmt(makeBlock(ifStmt.getThenStmt())); } if (ifStmt.hasElseStmt() && !isBlock(ifStmt.getElseStmt())) { ifStmt.setElseStmt(makeBlock(ifStmt.getElseStmt())); } } @Override public void visitForStmt(ForStmt forStmt) { super.visitForStmt(forStmt); handleLoopStmt(forStmt); } @Override public void visitWhileStmt(WhileStmt whileStmt) { super.visitWhileStmt(whileStmt); handleLoopStmt(whileStmt); } @Override public void visitDoStmt(DoStmt doStmt) { super.visitDoStmt(doStmt); handleLoopStmt(doStmt); } private boolean isBlock(Stmt stmt) { return stmt instanceof BlockStmt; } private Stmt makeBlock(Stmt stmt) { return new BlockStmt(Arrays.asList(stmt), true); } private void handleLoopStmt(LoopStmt loopStmt) { if (!isBlock(loopStmt.getBody())) { loopStmt.setBody(makeBlock(loopStmt.getBody())); } } public TranslationUnit addBraces(TranslationUnit tu) { visit(tu); return tu; } }.addBraces(tu.clone()); } private  AddBraces(); static TranslationUnit transform(TranslationUnit tu);  }
@Test public void danglingElse() throws Exception { String program = "void main() { if (a) if (b) s1; else s2; }"; TranslationUnit tu = ParseHelper.parse(program); TranslationUnit transformed = AddBraces.transform(tu); assertNotEquals(tu, transformed); String programAfter = "void main() { if (a) { if (b) { s1; } } else { s2; } }"; assertEquals(PrettyPrinterVisitor.prettyPrintAsString( ParseHelper.parse(programAfter)), PrettyPrinterVisitor.prettyPrintAsString(transformed)); }
@Test public void loops() throws Exception { String program = "void main() { for (a; b; c) while (d) do e; while (f); }"; String programAfter = "void main() { for (a; b; c) { while (d) { do { e; } while (f); } } }"; TranslationUnit tu = ParseHelper.parse(program); TranslationUnit transformed = AddBraces.transform(tu); assertNotEquals(tu, transformed); assertEquals(PrettyPrinterVisitor.prettyPrintAsString( ParseHelper.parse(programAfter)), PrettyPrinterVisitor.prettyPrintAsString(transformed)); }