module PostgresqlSyntax.Helpers.Gens where

import PostgresqlSyntax.Prelude
import Test.QuickCheck

downscale :: Gen a -> Gen a
downscale :: forall a. Gen a -> Gen a
downscale = (Int -> Int) -> Gen a -> Gen a
forall a. (Int -> Int) -> Gen a -> Gen a
scale (Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2)

recursive :: Gen a -> Gen a -> Gen a
recursive :: forall a. Gen a -> Gen a -> Gen a
recursive Gen a
nonRecursiveGen Gen a
recursiveGen = (Int -> Gen a) -> Gen a
forall a. (Int -> Gen a) -> Gen a
sized ((Int -> Gen a) -> Gen a) -> (Int -> Gen a) -> Gen a
forall a b. (a -> b) -> a -> b
$ \Int
size ->
  if Int
size Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1
    then Gen a
nonRecursiveGen
    else Gen a -> Gen a
forall a. Gen a -> Gen a
downscale Gen a
recursiveGen

oneofRec ::
  (Arbitrary a) =>
  [Gen a] ->
  [Gen a] ->
  Gen a
oneofRec :: forall a. Arbitrary a => [Gen a] -> [Gen a] -> Gen a
oneofRec [Gen a]
nonRecursiveGens [Gen a]
recursiveGens = (Int -> Gen a) -> Gen a
forall a. (Int -> Gen a) -> Gen a
sized ((Int -> Gen a) -> Gen a) -> (Int -> Gen a) -> Gen a
forall a b. (a -> b) -> a -> b
$ \Int
size ->
  if Int
size Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1
    then [Gen a] -> Gen a
forall a. HasCallStack => [Gen a] -> Gen a
oneof [Gen a]
nonRecursiveGens
    else
      [(Int, Gen a)] -> Gen a
forall a. HasCallStack => [(Int, Gen a)] -> Gen a
frequency
        [ (Int
1, [Gen a] -> Gen a
forall a. HasCallStack => [Gen a] -> Gen a
oneof [Gen a]
nonRecursiveGens),
          (Int
3, Gen a -> Gen a
forall a. Gen a -> Gen a
downscale ([Gen a] -> Gen a
forall a. HasCallStack => [Gen a] -> Gen a
oneof [Gen a]
recursiveGens))
        ]

-- | Generate a non-empty list of at most @n + 1@ elements, splitting the size
-- budget across them.
--
-- The split is what keeps growth bounded: generating every element at the
-- undiminished size would multiply the subtree's cost by the list length at no
-- size cost, and those multipliers compound through the AST.
nonEmptyUpTo :: Int -> Gen a -> Gen (NonEmpty a)
nonEmptyUpTo :: forall a. Int -> Gen a -> Gen (NonEmpty a)
nonEmptyUpTo Int
n Gen a
gen = (Int -> Gen (NonEmpty a)) -> Gen (NonEmpty a)
forall a. (Int -> Gen a) -> Gen a
sized ((Int -> Gen (NonEmpty a)) -> Gen (NonEmpty a))
-> (Int -> Gen (NonEmpty a)) -> Gen (NonEmpty a)
forall a b. (a -> b) -> a -> b
$ \Int
size -> do
  -- The 'max 0' matters: at size 0 the upper bound is negative, and 'choose'
  -- silently swaps inverted bounds instead of failing.
  Int
tailLen <- (Int, Int) -> Gen Int
forall a. Random a => (a, a) -> Gen a
choose (Int
0, Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
n (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)))
  let totalLen :: Int
totalLen = Int
tailLen Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
      subsize :: Int
subsize = Int
size Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
totalLen
      subgen :: Gen a
subgen = Int -> Gen a -> Gen a
forall a. HasCallStack => Int -> Gen a -> Gen a
resize Int
subsize Gen a
gen
  a
x <- Gen a
subgen
  [a]
xs <- Int -> Gen a -> Gen [a]
forall a. Int -> Gen a -> Gen [a]
vectorOf Int
tailLen Gen a
subgen
  pure (a
x a -> [a] -> NonEmpty a
forall a. a -> [a] -> NonEmpty a
:| [a]
xs)

terminatingMaybe :: Gen a -> Gen (Maybe a)
terminatingMaybe :: forall a. Gen a -> Gen (Maybe a)
terminatingMaybe Gen a
gen = (Int -> Gen (Maybe a)) -> Gen (Maybe a)
forall a. (Int -> Gen a) -> Gen a
sized ((Int -> Gen (Maybe a)) -> Gen (Maybe a))
-> (Int -> Gen (Maybe a)) -> Gen (Maybe a)
forall a b. (a -> b) -> a -> b
$ \Int
size ->
  if Int
size Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1
    then Maybe a -> Gen (Maybe a)
forall a. a -> Gen a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe a
forall a. Maybe a
Nothing
    else a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> Gen a -> Gen (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen a
gen