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

[maiz-an]

Generic Syntax Tree Modification for Adding Leading Comments

I came across your scenario about generically modifying a Ballerina syntax tree (ST) to inject leading comments at the start of a file. Since I couldn’t find an existing solution, I crafted one myself. Below, I’m sharing my approach, an analysis of challenges I encountered, and some suggestions for potential API enhancements. Hope this helps!

Solution

Here’s a working test case (paste it into SyntaxTreeModifierTest.java):

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

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

    // Locate the first token
    Token firstToken = firstToken(oldRoot);
    
    // Preserve existing minutiae
    for (Minutiae m : firstToken.leadingMinutiae()) {
        newLeadingMinutiae = newLeadingMinutiae.add(m);
    }

    // Modify the token
    Token modifiedFirstToken = firstToken.modify(newLeadingMinutiae, firstToken.trailingMinutiae());

    // Replace token in the tree
    TokenReplacer tokenReplacer = new TokenReplacer(firstToken, modifiedFirstToken);
    ModulePartNode newRoot = (ModulePartNode) oldRoot.apply(tokenReplacer);
    SyntaxTree newST = syntaxTree.modifyWith(newRoot);

    // Display results
    System.out.println("Original:\n" + syntaxTree.toSourceCode());
    System.out.println("===============");
    System.out.println("Modified:\n" + newST.toSourceCode());
}

// Helper to find the first token
private Token firstToken(Node node) {
    if (node instanceof Token) return (Token) node;
    return firstToken(((NonTerminalNode) node).children().get(0));
}

// Custom TreeModifier for token replacement
private static class TokenReplacer extends TreeModifier {
    private final Token target;
    private final 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;
    }
}

Output

Original:
// LINE
int foo;
===============
Modified:
// line1
// line2
// line3
// LINE
int foo;

This solution works generically, regardless of the file’s structure, unlike specific cases (e.g., modifying an import declaration).

Observations and Challenges

1. No Public API for First Token

The internal STNode has a firstToken() method, but SyntaxTree doesn’t expose it publicly. I wrote a custom firstToken(Node) helper, but it’s fragile (e.g., assumes non-empty children). Should SyntaxTree.firstToken() be added to the public API?

2. Why replaceNode() Failed

I initially tried:

SyntaxTree newST = syntaxTree.replaceNode(firstToken, modifiedFirstToken);

This replaced the root with the token, breaking the tree. It seems replaceNode() is designed for structural nodes (e.g., ModulePartNode), not leaf Tokens. Hence, I switched to a TokenReplacer extending TreeModifier.

3. Is This Optimal?

The TokenReplacer approach aligns with Ballerina’s design but feels verbose for a simple token swap. With hypothetical APIs like:

• SyntaxTree.firstToken()
• SyntaxTree.replaceToken(Token old, Token new)

The code could simplify to:

Token firstToken = syntaxTree.firstToken();
newLeadingMinutiae = newLeadingMinutiae.addAll(firstToken.leadingMinutiae());
Token modifiedFirstToken = firstToken.modify(newLeadingMinutiae, firstToken.trailingMinutiae());
SyntaxTree newST = syntaxTree.replaceToken(firstToken, modifiedFirstToken);

Suggestions

API Enhancement

I propose adding:

• SyntaxTree.firstToken(): Returns the first token in the tree.
• SyntaxTree.replaceToken(Token old, Token new): Swaps a token directly.

These would streamline use cases like adding comments or reformatting. Should I file a feature request for this?

Robust firstToken() Implementation

If added, it could look like:

public Token firstToken(Node node) {
    if (node == null) return null;
    if (node instanceof Token) return (Token) node;
    NonTerminalNode ntNode = (NonTerminalNode) node;
    return ntNode.children().isEmpty() ? null : firstToken(ntNode.children().get(0));
}

Closing Thoughts

This solution works, but the API could make it more elegant. Let me know what you think—happy to refine further or collaborate on a feature request!

Cheers,
🕷️