From 0b094dcda26967f2b7ab21aaf15161db3b5c4d96 Mon Sep 17 00:00:00 2001
From: Maximilian Algehed <m.algehed@gmail.com>
Date: Tue, 26 Mar 2024 09:37:23 +0100
Subject: [PATCH] add withMaxSize and withMaxShrinks modifiers

---
 src/Test/QuickCheck.hs          |  2 +
 src/Test/QuickCheck/Property.hs | 22 +++++++++++
 src/Test/QuickCheck/Test.hs     | 68 ++++++++++++++++++---------------
 tests/Misc.hs                   |  3 ++
 4 files changed, 64 insertions(+), 31 deletions(-)

diff --git a/src/Test/QuickCheck.hs b/src/Test/QuickCheck.hs
index 94cb6e81..12fb6822 100644
--- a/src/Test/QuickCheck.hs
+++ b/src/Test/QuickCheck.hs
@@ -305,6 +305,8 @@ module Test.QuickCheck
   , within
   , discardAfter
   , withDiscardRatio
+  , withMaxSize
+  , withMaxShrinks
   , once
   , again
   , mapSize
diff --git a/src/Test/QuickCheck/Property.hs b/src/Test/QuickCheck/Property.hs
index a092e2d5..cad7a861 100644
--- a/src/Test/QuickCheck/Property.hs
+++ b/src/Test/QuickCheck/Property.hs
@@ -273,6 +273,10 @@ data Result
     -- ^ required coverage confidence
   , maybeDiscardedRatio :: Maybe Int
     -- ^ maximum number of discarded tests per successful test
+  , maybeMaxShrinks     :: Maybe Int
+    -- ^ maximum number of shrinks
+  , maybeMaxTestSize    :: Maybe Int
+    -- ^ maximum test size
   , labels              :: [String]
     -- ^ test case labels
   , classes             :: [String]
@@ -317,6 +321,8 @@ succeeded, failed, rejected :: Result
       , maybeNumTests       = Nothing
       , maybeCheckCoverage  = Nothing
       , maybeDiscardedRatio = Nothing
+      , maybeMaxShrinks     = Nothing
+      , maybeMaxTestSize    = Nothing
       , labels              = []
       , classes             = []
       , tables              = []
@@ -482,6 +488,20 @@ withMaxSuccess n = n `seq` mapTotalResult (\res -> res{ maybeNumTests = Just n }
 withDiscardRatio :: Testable prop => Int -> prop -> Property
 withDiscardRatio n = n `seq` mapTotalResult (\res -> res{ maybeDiscardedRatio = Just n })
 
+-- | Configure the maximum number of times a property will be shrunk.
+--
+-- For example,
+--
+-- > quickCheck (withMaxShrinks 100 p)
+--
+-- will cause @p@ to only attempt 100 shrinks on failure.
+withMaxShrinks :: Testable prop => Int -> prop -> Property
+withMaxShrinks n = n `seq` mapTotalResult (\res -> res{ maybeMaxShrinks = Just n })
+
+-- | Configure the maximum size a property will be tested at.
+withMaxSize :: Testable prop => Int -> prop -> Property
+withMaxSize n = n `seq` mapTotalResult (\res -> res{ maybeMaxTestSize = Just n })
+
 -- | Check that all coverage requirements defined by 'cover' and 'coverTable'
 -- are met, using a statistically sound test, and fail if they are not met.
 --
@@ -932,6 +952,8 @@ disjoin ps =
                    maybeNumTests = Nothing,
                    maybeCheckCoverage = Nothing,
                    maybeDiscardedRatio = Nothing,
+                   maybeMaxShrinks = Nothing,
+                   maybeMaxTestSize = Nothing,
                    labels = [],
                    classes = [],
                    tables = [],
diff --git a/src/Test/QuickCheck/Test.hs b/src/Test/QuickCheck/Test.hs
index 1fbca6d3..4055d003 100644
--- a/src/Test/QuickCheck/Test.hs
+++ b/src/Test/QuickCheck/Test.hs
@@ -360,56 +360,62 @@ runATest st f =
      MkRose res ts <- protectRose (reduceRose (unProp (unGen (unProperty f_or_cov) rnd1 size)))
      res <- callbackPostTest st res
 
-     let continue break st' | abort res = break st'
-                            | otherwise = test st'
-
-     let st' = st{ coverageConfidence = maybeCheckCoverage res `mplus` coverageConfidence st
-                 , maxSuccessTests = fromMaybe (maxSuccessTests st) (maybeNumTests res)
-                 , maxDiscardedRatio = fromMaybe (maxDiscardedRatio st) (maybeDiscardedRatio res)
-                 , S.labels = Map.insertWith (+) (P.labels res) 1 (S.labels st)
-                 , S.classes = Map.unionWith (+) (S.classes st) (Map.fromList (zip (P.classes res) (repeat 1)))
-                 , S.tables =
-                   foldr (\(tab, x) -> Map.insertWith (Map.unionWith (+)) tab (Map.singleton x 1))
-                     (S.tables st) (P.tables res)
-                 , S.requiredCoverage =
-                   foldr (\(key, value, p) -> Map.insertWith max (key, value) p)
-                     (S.requiredCoverage st) (P.requiredCoverage res)
-                 , expected = expect res }
+     let continue break st' | abort res = break (addNewOptions st')
+                            | otherwise = test (addNewOptions st')
+
+         addNewOptions st0 = st0{ maxSuccessTests = fromMaybe (maxSuccessTests st0) (maybeNumTests res)
+                                , maxDiscardedRatio = fromMaybe (maxDiscardedRatio st0) (maybeDiscardedRatio res)
+                                , numTotMaxShrinks = fromMaybe (numTotMaxShrinks st0) (maybeMaxShrinks res)
+                                , maxTestSize = fromMaybe (maxTestSize st0) (maybeMaxTestSize res)
+                                }
+
+         addCoverageInfo st0 =
+           st0{ coverageConfidence = maybeCheckCoverage res `mplus` coverageConfidence st0
+              , S.labels = Map.insertWith (+) (P.labels res) 1 (S.labels st0)
+              , S.classes = Map.unionWith (+) (S.classes st0) (Map.fromList (zip (P.classes res) (repeat 1)))
+              , S.tables =
+                foldr (\(tab, x) -> Map.insertWith (Map.unionWith (+)) tab (Map.singleton x 1))
+                  (S.tables st0) (P.tables res)
+              , S.requiredCoverage =
+                foldr (\(key, value, p) -> Map.insertWith max (key, value) p)
+                  (S.requiredCoverage st0) (P.requiredCoverage res)
+              }
+
+         stC = addCoverageInfo st
 
      case res of
        MkResult{ok = Just True} -> -- successful test
          do continue doneTesting
-              st'{ numSuccessTests           = numSuccessTests st' + 1
+              stC{ numSuccessTests           = numSuccessTests st + 1
                  , numRecentlyDiscardedTests = 0
-                 , randomSeed = rnd2
+                 , randomSeed                = rnd2
+                 , expected                  = expect res
                  } f
 
        MkResult{ok = Nothing} -> -- discarded test
          do continue giveUp
               -- Don't add coverage info from this test
-              st{ numDiscardedTests         = numDiscardedTests st' + 1
-                , numRecentlyDiscardedTests = numRecentlyDiscardedTests st' + 1
-                , maxSuccessTests           = fromMaybe (maxSuccessTests st) (maybeNumTests res)
-                , maxDiscardedRatio         = fromMaybe (maxDiscardedRatio st) (maybeDiscardedRatio res)
+              st{ numDiscardedTests         = numDiscardedTests st + 1
+                , numRecentlyDiscardedTests = numRecentlyDiscardedTests st + 1
                 , randomSeed                = rnd2
                 } f
 
        MkResult{ok = Just False} -> -- failed test
-         do (numShrinks, totFailed, lastFailed, res) <- foundFailure st' res ts
-            theOutput <- terminalOutput (terminal st')
+         do (numShrinks, totFailed, lastFailed, res) <- foundFailure stC res ts
+            theOutput <- terminalOutput (terminal stC)
             if not (expect res) then
-              return Success{ labels = S.labels st',
-                              classes = S.classes st',
-                              tables = S.tables st',
-                              numTests = numSuccessTests st'+1,
-                              numDiscarded = numDiscardedTests st',
+              return Success{ labels = S.labels stC,
+                              classes = S.classes stC,
+                              tables = S.tables stC,
+                              numTests = numSuccessTests stC+1,
+                              numDiscarded = numDiscardedTests stC,
                               output = theOutput }
              else do
               testCase <- mapM showCounterexample (P.testCase res)
-              return Failure{ usedSeed        = randomSeed st' -- correct! (this will be split first)
+              return Failure{ usedSeed        = randomSeed stC -- correct! (this will be split first)
                             , usedSize        = size
-                            , numTests        = numSuccessTests st'+1
-                            , numDiscarded    = numDiscardedTests st'
+                            , numTests        = numSuccessTests stC + 1
+                            , numDiscarded    = numDiscardedTests stC
                             , numShrinks      = numShrinks
                             , numShrinkTries  = totFailed
                             , numShrinkFinal  = lastFailed
diff --git a/tests/Misc.hs b/tests/Misc.hs
index 016be569..93220af7 100644
--- a/tests/Misc.hs
+++ b/tests/Misc.hs
@@ -20,6 +20,9 @@ prop_failingTestCase (Blind p) = ioProperty $ do
   let [x, y, z] = failingTestCase res
   return (not (p (read x) (read y) (read z)))
 
+prop_maxSize :: Property
+prop_maxSize = withMaxSize 10 (forAll (arbitrary :: Gen Int) $ \ x -> abs x < 10)
+
 return []
 main = do
   True <- $quickCheckAll