{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeSynonymInstances #-}

module Test.Cardano.Ledger.Constrained.Combinators where

import Cardano.Ledger.Coin (Coin (..))
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Set (Set)
import qualified Data.Set as Set
import GHC.Stack (HasCallStack)
import Test.Cardano.Ledger.Binary.Random (QC (..))
import Test.Cardano.Ledger.Core.Arbitrary (uniformSubMap, uniformSubSet)
import Test.QuickCheck hiding (Fixed, total)

-- ==========================================================================
-- Tracking Gen-time errors

-- | Report a Gen-time error from the current message 'extra' and the
--   [messages] 'mess' describing the path to this call site.
errorMess :: HasCallStack => String -> [String] -> a
errorMess :: forall a. HasCallStack => String -> [String] -> a
errorMess String
extra [String]
mess = forall a. HasCallStack => String -> a
error ([String] -> String
unlines (String
"\nGen-time error" forall a. a -> [a] -> [a]
: forall a. [a] -> [a]
reverse (String
extra forall a. a -> [a] -> [a]
: [String]
mess)))

-- | suchThat version that tracks Gen-time errors
suchThatErr :: [String] -> Gen a -> (a -> Bool) -> Gen a
suchThatErr :: forall a. [String] -> Gen a -> (a -> Bool) -> Gen a
suchThatErr [String]
msgs Gen a
gen a -> Bool
p = forall a. (Int -> Gen a) -> Gen a
sized forall a b. (a -> b) -> a -> b
$ \Int
n -> Int -> Int -> Gen a
try (Int
10 :: Int) Int
n
  where
    try :: Int -> Int -> Gen a
try Int
0 Int
_ = forall a. HasCallStack => String -> [String] -> a
errorMess String
"SuchThat times out" [String]
msgs
    try Int
k Int
sz = do
      Maybe a
x <- forall a. HasCallStack => Int -> Gen a -> Gen a
resize Int
sz forall a b. (a -> b) -> a -> b
$ forall a. Gen a -> (a -> Bool) -> Gen (Maybe a)
suchThatMaybe Gen a
gen a -> Bool
p
      case Maybe a
x of
        Just a
y -> forall (f :: * -> *) a. Applicative f => a -> f a
pure a
y
        Maybe a
Nothing -> Int -> Int -> Gen a
try (Int
k forall a. Num a => a -> a -> a
- Int
1) (Int
sz forall a. Num a => a -> a -> a
+ Int
5)

-- =======================================================================

-- | add items from 'source' to 'base' until size 'n' is reached.
addUntilSize :: Ord a => [String] -> Set a -> Set a -> Int -> Gen (Set a)
addUntilSize :: forall a. Ord a => [String] -> Set a -> Set a -> Int -> Gen (Set a)
addUntilSize [String]
msgs Set a
base Set a
source Int
n = do
  let possible :: Set a
possible = forall a. Ord a => Set a -> Set a -> Set a
Set.difference Set a
source Set a
base
      p :: Int
p = forall a. Set a -> Int
Set.size Set a
possible
      m :: Int
m = forall a. Set a -> Int
Set.size Set a
base
      loop :: Set a -> Set a -> Gen (Set a)
loop Set a
result Set a
_ | forall a. Set a -> Int
Set.size Set a
result forall a. Ord a => a -> a -> Bool
>= Int
n = forall (f :: * -> *) a. Applicative f => a -> f a
pure Set a
result
      loop Set a
_ Set a
extra
        | forall a. Set a -> Bool
Set.null Set a
extra =
            forall a. HasCallStack => String -> [String] -> a
errorMess
              ( String
"There are not enough unused elements in 'source'("
                  forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
p
                  forall a. [a] -> [a] -> [a]
++ String
") to reach the size 'n'("
                  forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
n
                  forall a. [a] -> [a] -> [a]
++ String
")"
              )
              [String]
msgs
      loop Set a
result Set a
extra = do
        Int
i <- forall a. Random a => (a, a) -> Gen a
choose (Int
0, forall a. Set a -> Int
Set.size Set a
extra forall a. Num a => a -> a -> a
- Int
1)
        Set a -> Set a -> Gen (Set a)
loop (forall a. Ord a => a -> Set a -> Set a
Set.insert (forall a. Int -> Set a -> a
Set.elemAt Int
i Set a
extra) Set a
result) (forall a. Int -> Set a -> Set a
Set.deleteAt Int
i Set a
extra)
  case forall a. Ord a => a -> a -> Ordering
compare Int
m Int
n of
    Ordering
EQ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure Set a
base
    Ordering
GT ->
      forall a. HasCallStack => String -> [String] -> a
errorMess
        (String
"The size(" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
m forall a. [a] -> [a] -> [a]
++ String
") of the 'base' set exceeds the target size(" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
n forall a. [a] -> [a] -> [a]
++ String
")")
        [String]
msgs
    Ordering
LT -> Set a -> Set a -> Gen (Set a)
loop Set a
base Set a
possible

-- | Generate a random set of a fixed size 'size', use 'gen' to pick the elements.
setSized :: Ord a => [String] -> Int -> Gen a -> Gen (Set a)
setSized :: forall a. Ord a => [String] -> Int -> Gen a -> Gen (Set a)
setSized [String]
mess Int
size Gen a
gen = do
  Set a
set <- forall a. Ord a => [a] -> Set a
Set.fromList forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Int -> Gen a -> Gen [a]
vectorOf Int
size Gen a
gen
  forall a.
Ord a =>
[String] -> Int -> Int -> Gen a -> Set a -> Gen (Set a)
fixSet ((String
"setSized " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
size) forall a. a -> [a] -> [a]
: [String]
mess) Int
1000 Int
size Gen a
gen Set a
set

-- | Fix a a set 'source' that is supposed to have 'size' elements, but because of duplicates
--   may not. Fix it by adding random elements using 'genA', not currently in the set. If after 'numTries'
--   the set is still not fixed, report an error.
fixSet :: Ord a => [String] -> Int -> Int -> Gen a -> Set a -> Gen (Set a)
fixSet :: forall a.
Ord a =>
[String] -> Int -> Int -> Gen a -> Set a -> Gen (Set a)
fixSet [String]
mess Int
numTries Int
size Gen a
genA Set a
source
  | forall a. Set a -> Int
Set.size Set a
source forall a. Ord a => a -> a -> Bool
> Int
size = forall a. Ord a => [String] -> Set a -> Int -> Gen (Set a)
subsetFromSetWithSize (String
"fixSet" forall a. a -> [a] -> [a]
: [String]
mess) Set a
source Int
size
  | Bool
otherwise = Int -> Set a -> Gen a -> Gen (Set a)
go Int
numTries Set a
source Gen a
genA
  where
    go :: Int -> Set a -> Gen a -> Gen (Set a)
go Int
n !Set a
set Gen a
gen
      | Int
currentSize forall a. Eq a => a -> a -> Bool
== Int
size = forall (f :: * -> *) a. Applicative f => a -> f a
pure Set a
set
      | Int
n forall a. Ord a => a -> a -> Bool
< Int
0 =
          forall a. HasCallStack => String -> [String] -> a
errorMess
            ( String
"Ran out of tries("
                forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
numTries
                forall a. [a] -> [a] -> [a]
++ String
") in fixSet: need "
                forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
size
                forall a. [a] -> [a] -> [a]
++ String
" elements, have "
                forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
currentSize
            )
            [String]
mess
      | Bool
otherwise = do
          a
x <- Gen a
gen
          if a
x forall a. Ord a => a -> Set a -> Bool
`Set.member` Set a
set
            then Int -> Set a -> Gen a -> Gen (Set a)
go (Int
n forall a. Num a => a -> a -> a
- Int
1) Set a
set forall a b. (a -> b) -> a -> b
$ forall a. (Int -> Int) -> Gen a -> Gen a
scale (forall a. Num a => a -> a -> a
+ Int
5) Gen a
gen
            else Int -> Set a -> Gen a -> Gen (Set a)
go Int
n (forall a. Ord a => a -> Set a -> Set a
Set.insert a
x Set a
set) Gen a
gen
      where
        currentSize :: Int
currentSize = forall a. Set a -> Int
Set.size Set a
set

-- | Generate a random Map with a fixed size 'n'. use 'genA' to pick the
--   domain of the map, and 'genB' to pick the range.
mapSized :: Ord a => [String] -> Int -> Gen a -> Gen b -> Gen (Map a b)
mapSized :: forall a b.
Ord a =>
[String] -> Int -> Gen a -> Gen b -> Gen (Map a b)
mapSized [String]
mess Int
size Gen a
genA Gen b
genB = do
  Set a
keys <- forall a. Ord a => [String] -> Int -> Gen a -> Gen (Set a)
setSized ((String
"mapSized " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
size) forall a. a -> [a] -> [a]
: [String]
mess) Int
size Gen a
genA
  [b]
values <- forall a. Int -> Gen a -> Gen [a]
vectorOf Int
size Gen b
genB
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a b. Ord a => Set a -> [b] -> Map a b
mapFromDomRange Set a
keys [b]
values

coinSized :: Int -> Gen Coin
coinSized :: Int -> Gen Coin
coinSized Int
n = forall (f :: * -> *) a. Applicative f => a -> f a
pure (Integer -> Coin
Coin (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n))

-- | Generate a random subset of a set 'set' with a fixed size 'n'
subsetFromSetWithSize :: Ord a => [String] -> Set a -> Int -> Gen (Set a)
subsetFromSetWithSize :: forall a. Ord a => [String] -> Set a -> Int -> Gen (Set a)
subsetFromSetWithSize [String]
mess Set a
set Int
n
  | forall a. Set a -> Int
Set.size Set a
set forall a. Ord a => a -> a -> Bool
< Int
n =
      forall a. HasCallStack => String -> [String] -> a
errorMess
        ( String
"subsetFromSetWithSize: Can't make a subset of size "
            forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
n
            forall a. [a] -> [a] -> [a]
++ String
" from a smaller set of size "
            forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show (forall a. Set a -> Int
Set.size Set a
set)
        )
        [String]
mess
  | Bool
otherwise = forall g (m :: * -> *) k.
(StatefulGen g m, Ord k) =>
Maybe Int -> Set k -> g -> m (Set k)
uniformSubSet (forall a. a -> Maybe a
Just Int
n) Set a
set QC
QC

-- | Generate a larger map, from a smaller map 'subset'. The new larger map, should have all the
--   keys and values of the smaller map.
mapFromSubset :: Ord a => [String] -> Map a b -> Int -> Gen a -> Gen b -> Gen (Map a b)
mapFromSubset :: forall a b.
Ord a =>
[String] -> Map a b -> Int -> Gen a -> Gen b -> Gen (Map a b)
mapFromSubset [String]
mess Map a b
subset Int
n Gen a
genA Gen b
genB = do
  Map a b
additions <- forall a b.
Ord a =>
[String] -> Int -> Gen a -> Gen b -> Gen (Map a b)
mapSized ((String
"mapFromSubset " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
n) forall a. a -> [a] -> [a]
: [String]
mess) Int
n Gen a
genA Gen b
genB
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall k a. Ord k => Map k a -> Map k a -> Map k a
Map.union Map a b
subset Map a b
additions)

-- | Generate a map, from a set. The new map, should have all the
--   elements of 'set' in its keysSet.
mapFromSet :: Ord a => Set a -> Gen b -> Gen (Map a b)
mapFromSet :: forall a b. Ord a => Set a -> Gen b -> Gen (Map a b)
mapFromSet Set a
set Gen b
genB = Set a -> Gen (Map a b)
addRange Set a
set
  where
    addRange :: Set a -> Gen (Map a b)
addRange Set a
s = forall a b. (a -> b -> a) -> a -> Set b -> a
Set.foldl' Gen (Map a b) -> a -> Gen (Map a b)
accum (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall k a. Map k a
Map.empty) Set a
s
    accum :: Gen (Map a b) -> a -> Gen (Map a b)
accum Gen (Map a b)
ansGen a
dom = do Map a b
ans <- Gen (Map a b)
ansGen; b
rng <- Gen b
genB; forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert a
dom b
rng Map a b
ans)

-- | Generate a random element from a set. Also return the set with that element removed
itemFromSet :: [String] -> Set a -> Gen (a, Set a)
itemFromSet :: forall a. [String] -> Set a -> Gen (a, Set a)
itemFromSet [String]
mess Set a
set
  | forall a. Set a -> Bool
Set.null Set a
set = forall a. HasCallStack => String -> [String] -> a
errorMess String
"itemFromSet : Can't take an item from the empty set." [String]
mess
itemFromSet [String]
_ Set a
set =
  (\Int
ix -> (forall a. Int -> Set a -> a
Set.elemAt Int
ix Set a
set, forall a. Int -> Set a -> Set a
Set.deleteAt Int
ix Set a
set)) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int, Int) -> Gen Int
chooseInt (Int
0, forall a. Set a -> Int
Set.size Set a
set forall a. Num a => a -> a -> a
- Int
1)

-- Generate a map from a list 'bs'. The result should have the property (Map.elems result == bs)
mapFromRange :: forall a b. Ord a => [String] -> [b] -> Gen a -> Gen (Map a b)
mapFromRange :: forall a b. Ord a => [String] -> [b] -> Gen a -> Gen (Map a b)
mapFromRange [String]
msgs [b]
bs Gen a
genA = do
  Set a
keys <- forall a. Ord a => [String] -> Int -> Gen a -> Gen (Set a)
setSized [String]
msgs (forall (t :: * -> *) a. Foldable t => t a -> Int
length [b]
bs) Gen a
genA
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a b. Ord a => Set a -> [b] -> Map a b
mapFromDomRange Set a
keys [b]
bs

mapFromProj :: Ord a => [b] -> Gen a -> (b -> Gen c) -> Gen (Map a c)
mapFromProj :: forall a b c.
Ord a =>
[b] -> Gen a -> (b -> Gen c) -> Gen (Map a c)
mapFromProj [b]
bs Gen a
genA b -> Gen c
genC = do
  [c]
cs <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM b -> Gen c
genC [b]
bs
  forall a b. Ord a => [String] -> [b] -> Gen a -> Gen (Map a b)
mapFromRange [String
"mapFromProj"] [c]
cs Gen a
genA

-- | Generate a Map from a set 'dom' and a list 'bs'. They should have the same size
--   but this is not checked. The result should have the properties (Map.keysSet result == dom)
--   and (Map.elems result == bs)
mapFromDomRange :: Ord a => Set a -> [b] -> Map a b
mapFromDomRange :: forall a b. Ord a => Set a -> [b] -> Map a b
mapFromDomRange Set a
dom [b]
bs = forall k a. Eq k => [(k, a)] -> Map k a
Map.fromAscList forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip (forall a. Set a -> [a]
Set.toAscList Set a
dom) [b]
bs

-- | Generate an Int, that would be a valid size for a subset of a set 's'
--   it should return 0 or the size of the set infrequently.
subsetSize :: Set a -> Gen Int
subsetSize :: forall a. Set a -> Gen Int
subsetSize Set a
s | forall a. Set a -> Bool
Set.null Set a
s = forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
0
subsetSize Set a
s = forall a. HasCallStack => [(Int, Gen a)] -> Gen a
frequency [(Int
1, forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
0), (Int
20, forall a. Random a => (a, a) -> Gen a
choose (Int
1, Int
n forall a. Num a => a -> a -> a
- Int
1)), (Int
1, forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
n)]
  where
    n :: Int
n = forall a. Set a -> Int
Set.size Set a
s

-- | Generate a random subset of a set 'set'
subsetFromSet :: Ord a => [String] -> Set a -> Gen (Set a)
subsetFromSet :: forall a. Ord a => [String] -> Set a -> Gen (Set a)
subsetFromSet [String]
mess Set a
set = do
  Int
n <- forall a. Set a -> Gen Int
subsetSize Set a
set
  forall a. Ord a => [String] -> Set a -> Int -> Gen (Set a)
subsetFromSetWithSize (String
"From subsetFromSet" forall a. a -> [a] -> [a]
: [String]
mess) Set a
set Int
n

-- | Generate a random superset of a set 'set'. Use 'genA' to pick random additional elements
superSetFromSet :: Ord a => Gen a -> Set a -> Gen (Set a)
superSetFromSet :: forall a. Ord a => Gen a -> Set a -> Gen (Set a)
superSetFromSet Gen a
genA Set a
setA = do
  Int
n <- forall a. Random a => (a, a) -> Gen a
choose (Int
0, Int
4)
  if Int
n forall a. Eq a => a -> a -> Bool
== Int
0
    then forall (f :: * -> *) a. Applicative f => a -> f a
pure Set a
setA
    else forall a. Ord a => [String] -> Int -> Gen a -> Set a -> Gen (Set a)
superSetFromSetWithSize [String
"supersetFromSet " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
n] (Int
n forall a. Num a => a -> a -> a
+ forall a. Set a -> Int
Set.size Set a
setA) Gen a
genA Set a
setA

-- | Generate a random superset of a set 'set' with s fixed size 'n'.
--   Use 'genA' to pick random additional elements.
superSetFromSetWithSize :: Ord a => [String] -> Int -> Gen a -> Set a -> Gen (Set a)
superSetFromSetWithSize :: forall a. Ord a => [String] -> Int -> Gen a -> Set a -> Gen (Set a)
superSetFromSetWithSize [String]
mess Int
n Gen a
_ Set a
setA
  | Int
n forall a. Ord a => a -> a -> Bool
< forall a. Set a -> Int
Set.size Set a
setA =
      forall a. HasCallStack => String -> [String] -> a
errorMess
        ( String
"superSetFromSetWithSize: Size of the superset ("
            forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
n
            forall a. [a] -> [a] -> [a]
++ String
") is smaller than "
            forall a. [a] -> [a] -> [a]
++ String
"the size of the given set ("
            forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show (forall a. Set a -> Int
Set.size Set a
setA)
            forall a. [a] -> [a] -> [a]
++ String
")."
        )
        [String]
mess
superSetFromSetWithSize [String]
msgs Int
size Gen a
genA Set a
setA =
  let tries :: Int
tries = Int
10 forall a. Num a => a -> a -> a
* (Int
size forall a. Num a => a -> a -> a
- forall a. Set a -> Int
Set.size Set a
setA) -- at most 10 times per new element
   in forall a.
Ord a =>
[String] -> Int -> Int -> Gen a -> Set a -> Gen (Set a)
fixSet (String
"superSetFromSetWithSize" forall a. a -> [a] -> [a]
: [String]
msgs) Int
tries Int
size Gen a
genA Set a
setA

-- | Generates things both in (Set a) and not in (Set a) with differing frequencies.
superItemFromSet :: Ord a => Gen a -> Set a -> Gen a
superItemFromSet :: forall a. Ord a => Gen a -> Set a -> Gen a
superItemFromSet Gen a
genA Set a
set | forall a. Set a -> Bool
Set.null Set a
set = Gen a
genA
superItemFromSet Gen a
genA Set a
set =
  forall a. HasCallStack => [(Int, Gen a)] -> Gen a
frequency
    [ (Int
3, forall a b. (a, b) -> a
fst forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. [String] -> Set a -> Gen (a, Set a)
itemFromSet [String
"Not possible since set is not empty"] Set a
set)
    , (Int
1, forall a. Gen a -> (a -> Bool) -> Gen a
suchThat Gen a
genA (forall a. Ord a => a -> Set a -> Bool
`Set.notMember` Set a
set))
    ]

-- | Pick a random (key,value) pair from a Map
genFromMap :: [String] -> Map k a -> Gen (k, a)
genFromMap :: forall k a. [String] -> Map k a -> Gen (k, a)
genFromMap [String]
msgs Map k a
m
  | Int
n forall a. Eq a => a -> a -> Bool
== Int
0 = forall a. HasCallStack => String -> [String] -> a
errorMess String
"The map is empty in genFromMap" [String]
msgs
  | Bool
otherwise = do
      Int
i <- forall a. Random a => (a, a) -> Gen a
choose (Int
0, Int
n forall a. Num a => a -> a -> a
- Int
1)
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall k a. Int -> Map k a -> (k, a)
Map.elemAt Int
i Map k a
m
  where
    n :: Int
n = forall k a. Map k a -> Int
Map.size Map k a
m

subMapFromMapWithSize :: Ord k => Int -> Map k a -> Gen (Map k a)
subMapFromMapWithSize :: forall k a. Ord k => Int -> Map k a -> Gen (Map k a)
subMapFromMapWithSize Int
n Map k a
m = forall g (m :: * -> *) k v.
(StatefulGen g m, Ord k) =>
Maybe Int -> Map k v -> g -> m (Map k v)
uniformSubMap (forall a. a -> Maybe a
Just Int
n) Map k a
m QC
QC