[Valid-data] Support channel attributes in formatter.

xls/dslx/fmt/ast_fmt.cc

@@ -2306,9 +2306,21 @@ DocRef Formatter::FormatFunction(const Function& n, bool is_test) {
 }
 
 DocRef Formatter::FormatProcMember(const ProcMember& n) {
-  return ConcatNGroup(
-      arena_, {FormatNameDef(*n.name_def()), arena_.colon(), arena_.break1(),
-               FormatTypeAnnotation(*n.type_annotation())});
+  std::vector<DocRef> member_pieces;
+  if (n.strictness().has_value()) {
+    member_pieces.push_back(arena_.MakeText(n.StrictnessAttributeString()));
+    member_pieces.push_back(arena_.hard_line());
+  }
+  if (n.flow_control().has_value()) {
+    member_pieces.push_back(arena_.MakeText(n.FlowControlAttributeString()));
+    member_pieces.push_back(arena_.hard_line());
+  }
+
+  member_pieces.push_back(FormatNameDef(*n.name_def()));
+  member_pieces.push_back(arena_.colon());
+  member_pieces.push_back(arena_.space());
+  member_pieces.push_back(FormatTypeAnnotation(*n.type_annotation()));
+  return ConcatNGroup(arena_, member_pieces);
 }
 
 DocRef Formatter::FormatProc(const Proc& n, bool is_test) {

xls/dslx/fmt/ast_fmt_test.cc

@@ -2566,6 +2566,61 @@ TEST_F(ModuleFmtTest, SimpleProcWithMembers) {
 )");
 }
 
+TEST_F(ModuleFmtTest, SimpleProcWithStrictness) {
+  DoFmt(
+      R"(#![feature(channel_attributes)]
+
+pub proc p {
+    #[channel_strictness("runtime_ordered")]
+    cin: chan<u32> in;
+    cout: chan<u32> out;
+
+    config(cin: chan<u32> in, cout: chan<u32> out) { (cin, cout) }
+
+    init { () }
+
+    next(state: ()) { () }
+}
+)");
+}
+
+TEST_F(ModuleFmtTest, SimpleProcWithFlowControl) {
+  DoFmt(
+      R"(#![feature(channel_attributes)]
+
+pub proc p {
+    cin: chan<u32> in;
+    #[channel_flow_control("valid_data")]
+    cout: chan<u32> out;
+
+    config(cin: chan<u32> in, cout: chan<u32> out) { (cin, cout) }
+
+    init { () }
+
+    next(state: ()) { () }
+}
+)");
+}
+
+TEST_F(ModuleFmtTest, SimpleProcWithFlowControlAndStrictness) {
+  DoFmt(
+      R"(#![feature(channel_attributes)]
+
+pub proc p {
+    #[channel_strictness("runtime_ordered")]
+    #[channel_flow_control("valid_data")]
+    cin: chan<u32> in;
+    cout: chan<u32> out;
+
+    config(cin: chan<u32> in, cout: chan<u32> out) { (cin, cout) }
+
+    init { () }
+
+    next(state: ()) { () }
+}
+)");
+}
+
 TEST_F(ModuleFmtTest, SimpleProcWithMembersImplStyle) {
   DoFmt(
       R"(pub proc P {

xls/dslx/frontend/ast_cloner.cc

@@ -723,7 +723,8 @@ class AstCloner : public AstNodeVisitor {
     XLS_RETURN_IF_ERROR(VisitChildren(n));
     old_to_new_[n] = module(n)->Make<ProcMember>(
         absl::down_cast<NameDef*>(old_to_new_.at(n->name_def())),
-        absl::down_cast<TypeAnnotation*>(old_to_new_.at(n->type_annotation())));
+        absl::down_cast<TypeAnnotation*>(old_to_new_.at(n->type_annotation())),
+        n->strictness(), n->flow_control());
     return absl::OkStatus();
   }
 

xls/dslx/frontend/parser_test.cc

@@ -4575,16 +4575,22 @@ TEST_F(ParserTest, CannotSpecifyParametricsForNumber) {
 }
 
 TEST_F(ParserTest, ChannelStrictnessAttribute) {
-  constexpr std::string_view kProgram = R"(
-#![feature(channel_attributes)]
+  constexpr std::string_view kProgram = R"(#![feature(channel_attributes)]
+
 proc p {
-  #[channel_strictness("runtime_ordered")]
-  c: chan<u32> in;
-  config(c: chan<u32> in) { (c,) }
-  init { () }
-  next(tok: token) { tok }
+    #[channel_strictness("runtime_ordered")]
+    c: chan<u32> in;
+    config(c: chan<u32> in) {
+        (c,)
+    }
+    init {
+        ()
+    }
+    next(tok: token) {
+        tok
+    }
 })";
-  std::unique_ptr<Module> module = ExpectParsesSuccessfully(kProgram);
+  std::unique_ptr<Module> module = RoundTrip(kProgram);
   ASSERT_NE(module, nullptr);
   ModuleMember* member = module->FindMemberWithName("p").value();
   Proc* p = std::get<Proc*>(*member);
@@ -4595,16 +4601,22 @@ proc p {
 }
 
 TEST_F(ParserTest, ChannelFlowControlAttribute) {
-  constexpr std::string_view kProgram = R"(
-#![feature(channel_attributes)]
+  constexpr std::string_view kProgram = R"(#![feature(channel_attributes)]
+
 proc p {
-  #[channel_flow_control("valid_data")]
-  c: chan<u32> in;
-  config(c: chan<u32> in) { (c,) }
-  init { () }
-  next(tok: token) { tok }
+    #[channel_flow_control("valid_data")]
+    c: chan<u32> in;
+    config(c: chan<u32> in) {
+        (c,)
+    }
+    init {
+        ()
+    }
+    next(tok: token) {
+        tok
+    }
 })";
-  std::unique_ptr<Module> module = ExpectParsesSuccessfully(kProgram);
+  std::unique_ptr<Module> module = RoundTrip(kProgram);
   ASSERT_NE(module, nullptr);
   ModuleMember* member = module->FindMemberWithName("p").value();
   Proc* p = std::get<Proc*>(*member);

xls/dslx/frontend/proc.cc

@@ -18,6 +18,7 @@
 #include <string>
 #include <string_view>
 #include <utility>
+#include <variant>
 #include <vector>
 
 #include "absl/log/check.h"
@@ -182,4 +183,46 @@ ProcMember::ProcMember(Module* owner, NameDef* name_def,
 
 ProcMember::~ProcMember() = default;
 
+using ProcMemberAttribute = std::variant<ChannelStrictness, FlowControl>;
+
+std::string AttributeToString(ProcMemberAttribute attribute) {
+  const auto fmt_attribute = [](std::string name, std::string value) {
+    return absl::StrFormat("#[%s(\"%s\")]", name, value);
+  };
+  if (std::holds_alternative<ChannelStrictness>(attribute)) {
+    return fmt_attribute(
+        "channel_strictness",
+        ChannelStrictnessToString(std::get<ChannelStrictness>(attribute)));
+  }
+  return fmt_attribute("channel_flow_control",
+                       FlowControlToString(std::get<FlowControl>(attribute)));
+}
+
+std::string ProcMember::StrictnessAttributeString() const {
+  if (!strictness_.has_value()) {
+    return "";
+  }
+  return AttributeToString(*strictness_);
+}
+
+std::string ProcMember::FlowControlAttributeString() const {
+  if (!flow_control_.has_value()) {
+    return "";
+  }
+  return AttributeToString(*flow_control_);
+}
+
+std::string ProcMember::ToString() const {
+  std::string result;
+  if (strictness_.has_value()) {
+    result += AttributeToString(*strictness_) + "\n";
+  }
+  if (flow_control_.has_value()) {
+    result += AttributeToString(*flow_control_) + "\n";
+  }
+  result += absl::StrFormat("%s: %s;", name_def_->ToString(),
+                            type_annotation_->ToString());
+  return result;
+}
+
 }  // namespace xls::dslx

xls/dslx/frontend/proc.h

@@ -25,7 +25,6 @@
 #include "absl/log/check.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
-#include "absl/strings/str_format.h"
 #include "absl/types/span.h"
 #include "xls/dslx/frontend/ast.h"
 #include "xls/dslx/frontend/pos.h"
@@ -39,6 +38,7 @@ namespace xls::dslx {
 // This is very similar to a `Param` at the moment, but we make them distinct
 // types for structural clarity in the AST. Params are really "parameters to
 // functions".
+
 class ProcMember : public AstNode {
  public:
   ProcMember(Module* owner, NameDef* name_def, TypeAnnotation* type,
@@ -54,10 +54,8 @@ class ProcMember : public AstNode {
   }
 
   std::string_view GetNodeTypeName() const override { return "ProcMember"; }
-  std::string ToString() const override {
-    return absl::StrFormat("%s: %s;", name_def_->ToString(),
-                           type_annotation_->ToString());
-  }
+
+  std::string ToString() const override;
 
   std::vector<AstNode*> GetChildren(bool want_types) const override {
     return {name_def_, type_annotation_};
@@ -77,6 +75,9 @@ class ProcMember : public AstNode {
     return flow_control_;
   }
 
+  std::string StrictnessAttributeString() const;
+  std::string FlowControlAttributeString() const;
+
  private:
   NameDef* name_def_;
   TypeAnnotation* type_annotation_;