-- | Test case generation.
module Test.QuickCheck.Gen where

--------------------------------------------------------------------------
-- imports

import System.Random
( Random
, StdGen
, randomR
, split
, newStdGen
)

( liftM
, ap
)

import Control.Applicative
( Applicative(..)
)

--------------------------------------------------------------------------
-- ** Generator type

newtype Gen a = MkGen{ unGen :: StdGen -> Int -> a }

instance Functor Gen where
fmap f (MkGen h) =
MkGen (\r n -> f (h r n))

instance Applicative Gen where
pure  = return
(<*>) = ap

return x =
MkGen (\_ _ -> x)

MkGen m >>= k =
MkGen (\r n ->
let (r1,r2)  = split r
MkGen m' = k (m r1 n)
in m' r2 n
)

--------------------------------------------------------------------------
-- ** Primitive generator combinators

-- | Modifies a generator using an integer seed.
variant :: Integral n => n -> Gen a -> Gen a
variant k0 (MkGen m) = MkGen (\r n -> m (var k0 r) n)
where
var k = (if k == k' then id  else var k')
. (if even k  then fst else snd)
. split
where
k' = k `div` 2

-- | Used to construct generators that depend on the size parameter.
sized :: (Int -> Gen a) -> Gen a
sized f = MkGen (\r n -> let MkGen m = f n in m r n)

-- | Overrides the size parameter. Returns a generator which uses
-- the given size instead of the runtime-size parameter.
resize :: Int -> Gen a -> Gen a
resize n (MkGen m) = MkGen (\r _ -> m r n)

-- | Generates a random element in the given inclusive range.
choose :: Random a => (a,a) -> Gen a
choose rng = MkGen (\r _ -> let (x,_) = randomR rng r in x)

-- | Promotes a monadic generator to a generator of monadic values.
promote :: Monad m => m (Gen a) -> Gen (m a)
promote m = MkGen (\r n -> liftM (\(MkGen m') -> m' r n) m)

-- | Generates some example values.
sample' :: Gen a -> IO [a]
sample' (MkGen m) =
do rnd0 <- newStdGen
let rnds rnd = rnd1 : rnds rnd2 where (rnd1,rnd2) = split rnd
return [(m r n) | (r,n) <- rnds rnd0 `zip` [0,2..20] ]

-- | Generates some example values and prints them to 'stdout'.
sample :: Show a => Gen a -> IO ()
sample g =
do cases <- sample' g
sequence_ (map print cases)

--------------------------------------------------------------------------
-- ** Common generator combinators

-- | Generates a value that satisfies a predicate.
suchThat :: Gen a -> (a -> Bool) -> Gen a
gen `suchThat` p =
do mx <- gen `suchThatMaybe` p
case mx of
Just x  -> return x
Nothing -> sized (\n -> resize (n+1) (gen `suchThat` p))

-- | Tries to generate a value that satisfies a predicate.
suchThatMaybe :: Gen a -> (a -> Bool) -> Gen (Maybe a)
gen `suchThatMaybe` p = sized (try 0 . max 1)
where
try _ 0 = return Nothing
try k n = do x <- resize (2*k+n) gen
if p x then return (Just x) else try (k+1) (n-1)

-- | Randomly uses one of the given generators. The input list
-- must be non-empty.
oneof :: [Gen a] -> Gen a
oneof [] = error "QuickCheck.oneof used with empty list"
oneof gs = choose (0,length gs - 1) >>= (gs !!)

-- | Chooses one of the given generators, with a weighted random distribution.
-- The input list must be non-empty.
frequency :: [(Int, Gen a)] -> Gen a
frequency [] = error "QuickCheck.frequency used with empty list"
frequency xs0 = choose (1, tot) >>= (`pick` xs0)
where
tot = sum (map fst xs0)

pick n ((k,x):xs)
| n <= k    = x
| otherwise = pick (n-k) xs
pick _ _  = error "QuickCheck.pick used with empty list"

-- | Generates one of the given values. The input list must be non-empty.
elements :: [a] -> Gen a
elements [] = error "QuickCheck.elements used with empty list"
elements xs = (xs !!) `fmap` choose (0, length xs - 1)

-- | Takes a list of elements of increasing size, and chooses
-- among an initial segment of the list. The size of this initial
-- segment increases with the size parameter.
-- The input list must be non-empty.
growingElements :: [a] -> Gen a
growingElements [] = error "QuickCheck.growingElements used with empty list"
growingElements xs = sized \$ \n -> elements (take (1 `max` size n) xs)
where
k      = length xs
mx     = 100
log'   = round . log . fromIntegral
size n = (log' n + 1) * k `div` log' mx

{- WAS:
growingElements xs = sized \$ \n -> elements (take (1 `max` (n * k `div` 100)) xs)
where
k = length xs
-}

-- | Generates a list of random length. The maximum length depends on the
-- size parameter.
listOf :: Gen a -> Gen [a]
listOf gen = sized \$ \n ->
do k <- choose (0,n)
vectorOf k gen

-- | Generates a non-empty list of random length. The maximum length
-- depends on the size parameter.
listOf1 :: Gen a -> Gen [a]
listOf1 gen = sized \$ \n ->
do k <- choose (1,1 `max` n)
vectorOf k gen

-- | Generates a list of the given length.
vectorOf :: Int -> Gen a -> Gen [a]
vectorOf k gen = sequence [ gen | _ <- [1..k] ]

--------------------------------------------------------------------------
-- the end.