Regenerate source from Antlr4 ParseTree preserving whitespaces - antlr

Question

I am able to use Java.g4 grammar and generate lexers and parsers. This grammar is used to get started. I have a ParseTree and I walk it and change whatever I want and write it back to a Java source code file. The ParseTree is not directly changed though.
So for example this is the code. But I also modify method names and add annotations to methods. In some cases I also remove method parameters
#Override
public void enterClassDeclaration(JavaParser.ClassDeclarationContext ctx) {
printer.write( intervals );
printer.writeList(importFilter );
ParseTree pt = ctx.getChild(1);
/*
The class name is changed by a visitor because
we get a ParseTree back
*/
pt.accept(new ParseTreeVisitor<Object>() {
#Override
public Object visitChildren(RuleNode ruleNode) {
return null;
}
#Override
public Object visitErrorNode(ErrorNode errorNode) {
return null;
}
#Override
public Object visit(ParseTree parseTree) {
return null;
}
#Override
public Object visitTerminal(TerminalNode terminalNode) {
String className = terminalNode.getText();
System.out.println("Name of the class is [ " + className + "]");
printer.writeText( classModifier.get() + " class " + NEW_CLASS_IDENTIFIER );
return null;
}
});
}
But I am not sure how to print the changed Java code while preserving all the original whitespaces.
How is that done ?
Update : It seems that the whitespaces and comments are there but not accessible easily. So it looks like I need to specifically keep track of them and write them along with the code. Not sure though.
So more specifically the code is this.
package x;
import java.util.Enumeration;
import java.util.*;
As I hit the first ImportDeclarationContext I need to store all the hidden space tokens. When I write this code back I want to include those spaces too.
Solution :
Don't skip but add to a HIDDEN channel
//
// Whitespace and comments
//
WS : [ \t\r\n\u000C]+ -> channel(HIDDEN)
;
COMMENT
: '/*' .*? '*/' -> channel(HIDDEN)
;
Use code like this to get them back
I use this to get the whitespaces and comments before each method. But it should be possible to get whitespaces from other places too. I think.
((CommonTokenStream) tokens).getHiddenTokensToLeft( classOrInterfaceModifierContext.getStart().getTokenIndex(),
Token.HIDDEN_CHANNEL);