Best way to add comments at the top of a bal file by modifying syntax tree?

[lochana-chathura]

Playing around with modifying the syntax tree after a while. :) Today I encountered the scenario mentioned in the title. (Raised by @kalaiyarasiganeshalingam) As I couldn't find an already written similar code, thought of writing one myself.

One way of modifying ST is by knowing the internal structure. e.g. If we know the first line is an import declaration.
We could simply modify it by,

// Here modifiedFirstToken = modified import keyword by injecting new comments as leading minutiae
importDeclarationNode newImportDeclNode = importDeclNode.modify().withImportKeyword(modifiedFirstToken).apply();
SyntaxTree newST = syntaxTree.replaceNode(importDeclNode, newImportDeclNode);

However, when the structure is unknown (generic scenario), I had to write one new method (Token firstToken(Node node)) and a TokenReplacer by extending the TreeModifier.

My code is as follows (can run it by copying it to SyntaxTreeModifierTest.java class),

    @Test
    public void testAddingCommentAtTheBeginningOfBalFile() {
        // Generate minutiae to be injected
        MinutiaeList newLeadingMinutiae = NodeParser.parseModulePart(
                        "// line1" + System.lineSeparator() +
                                "// line2" + System.lineSeparator() +
                                "// line3" + System.lineSeparator())
                .eofToken().leadingMinutiae();

        // Get syntax tree
        SyntaxTree syntaxTree = SyntaxTree.from(TextDocuments.from("// LINE" + System.lineSeparator() + "int foo;"));
        ModulePartNode oldRoot = syntaxTree.rootNode();

        // Get the very first token of the bal file
        Token firstToken = firstToken(oldRoot);
        
        // Add the existing leading minutiae to the newLeadingMinutiae
        for (Minutiae m : firstToken.leadingMinutiae()) {
            newLeadingMinutiae = newLeadingMinutiae.add(m);
        }

        // Modify first token with new leading minutiae.
        Token modifiedFirstToken = firstToken.modify(newLeadingMinutiae, firstToken.trailingMinutiae());

        // Modify the syntax tree with the modifiedFirstToken.
        // SyntaxTree newST = syntaxTree.replaceNode(firstToken, modifiedFirstToken); <----------- [1]
        // TODO: check the reason why we cannot replace using the above single line
        TokenReplacer tokenReplacer = new TokenReplacer(firstToken, modifiedFirstToken);
        ModulePartNode newRoot = (ModulePartNode) oldRoot.apply(tokenReplacer);
        SyntaxTree newST = syntaxTree.modifyWith(newRoot);

        System.out.println("original code: \n" + syntaxTree.toSourceCode());
        System.out.println("===============");
        System.out.println("modified code: \n" + newST.toSourceCode());
    }

    Token firstToken(Node node) {
        if (node instanceof Token) {
            return (Token) node;
        }
        return firstToken(((NonTerminalNode) node).children().get(0));
    }

    private static class TokenReplacer extends TreeModifier {
        private Token target;
        private Token replacement;

        public TokenReplacer(Token target, Token replacement) {
            this.target = target;
            this.replacement = replacement;
        }

        @Override
        public Token transform(Token token) {
            return token == target ? replacement : token;
        }
    }

1. We have an API(firstToken()) to get the first token of a given STNode. However, no API to achieve the same in the external tree. So had to write Token firstToken(Node node) method. Should we add such API?
2. I expected [1] line marked in the code to work. However, for "Token"s it does not work as I intended. For tokens, syntaxTree.replaceNode() returns a new syntax with the root being that token. I'm not sure why. So I ended up writing a TokenReplacer to achieve the token replacement in the syntax tree.

Want to make sure if the above is the best way to achieve it?

[lochana-chathura]

looping @sameerajayasoma @rdulmina @IMS94 @malinthar @hasithaa