From d9f761036a008ffea70087ab8cc0519981f0e8f9 Mon Sep 17 00:00:00 2001
From: Ryan Roelke <ryan.roelke@tiledb.com>
Date: Thu, 6 Feb 2025 15:24:35 -0500
Subject: [PATCH] Fix rapidcheck allow_dups issues

---
 test/src/unit-sparse-global-order-reader.cc | 24 +++++++++++++++------
 test/support/rapidcheck/array_templates.h   | 10 +++++++--
 test/support/src/array_templates.h          | 11 ++++++++++
 3 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/test/src/unit-sparse-global-order-reader.cc b/test/src/unit-sparse-global-order-reader.cc
index 2fc55a74ff5..03007cf0a00 100644
--- a/test/src/unit-sparse-global-order-reader.cc
+++ b/test/src/unit-sparse-global-order-reader.cc
@@ -3169,8 +3169,7 @@ struct Arbitrary<FxRun1D> {
                templates::Dimension<DIMENSION_TYPE>,
                std::vector<templates::Domain<CoordType>>,
                std::vector<templates::Fragment1D<int, int>>> fragments,
-           int num_user_cells,
-           bool allow_dups) {
+           int num_user_cells) {
           FxRun1D instance;
           std::tie(
               instance.array.allow_dups_,
@@ -3179,13 +3178,11 @@ struct Arbitrary<FxRun1D> {
               instance.fragments) = fragments;
 
           instance.num_user_cells = num_user_cells;
-          instance.array.allow_dups_ = allow_dups;
 
           return instance;
         },
         fragments,
-        num_user_cells,
-        gen::arbitrary<bool>());
+        num_user_cells);
   }
 };
 
@@ -3232,7 +3229,22 @@ struct Arbitrary<FxRun2D> {
     auto d0 = gen::arbitrary<templates::Dimension<Dim0Type>>();
     auto d1 = gen::arbitrary<templates::Dimension<Dim1Type>>();
 
-    auto fragments = gen::mapcat(gen::tuple(allow_dups, d0, d1), [](auto arg) {
+    auto domain =
+        gen::suchThat(gen::tuple(allow_dups, d0, d1), [](const auto& arg) {
+          bool allow_dups;
+          templates::Dimension<Dim0Type> d0;
+          templates::Dimension<Dim1Type> d1;
+          std::tie(allow_dups, d0, d1) = arg;
+          if (allow_dups) {
+            return true;
+          } else {
+            // need to ensure that rapidcheck uniqueness can generate enough
+            // cases
+            return (d0.domain.num_cells() + d1.domain.num_cells()) >= 12;
+          }
+        });
+
+    auto fragments = gen::mapcat(domain, [](auto arg) {
       bool allow_dups;
       templates::Dimension<Dim0Type> d0;
       templates::Dimension<Dim1Type> d1;
diff --git a/test/support/rapidcheck/array_templates.h b/test/support/rapidcheck/array_templates.h
index 2ec421ff69f..95f8b671030 100644
--- a/test/support/rapidcheck/array_templates.h
+++ b/test/support/rapidcheck/array_templates.h
@@ -173,9 +173,11 @@ Gen<Fragment1D<D, Att...>> make_fragment_1d(
 
   using Cell = std::tuple<D, Att...>;
 
+  auto uniqueCoords = [](const Cell& cell) { return std::get<0>(cell); };
+
   auto cells = gen::nonEmpty(
       allow_duplicates ? gen::container<std::vector<Cell>>(cell) :
-                         gen::unique<std::vector<Cell>>(cell));
+                         gen::uniqueBy<std::vector<Cell>>(cell, uniqueCoords));
 
   return gen::map(cells, [](std::vector<Cell> cells) {
     std::vector<D> coords;
@@ -204,9 +206,13 @@ Gen<Fragment2D<D1, D2, Att...>> make_fragment_2d(
 
   auto cell = gen::tuple(coord_d1, coord_d2, gen::arbitrary<Att>()...);
 
+  auto uniqueCoords = [](const Cell& cell) {
+    return std::make_pair(std::get<0>(cell), std::get<1>(cell));
+  };
+
   auto cells = gen::nonEmpty(
       allow_duplicates ? gen::container<std::vector<Cell>>(cell) :
-                         gen::unique<std::vector<Cell>>(cell));
+                         gen::uniqueBy<std::vector<Cell>>(cell, uniqueCoords));
 
   return gen::map(cells, [](std::vector<Cell> cells) {
     std::vector<D1> coords_d1;
diff --git a/test/support/src/array_templates.h b/test/support/src/array_templates.h
index 024cc6176c7..d70d2217b20 100644
--- a/test/support/src/array_templates.h
+++ b/test/support/src/array_templates.h
@@ -142,6 +142,17 @@ struct Domain {
       , upper_bound(std::max(d1, d2)) {
   }
 
+  uint64_t num_cells() const {
+    // FIXME: this is incorrect for 64-bit domains which need to check overflow
+    if (std::is_signed<D>::value) {
+      return static_cast<int64_t>(upper_bound) -
+             static_cast<int64_t>(lower_bound) + 1;
+    } else {
+      return static_cast<uint64_t>(upper_bound) -
+             static_cast<uint64_t>(lower_bound) + 1;
+    }
+  }
+
   bool contains(D point) const {
     return lower_bound <= point && point <= upper_bound;
   }