[WIP]: comment formatting for structs and typedefs

Signed-off-by: Nitish <snapdgnn@proton.me>
p4lang · Aug 27, 2024 · 59ff847 · 59ff847
1 parent 0dfc44a
commit 59ff847
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 7 deletions.
diff --git a/backends/p4fmt/p4fmt.cpp b/backends/p4fmt/p4fmt.cpp
@@ -38,10 +38,10 @@ std::stringstream getFormattedOutput(std::filesystem::path inputFile) {
         }
     }
 
-    auto top4 = P4Fmt::P4Formatter(&formattedOutput);
     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,18 @@ 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;
+        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 +184,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 +401,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 +432,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()) {
@@ -1456,7 +1482,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,9 @@
 #ifndef BACKENDS_P4FMT_P4FORMATTER_H_
 #define BACKENDS_P4FMT_P4FORMATTER_H_
 
+#include <utility>
+
+#include "backends/p4fmt/attach.h"
 #include "frontends/common/resolveReferences/resolveReferences.h"
 #include "ir/ir.h"
 #include "ir/visitor.h"
@@ -123,13 +126,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 +150,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 +296,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);