[DSLX Fuzz testing] Add a FuzzTestFunction AST node in addition to the

`#[fuzz_test]` attribute. This is needed to unlock parsing and type-
checking the fuzz test domains as a tuple.

PiperOrigin-RevId: 896502488

Author: dplassgit

xls/dslx/frontend/ast.cc

@@ -311,6 +311,8 @@ std::string_view AstNodeKindToString(AstNodeKind kind) {
       return "for";
     case AstNodeKind::kFunctionRef:
       return "function-ref";
+    case AstNodeKind::kFuzzTestFunction:
+      return "fuzz test function";
     case AstNodeKind::kStatementBlock:
       return "statement-block";
     case AstNodeKind::kTrait:
@@ -2419,6 +2421,10 @@ Function::LambdaReturnTypeParametrics() const {
 
 TestFunction::~TestFunction() = default;
 
+// -- class FuzzTestFunction
+
+FuzzTestFunction::~FuzzTestFunction() = default;
+
 // -- class Lambda
 
 Lambda::Lambda(Module* owner, Span span, Function* function, bool in_parens)

xls/dslx/frontend/ast.h

@@ -94,6 +94,7 @@
   X(ConstantDef)                    \
   X(EnumDef)                        \
   X(Function)                       \
+  X(FuzzTestFunction)               \
   X(Impl)                           \
   X(Import)                         \
   X(Let)                            \
@@ -3843,6 +3844,8 @@ class Range : public Expr {
 // #[test]
 // fn test_foo() { ... }
 // ```
+// TODO: davidplass - Add an optional parameter Expr to allow re-using as a
+// FuzzTestFunction or QuickCheck.
 class TestFunction : public AstNode {
  public:
   static std::string_view GetDebugTypeName() { return "test function"; }
@@ -3879,6 +3882,72 @@ class TestFunction : public AstNode {
   Function& fn_;
 };
 
+// Represents a fuzz test construct.
+//
+// These are specified with an annotation as follows:
+//
+// ```dslx
+// #[fuzz_test]
+// fn test_foo() { ... }
+// ```
+//
+// or with a domains argument:
+//
+// ```dslx
+// #[fuzz_test(domains=`...`)]
+// fn test_foo() { ... }
+// ```
+// TODO: davidplass - Migrate this into TestFunction, after adding an optional
+// parameter to TestFunction to capture the domain tuple.
+class FuzzTestFunction : public AstNode {
+ public:
+  static std::string_view GetDebugTypeName() { return "fuzz test function"; }
+
+  FuzzTestFunction(Module* owner, Span span, Function& fn,
+                   std::optional<Expr*> domains)
+      : AstNode(owner),
+        span_(std::move(span)),
+        fn_(fn),
+        domains_(std::move(domains)) {}
+
+  ~FuzzTestFunction() override;
+
+  AstNodeKind kind() const override { return AstNodeKind::kFuzzTestFunction; }
+  NameDef* name_def() const { return fn_.name_def(); }
+
+  absl::Status Accept(AstNodeVisitor* v) const override {
+    return v->HandleFuzzTestFunction(this);
+  }
+
+  std::vector<AstNode*> GetChildren(bool want_types) const override {
+    return {&fn_};
+  }
+
+  std::string_view GetNodeTypeName() const override {
+    return "FuzzTestFunction";
+  }
+  std::string ToString() const override {
+    if (domains_.has_value()) {
+      return absl::StrFormat("#[fuzz_test(domains=`%s`)]\n%s",
+                             (*domains_)->ToString(), fn_.ToString());
+    }
+    return absl::StrFormat("#[fuzz_test]\n%s", fn_.ToString());
+  }
+
+  Function& fn() const { return fn_; }
+  std::optional<Span> GetSpan() const override { return span(); }
+  const Span& span() const { return span_; }
+
+  const std::string& identifier() const { return fn_.name_def()->identifier(); }
+
+  const std::optional<Expr*>& domains() const { return domains_; }
+
+ private:
+  const Span span_;
+  Function& fn_;
+  const std::optional<Expr*> domains_;
+};
+
 enum class QuickCheckTestCasesTag {
   kExhaustive,
   kCounted,
@@ -3914,6 +3983,8 @@ class QuickCheckTestCases {
 };
 
 // Represents a function to be quick-check'd.
+// TODO: davidplass - Migrate this into TestFunction, since QuickCheck is an
+// attribute which can capture the quick check test cases.
 class QuickCheck : public AstNode {
  public:
   static std::string_view GetDebugTypeName() { return "quickcheck"; }

xls/dslx/frontend/ast_cloner.cc

@@ -1047,6 +1047,17 @@ class AstCloner : public AstNodeVisitor {
     return absl::OkStatus();
   }
 
+  absl::Status HandleFuzzTestFunction(const FuzzTestFunction* n) override {
+    XLS_RETURN_IF_ERROR(VisitChildren(n));
+
+    XLS_RETURN_IF_ERROR(ReplaceOrVisit(&n->fn()));
+    XLS_ASSIGN_OR_RETURN(Function * new_fn, CastIfNotVerbatim<Function*>(
+                                                old_to_new_.at(&n->fn())));
+    old_to_new_[n] =
+        module(n)->Make<FuzzTestFunction>(n->span(), *new_fn, n->domains());
+    return absl::OkStatus();
+  }
+
   absl::Status HandleTestProc(const TestProc* n) override {
     XLS_RETURN_IF_ERROR(VisitChildren(n));
 

xls/dslx/frontend/ast_node.h

@@ -52,6 +52,7 @@ enum class AstNodeKind : uint8_t {
   kFormatMacro,
   kFunction,
   kFunctionRef,
+  kFuzzTestFunction,
   kImpl,
   kImport,
   kIndex,

xls/dslx/ir_convert/function_converter.cc

@@ -399,6 +399,7 @@ class FunctionConverterVisitor : public AstNodeVisitor {
   // keep-sorted start
   INVALID(Attribute)
   INVALID(FunctionRef)
+  INVALID(FuzzTestFunction)
   INVALID(MatchArm)
   INVALID(NameDef)
   INVALID(NameDefTree)

xls/dslx/type_system/type_info.proto

@@ -97,6 +97,7 @@ enum AstNodeKindProto {
   AST_NODE_KIND_PROC_ALIAS = 72;
   AST_NODE_KIND_TRAIT = 73;
   AST_NODE_KIND_ATTRIBUTE = 74;
+  AST_NODE_KIND_FUZZ_TEST_FUNCTION = 75;
 }
 
 message BitsValueProto {

xls/dslx/type_system/type_info_to_proto.cc

@@ -124,6 +124,8 @@ AstNodeKindProto ToProto(AstNodeKind kind) {
       return AST_NODE_KIND_SEND_IF;
     case AstNodeKind::kTestFunction:
       return AST_NODE_KIND_TEST_FUNCTION;
+    case AstNodeKind::kFuzzTestFunction:
+      return AST_NODE_KIND_FUZZ_TEST_FUNCTION;
     case AstNodeKind::kTestProc:
       return AST_NODE_KIND_TEST_PROC;
     case AstNodeKind::kWildcardPattern:
@@ -738,6 +740,8 @@ absl::StatusOr<AstNodeKind> FromProto(AstNodeKindProto p) {
       return AstNodeKind::kSendIf;
     case AST_NODE_KIND_TEST_FUNCTION:
       return AstNodeKind::kTestFunction;
+    case AST_NODE_KIND_FUZZ_TEST_FUNCTION:
+      return AstNodeKind::kFuzzTestFunction;
     case AST_NODE_KIND_TEST_PROC:
       return AstNodeKind::kTestProc;
     case AST_NODE_KIND_WILDCARD_PATTERN: