test: Don't generate unreasonably large replicated arraysHEADmaster

2 files changed, 20 insertions, 6 deletions

diff --git a/test/Gen.hs b/test/Gen.hs
index 9ff206c..044de14 100644
--- a/test/Gen.hs
+++ b/test/Gen.hs
@@ -59,9 +59,12 @@ shuffleShR = \sh -> go (length sh) (toList sh) sh
       (dim :$:) <$> go (nbag - 1) bag' sh
 
 genShR :: SNat n -> Gen (IShR n)
-genShR sn = do
+genShR = genShRwithTarget 100_000
+
+genShRwithTarget :: Int -> SNat n -> Gen (IShR n)
+genShRwithTarget targetMax sn = do
   let n = fromSNat' sn
-  targetSize <- Gen.int (Range.linear 0 100_000)
+  targetSize <- Gen.int (Range.linear 0 targetMax)
   let genDims :: SNat m -> Int -> Gen (IShR m)
       genDims SZ _ = return ZSR
       genDims (SS m) 0 = do
@@ -93,10 +96,14 @@ genShR sn = do
 -- other dimensions might be zero.
 genReplicatedShR :: m <= n => SNat m -> SNat n -> Gen (IShR m, IShR n, IShR n)
 genReplicatedShR = \m n -> do
-  sh1 <- genShR m
+  let expectedSizeIncrease = round (repvalavg ^ (fromSNat' n - fromSNat' m))
+  sh1 <- genShRwithTarget (1_000_000 `div` expectedSizeIncrease) m
   (sh2, sh3) <- injectOnes n sh1 sh1
   return (sh1, sh2, sh3)
   where
+    repvalrange = (1::Int, 5)
+    repvalavg = let (lo, hi) = repvalrange in fromIntegral (lo + hi) / 2 :: Double
+
     injectOnes :: m <= n => SNat n -> IShR m -> IShR m -> Gen (IShR n, IShR n)
     injectOnes n@SNat shOnes sh
       | m@SNat <- shrRank sh
@@ -105,7 +112,7 @@ genReplicatedShR = \m n -> do
           EQI -> return (shOnes, sh)
           GTI -> do
             index <- Gen.int (Range.linear 0 (fromSNat' m))
-            value <- Gen.int (Range.linear 1 5)
+            value <- Gen.int (uncurry Range.linear repvalrange)
             Refl <- return (lem n m)
             injectOnes n (inject index 1 shOnes) (inject index value sh)
 
@@ -115,7 +122,7 @@ genReplicatedShR = \m n -> do
     inject :: Int -> Int -> IShR m -> IShR (m + 1)
     inject 0 v sh = v :$: sh
     inject i v (w :$: sh) = w :$: inject (i - 1) v sh
-    inject _ v ZSR = v :$: ZSR  -- invalid input, but meh
+    inject _ _ ZSR = error "unreachable"
 
 genStorables :: forall a. Storable a => Range Int -> (Word64 -> a) -> GenT IO (VS.Vector a)
 genStorables rng f = do
diff --git a/test/Tests/C.hs b/test/Tests/C.hs
index 50e0e63..6881cd4 100644
--- a/test/Tests/C.hs
+++ b/test/Tests/C.hs
@@ -24,7 +24,7 @@ import Data.Array.Nested.Types (fromSNat')
 
 import Hedgehog
 import Hedgehog.Gen qualified as Gen
-import Hedgehog.Internal.Property (forAllT)
+import Hedgehog.Internal.Property (forAllT, LabelName (..))
 import Hedgehog.Range qualified as Range
 import Test.Tasty
 import Test.Tasty.Hedgehog
@@ -39,6 +39,9 @@ import Util
 fineTol :: Double
 fineTol = 1e-8
 
+debugCoverage :: Bool
+debugCoverage = False
+
 prop_sum_nonempty :: Property
 prop_sum_nonempty = property $ genRank $ \outrank@(SNat @n) -> do
   -- Test nonempty _results_. The first dimension of the input is allowed to be 0, because then OR.rerank doesn't fail yet.
@@ -89,6 +92,10 @@ prop_sum_replicated doTranspose = property $
       LTI -> return Refl  -- actually we only continue if m < n
       _ -> discard
     (sh1, sh2, sh3) <- forAll $ genReplicatedShR inrank1 inrank2
+    when debugCoverage $ do
+      label (LabelName ("rankdiff " ++ show (fromSNat' inrank2 - fromSNat' inrank1)))
+      label (LabelName ("size sh1 10^" ++ show (floor (logBase 10 (fromIntegral (shrSize sh1) :: Double)) :: Int)))
+      label (LabelName ("size sh3 10^" ++ show (floor (logBase 10 (fromIntegral (shrSize sh3) :: Double)) :: Int)))
     guard (all (> 0) sh3)
     arr <- forAllT $
       OR.stretch (toList sh3)