p4lang · snapdgn · Aug 1, 2024 · Aug 26, 2024 · Aug 26, 2024 · Aug 27, 2024
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,78 @@
+#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() || processedComments.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();
+
+    for (auto &[comment, isAttached] : processedComments) {
+        // Skip if already attached
+        if (isAttached) {
+            continue;
+        }
+
+        const auto &commentEnd = comment->getEndPosition();
+
+        switch (ttype) {
+            case TraversalType::Preorder:
+                if (commentEnd.getLineNumber() == nodeStart.getLineNumber() - 1) {
+                    addPrefixComments(node->id, comment);
+                    isAttached = true;  // Mark the comment as attached
+                }
+                break;
+
+            case TraversalType::Postorder:
+                if (commentEnd.getLineNumber() == nodeStart.getLineNumber()) {
+                    addSuffixComments(node->id, comment);
+                    isAttached = true;
+                }
+                break;
+
+            default:
+                ::P4::error(ErrorType::ERR_INVALID, "traversal type unknown/unsupported.");
+                return node;
+        }
+    }
+
+    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,51 @@
+#ifndef BACKENDS_P4FMT_ATTACH_H_
+#define BACKENDS_P4FMT_ATTACH_H_
+
+#include <filesystem>
+#include <unordered_map>
+#include <vector>
+
+#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_map<const Util::Comment *, bool> &processedComments)
+        : processedComments(processedComments){};
+    ~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:
+    // This Hashmap tracks each comment’s attachment status to IR nodes. Initially, all comments are
+    // set to 'false'.
+    std::unordered_map<const Util::Comment *, bool> processedComments;
+
+    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,8 +25,23 @@ std::stringstream getFormattedOutput(std::filesystem::path inputFile) {
         return formattedOutput;
     }
 
-    auto top4 = P4Fmt::P4Formatter(&formattedOutput);
+    std::unordered_map<const Util::Comment *, bool> globalCommentsMap;
+
+    // Initialize the global comments map from the list of comments in the program.
+    if (!program->objects.empty()) {
+        const auto *firstNode = program->objects.front();
+        if (firstNode->srcInfo.isValid()) {
+            for (const auto *comment : firstNode->srcInfo.getAllFileComments()) {
+                globalCommentsMap[comment] =
+                    false;  // Initialize all comments as not yet attached to nodes
+            }
+        }
+    }
+
+    auto attach = P4::P4Fmt::Attach(globalCommentsMap);
+    program = program->apply(attach);
     // Print the program before running front end passes.
+    auto top4 = P4Fmt::P4Formatter(&formattedOutput, attach.getCommentsMap());
     program->apply(top4);
 
     if (::P4::errorCount() > 0) {

diff --git a/backends/p4fmt/p4formatter.cpp b/backends/p4fmt/p4formatter.cpp
@@ -1,10 +1,8 @@
 #include "p4formatter.h"
 
-#include <deque>
 #include <sstream>
 #include <string>
 
-#include "frontends/common/options.h"
 #include "frontends/p4/fromv1.0/v1model.h"
 #include "frontends/parsers/p4/p4parser.hpp"
 #include "ir/dump.h"
@@ -43,6 +41,19 @@ std::filesystem::path P4Formatter::ifSystemFile(const IR::Node *node) {
     return {};
 }
 
+std::pair<cstring, cstring> P4Formatter::extractNodeComments(int nodeId) {
+    auto commsIt = comMap.find(nodeId);
+    if (commsIt != comMap.end()) {
+        const Attach::Comments &comms = commsIt->second;
+        // Since we have only one prefix & suffix comment per node for now
+        cstring prefixComm = comms.prefix.empty() ? cstring("") : comms.prefix.front()->toString();
+        cstring suffixComm = comms.suffix.empty() ? cstring("") : comms.suffix.front()->toString();
+        return std::make_pair(prefixComm, suffixComm);
+    }
+    // Return empty strings if node ID not found
+    return {cstring(""), cstring("")};
+}
+
 bool P4Formatter::preorder(const IR::Node *node) {
     P4C_UNIMPLEMENTED("Unhandled IR node type: ", node->node_type_name());
     return false;
@@ -174,6 +185,11 @@ bool P4Formatter::preorder(const IR::Argument *arg) {
 }
 
 bool P4Formatter::preorder(const IR::Type_Typedef *t) {
+    auto [prefixc, suffixc] = extractNodeComments(t->id);
+    if (*prefixc != 0) {
+        builder.append(prefixc);
+        builder.append("\n");
+    }
     if (!t->annotations->annotations.empty()) {
         visit(t->annotations);
         builder.spc();
@@ -386,6 +402,10 @@ bool P4Formatter::preorder(const IR::Type_Package *package) {
 }
 
 bool P4Formatter::process(const IR::Type_StructLike *t, const char *name) {
+    auto [prefixc, suffixc] = extractNodeComments(t->id);
+
+    builder.append(prefixc);
+    builder.append("\n");
     if (isDeclaration) {
         builder.emitIndent();
         if (!t->annotations->annotations.empty()) {
@@ -413,6 +433,13 @@ bool P4Formatter::process(const IR::Type_StructLike *t, const char *name) {
     }
 
     for (auto f : t->fields) {
+        auto [fieldPrefixc, fieldSuffixc] = extractNodeComments(f->id);
+        if (*fieldPrefixc != 0) {
+            builder.emitIndent();
+            builder.append(fieldPrefixc);
+            builder.newline();
+        }
+
         if (f->annotations->size() > 0) {
             builder.emitIndent();
             if (!f->annotations->annotations.empty()) {
@@ -453,6 +480,11 @@ bool P4Formatter::preorder(const IR::Type_Control *t) {
         visit(t->annotations);
         builder.spc();
     }
+    auto [prefixc, suffixc] = extractNodeComments(t->id);
+    if (*prefixc != 0) {
+        builder.append(prefixc);
+        builder.append("\n");
+    }
     builder.append("control ");
     builder.append(t->name);
     visit(t->typeParameters);
@@ -1227,6 +1259,12 @@ bool P4Formatter::preorder(const IR::Parameter *p) {
 }
 
 bool P4Formatter::preorder(const IR::P4Control *c) {
+    auto [prefixc, suffixc] = extractNodeComments(c->id);
+    if (*prefixc != 0) {
+        builder.append(prefixc);
+        builder.append("\n");
+        builder.emitIndent();
+    }
     bool decl = isDeclaration;
     isDeclaration = false;
     visit(c->type);
@@ -1239,7 +1277,6 @@ bool P4Formatter::preorder(const IR::P4Control *c) {
         visit(s);
         builder.newline();
     }
-
     builder.emitIndent();
     builder.append("apply ");
     visit(c->body);
@@ -1256,6 +1293,12 @@ bool P4Formatter::preorder(const IR::ParameterList *p) {
 }
 
 bool P4Formatter::preorder(const IR::P4Action *c) {
+    auto [prefixc, suffixc] = extractNodeComments(c->id);
+    if (*prefixc != 0) {
+        builder.append(prefixc);
+        builder.append("\n");
+        builder.emitIndent();
+    }
     if (!c->annotations->annotations.empty()) {
         visit(c->annotations);
         builder.spc();
@@ -1378,6 +1421,12 @@ bool P4Formatter::preorder(const IR::Key *v) {
 }
 
 bool P4Formatter::preorder(const IR::Property *p) {
+    auto [prefixc, suffixc] = extractNodeComments(p->id);
+    if (*prefixc != 0) {
+        builder.append(prefixc);
+        builder.append("\n");
+        builder.emitIndent();
+    }
     if (!p->annotations->annotations.empty()) {
         visit(p->annotations);
         builder.spc();
@@ -1456,7 +1505,8 @@ bool P4Formatter::preorder(const IR::Path *p) {
 
 std::string toP4(const IR::INode *node) {
     std::stringstream stream;
-    P4Fmt::P4Formatter toP4(&stream);
+    Attach::CommentsMap comMap;
+    P4Fmt::P4Formatter toP4(&stream, comMap);
     node->getNode()->apply(toP4);
     return stream.str();
 }

diff --git a/backends/p4fmt/p4formatter.h b/backends/p4fmt/p4formatter.h
@@ -1,6 +1,7 @@
 #ifndef BACKENDS_P4FMT_P4FORMATTER_H_
 #define BACKENDS_P4FMT_P4FORMATTER_H_
 
+#include "backends/p4fmt/attach.h"
 #include "frontends/common/resolveReferences/resolveReferences.h"
 #include "ir/ir.h"
 #include "ir/visitor.h"
@@ -123,13 +124,15 @@ class P4Formatter : public Inspector, ::P4::ResolutionContext {
         visitDagOnce = false;
         setName("P4Formatter");
     }
-    explicit P4Formatter(std::ostream *outStream, cstring mainFile = nullptr)
+    explicit P4Formatter(std::ostream *outStream, Attach::Attach::CommentsMap comMap,
+                         cstring mainFile = nullptr)
         : expressionPrecedence(DBPrint::Prec_Low),
           isDeclaration(true),
           withinArgument(false),
           builder(*new Util::SourceCodeBuilder()),
           outStream(outStream),
-          mainFile(mainFile) {
+          mainFile(mainFile),
+          comMap(std::move(comMap)) {
         visitDagOnce = false;
         setName("P4Formatter");
     }
@@ -145,10 +148,11 @@ class P4Formatter : public Inspector, ::P4::ResolutionContext {
         setName("P4Formatter");
     }
 
+    std::pair<cstring, cstring> extractNodeComments(int nodeId);
     void setnoIncludesArg(bool condition) { noIncludes = condition; }
 
     void setListTerm(const char *start, const char *end) {
-        listTerminators.push_back(ListPrint(start, end));
+        listTerminators.emplace_back(start, end);
     }
     Visitor::profile_t init_apply(const IR::Node *node) override;
     void end_apply(const IR::Node *node) override;
@@ -290,6 +294,9 @@ class P4Formatter : public Inspector, ::P4::ResolutionContext {
 
     // in case it is accidentally called on a V1Program
     bool preorder(const IR::V1Program *) override { return false; }
+
+ private:
+    Attach::Attach::CommentsMap comMap;
 };
 
 std::string toP4(const IR::INode *node);