[P4fmt]: attach comments

Signed-off-by: Nitish <snapdgnn@proton.me>
p4lang · snapdgn · Aug 1, 2024 · Aug 26, 2024 · Aug 26, 2024 · Aug 27, 2024
commit b424c0628bf6852823def7be57629c130af80ebc
diff --git a/backends/p4fmt/CMakeLists.txt b/backends/p4fmt/CMakeLists.txt
@@ -3,13 +3,15 @@ set(FMT_SRCS
     options.cpp
     main.cpp
     p4formatter.cpp
+    attach.cpp
 )
 
 set(REFCHECK_SRCS
     refcheck.cpp
     options.cpp
     p4fmt.cpp
     p4formatter.cpp
+    attach.cpp
 )
 
 # p4fmt

diff --git a/backends/p4fmt/attach.cpp b/backends/p4fmt/attach.cpp
@@ -0,0 +1,81 @@
+#include "backends/p4fmt/attach.h"
+
+#include "frontends/common/parser_options.h"
+#include "lib/source_file.h"
+
+namespace P4::P4Fmt {
+
+Attach::~Attach() = default;
+
+void Attach::addPrefixComments(NodeId node, const Util::Comment *prefix) {
+    commentsMap[node].prefix.push_back(prefix);
+}
+
+void Attach::addSuffixComments(NodeId node, const Util::Comment *suffix) {
+    commentsMap[node].suffix.push_back(suffix);
+}
+
+bool Attach::isSystemFile(const std::filesystem::path &file) {
+    const std::filesystem::path p4include(p4includePath);
+    return file.parent_path() == p4include;
+}
+
+const Attach::CommentsMap &Attach::getCommentsMap() const { return commentsMap; }
+
+const IR::Node *Attach::attachCommentsToNode(IR::Node *node, TraversalType ttype) {
+    if (node == nullptr || !node->srcInfo.isValid() || clist.empty()) {
+        return node;
+    }
+    std::filesystem::path sourceFile(node->srcInfo.getSourceFile().c_str());
+    if (isSystemFile(sourceFile)) {
+        // Skip attachment for system files
+        return node;
+    }
+
+    const auto nodeStart = node->srcInfo.getStart();
+
+    auto itr = clist.begin();
+    while (itr != clist.end()) {
+        const Util::Comment *comment = *itr;
+        const auto &commentEnd = comment->getEndPosition();
+
+        bool shouldRemove = false;
+
+        switch (ttype) {
+            {
+                case TraversalType::Preorder:
+                    if (commentEnd.getLineNumber() < nodeStart.getLineNumber()) {
+                        addPrefixComments(node->id, comment);
+                        shouldRemove = true;
+                    }
+                    break;
+            }
+
+            {
+                case TraversalType::Postorder:
+                    if (commentEnd.getLineNumber() == nodeStart.getLineNumber()) {
+                        addSuffixComments(node->id, comment);
+                        shouldRemove = true;
+                    }
+                    break;
+            }
+            default:
+                ::P4::error(ErrorType::ERR_INVALID, "traversal type unknown/unsupported.");
+                return node;
+        }
+
+        itr = shouldRemove ? clist.erase(itr) : std::next(itr);
+    }
+
+    return node;
+}
+
+const IR::Node *Attach::preorder(IR::Node *node) {
+    return attachCommentsToNode(node, TraversalType::Preorder);
+}
+
+const IR::Node *Attach::postorder(IR::Node *node) {
+    return attachCommentsToNode(node, TraversalType::Postorder);
+}
+
+}  // namespace P4::P4Fmt
diff --git a/backends/p4fmt/attach.h b/backends/p4fmt/attach.h
@@ -0,0 +1,49 @@
+#ifndef BACKENDS_P4FMT_ATTACH_H_
+#define BACKENDS_P4FMT_ATTACH_H_
+
+#include <unordered_map>
+#include <unordered_set>
+#include <vector>
+
+#include "backends/p4fmt/p4fmt.h"
+#include "ir/ir.h"
+#include "ir/visitor.h"
+#include "lib/source_file.h"
+
+namespace P4::P4Fmt {
+
+class Attach : public Transform {
+ public:
+    using NodeId = int;
+    struct Comments {
+        std::vector<const Util::Comment *> prefix;
+        std::vector<const Util::Comment *> suffix;
+    };
+    using CommentsMap = std::unordered_map<NodeId, Comments>;
+    enum class TraversalType { Preorder, Postorder };
+
+    explicit Attach(const std::unordered_set<const Util::Comment *> &clist) : clist(clist){};
+    ~Attach() override;
+
+    const IR::Node *attachCommentsToNode(IR::Node *, TraversalType);
+
+    using Transform::postorder;
+    using Transform::preorder;
+
+    const IR::Node *preorder(IR::Node *node) override;
+    const IR::Node *postorder(IR::Node *node) override;
+
+    static bool isSystemFile(const std::filesystem::path &file);
+
+    void addPrefixComments(NodeId, const Util::Comment *);
+    void addSuffixComments(NodeId, const Util::Comment *);
+    const CommentsMap &getCommentsMap() const;
+
+ private:
+    std::unordered_set<const Util::Comment *> clist;
+    CommentsMap commentsMap;
+};
+
+}  // namespace P4::P4Fmt
+
+#endif /* BACKENDS_P4FMT_ATTACH_H_ */
diff --git a/backends/p4fmt/p4fmt.cpp b/backends/p4fmt/p4fmt.cpp
@@ -1,5 +1,6 @@
 #include "backends/p4fmt/p4fmt.h"
 
+#include "backends/p4fmt/attach.h"
 #include "frontends/common/parseInput.h"
 #include "frontends/common/parser_options.h"
 #include "ir/ir.h"
@@ -24,7 +25,20 @@ std::stringstream getFormattedOutput(std::filesystem::path inputFile) {
         return formattedOutput;
     }
 
+    std::unordered_set<const Util::Comment *> globalCommentsList;
+    // get the global list of comments
+    if (!program->objects.empty()) {
+        const auto *firstNode = program->objects.front();
+        if (firstNode->srcInfo.isValid()) {
+            for (const auto *comment : firstNode->srcInfo.getAllFileComments()) {
+                globalCommentsList.insert(comment);
+            }
+        }
+    }
+
     auto top4 = P4Fmt::P4Formatter(&formattedOutput);
+    auto attach = P4::P4Fmt::Attach(globalCommentsList);
+    program = program->apply(attach);
     // Print the program before running front end passes.
     program->apply(top4);
 

diff --git a/lib/source_file.cpp b/lib/source_file.cpp
@@ -69,6 +69,13 @@ std::ostream &operator<<(std::ostream &os, const SourceInfo &info) {
     return os;
 }
 
+const std::vector<Comment *> SourceInfo::getAllFileComments() const {
+    if (sources == nullptr) {
+        return {};
+    }
+    return this->sources->getAllComments();
+}
+
 //////////////////////////////////////////////////////////////////////////////////////////
 
 InputSources::InputSources() : sealed(false) {
@@ -83,6 +90,8 @@ void InputSources::addComment(SourceInfo srcInfo, bool singleLine, cstring body)
     comments.push_back(new Comment(srcInfo, singleLine, body));
 }
 
+const std::vector<Comment *> &InputSources::getAllComments() const { return comments; }
+
 /// prevent further changes
 void InputSources::seal() {
     LOG4(toDebugString());

diff --git a/lib/source_file.h b/lib/source_file.h
@@ -110,6 +110,7 @@ class SourcePosition final {
 };
 
 class InputSources;
+class Comment;
 
 /**
 Information about the source position of a language element -
@@ -188,6 +189,8 @@ class SourceInfo final {
 
     const SourcePosition &getEnd() const { return this->end; }
 
+    [[nodiscard]] const std::vector<Comment *> getAllFileComments() const;
+
     /**
        True if this comes 'before' this source position.
        'invalid' source positions come first.
@@ -261,6 +264,10 @@ class Comment final : IHasDbPrint {
         out << body;
         if (!singleLine) out << "*/";
     }
+
+    // Retrieve the source position associated with this comment.
+    [[nodiscard]] const SourcePosition &getStartPosition() const { return srcInfo.getStart(); }
+    [[nodiscard]] const SourcePosition &getEndPosition() const { return srcInfo.getEnd(); }
 };
 
 /**
@@ -280,7 +287,6 @@ class InputSources final {
 
  public:
     InputSources();
-
     std::string_view getLine(unsigned lineNumber) const;
     /// Original source line that produced the line with the specified number
     SourceFileLine getSourceLine(unsigned line) const;
@@ -312,6 +318,9 @@ class InputSources final {
     cstring toDebugString() const;
     void addComment(SourceInfo srcInfo, bool singleLine, cstring body);
 
+    // Returns a list of all the comments found in the file, stored as a part of `InputSources`
+    const std::vector<Comment *> &getAllComments() const;
+
  private:
     /// Append this text to the last line; must not contain newlines
     void appendToLastLine(std::string_view text);