Merge pull request #43376 from KavinduZoysa/comment-as-a-node

Model comment as a Ballerina node

compiler/ballerina-parser/src/main/java/io/ballerina/compiler/syntax/tree/BlockStatementNode.java

@@ -40,6 +40,10 @@ public NodeList<StatementNode> statements() {
         return new NodeList<>(childInBucket(1));
     }
 
+    public NodeAndCommentList<Node> statementsWithComments() {
+        return new NodeAndCommentList<>(childInBucket(1), childInBucket(2));
+    }
+
     public Token closeBraceToken() {
         return childInBucket(2);
     }

compiler/ballerina-parser/src/main/java/io/ballerina/compiler/syntax/tree/CommentNode.java

@@ -0,0 +1,80 @@
+/*
+ * Copyright (c) 2024, WSO2 LLC. (http://wso2.com)
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.ballerina.compiler.syntax.tree;
+
+import io.ballerina.compiler.internal.parser.tree.STNode;
+import io.ballerina.compiler.internal.parser.tree.STNodeDiagnostic;
+import io.ballerina.compiler.internal.parser.tree.STToken;
+import io.ballerina.tools.diagnostics.Diagnostic;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+
+/**
+ * Represents a comment. This is not a part of the Ballerina syntax tree.
+ *
+ * @since 2201.10.0
+ */
+public class CommentNode extends NonTerminalNode {
+    private Node commentAttachedNode;
+    private Minutiae lastMinutiae;
+    private List<String> commentLines;
+
+    public CommentNode(STNode commentAttachedSTNode, int position, NonTerminalNode commentAttachedNode) {
+        super(commentAttachedSTNode, position, commentAttachedNode);
+    }
+
+    public Node getCommentAttachedNode() {
+        return this.commentAttachedNode;
+    }
+
+    public void setCommentAttachedNode(Node commentAttachedNode) {
+        this.commentAttachedNode = commentAttachedNode;
+    }
+
+    public Minutiae getLastMinutiae() {
+        return this.lastMinutiae;
+    }
+
+    public void setLastMinutiae(Minutiae lastMinutiae) {
+        this.lastMinutiae = lastMinutiae;
+    }
+
+    public List<String> getCommentLines() {
+        return this.commentLines;
+    }
+
+    public void setCommentLines(List<String> commentLines) {
+        this.commentLines = commentLines;
+    }
+
+    @Override
+    protected String[] childNames() {
+        return new String[0];
+    }
+
+    @Override
+    public void accept(NodeVisitor visitor) {
+        visitor.visit(this);
+    }
+
+    @Override
+    public <T> T apply(NodeTransformer<T> visitor) {
+        return visitor.transform(this);
+    }
+}

compiler/ballerina-parser/src/main/java/io/ballerina/compiler/syntax/tree/FunctionBodyBlockNode.java

@@ -45,6 +45,10 @@ public NodeList<StatementNode> statements() {
         return new NodeList<>(childInBucket(2));
     }
 
+    public NodeAndCommentList<Node> statementsWithComments() {
+        return new NodeAndCommentList<>(childInBucket(2), childInBucket(3));
+    }
+
     public Token closeBraceToken() {
         return childInBucket(3);
     }

compiler/ballerina-parser/src/main/java/io/ballerina/compiler/syntax/tree/NodeAndCommentList.java

@@ -0,0 +1,210 @@
+/*
+ * Copyright (c) 2024, WSO2 LLC. (http://wso2.com)
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.ballerina.compiler.syntax.tree;
+
+import io.ballerina.compiler.internal.parser.tree.STNode;
+import io.ballerina.compiler.internal.parser.tree.STNodeList;
+import io.ballerina.compiler.internal.syntax.NodeListUtils;
+
+import java.util.*;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+import java.util.stream.StreamSupport;
+
+import static io.ballerina.compiler.internal.syntax.NodeListUtils.rangeCheck;
+import static io.ballerina.compiler.internal.syntax.NodeListUtils.rangeCheckForAdd;
+
+/**
+ * Represent both nodes and attached comments to each node.
+ *
+ * @param <T> The type of Node
+ */
+public class NodeAndCommentList<T extends Node> implements Iterable<T> {
+    protected final STNodeList internalListNode;
+    protected final NonTerminalNode nonTerminalNode;
+    protected final int size;
+    protected final Node[] nodes;
+
+    NodeAndCommentList(NonTerminalNode nonTerminalNode, Token semicolon) {
+        this(nonTerminalNode, semicolon, nonTerminalNode.bucketCount() * 2 + 1);
+    }
+
+    protected NodeAndCommentList(NonTerminalNode nonTerminalNode, Token semicolon, int size) {
+        if (!NodeListUtils.isSTNodeList(nonTerminalNode.internalNode())) {
+            throw new IllegalArgumentException("An STNodeList instance is expected");
+        }
+
+        this.internalListNode = (STNodeList) nonTerminalNode.internalNode();
+        this.nonTerminalNode = nonTerminalNode;
+        this.nodes = new Node[size];
+        int nodeIndex = 0;
+        for (int i = 0; i < nonTerminalNode.bucketCount(); i++) {
+            Node node = nonTerminalNode.childInBucket(i);
+            CommentNode commentNode = getCommentNode(node);
+            if (commentNode != null) {
+                this.nodes[nodeIndex++] = commentNode;
+            }
+            this.nodes[nodeIndex++] = node;
+        }
+
+        CommentNode commentNodeBeforeEnd = getCommentNode(semicolon);
+        if (commentNodeBeforeEnd != null) {
+            this.nodes[nodeIndex++] = commentNodeBeforeEnd;
+        }
+        this.size = nodeIndex;
+    }
+
+    private CommentNode getCommentNode(Node node) {
+        List<String> commentLines = new ArrayList<>();
+        Minutiae lastMinutiae = null;
+        for (Minutiae minutiae : node.leadingMinutiae()) {
+            String[] splits = minutiae.text().split("// ");
+            if (splits.length >= 2) {
+                commentLines.add(splits[1]);
+                lastMinutiae = minutiae;
+            } else if (splits.length == 1 && splits[0].contains("//")) {
+                commentLines.add("");
+                lastMinutiae = minutiae;
+            }
+        }
+        if (commentLines.isEmpty()) {
+            return null;
+        }
+        CommentNode commentNode = new CommentNode(node.internalNode(), 0, null);
+        commentNode.setCommentAttachedNode(node);
+        commentNode.setLastMinutiae(lastMinutiae);
+        commentNode.setCommentLines(commentLines);
+        return commentNode;
+    }
+
+    // Positional access methods
+
+    public T get(int index) {
+        rangeCheck(index, size);
+        return (T) this.nodes[index];
+    }
+
+    // Modification methods
+
+    public NodeAndCommentList<T> add(T node) {
+        Objects.requireNonNull(node, "node should not be null");
+        return new NodeAndCommentList<>(internalListNode.add(node.internalNode()).createUnlinkedFacade(), null);
+    }
+
+    public NodeAndCommentList<T> add(int index, T node) {
+        Objects.requireNonNull(node, "node should not be null");
+        rangeCheckForAdd(index, size);
+        return new NodeAndCommentList<>(internalListNode.add(index, node.internalNode()).createUnlinkedFacade(), null);
+    }
+
+    public NodeAndCommentList<T> addAll(Collection<T> c) {
+        if (c.isEmpty()) {
+            return this;
+        }
+
+        List<STNode> stNodesToBeAdded = c.stream()
+                .map(node -> Objects.requireNonNull(node, "node should not be null"))
+                .map(Node::internalNode)
+                .collect(Collectors.toList());
+        return new NodeAndCommentList<>(internalListNode.addAll(stNodesToBeAdded).createUnlinkedFacade(), null);
+    }
+
+    public NodeAndCommentList<T> set(int index, T node) {
+        Objects.requireNonNull(node, "node should not be null");
+        rangeCheck(index, size);
+        if (nonTerminalNode.checkForReferenceEquality(index, node)) {
+            return this;
+        }
+
+        return new NodeAndCommentList<>(internalListNode.set(index, node.internalNode()).createUnlinkedFacade(), null);
+    }
+
+    public NodeAndCommentList<T> remove(int index) {
+        rangeCheck(index, size);
+        return new NodeAndCommentList<>(internalListNode.remove(index).createUnlinkedFacade(), null);
+    }
+
+    public NodeAndCommentList<T> remove(T node) {
+        Objects.requireNonNull(node, "node should not be null");
+        for (int bucket = 0; bucket < nonTerminalNode.bucketCount(); bucket++) {
+            if (nonTerminalNode.checkForReferenceEquality(bucket, node)) {
+                return remove(bucket);
+            }
+        }
+        return this;
+    }
+
+    @SuppressWarnings("SuspiciousMethodCalls")
+    public NodeAndCommentList<T> removeAll(Collection<T> c) {
+        if (c.isEmpty()) {
+            return this;
+        }
+        c.forEach(node -> Objects.requireNonNull(node, "node should not be null"));
+
+        List<STNode> toBeDeletedList = new ArrayList<>();
+        for (int bucket = 0; bucket < nonTerminalNode.bucketCount(); bucket++) {
+            Node childNode = nonTerminalNode.childBuckets[bucket];
+            if (c.contains(childNode)) {
+                toBeDeletedList.add(childNode.internalNode());
+            }
+        }
+
+        return new NodeAndCommentList<>(internalListNode.removeAll(toBeDeletedList).createUnlinkedFacade(), null);
+    }
+
+    //query methods
+
+    public int size() {
+        return this.size;
+    }
+
+    public boolean isEmpty() {
+        return this.size == 0;
+    }
+
+    @Override
+    public Iterator<T> iterator() {
+        return new NodeAndCommentListIterator();
+    }
+
+    public Stream<T> stream() {
+        return StreamSupport.stream(spliterator(), false);
+    }
+
+    NonTerminalNode underlyingListNode() {
+        return this.nonTerminalNode;
+    }
+
+    /**
+     * An iterator for this list of nodes.
+     *
+     * @since 2201.10.0
+     */
+    protected class NodeAndCommentListIterator implements Iterator<T> {
+        private int currentIndex = 0;
+
+        @Override
+        public boolean hasNext() {
+            return this.currentIndex < size;
+        }
+
+        @Override
+        public T next() {
+            return get(currentIndex++);
+        }
+    }
+}

compiler/ballerina-parser/src/main/java/io/ballerina/compiler/syntax/tree/NodeTransformer.java

@@ -956,6 +956,10 @@ public T transform(ReceiveFieldNode receiveFieldNode) {
         return transformSyntaxNode(receiveFieldNode);
     }
 
+    public T transform(CommentNode commentNode) {
+        return transformSyntaxNode(commentNode);
+    }
+
     // Tokens
 
     public T transform(Token token) {

compiler/ballerina-parser/src/main/java/io/ballerina/compiler/syntax/tree/NodeVisitor.java

@@ -954,6 +954,9 @@ public void visit(MemberTypeDescriptorNode memberTypeDescriptorNode) {
     public void visit(ReceiveFieldNode receiveFieldNode) {
         visitSyntaxNode(receiveFieldNode);
     }
+    public void visit(CommentNode commentNode) {
+        visitSyntaxNode(commentNode);
+    }
 
     // Tokens
 

compiler/ballerina-parser/src/test/java/io/ballerinalang/compiler/parser/test/ParserTestUtils.java

@@ -32,6 +32,7 @@
 import io.ballerina.compiler.internal.parser.tree.STNodeList;
 import io.ballerina.compiler.internal.parser.tree.STToken;
 import io.ballerina.compiler.internal.syntax.SyntaxUtils;
+import io.ballerina.compiler.syntax.tree.CommentNode;
 import io.ballerina.compiler.syntax.tree.Node;
 import io.ballerina.compiler.syntax.tree.SyntaxKind;
 import io.ballerina.compiler.syntax.tree.SyntaxTree;
@@ -47,6 +48,7 @@
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.Collection;
+import java.util.List;
 
 import static io.ballerina.compiler.internal.syntax.SyntaxUtils.isSTNodePresent;
 import static io.ballerinalang.compiler.parser.test.ParserTestConstants.CHILDREN_FIELD;
@@ -1476,4 +1478,14 @@ private static SyntaxKind getDocumentationKind(String kind) {
                 throw new UnsupportedOperationException("cannot find syntax kind: " + kind);
         }
     }
+
+    public static void assertCommentNode(Node node, List<String> comments) {
+        Assert.assertTrue(node instanceof CommentNode);
+        CommentNode commentNode = (CommentNode) node;
+        List<String> commentLines = commentNode.getCommentLines();
+        Assert.assertEquals(commentLines.size(), comments.size());
+        for (int i = 0; i < comments.size(); i++) {
+            Assert.assertEquals(commentLines.get(i), comments.get(i));
+        }
+    }
 }