{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TupleSections #-}

module Test.QuickCheck.SafeGen.Internal where

import Control.Applicative (liftA2)
import Data.Foldable (toList)
import Data.List.NonEmpty (NonEmpty (..))
import Test.QuickCheck (Arbitrary (..), Gen)
import qualified Test.QuickCheck as QuickCheck

-- | 'SafeGen' is a tree of 'Gen' values, that automatically distributes the size parameter over its branches.
-- To run/consume a tree, use 'runSafeGen'.
-- Branches are either product types (composed through the 'Applicative' interface), or sum types composed using 'oneof' or 'frequency'.
data SafeGen a
  = Gen (Gen a)
  | -- 'Pure' could be encoded as @Gen . pure@, but by special-casing it we can maintain the Applicative laws.
    -- Specifically, when dividing the size parameter, we don't count 'Pure' branches.
    Pure a
  | forall i.
    Ap
      (SafeGen (i -> a))
      (SafeGen i)
  | Choice (NonEmpty (Int, SafeGen a))

deriving instance Functor SafeGen

instance Applicative SafeGen where
  pure :: forall a. a -> SafeGen a
pure = forall a. a -> SafeGen a
Pure
  Pure a -> b
a <*> :: forall a b. SafeGen (a -> b) -> SafeGen a -> SafeGen b
<*> SafeGen a
b = a -> b
a forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SafeGen a
b
  SafeGen (a -> b)
a <*> Pure a
b = (forall a b. (a -> b) -> a -> b
$ a
b) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SafeGen (a -> b)
a
  SafeGen (a -> b)
a <*> SafeGen a
b = forall a i. SafeGen (i -> a) -> SafeGen i -> SafeGen a
Ap SafeGen (a -> b)
a SafeGen a
b

instance Arbitrary a => Arbitrary (SafeGen a) where
  arbitrary :: Gen (SafeGen a)
arbitrary =
    forall a. [(Int, Gen a)] -> Gen a
QuickCheck.frequency
      [ (Int
1, forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. Gen a -> SafeGen a
Gen forall a. Arbitrary a => Gen a
arbitrary)),
        (Int
1, forall a. a -> SafeGen a
Pure forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Arbitrary a => Gen a
arbitrary),
        (Int
2, forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 forall a i. SafeGen (i -> a) -> SafeGen i -> SafeGen a
Ap (forall a. Arbitrary a => Gen a
arbitrary :: Gen (SafeGen (Int -> a))) forall a. Arbitrary a => Gen a
arbitrary),
        (Int
2, forall a. NonEmpty (Int, SafeGen a) -> SafeGen a
Choice forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(QuickCheck.Positive (QuickCheck.Small Int
w), SafeGen a
g) -> (Int
w, SafeGen a
g)) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall b. Arbitrary b => Gen (NonEmpty b)
nonEmpty)
      ]
    where
      nonEmpty :: Arbitrary b => Gen (NonEmpty b)
      nonEmpty :: forall b. Arbitrary b => Gen (NonEmpty b)
nonEmpty = forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 forall a. a -> [a] -> NonEmpty a
(:|) forall a. Arbitrary a => Gen a
arbitrary forall a. Arbitrary a => Gen a
arbitrary

data Nat = Zero | Succ Nat
  deriving (Int -> Nat -> ShowS
[Nat] -> ShowS
Nat -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Nat] -> ShowS
$cshowList :: [Nat] -> ShowS
show :: Nat -> String
$cshow :: Nat -> String
showsPrec :: Int -> Nat -> ShowS
$cshowsPrec :: Int -> Nat -> ShowS
Show)

-- | Run a 'SafeGen' using the current context's size parameter.
-- If the 'SafeGen' value does not have a leaf within 20 layers, assume it has infinite recursion, and throw an exception.
runSafeGen :: SafeGen a -> Gen a
runSafeGen :: forall a. SafeGen a -> Gen a
runSafeGen SafeGen a
sg
  | Bool -> Bool
not (Nat -> Int -> Bool
leqInt (forall a. SafeGen a -> Nat
shallowness SafeGen a
sg) Int
20) = forall a. HasCallStack => String -> a
error String
"runSafeGen: Minimum depth more than 20, likely because all paths have infinite recursion!"
  | Bool
otherwise = forall a. SafeGen a -> Gen a
runSafeGenNoCheck SafeGen a
sg

-- | like 'runSafeGen', but doesn't first check if this generator can terminate.
runSafeGenNoCheck :: SafeGen a -> Gen a
runSafeGenNoCheck :: forall a. SafeGen a -> Gen a
runSafeGenNoCheck SafeGen a
sg0 = forall a. (Int -> Gen a) -> Gen a
QuickCheck.sized (\Int
size -> forall a. Int -> SafeGen a -> Gen a
go Int
size SafeGen a
sg0)
  where
    go :: Int -> SafeGen a -> Gen a
    go :: forall a. Int -> SafeGen a -> Gen a
go Int
_ (Pure a
a) = forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a
    go !Int
size (Gen Gen a
g) = forall a. Int -> Gen a -> Gen a
QuickCheck.resize Int
size Gen a
g
    go !Int
size p :: SafeGen a
p@Ap {} = forall a. Int -> SafeGen a -> Gen a
goProduct (Int
size forall a. Integral a => a -> a -> a
`div` forall a. Ord a => a -> a -> a
max Int
1 (forall a. SafeGen a -> Int
arity SafeGen a
p)) SafeGen a
p
    go !Int
size (Choice ((Int
_, SafeGen a
a) :| [])) = forall a. Int -> SafeGen a -> Gen a
go Int
size SafeGen a
a
    go !Int
size (Choice NonEmpty (Int, SafeGen a)
as) =
      case forall a. (a -> Bool) -> [a] -> [a]
filter (forall a b c. (a -> b -> c) -> b -> a -> c
flip Nat -> Int -> Bool
leqInt Int
size forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. SafeGen a -> Nat
shallowness forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) (forall (t :: * -> *) a. Foldable t => t a -> [a]
toList NonEmpty (Int, SafeGen a)
as) of
        [] -> forall a. [(Int, Gen a)] -> Gen a
QuickCheck.frequency ((forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap) (forall a. Int -> SafeGen a -> Gen a
go Int
size) (forall (t :: * -> *) a. Foldable t => t a -> [a]
toList (forall (t :: * -> *) a.
Traversable t =>
(a -> Nat) -> t a -> NonEmpty a
safeMinBy (forall a. SafeGen a -> Nat
shallowness forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) NonEmpty (Int, SafeGen a)
as)))
        [(Int, SafeGen a)]
as' -> forall a. [(Int, Gen a)] -> Gen a
QuickCheck.frequency ((forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap) (forall a. Int -> SafeGen a -> Gen a
go Int
size) [(Int, SafeGen a)]
as')

    goProduct :: Int -> SafeGen a -> Gen a
    goProduct :: forall a. Int -> SafeGen a -> Gen a
goProduct !Int
size (Ap SafeGen (i -> a)
l SafeGen i
r) = forall a. Int -> SafeGen a -> Gen a
goProduct Int
size SafeGen (i -> a)
l forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. Int -> SafeGen a -> Gen a
goProduct Int
size SafeGen i
r
    goProduct !Int
size SafeGen a
sg = forall a. Int -> SafeGen a -> Gen a
go Int
size SafeGen a
sg

    arity :: SafeGen a -> Int
    arity :: forall a. SafeGen a -> Int
arity (Ap SafeGen (i -> a)
l SafeGen i
r) = forall a. SafeGen a -> Int
arity SafeGen (i -> a)
l forall a. Num a => a -> a -> a
+ forall a. SafeGen a -> Int
arity SafeGen i
r
    arity (Pure a
_) = Int
0
    arity (Gen Gen a
_) = Int
1
    arity (Choice NonEmpty (Int, SafeGen a)
_) = Int
1

    safeMinBy :: Traversable t => (a -> Nat) -> t a -> NonEmpty a
    safeMinBy :: forall (t :: * -> *) a.
Traversable t =>
(a -> Nat) -> t a -> NonEmpty a
safeMinBy a -> Nat
fdepth = forall {t :: * -> *} {a}. Traversable t => t (a, Nat) -> NonEmpty a
goMin forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\a
x -> (a
x, a -> Nat
fdepth a
x))
      where
        unpeel :: (a, Nat) -> Validation a (a, Nat)
unpeel (a
x, Nat
Zero) = forall e a. NonEmpty e -> Validation e a
VLeft (forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x)
        unpeel (a
x, Succ Nat
n) = forall e a. a -> Validation e a
VRight (a
x, Nat
n)
        goMin :: t (a, Nat) -> NonEmpty a
goMin t (a, Nat)
xs = case forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse forall {a}. (a, Nat) -> Validation a (a, Nat)
unpeel t (a, Nat)
xs of
          VLeft NonEmpty a
a -> NonEmpty a
a
          VRight t (a, Nat)
xs' -> t (a, Nat) -> NonEmpty a
goMin t (a, Nat)
xs'

leqInt :: Nat -> Int -> Bool
leqInt :: Nat -> Int -> Bool
leqInt Nat
Zero !Int
x = Int
x forall a. Ord a => a -> a -> Bool
> Int
0
leqInt (Succ Nat
n) !Int
x = Int
x forall a. Ord a => a -> a -> Bool
> Int
1 Bool -> Bool -> Bool
&& Nat -> Int -> Bool
leqInt Nat
n (Int
x forall a. Num a => a -> a -> a
- Int
1)

-- | Lift a 'Gen' generator into 'SafeGen'.
gen :: Gen a -> SafeGen a
gen :: forall a. Gen a -> SafeGen a
gen = forall a. Gen a -> SafeGen a
Gen

-- | Convenient synonym for 'gen arbitrary'.
arb :: Arbitrary a => SafeGen a
arb :: forall a. Arbitrary a => SafeGen a
arb = forall a. Gen a -> SafeGen a
gen forall a. Arbitrary a => Gen a
arbitrary

-- | Pick one of these branches, with equal probability.
-- Only branches shallower than the current size are considered.
oneof :: [SafeGen a] -> SafeGen a
oneof :: forall a. [SafeGen a] -> SafeGen a
oneof [] = forall a. HasCallStack => String -> a
error String
"SafeGen.oneof: empty list"
oneof (SafeGen a
a : [SafeGen a]
as) = forall a. NonEmpty (Int, SafeGen a) -> SafeGen a
Choice forall a b. (a -> b) -> a -> b
$ (Int
1,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SafeGen a
a forall a. a -> [a] -> NonEmpty a
:| [SafeGen a]
as

-- | Pick one of these branches, with weighted probability.
-- Only branches shallower than the current size are considered.
frequency :: [(Int, SafeGen a)] -> SafeGen a
frequency :: forall a. [(Int, SafeGen a)] -> SafeGen a
frequency [] = forall a. HasCallStack => String -> a
error String
"SafeGen.frequency: empty list"
frequency ((Int, SafeGen a)
a : [(Int, SafeGen a)]
as) = forall a. NonEmpty (Int, SafeGen a) -> SafeGen a
Choice ((Int, SafeGen a)
a forall a. a -> [a] -> NonEmpty a
:| [(Int, SafeGen a)]
as)

-- TODO memoize this into 'SafeGen' directly
shallowness :: SafeGen a -> Nat
shallowness :: forall a. SafeGen a -> Nat
shallowness = forall a. SafeGen a -> Nat
go
  where
    go :: SafeGen a -> Nat
    go :: forall a. SafeGen a -> Nat
go (Gen Gen a
_) = Nat
Zero
    go (Pure a
_) = Nat
Zero
    go (Choice NonEmpty (Int, SafeGen a)
as) = Nat -> Nat
Succ forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *). Traversable t => t Nat -> Nat
safeMin (forall a. SafeGen a -> Nat
go forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> NonEmpty (Int, SafeGen a)
as)
    go p :: SafeGen a
p@Ap {} = Nat -> Nat
Succ forall a b. (a -> b) -> a -> b
$ forall a. SafeGen a -> Nat
goProduct SafeGen a
p

    goProduct :: SafeGen a -> Nat
    goProduct :: forall a. SafeGen a -> Nat
goProduct (Ap SafeGen (i -> a)
l SafeGen i
r) = Nat -> Nat -> Nat
safeMax (forall a. SafeGen a -> Nat
goProduct SafeGen (i -> a)
l) (forall a. SafeGen a -> Nat
goProduct SafeGen i
r)
    goProduct SafeGen a
sg = forall a. SafeGen a -> Nat
go SafeGen a
sg

    safeMax :: Nat -> Nat -> Nat
    safeMax :: Nat -> Nat -> Nat
safeMax Nat
Zero Nat
b = Nat
b
    safeMax Nat
a Nat
Zero = Nat
a
    safeMax (Succ Nat
a) (Succ Nat
b) = Nat -> Nat
Succ (Nat -> Nat -> Nat
safeMax Nat
a Nat
b)

    safeMin :: Traversable t => t Nat -> Nat
    safeMin :: forall (t :: * -> *). Traversable t => t Nat -> Nat
safeMin t Nat
xs = case forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Nat -> Maybe Nat
unsucc t Nat
xs of
      Maybe (t Nat)
Nothing -> Nat
Zero
      Just t Nat
xs' -> Nat -> Nat
Succ forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *). Traversable t => t Nat -> Nat
safeMin t Nat
xs'
      where
        unsucc :: Nat -> Maybe Nat
        unsucc :: Nat -> Maybe Nat
unsucc Nat
Zero = forall a. Maybe a
Nothing
        unsucc (Succ Nat
l) = forall a. a -> Maybe a
Just Nat
l

-- | 'Either' that collects _all_ its failures in a list
data Validation e a = VLeft (NonEmpty e) | VRight a
  deriving (forall a b. a -> Validation e b -> Validation e a
forall a b. (a -> b) -> Validation e a -> Validation e b
forall e a b. a -> Validation e b -> Validation e a
forall e a b. (a -> b) -> Validation e a -> Validation e b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> Validation e b -> Validation e a
$c<$ :: forall e a b. a -> Validation e b -> Validation e a
fmap :: forall a b. (a -> b) -> Validation e a -> Validation e b
$cfmap :: forall e a b. (a -> b) -> Validation e a -> Validation e b
Functor)

instance Applicative (Validation e) where
  pure :: forall a. a -> Validation e a
pure = forall e a. a -> Validation e a
VRight
  VRight a -> b
f <*> :: forall a b.
Validation e (a -> b) -> Validation e a -> Validation e b
<*> VRight a
a = forall e a. a -> Validation e a
VRight (a -> b
f a
a)
  VLeft NonEmpty e
e1 <*> VLeft NonEmpty e
e2 = forall e a. NonEmpty e -> Validation e a
VLeft (NonEmpty e
e1 forall a. Semigroup a => a -> a -> a
<> NonEmpty e
e2)
  VLeft NonEmpty e
e1 <*> Validation e a
_ = forall e a. NonEmpty e -> Validation e a
VLeft NonEmpty e
e1
  Validation e (a -> b)
_ <*> VLeft NonEmpty e
e2 = forall e a. NonEmpty e -> Validation e a
VLeft NonEmpty e
e2