Add leading comment support (#1550)

Author: AndrewNolte

include/slang/syntax/SyntaxNode.h

@@ -146,6 +146,36 @@ class SLANG_EXPORT SyntaxNode {
     /// Gets the number of (direct) children underneath this node in the tree.
     size_t getChildCount() const; // Note: implemented in AllSyntax.cpp
 
+    /// An iterator that walks through all tokens in a syntax node in order.
+    class SLANG_EXPORT token_iterator : public iterator_facade<token_iterator> {
+    public:
+        token_iterator() : node(nullptr), childIndex(0), currentToken() {}
+
+        parsing::Token dereference() const { return currentToken; }
+
+        void increment();
+
+        bool equals(const token_iterator& other) const;
+
+    private:
+        friend class SyntaxNode;
+
+        explicit token_iterator(const SyntaxNode* node);
+
+        void findNextToken();
+
+        const SyntaxNode* node;
+        size_t childIndex;
+        parsing::Token currentToken;
+        SmallVector<std::pair<const SyntaxNode*, size_t>, 4> stack;
+    };
+
+    /// Gets an iterator to the first token in this subtree.
+    token_iterator tokens_begin() const { return token_iterator(this); }
+
+    /// Gets an iterator representing the end of tokens in this subtree.
+    token_iterator tokens_end() const { return token_iterator(); }
+
     /// Returns true if this syntax node is "equivalent" to the other provided
     /// syntax node. Equivalence here is determined by the entire subtrees having
     /// the same kinds of syntax nodes in the same order and all leaf tokens

include/slang/syntax/SyntaxPrinter.h

@@ -43,6 +43,23 @@ class SLANG_EXPORT SyntaxPrinter {
     /// @return a reference to this object, to allow chaining additional method calls.
     SyntaxPrinter& print(const SyntaxNode& node);
 
+    /// Print the leading comments and trivia for the provided @a node.
+    /// This will attempt to grab the leading trivia that is intended to provide info for this
+    /// syntax node. This includes doc comments, but also other typical annotation styles. See
+    /// SyntaxPrinterTests.cpp for more detail.
+    /// @return a reference to this object, to allow chaining additional method calls.
+    SyntaxPrinter& printLeadingComments(const SyntaxNode& node);
+
+    /// Print the provided @a node to the internal buffer, with the leading comments and trivia on
+    /// the primary node. See @a printLeadingComments for detail on what is considered a leading
+    /// comment.
+    /// @return a reference to this object, to allow chaining additional method calls.
+    SyntaxPrinter& printWithLeadingComments(const SyntaxNode& node);
+
+    /// Print the provided @a node to the internal buffer, excluding any leading trivia.
+    /// @return a reference to this object, to allow chaining additional method calls.
+    SyntaxPrinter& printExcludingLeadingComments(const SyntaxNode& node);
+
     /// Print the provided @a tree to the internal buffer.
     /// @return a reference to this object, to allow chaining additional method calls.
     SyntaxPrinter& print(const SyntaxTree& tree);

source/syntax/SyntaxNode.cpp

@@ -148,6 +148,63 @@ TokenOrSyntax SyntaxNode::getChild(size_t index) {
     return visit(visitor, index);
 }
 
+SyntaxNode::token_iterator::token_iterator(const SyntaxNode* node) :
+    node(node), childIndex(0), currentToken() {
+    if (node)
+        findNextToken();
+}
+
+void SyntaxNode::token_iterator::increment() {
+    if (!node)
+        return;
+
+    // Move to the next child in the current node
+    childIndex++;
+    findNextToken();
+}
+
+bool SyntaxNode::token_iterator::equals(const token_iterator& other) const {
+    // begin iterator will have node set, end will have nullptr node
+    return node == other.node && childIndex == other.childIndex;
+}
+
+void SyntaxNode::token_iterator::findNextToken() {
+    while (true) {
+        // If we're out of children in the current node, pop from the stack
+        if (childIndex >= node->getChildCount()) {
+            if (stack.empty()) {
+                // We're done iterating
+                node = nullptr;
+                childIndex = 0;
+                return;
+            }
+
+            // Pop the last saved state
+            auto [parentNode, parentIndex] = stack.back();
+            stack.pop_back();
+            node = parentNode;
+            childIndex = parentIndex + 1;
+            continue;
+        }
+
+        auto token = node->childToken(childIndex);
+        if (token) {
+            currentToken = token;
+            return;
+        }
+
+        auto childNode = node->childNode(childIndex);
+        if (childNode) {
+            stack.push_back({node, childIndex});
+            node = childNode;
+            childIndex = 0;
+        }
+        else {
+            childIndex++;
+        }
+    }
+}
+
 const SyntaxNode* SyntaxNode::childNode(size_t index) const {
     auto child = getChild(index);
     if (child.isToken())

source/syntax/SyntaxPrinter.cpp

@@ -8,10 +8,12 @@
 #include "slang/syntax/SyntaxPrinter.h"
 
 #include "slang/parsing/ParserMetadata.h"
+#include "slang/parsing/Token.h"
 #include "slang/syntax/SyntaxNode.h"
 #include "slang/syntax/SyntaxTree.h"
 #include "slang/text/SourceLocation.h"
 #include "slang/text/SourceManager.h"
+#include "slang/util/Util.h"
 
 namespace slang::syntax {
 
@@ -118,6 +120,82 @@ SyntaxPrinter& SyntaxPrinter::print(const SyntaxNode& node) {
     return *this;
 }
 
+SyntaxPrinter& SyntaxPrinter::printLeadingComments(const SyntaxNode& node) {
+    auto triviaSpan = node.getFirstToken().trivia();
+    using Iterator = std::span<const Trivia>::iterator;
+    std::optional<Iterator> lastComment;
+    std::optional<Iterator> leadingCommentStart;
+
+    // Walk backwards through trivia until
+    // - block comment
+    // - double new line after seeing a comment
+    // This misses leading trivia at first line, although that's typically for license/file
+    auto findDocBoundary = [&]() {
+        bool lastIsNewline = false;
+        for (auto it = triviaSpan.rbegin(); it != triviaSpan.rend(); it++) {
+            const auto& trivia = *it;
+            switch (trivia.kind) {
+                case TriviaKind::EndOfLine:
+                    if (lastIsNewline && lastComment) {
+                        // found a double newline after a comment, stop here
+                        return;
+                    }
+                    leadingCommentStart = lastComment;
+                    lastIsNewline = true;
+                    break;
+                case TriviaKind::BlockComment:
+                    // the first block comment is the start
+                    leadingCommentStart = it.base() - 1;
+                    return;
+                case TriviaKind::LineComment:
+                    lastComment = it.base() - 1;
+                    [[fallthrough]];
+                default:
+                    lastIsNewline = false;
+            }
+        }
+    };
+    findDocBoundary();
+
+    if (leadingCommentStart) {
+        for (auto it = *leadingCommentStart; it != triviaSpan.end(); it++) {
+            print(*it);
+        }
+    }
+
+    return *this;
+}
+
+SyntaxPrinter& SyntaxPrinter::printExcludingLeadingComments(const SyntaxNode& node) {
+    if (!includeTrivia) {
+        print(node);
+        return *this;
+    }
+
+    auto it = node.tokens_begin();
+    auto end = node.tokens_end();
+
+    if (it == end)
+        return *this;
+
+    // Print first token without leading trivia
+    includeTrivia = false;
+    print(*it);
+    includeTrivia = true;
+
+    for (++it; it != end; ++it) {
+        print(*it);
+    }
+
+    return *this;
+}
+
+SyntaxPrinter& SyntaxPrinter::printWithLeadingComments(const SyntaxNode& node) {
+    printLeadingComments(node);
+    printExcludingLeadingComments(node);
+    return *this;
+}
+
 SyntaxPrinter& SyntaxPrinter::print(const SyntaxTree& tree) {
     print(tree.root());
     if (tree.root().kind != SyntaxKind::CompilationUnit && tree.getMetadata().eofToken)

tests/unittests/CMakeLists.txt

@@ -38,6 +38,7 @@ add_executable(
   parsing/RewriterExpandTests.cpp
   parsing/VisitorTests.cpp
   parsing/StatementParsingTests.cpp
+  parsing/SyntaxPrinterTests.cpp
   util/CommandLineTests.cpp
   util/IntervalMapTests.cpp
   util/NumericTests.cpp