[P4fmt]: attaching comments to IR Nodes

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

backends/p4fmt/attach.cpp

@@ -0,0 +1,71 @@
+#include "backends/p4fmt/attach.h"
+
+#include "frontends/common/parser_options.h"
+#include "lib/source_file.h"
+
+namespace P4::P4Fmt {
+
+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);
+}
+
+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().string_view());
+    if (isSystemFile(sourceFile.string())) {
+        // 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->getSourceInfo().getEnd();
+
+        switch (ttype) {
+            case TraversalType::Preorder:
+                if (commentEnd.getLineNumber() == nodeStart.getLineNumber() - 1) {
+                    addPrefixComments(node->clone_id, comment);
+                    isAttached = true;  // Mark the comment as attached
+                }
+                break;
+
+            case TraversalType::Postorder:
+                if (commentEnd.getLineNumber() == nodeStart.getLineNumber()) {
+                    addSuffixComments(node->clone_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

backends/p4fmt/attach.h

@@ -0,0 +1,45 @@
+#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){};
+
+    const IR::Node *attachCommentsToNode(IR::Node *, TraversalType);
+
+    const IR::Node *preorder(IR::Node *node) override;
+    const IR::Node *postorder(IR::Node *node) override;
+
+    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_ */

backends/p4fmt/p4fmt.cpp

@@ -1,15 +1,46 @@
 #include "backends/p4fmt/p4fmt.h"
 
+#include "backends/p4fmt/attach.h"
+#include "backends/p4fmt/p4formatter.h"
 #include "frontends/common/parseInput.h"
 #include "frontends/common/parser_options.h"
+#include "frontends/parsers/parserDriver.h"
 #include "ir/ir.h"
 #include "lib/compile_context.h"
 #include "lib/error.h"
+#include "lib/source_file.h"
 #include "options.h"
-#include "p4formatter.h"
 
 namespace P4::P4Fmt {
 
+std::optional<std::pair<const IR::P4Program *, const Util::InputSources *>> parseProgram(
+    const ParserOptions &options) {
+    auto *file = fopen(options.file.c_str(), "r");
+    if (file == nullptr) {
+        ::P4::error(ErrorType::ERR_NOT_FOUND, "%1%: No such file or directory.", options.file);
+        return std::nullopt;
+    }
+    if (options.isv1()) {
+        ::P4::error(ErrorType::ERR_UNKNOWN, "p4fmt cannot deal with p4-14 programs");
+        fclose(file);
+        return std::nullopt;
+    }
+    auto preprocessorResult = options.preprocess();
+    auto result =
+        P4ParserDriver::parseProgramSources(preprocessorResult.value().get(), options.file.c_str());
+    fclose(file);
+
+    if (::P4::errorCount() > 0) {
+        ::P4::error(ErrorType::ERR_OVERLIMIT, "%1% errors encountered, aborting compilation",
+                    ::P4::errorCount());
+        return std::nullopt;
+    }
+
+    BUG_CHECK(result.first != nullptr, "Parsing failed, but no error was reported");
+
+    return result;
+}
+
 std::stringstream getFormattedOutput(std::filesystem::path inputFile) {
     AutoCompileContext autoP4FmtContext(new P4Fmt::P4FmtContext);
     auto &options = P4Fmt::P4FmtContext::get().options();
@@ -18,13 +49,28 @@ std::stringstream getFormattedOutput(std::filesystem::path inputFile) {
 
     std::stringstream formattedOutput;
 
-    const IR::P4Program *program = P4::parseP4File(options);
-    if (program == nullptr && ::P4::errorCount() != 0) {
-        ::P4::error("Failed to parse P4 file.");
+    auto parseResult = parseProgram(options);
+    if (!parseResult) {
+        if (::P4::errorCount() > 0) {
+            ::P4::error("Failed to parse P4 file.");
+        }
         return formattedOutput;
     }
+    const auto &[program, sources] = *parseResult;
+
+    std::unordered_map<const Util::Comment *, bool> globalCommentsMap;
+
+    // Initialize the global comments map from the list of comments in the program.
+    if (sources != nullptr) {
+        for (const auto *comment : sources->getAllComments()) {
+            globalCommentsMap[comment] =
+                false;  // Initialize all comments as not yet attached to nodes
+        }
+    }
 
     auto top4 = P4Fmt::P4Formatter(&formattedOutput);
+    auto attach = P4::P4Fmt::Attach(globalCommentsMap);
+    program->apply(attach);
     // Print the program before running front end passes.
     program->apply(top4);
 

frontends/parsers/parserDriver.cpp

@@ -153,6 +153,28 @@ bool P4ParserDriver::parse(AbstractP4Lexer &lexer, std::string_view sourceFile,
     return parse(inputStream.get(), sourceFile, sourceLine);
 }
 
+/* static */ std::pair<const IR::P4Program *, const Util::InputSources *>
+P4ParserDriver::parseProgramSources(std::istream &in, std::string_view sourceFile,
+                                    unsigned sourceLine /* = 1 */) {
+    P4ParserDriver driver;
+    P4Lexer lexer(in);
+    if (!driver.parse(lexer, sourceFile, sourceLine)) {
+        return {nullptr, nullptr};
+    }
+
+    auto *program = new IR::P4Program(driver.nodes->srcInfo, *driver.nodes);
+    const Util::InputSources *sources = driver.sources;
+
+    return {program, sources};
+}
+
+/*static */ std::pair<const IR::P4Program *, const Util::InputSources *>
+P4ParserDriver::parseProgramSources(FILE *in, std::string_view sourceFile,
+                                    unsigned sourceLine /* = 1 */) {
+    AutoStdioInputStream inputStream(in);
+    return parseProgramSources(inputStream.get(), sourceFile, sourceLine);
+}
+
 template <typename T>
 const T *P4ParserDriver::parse(P4AnnotationLexer::Type type, const Util::SourceInfo &srcInfo,
                                const IR::Vector<IR::AnnotationToken> &body) {

frontends/parsers/parserDriver.h

@@ -115,6 +115,15 @@ class P4ParserDriver final : public AbstractParserDriver {
     static const IR::P4Program *parse(FILE *in, std::string_view sourceFile,
                                       unsigned sourceLine = 1);
 
+    /// Parses the input and returns a pair with the P4Program and InputSources.
+    /// Use this when both the parsed P4Program and InputSources are required,
+    /// as opposed to the `parse` method, which only returns the P4Program.
+    static std::pair<const IR::P4Program *, const Util::InputSources *> parseProgramSources(
+        std::istream &in, std::string_view sourceFile, unsigned sourceLine = 1);
+
+    static std::pair<const IR::P4Program *, const Util::InputSources *> parseProgramSources(
+        FILE *in, std::string_view sourceFile, unsigned sourceLine = 1);
+
     /**
      * Parses a P4-16 annotation body.
      *

lib/source_file.cpp

@@ -83,6 +83,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());

lib/source_file.h

@@ -110,6 +110,7 @@ class SourcePosition final {
 };
 
 class InputSources;
+class Comment;
 
 /**
 Information about the source position of a language element -
@@ -246,7 +247,7 @@ struct SourceFileLine {
     cstring toString() const;
 };
 
-class Comment final : IHasDbPrint {
+class Comment final : IHasDbPrint, IHasSourceInfo {
  private:
     SourceInfo srcInfo;
     bool singleLine;
@@ -255,7 +256,7 @@ class Comment final : IHasDbPrint {
  public:
     Comment(SourceInfo srcInfo, bool singleLine, cstring body)
         : srcInfo(srcInfo), singleLine(singleLine), body(body) {}
-    cstring toString() const {
+    cstring toString() const override {
         std::stringstream str;
         dbprint(str);
         return str.str();
@@ -268,6 +269,9 @@ class Comment final : IHasDbPrint {
         out << body;
         if (!singleLine) out << "*/";
     }
+
+    /// Retrieve the source position associated with this comment.
+    [[nodiscard]] SourceInfo getSourceInfo() const override { return srcInfo; }
 };
 
 /**
@@ -287,7 +291,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;
@@ -319,6 +322,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);