Add VerbatimNode to the ImplMember variant, so we can disable formatt…

…ing for impl functions.

Improved format disabler to keep track of which nodes it's already considered, to reduce infinite loops and hopefully other weirdness.

Addresses most of #1781

PiperOrigin-RevId: 705543728

xls/dslx/fmt/BUILD

@@ -166,6 +166,7 @@ cc_library(
         "//xls/dslx/frontend:comment_data",
         "//xls/dslx/frontend:module",
         "//xls/dslx/frontend:pos",
+        "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",

xls/dslx/fmt/ast_fmt.cc

@@ -2414,6 +2414,7 @@ DocRef Formatter::Format(const ImplMember& n) {
   return absl::visit(Visitor{
                          [&](const Function* n) { return Format(*n); },
                          [&](const ConstantDef* n) { return Format(*n); },
+                         [&](const VerbatimNode* n) { return Format(*n); },
                      },
                      n);
 }

xls/dslx/fmt/ast_fmt_test.cc

@@ -1603,6 +1603,54 @@ impl MyStruct {
 )");
 }
 
+TEST_F(ModuleFmtTest, ImplWithDisabledFmtFn) {
+  DoFmt(
+      R"(struct MyStruct {}
+
+impl MyStruct {
+    // dslx-fmt::off
+fn my_function() -> u32 {
+// Do something important.
+u32:3
+}
+// dslx-fmt::on
+
+}
+)");
+}
+
+TEST_F(ModuleFmtTest, ImplWithDisabledFmtExpr) {
+  DoFmt(
+      R"(struct MyStruct {}
+
+impl MyStruct {
+    fn my_function() -> u32 {
+        // dslx-fmt::off
+// Do something important.
+u32:3
+// dslx-fmt::on
+
+    }
+}
+)");
+}
+
+TEST_F(ModuleFmtTest, DISABLED_ImplWithDisabledFmtImpl) {
+  // This test is disabled because the impl becomes a VerbatimNode. Then the
+  // cloner tries to set it into the struct as its impl, but structs cannot
+  // take VerbatimNodes as their impl yet.
+  DoFmt(
+      R"(struct MyStruct {}
+
+// dslx-fmt::off
+impl MyStruct { fn my_function() -> u32 {
+// Do something important.
+u32:3
+}}
+// dslx-fmt::on
+)");
+}
+
 TEST_F(ModuleFmtTest, ImplMethodWithVars) {
   DoFmt(R"(struct MyStruct { x: u32 }
 

xls/dslx/fmt/format_disabler.cc

@@ -55,33 +55,44 @@ bool IsInDesugaredFn(const AstNode *node) {
 
 absl::StatusOr<std::optional<AstNode *>> FormatDisabler::operator()(
     const AstNode *node) {
+  if (seen_nodes_.contains(node)) {
+    return std::nullopt;
+  }
+  seen_nodes_.insert(node);
+
   if (node == nullptr || !node->GetSpan().has_value()) {
     // If there's no node, or no span, we can't know if it's in the unformatted
     // range, so just return nullopt to indicate it should be unchanged.
     return std::nullopt;
   }
+  VLOG(5) << "FormatDisabler looking at " << node->ToString() << " @ "
+          << (*node->GetSpan()).ToString(node->owner()->file_name());
 
   // If this node is part of a desugared proc function, we skip it, because it
   // won't be formatted anyway.
   // TODO: https://github.com/google/xls/issues/1029 remove desugared proc
   // functions.
   if (IsInDesugaredFn(node)) {
+    VLOG(5) << "In desugared function, stopping";
     return std::nullopt;
   }
 
   if (unformatted_end_.has_value()) {
+    VLOG(6) << "In format disabled mode";
     // We are in "format disabled" mode.
 
     if (node->GetSpan()->start() < *unformatted_end_) {
       // This node is within the unformatted range; delete it by returning
       // an empty VerbatimNode. This is safe because its text has already been
       // extracted in a VerbatimNode.
+      VLOG(5) << "Setting previous node to " << node->ToString();
       previous_node_ = node;
 
       return node->owner()->Make<VerbatimNode>(*node->GetSpan());
     }
 
     // We're past the end of the unformatted range; clear the setting.
+    VLOG(5) << "Past end of unformatted range";
     unformatted_end_ = std::nullopt;
 
     // Note: continue and process this node normally now.
@@ -91,6 +102,8 @@ absl::StatusOr<std::optional<AstNode *>> FormatDisabler::operator()(
   std::vector<const CommentData *> disable_comments =
       FindDisablesBetween(previous_node_, node);
   if (disable_comments.empty()) {
+    VLOG(5) << "No comments between previous node and this node";
+    VLOG(5) << "Setting previous node to " << node->ToString();
     previous_node_ = node;
     // Node should be unchanged.
     return std::nullopt;
@@ -105,6 +118,7 @@ absl::StatusOr<std::optional<AstNode *>> FormatDisabler::operator()(
         absl::StrCat("Multiple dslx-fmt::off commands between ",
                      previous_node_->ToString(), " and ", node->ToString()));
   }
+  VLOG(5) << "Found disable comment between previous node and this node.";
   const CommentData *disable_comment = disable_comments[0];
 
   // If there's a disable comment between the previous node and this node:
@@ -134,11 +148,14 @@ absl::StatusOr<std::optional<AstNode *>> FormatDisabler::operator()(
 
   // d. Set a field that indicates the ending position of the last node that
   // should be unformatted.
+  VLOG(6) << "Setting unformatted end to " << limit;
   unformatted_end_ = limit;
 
   previous_node_ = node;
 
   // e. Replace the current node with the verbatim node from step c
+  VLOG(5) << "Setting previous node to, and replacing " << node->ToString()
+          << " with VerbatimNode";
   return verbatim_node;
 }
 

xls/dslx/fmt/format_disabler.h

@@ -21,6 +21,7 @@
 #include <string_view>
 #include <vector>
 
+#include "absl/container/flat_hash_set.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "xls/dslx/fmt/comments.h"
@@ -87,6 +88,7 @@ class FormatDisabler {
 
   // Parsed lines of the module.
   std::vector<std::string> lines_;
+  absl::flat_hash_set<const AstNode *> seen_nodes_;
 };
 
 }  // namespace xls::dslx

xls/dslx/frontend/ast.cc

@@ -1538,7 +1538,12 @@ std::vector<ConstantDef*> Impl::GetConstants() const {
 
 static std::string ImplMemberIdentifier(const ImplMember member) {
   return absl::visit(
-      Visitor{[](auto* n) { return n->name_def()->identifier(); }}, member);
+      Visitor{
+          [](const ConstantDef* n) { return n->name_def()->identifier(); },
+          [](const Function* n) { return n->name_def()->identifier(); },
+          [](const VerbatimNode* n) { return std::string("Verbatim"); },
+      },
+      member);
 }
 
 std::optional<ImplMember> Impl::GetMember(std::string_view name) const {

xls/dslx/frontend/ast.h

@@ -2611,7 +2611,7 @@ inline std::optional<StructDefBase*> TypeDefinitionToStructDefBase(
   return result == nullptr ? std::nullopt : std::make_optional(result);
 }
 
-using ImplMember = std::variant<ConstantDef*, Function*>;
+using ImplMember = std::variant<ConstantDef*, Function*, VerbatimNode*>;
 
 // Represents an impl for a struct.
 class Impl : public AstNode {

xls/dslx/frontend/ast_cloner.cc

@@ -785,10 +785,12 @@ class AstCloner : public AstNodeVisitor {
         XLS_RETURN_IF_ERROR(ReplaceOrVisit(member_node));
       }
       AstNode* new_node = old_to_new_.at(member_node);
-      if (std::holds_alternative<ConstantDef*>(member)) {
-        new_members.push_back(down_cast<ConstantDef*>(new_node));
+      if (auto* new_constant_def = dynamic_cast<ConstantDef*>(new_node)) {
+        new_members.push_back(new_constant_def);
+      } else if (auto* new_function = dynamic_cast<Function*>(new_node)) {
+        new_members.push_back(new_function);
       } else {
-        new_members.push_back(down_cast<Function*>(new_node));
+        new_members.push_back(down_cast<VerbatimNode*>(new_node));
       }
     }
     new_impl->set_members(new_members);