How to manage multiple modifications on the same AST

Question

I need to visit each method declaration within a class in order to add some comments to the methods in this class.

    void parseMethodDeclaration() {

        ICompilationUnit unit = getCompilationUnitByName(unitName);
        ASTParser parser = ASTParser.newParser(AST.JLS4);
        parser.setKind(ASTParser.K_COMPILATION_UNIT);
        parser.setSource(unit);
        parser.setResolveBindings(true);

        CompilationUnit parse = (CompilationUnit) parser.createAST(null);

         parse.accept(new ASTVisitor() {

            @Override
            public boolean visit(MethodDeclaration node) {

            String unitName = node.resolveBinding().getDeclaringClass().getName();

            // .. if some conditions are satisfied
            addComment(unitName, node);

            return super.visit(node);
        }
    });
}

And the addComment method is defined as follows to add comments to the given method:

void addComment(String unitName, MethodDeclaration method){

    ICompilationUnit unit = getCompilationUnitByName(unitName);
    AST ast = method.getAST();
    ASTRewrite rewriter = ASTRewrite.create(ast);
    ListRewrite listRewrite = rewriter.getListRewrite(method.getBody(), 
        Block.STATEMENTS_PROPERTY);

    Statement placeHolder = (Statement)rewriter.createStringPlaceholder("//MyComment",
        ASTNode.EMPTY_STATEMENT);

    listRewrite.insertFirst(placeHolder, null);

    TextEdit edits = rewriter.rewriteAST();
    Document document = new Document(unit.getSource());
    edits.apply(document);
    unit.getBuffer().setContents(document.get());
}

This code works properly to add comments to the first visited method, but it doesn't apply comments in the right place in other methods. The problem originates from the fact that I'm using the original AST to visit each method while its corresponding sourcecode (unit) will be changed after the comments are added to the first method. Then, listRewrite.insertFirst(placeHolder, null) will put the comment in a wrong place.

How can I resolve this issue. Is there anyway to update the original AST as well?

Answer 1

Finally, I found the answer by studying different open-source projects. Unfortunately, most examples on the web are dealing with simple cases in which they apply a series of changes immediately.

If you need to visit different entities of a unit, and accordingly change different places of that unit, the only thing you need is to keep the same instance of an ASTRewrite for the whole unit. IN other words, the creation, rewriting and applying of ASTRewrite, should be only once for each CompilationUnit.

For each Method declaration we break-down the addComment method:

    ListRewrite listRewrite = rewriter.getListRewrite(method.getBody(), 
        Block.STATEMENTS_PROPERTY);
    Statement placeHolder = (Statement)rewriter.createStringPlaceholder("//MyComment",
        ASTNode.EMPTY_STATEMENT);
    listRewrite.insertFirst(placeHolder, null);

Lastly, whenever the whole unit has been visited, we should apply changes once:

    TextEdit edits = rewriter.rewriteAST();
    Document document = new Document(unit.getSource());
    edits.apply(document);
    unit.getBuffer().setContents(document.get());

Thanks to open-source community :)