Merge pull request #1829 from xlsynth:cdleary/2025-01-04-quickcheck-e…

…xhaustive

PiperOrigin-RevId: 712670235

.bazelrc

@@ -25,6 +25,7 @@ build --copt "-Wextra"                 --host_copt "-Wextra"
 
 # Turn some specific warnings into errors.
 build --copt "-Werror=switch" --host_copt "-Werror=switch"
+build --copt "-Werror=return-type" --host_copt "-Werror=return-type"
 
 # ... and disable the warnings we're not interested in.
 build --copt "-Wno-sign-compare"       --host_copt "-Wno-sign-compare"

docs_src/dslx_reference.md

@@ -2052,7 +2052,7 @@ implementation from above:
 // Reversing a value twice gets you the original value.
 
 #[quickcheck]
-fn prop_double_reverse(x: u32) -> bool {
+fn prop_double_bitreverse(x: u32) -> bool {
   x == rev(rev(x))
 }
 ```
@@ -2078,6 +2078,24 @@ for production can be found in the execution log.
 For determinism, the DSLX interpreter should be run with the `seed` flag:
 `./interpreter_main --seed=1234 <DSLX source file>`
 
+#### Exhaustive QuickCheck
+
+For small domains the quickcheck directive can also be placed in exhaustive mode via the `exhaustive` directive:
+
+```dslx
+// `u8` space is small enough to check exhaustively.
+#[quickcheck(exhaustive)]
+fn prop_double_bitreverse(x: u8) -> bool {
+  x == rev(rev(x))
+}
+```
+
+This is useful when there are small domains where we would prefer exhaustive
+stimulus over randomized stimulus. Note that as the space becomes large,
+exhaustive concrete-stimulus-based testing becomes implausible, and users should
+consider attempting to prove the QuickCheck formally via the
+`prove_quickcheck_main` tool.
+
 [hughes-paper]: https://www.cs.tufts.edu/~nr/cs257/archive/john-hughes/quick.pdf
 
 ## Communicating Sequential Processes (AKA procs)

xls/data_structures/inline_bitmap.h

@@ -352,6 +352,42 @@ class InlineBitmap {
   absl::InlinedVector<uint64_t, 1> data_;
 };
 
+class BitmapView {
+ public:
+  explicit BitmapView(const InlineBitmap& bitmap,
+                      std::optional<int64_t> start_bit = std::nullopt,
+                      std::optional<int64_t> bit_count = std::nullopt)
+      : bitmap_(bitmap) {
+    start_bit_ = start_bit.value_or(0);
+    CHECK_LE(start_bit_, bitmap_.bit_count());
+    bit_count_ = bit_count.value_or(bitmap_.bit_count() - start_bit_);
+  }
+
+  bool Get(int64_t bit_index) const {
+    CHECK_LT(bit_index, bit_count_);
+    return bitmap_.Get(start_bit_ + bit_index);
+  }
+
+  BitmapView Slice(int64_t start_bit, int64_t bit_count) const {
+    return BitmapView(bitmap_, start_bit_ + start_bit, bit_count);
+  }
+
+  InlineBitmap ToBitmap() const {
+    InlineBitmap result(bit_count_, false);
+    for (int64_t i = 0; i < bit_count_; ++i) {
+      result.Set(i, Get(i));
+    }
+    return result;
+  }
+
+  int64_t bit_count() const { return bit_count_; }
+
+ private:
+  const InlineBitmap& bitmap_;
+  int64_t start_bit_;
+  int64_t bit_count_;
+};
+
 }  // namespace xls
 
 #endif  // XLS_DATA_STRUCTURES_INLINE_BITMAP_H_

xls/dslx/fmt/ast_fmt.cc

@@ -2373,11 +2373,18 @@ DocRef Formatter::Format(const TestProc& n) {
 
 DocRef Formatter::Format(const QuickCheck& n) {
   std::vector<DocRef> pieces;
-  if (std::optional<int64_t> test_count = n.test_count()) {
-    pieces.push_back(arena_.MakeText(
-        absl::StrFormat("#[quickcheck(test_count=%d)]", test_count.value())));
-  } else {
-    pieces.push_back(arena_.MakeText("#[quickcheck]"));
+  switch (n.test_cases().tag()) {
+    case QuickCheckTestCasesTag::kExhaustive:
+      pieces.push_back(arena_.MakeText("#[quickcheck(exhaustive)]"));
+      break;
+    case QuickCheckTestCasesTag::kCounted:
+      if (n.test_cases().count().has_value()) {
+        pieces.push_back(arena_.MakeText(absl::StrFormat(
+            "#[quickcheck(test_count=%d)]", *n.test_cases().count())));
+      } else {
+        pieces.push_back(arena_.MakeText("#[quickcheck]"));
+      }
+      break;
   }
   pieces.push_back(arena_.hard_line());
   pieces.push_back(Format(*n.fn()));

xls/dslx/frontend/ast.cc

@@ -2182,22 +2182,40 @@ VerbatimNode::~VerbatimNode() = default;
 
 // -- class QuickCheck
 
+std::string QuickCheckTestCases::ToString() const {
+  switch (tag()) {
+    case QuickCheckTestCasesTag::kExhaustive:
+      return "exhaustive";
+    case QuickCheckTestCasesTag::kCounted:
+      if (count().has_value()) {
+        return absl::StrFormat("test_count=%d", count().value());
+      }
+      return absl::StrFormat("test_count=default=%d", kDefaultTestCount);
+  }
+}
+
 QuickCheck::QuickCheck(Module* owner, Span span, Function* fn,
-                       std::optional<int64_t> test_count)
+                       QuickCheckTestCases test_cases)
     : AstNode(owner),
       span_(std::move(span)),
       fn_(fn),
-      test_count_(test_count) {}
+      test_cases_(test_cases) {}
 
 QuickCheck::~QuickCheck() = default;
 
 std::string QuickCheck::ToString() const {
-  std::string test_count_str;
-  if (test_count_.has_value()) {
-    test_count_str = absl::StrFormat("(test_count=%d)", *test_count_);
+  std::string spec_str;
+  switch (test_cases_.tag()) {
+    case QuickCheckTestCasesTag::kExhaustive:
+      spec_str = "(exhaustive)";
+      break;
+    case QuickCheckTestCasesTag::kCounted:
+      if (test_cases_.count().has_value()) {
+        spec_str = absl::StrFormat("(test_count=%d)", *test_cases_.count());
+      }
+      break;
   }
-  return absl::StrFormat("#[quickcheck%s]\n%s", test_count_str,
-                         fn_->ToString());
+  return absl::StrFormat("#[quickcheck%s]\n%s", spec_str, fn_->ToString());
 }
 
 // -- class TupleIndex

xls/dslx/frontend/ast.h

@@ -3037,15 +3037,47 @@ class TestFunction : public AstNode {
   Function& fn_;
 };
 
+enum class QuickCheckTestCasesTag {
+  kExhaustive,
+  kCounted,
+};
+
+// Describes the test cases that should be run for a quickcheck test function --
+// they can be either counted or exhaustive, and for counted we have a default
+// count if the user doesn't explicitly specify a count.
+class QuickCheckTestCases {
+ public:
+  // The number of test cases we run if a count is not explicitly specified.
+  static constexpr int64_t kDefaultTestCount = 1000;
+
+  static QuickCheckTestCases Exhaustive() {
+    return QuickCheckTestCases(QuickCheckTestCasesTag::kExhaustive);
+  }
+  static QuickCheckTestCases Counted(std::optional<int64_t> count) {
+    return QuickCheckTestCases(QuickCheckTestCasesTag::kCounted, count);
+  }
+
+  std::string ToString() const;
+
+  QuickCheckTestCasesTag tag() const { return tag_; }
+  std::optional<int64_t> count() const { return count_; }
+
+ private:
+  explicit QuickCheckTestCases(QuickCheckTestCasesTag tag,
+                               std::optional<int64_t> count = std::nullopt)
+      : tag_(tag), count_(count) {}
+
+  QuickCheckTestCasesTag tag_;
+  std::optional<int64_t> count_;
+};
+
 // Represents a function to be quick-check'd.
 class QuickCheck : public AstNode {
  public:
   static std::string_view GetDebugTypeName() { return "quickcheck"; }
 
-  static constexpr int64_t kDefaultTestCount = 1000;
-
   QuickCheck(Module* owner, Span span, Function* fn,
-             std::optional<int64_t> test_count = std::nullopt);
+             QuickCheckTestCases test_cases);
 
   ~QuickCheck() override;
 
@@ -3066,17 +3098,14 @@ class QuickCheck : public AstNode {
   const std::string& identifier() const { return fn_->identifier(); }
 
   Function* fn() const { return fn_; }
-  int64_t GetTestCountOrDefault() const {
-    return test_count_ ? *test_count_ : kDefaultTestCount;
-  }
-  std::optional<int64_t> test_count() const { return test_count_; }
+  QuickCheckTestCases test_cases() const { return test_cases_; }
   std::optional<Span> GetSpan() const override { return span_; }
   const Span& span() const { return span_; }
 
  private:
   Span span_;
   Function* fn_;
-  std::optional<int64_t> test_count_;
+  QuickCheckTestCases test_cases_;
 };
 
 // Represents an index into a tuple, e.g., "(u32:7, u32:8).1".

xls/dslx/frontend/ast_cloner.cc

@@ -684,7 +684,7 @@ class AstCloner : public AstNodeVisitor {
     XLS_RETURN_IF_ERROR(VisitChildren(n));
     old_to_new_[n] = module_->Make<QuickCheck>(
         n->GetSpan().value(), down_cast<Function*>(old_to_new_.at(n->fn())),
-        n->test_count());
+        n->test_cases());
     return absl::OkStatus();
   }
 

xls/dslx/frontend/parser.cc

@@ -1701,42 +1701,42 @@ absl::StatusOr<Function*> Parser::ParseFunctionInternal(
   return f;
 }
 
+absl::StatusOr<QuickCheckTestCases> Parser::ParseQuickCheckConfig() {
+  XLS_RETURN_IF_ERROR(DropTokenOrError(TokenKind::kOParen));
+  XLS_ASSIGN_OR_RETURN(Token tok, PopTokenOrError(TokenKind::kIdentifier));
+  if (tok.GetValue() == "exhaustive") {
+    XLS_RETURN_IF_ERROR(DropTokenOrError(TokenKind::kCParen));
+    return QuickCheckTestCases::Exhaustive();
+  }
+  if (tok.GetValue() == "test_count") {
+    XLS_RETURN_IF_ERROR(DropTokenOrError(TokenKind::kEquals));
+    XLS_ASSIGN_OR_RETURN(Token tok, PopTokenOrError(TokenKind::kNumber));
+    XLS_ASSIGN_OR_RETURN(int64_t count, tok.GetValueAsInt64());
+    XLS_RETURN_IF_ERROR(DropTokenOrError(TokenKind::kCParen));
+    return QuickCheckTestCases::Counted(count);
+  }
+  return ParseErrorStatus(tok.span(),
+                          "Expected 'exhaustive' or 'test_count' in "
+                          "quickcheck directive");
+}
+
 absl::StatusOr<QuickCheck*> Parser::ParseQuickCheck(
     absl::flat_hash_map<std::string, Function*>* name_to_fn, Bindings& bindings,
     const Pos& hash_pos) {
-  std::optional<int64_t> test_count;
-  XLS_ASSIGN_OR_RETURN(bool peek_is_paren, PeekTokenIs(TokenKind::kOParen));
-  if (peek_is_paren) {  // Config is specified.
-    DropTokenOrDie();
-    Span config_name_span;
-    XLS_ASSIGN_OR_RETURN(std::string config_name,
-                         PopIdentifierOrError(&config_name_span));
-    XLS_RETURN_IF_ERROR(DropTokenOrError(TokenKind::kEquals));
-    if (config_name == "test_count") {
-      XLS_ASSIGN_OR_RETURN(Token count_token,
-                           PopTokenOrError(TokenKind::kNumber));
-      XLS_ASSIGN_OR_RETURN(test_count, count_token.GetValueAsInt64());
-      if (test_count <= 0) {
-        return ParseErrorStatus(
-            count_token.span(),
-            absl::StrFormat("Number of tests should be > 0, got %d",
-                            *test_count));
-      }
-      XLS_RETURN_IF_ERROR(DropTokenOrError(TokenKind::kCParen));
-    } else {
-      return ParseErrorStatus(
-          config_name_span,
-          absl::StrFormat("Unknown configuration key in directive: '%s'",
-                          config_name));
-    }
+  std::optional<QuickCheckTestCases> test_cases;
+  XLS_ASSIGN_OR_RETURN(bool peek_is_oparen, PeekTokenIs(TokenKind::kOParen));
+  if (peek_is_oparen) {  // Config is specified.
+    XLS_ASSIGN_OR_RETURN(test_cases, ParseQuickCheckConfig());
+  } else {
+    test_cases = QuickCheckTestCases::Counted(std::nullopt);
   }
 
   XLS_RETURN_IF_ERROR(DropTokenOrError(TokenKind::kCBrack));
   XLS_ASSIGN_OR_RETURN(
       Function * fn,
       ParseFunction(GetPos(), /*is_public=*/false, bindings, name_to_fn));
   const Span quickcheck_span(hash_pos, fn->span().limit());
-  return module_->Make<QuickCheck>(quickcheck_span, fn, test_count);
+  return module_->Make<QuickCheck>(quickcheck_span, fn, test_cases.value());
 }
 
 absl::StatusOr<XlsTuple*> Parser::ParseTupleRemainder(const Pos& start_pos,

xls/dslx/frontend/parser.h

@@ -611,6 +611,9 @@ class Parser : public TokenParser {
       absl::flat_hash_map<std::string, Function*>* name_to_fn,
       Bindings& bindings, const Pos& hash_pos);
 
+  // Parses the test count configuration for a quickcheck directive.
+  absl::StatusOr<QuickCheckTestCases> ParseQuickCheckConfig();
+
   // Parses a module-level attribute -- cursor should be over the open bracket.
   //
   // Side-effect: module_ is tagged with the parsed attribute on success.

xls/dslx/run_routines/run_comparator.h

@@ -75,6 +75,7 @@ class RunComparator : public AbstractRunComparator {
  private:
   XLS_FRIEND_TEST(RunRoutinesTest, TestInvokedFunctionDoesJit);
   XLS_FRIEND_TEST(RunRoutinesTest, QuickcheckInvokedFunctionDoesJit);
+  XLS_FRIEND_TEST(RunRoutinesTest, QuickcheckExhaustive);
   XLS_FRIEND_TEST(RunRoutinesTest, NoSeedStillQuickChecks);
 
   absl::flat_hash_map<std::string, std::unique_ptr<FunctionJit>> jit_cache_;