From 428152ffb3903e44588f80826864d52de4b96ebf Mon Sep 17 00:00:00 2001
From: Maximilian Algehed <m.algehed@gmail.com>
Date: Thu, 21 Mar 2024 13:49:57 +0100
Subject: [PATCH] introduce `suchThatDiscard`

---
 src/Test/QuickCheck.hs     |  1 +
 src/Test/QuickCheck/Gen.hs | 11 +++++++++++
 tests/DiscardRatio.hs      |  4 ++++
 3 files changed, 16 insertions(+)

diff --git a/src/Test/QuickCheck.hs b/src/Test/QuickCheck.hs
index 0ce960f1..3c0eb24c 100644
--- a/src/Test/QuickCheck.hs
+++ b/src/Test/QuickCheck.hs
@@ -112,6 +112,7 @@ module Test.QuickCheck
   , suchThat
   , suchThatMap
   , suchThatMaybe
+  , suchThatDiscard
   , applyArbitrary2
   , applyArbitrary3
   , applyArbitrary4
diff --git a/src/Test/QuickCheck/Gen.hs b/src/Test/QuickCheck/Gen.hs
index aca58183..79e6bee6 100644
--- a/src/Test/QuickCheck/Gen.hs
+++ b/src/Test/QuickCheck/Gen.hs
@@ -35,6 +35,7 @@ import Control.Applicative
   ( Applicative(..) )
 
 import Test.QuickCheck.Random
+import Test.QuickCheck.Exception
 import Data.List (sortBy)
 import Data.Ord
 import Data.Maybe
@@ -293,6 +294,16 @@ gen `suchThatMaybe` p = sized (\n -> try n (2*n))
         x <- resize m gen
         if p x then return (Just x) else try (m+1) n
 
+-- | Tries to generate a value that satisfies a predicate.
+-- If it fails to do so it discards the test case if the result
+-- is used in the test.
+suchThatDiscard :: Gen a -> (a -> Bool) -> Gen a
+suchThatDiscard g p = do
+  a <- g
+  if p a
+  then pure a
+  else discard
+
 -- | Randomly uses one of the given generators. The input list
 -- must be non-empty.
 oneof :: WITHCALLSTACK([Gen a] -> Gen a)
diff --git a/tests/DiscardRatio.hs b/tests/DiscardRatio.hs
index db6f84b5..83a0e53d 100644
--- a/tests/DiscardRatio.hs
+++ b/tests/DiscardRatio.hs
@@ -49,3 +49,7 @@ main = do
   putStrLn "\nExpecting success (discard ratio 40): x < 50 ==> True"
   quickCheckYes $ withDiscardRatio 40 p50
   quickCheckYesWith stdArgs{maxDiscardRatio = 40} p50
+
+  quickCheckNo $ forAll  (choose (1 :: Int, 10) `suchThatDiscard` const False) $ \ x -> x == x
+  quickCheckYes $ forAll (choose (1 :: Int, 10) `suchThatDiscard` const False) $ \ _ -> True
+  quickCheckYes $ forAll (choose (1 :: Int, 10) `suchThatDiscard` (> 3))       $ \ x -> x == x