{-# LANGUAGE TypeApplications #-}
module PlutusCore.Generators.QuickCheck.Split where
import Control.Monad
import Data.Bifunctor
import Data.Coerce (coerce)
import Data.List (sortBy)
import Data.Ord (comparing)
import Test.QuickCheck
smallLength :: Int
smallLength :: Int
smallLength = Int
6
sublistN :: Int -> [a] -> Gen [a]
sublistN :: forall a. Int -> [a] -> Gen [a]
sublistN Int
lenRes
= ([(Int, a)] -> [a]) -> Gen [(Int, a)] -> Gen [a]
forall a b. (a -> b) -> Gen a -> Gen b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (((Int, a) -> a) -> [(Int, a)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (Int, a) -> a
forall a b. (a, b) -> b
snd ([(Int, a)] -> [a])
-> ([(Int, a)] -> [(Int, a)]) -> [(Int, a)] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Int, a) -> (Int, a) -> Ordering) -> [(Int, a)] -> [(Int, a)]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy (((Int, a) -> Int) -> (Int, a) -> (Int, a) -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing (Int, a) -> Int
forall a b. (a, b) -> a
fst) ([(Int, a)] -> [(Int, a)])
-> ([(Int, a)] -> [(Int, a)]) -> [(Int, a)] -> [(Int, a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [(Int, a)] -> [(Int, a)]
forall a. Int -> [a] -> [a]
take Int
lenRes)
(Gen [(Int, a)] -> Gen [a])
-> ([a] -> Gen [(Int, a)]) -> [a] -> Gen [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Int, a)] -> Gen [(Int, a)]
forall a. [a] -> Gen [a]
shuffle
([(Int, a)] -> Gen [(Int, a)])
-> ([a] -> [(Int, a)]) -> [a] -> Gen [(Int, a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 :: Int ..]
toMaxChunkNumber :: Int -> Int
toMaxChunkNumber :: Int -> Int
toMaxChunkNumber Int
len
| Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
smallLength = Int
len
| Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
smallLength Int -> Int -> Int
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int) = Int
smallLength Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
smallLength
| Bool
otherwise = Int
smallLength Int -> Int -> Int
forall a. Num a => a -> a -> a
+ forall a b. (RealFrac a, Integral b) => a -> b
round @Double (Double -> Double
forall a. Floating a => a -> a
sqrt (Double -> Double) -> Double -> Double
forall a b. (a -> b) -> a -> b
$ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len)
toChunkNumber :: Int -> Int -> Int
toChunkNumber :: Int -> Int -> Int
toChunkNumber Int
len Int
chunkNum =
[Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1, Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2 .. Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
chunkNum Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1] Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div`
[Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int
chunkNum Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1, Int
chunkNum Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2 .. Int
2]
toChunkFrequencies :: Int -> [(Int, Int)]
toChunkFrequencies :: Int -> [(Int, Int)]
toChunkFrequencies Int
len
| Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
smallLength = (Int -> (Int, Int)) -> [Int] -> [(Int, Int)]
forall a b. (a -> b) -> [a] -> [b]
map (\Int
num -> (Int -> Int -> Int
toChunkNumber Int
len Int
num, Int
num)) [Int]
chunks
| Bool
otherwise =
let
singleElemProb :: Int
singleElemProb = Int
3
deltaN :: Int
deltaN = Int
chunkMax Int -> Int -> Int
forall a. Num a => a -> a -> a
* (Int
chunkMax Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
delta :: Int
delta = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 (Int -> Int) -> Int -> Int
forall a b. (a -> b) -> a -> b
$ (Int
100 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
chunkMax Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
singleElemProb) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
deltaN
in [Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Int -> Int) -> Int -> [Int]
forall a. (a -> a) -> a -> [a]
iterate (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
delta) Int
singleElemProb) [Int]
chunks
where
chunkMax :: Int
chunkMax = Int -> Int
toMaxChunkNumber Int
len
chunks :: [Int]
chunks = [Int
1 .. Int
chunkMax]
toChunks :: [Int] -> [a] -> [[a]]
toChunks :: forall a. [Int] -> [a] -> [[a]]
toChunks [] [a]
xs = [[a]
xs]
toChunks (Int
n : [Int]
ns) [a]
xs = [a]
chunk [a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
: [Int] -> [a] -> [[a]]
forall a. [Int] -> [a] -> [[a]]
toChunks [Int]
ns [a]
xs' where
([a]
chunk, [a]
xs') = Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
n [a]
xs
multiSplit1In :: Int -> [a] -> Gen [NonEmptyList a]
multiSplit1In :: forall a. Int -> [a] -> Gen [NonEmptyList a]
multiSplit1In Int
_ [] = [NonEmptyList a] -> Gen [NonEmptyList a]
forall a. a -> Gen a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
multiSplit1In Int
chunkNum [a]
xs = do
let len :: Int
len = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs
[Int]
breakpoints <- Int -> [Int] -> Gen [Int]
forall a. Int -> [a] -> Gen [a]
sublistN (Int
chunkNum Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [Int
1 .. Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
let chunkLens :: [Int]
chunkLens = (Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (-) [Int]
breakpoints (Int
0 Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
breakpoints)
[NonEmptyList a] -> Gen [NonEmptyList a]
forall a. a -> Gen a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([NonEmptyList a] -> Gen [NonEmptyList a])
-> ([[a]] -> [NonEmptyList a]) -> [[a]] -> Gen [NonEmptyList a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[a]] -> [NonEmptyList a]
forall a b. Coercible a b => a -> b
coerce ([[a]] -> Gen [NonEmptyList a]) -> [[a]] -> Gen [NonEmptyList a]
forall a b. (a -> b) -> a -> b
$ [Int] -> [a] -> [[a]]
forall a. [Int] -> [a] -> [[a]]
toChunks [Int]
chunkLens [a]
xs
multiSplit1 :: [a] -> Gen [NonEmptyList a]
multiSplit1 :: forall a. [a] -> Gen [NonEmptyList a]
multiSplit1 [a]
xs = do
Int
chunkNum <- [(Int, Gen Int)] -> Gen Int
forall a. [(Int, Gen a)] -> Gen a
frequency ([(Int, Gen Int)] -> Gen Int)
-> (Int -> [(Int, Gen Int)]) -> Int -> Gen Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Int, Int) -> (Int, Gen Int)) -> [(Int, Int)] -> [(Int, Gen Int)]
forall a b. (a -> b) -> [a] -> [b]
map ((Int -> Gen Int) -> (Int, Int) -> (Int, Gen Int)
forall a b. (a -> b) -> (Int, a) -> (Int, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Int -> Gen Int
forall a. a -> Gen a
forall (f :: * -> *) a. Applicative f => a -> f a
pure) ([(Int, Int)] -> [(Int, Gen Int)])
-> (Int -> [(Int, Int)]) -> Int -> [(Int, Gen Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [(Int, Int)]
toChunkFrequencies (Int -> Gen Int) -> Int -> Gen Int
forall a b. (a -> b) -> a -> b
$ [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs
Int -> [a] -> Gen [NonEmptyList a]
forall a. Int -> [a] -> Gen [NonEmptyList a]
multiSplit1In Int
chunkNum [a]
xs
halve :: Bool -> [a] -> ([a], [a])
halve :: forall a. Bool -> [a] -> ([a], [a])
halve Bool
isOddToLeft [a]
xs0 = [a] -> [a] -> ([a], [a])
forall {a} {a}. [a] -> [a] -> ([a], [a])
go [a]
xs0 [a]
xs0 where
go :: [a] -> [a] -> ([a], [a])
go (a
_ : a
_ : [a]
xsFast) (a
x : [a]
xsSlow) = ([a] -> [a]) -> ([a], [a]) -> ([a], [a])
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
:) (([a], [a]) -> ([a], [a])) -> ([a], [a]) -> ([a], [a])
forall a b. (a -> b) -> a -> b
$ [a] -> [a] -> ([a], [a])
go [a]
xsFast [a]
xsSlow
go [a
_] (a
x : [a]
xsSlow) | Bool
isOddToLeft = ([a
x], [a]
xsSlow)
go [a]
_ [a]
xsSlow = ([], [a]
xsSlow)
insertManyPreferRight :: forall a. Bool -> Double -> a -> [a] -> Gen [a]
insertManyPreferRight :: forall a. Bool -> Double -> a -> [a] -> Gen [a]
insertManyPreferRight Bool
keepPrefix Double
lastProb a
y [a]
xs0 = Bool -> Double -> [a] -> Gen [a]
go Bool
keepPrefix Double
initWeight [a]
xs0 where
initWeight :: Double
initWeight = Double
10
scaling :: Double
scaling = Double
1.1
topWeight :: Double
topWeight = Double
scaling Double -> Double -> Double
forall a. Floating a => a -> a -> a
** Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs0) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
initWeight
noopWeight :: Int
noopWeight = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor (Double -> Int) -> Double -> Int
forall a b. (a -> b) -> a -> b
$ Double
topWeight Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
lastProb Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1)
go :: Bool -> Double -> [a] -> Gen [a]
go :: Bool -> Double -> [a] -> Gen [a]
go Bool
keep Double
weight [a]
xs = do
Bool
doCons <- [(Int, Gen Bool)] -> Gen Bool
forall a. [(Int, Gen a)] -> Gen a
frequency [(Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor Double
weight, Bool -> Gen Bool
forall a. a -> Gen a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True), (Int
noopWeight, Bool -> Gen Bool
forall a. a -> Gen a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False)]
if Bool
doCons
then ([a
y | Bool
keep] [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++) ([a] -> [a]) -> Gen [a] -> Gen [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bool -> Double -> [a] -> Gen [a]
go Bool
keep Double
weight [a]
xs
else case [a]
xs of
[] -> [a] -> Gen [a]
forall a. a -> Gen a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
a
x : [a]
xs' -> (a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
:) ([a] -> [a]) -> Gen [a] -> Gen [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bool -> Double -> [a] -> Gen [a]
go Bool
True (Double
weight Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
scaling) [a]
xs'
insertManyPreferLeft :: Bool -> Double -> a -> [a] -> Gen [a]
insertManyPreferLeft :: forall a. Bool -> Double -> a -> [a] -> Gen [a]
insertManyPreferLeft Bool
keepSuffix Double
headProb a
y =
([a] -> [a]) -> Gen [a] -> Gen [a]
forall a b. (a -> b) -> Gen a -> Gen b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [a] -> [a]
forall a. [a] -> [a]
reverse (Gen [a] -> Gen [a]) -> ([a] -> Gen [a]) -> [a] -> Gen [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> Double -> a -> [a] -> Gen [a]
forall a. Bool -> Double -> a -> [a] -> Gen [a]
insertManyPreferRight Bool
keepSuffix Double
headProb a
y ([a] -> Gen [a]) -> ([a] -> [a]) -> [a] -> Gen [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> [a]
forall a. [a] -> [a]
reverse
insertManyPreferEnds :: Double -> a -> [a] -> Gen [a]
insertManyPreferEnds :: forall a. Double -> a -> [a] -> Gen [a]
insertManyPreferEnds Double
endProb a
y [a]
xs = do
Bool
isOddToLeft <- Gen Bool
forall a. Arbitrary a => Gen a
arbitrary
let ([a]
xsL, [a]
xsR) = Bool -> [a] -> ([a], [a])
forall a. Bool -> [a] -> ([a], [a])
halve Bool
isOddToLeft [a]
xs
[a]
xsL' <- Bool -> Double -> a -> [a] -> Gen [a]
forall a. Bool -> Double -> a -> [a] -> Gen [a]
insertManyPreferLeft ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xsL Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xsR) Double
endProb a
y [a]
xsL
[a]
xsR' <- Bool -> Double -> a -> [a] -> Gen [a]
forall a. Bool -> Double -> a -> [a] -> Gen [a]
insertManyPreferRight Bool
True Double
endProb a
y [a]
xsR
[a] -> Gen [a]
forall a. a -> Gen a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([a] -> Gen [a]) -> [a] -> Gen [a]
forall a b. (a -> b) -> a -> b
$ [a]
xsL' [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
xsR'
multiSplit0 :: Double -> [a] -> Gen [[a]]
multiSplit0 :: forall a. Double -> [a] -> Gen [[a]]
multiSplit0 Double
endProb = [a] -> Gen [NonEmptyList a]
forall a. [a] -> Gen [NonEmptyList a]
multiSplit1 ([a] -> Gen [NonEmptyList a])
-> ([NonEmptyList a] -> Gen [[a]]) -> [a] -> Gen [[a]]
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> Double -> [a] -> [[a]] -> Gen [[a]]
forall a. Double -> a -> [a] -> Gen [a]
insertManyPreferEnds Double
endProb [] ([[a]] -> Gen [[a]])
-> ([NonEmptyList a] -> [[a]]) -> [NonEmptyList a] -> Gen [[a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [NonEmptyList a] -> [[a]]
forall a b. Coercible a b => a -> b
coerce