-- editorconfig-checker-disable-file
{- |  Generators for various types of data for benchmarking built-in functions -}

module Generators where

import PlutusCore.Data
import PlutusCore.Evaluation.Machine.ExMemory (ExMemoryUsage (..))

import Control.Monad
import Data.Bits
import Data.ByteString (ByteString)
import Data.Int (Int64)
import Data.List (foldl')
import Data.Text (Text)
import Data.Word (Word64)

import Hedgehog qualified as H
import Hedgehog.Internal.Gen qualified as G
import Hedgehog.Internal.Range qualified as R
import Hedgehog.Internal.Tree qualified as T
import System.IO.Unsafe (unsafePerformIO)
import System.Random (StdGen, randomR)
import Test.QuickCheck
import Test.QuickCheck.Instances.ByteString ()


  TODO : we 're using Hedgehog for some things , QuickCheck for others , and 
    System . Random for others .   We should rationalise this .     is 
    more future - proof since it can produce random instances of a wide variety of 
    types .   On the other hand , you have to be careful with Hedgehog Ranges since 
    they do n't necessarily do what you might expect .   It might be a bit tricky to 
    use Hedgehog everywhere because we 'd need a lot of monadic code to take care 
    of generator states .
   System.Random for others.  We should rationalise this.  Pehaps Hedgehog is
   more future-proof since it can produce random instances of a wide variety of
   types.  On the other hand, you have to be careful with Hedgehog Ranges since
   they don't necessarily do what you might expect.  It might be a bit tricky to
   use Hedgehog everywhere because we'd need a lot of monadic code to take care
   of generator states. -}


---------------- Creating things of a given size ----------------

  Generate a random n - word ( ie , 64n - bit ) integer
  In principle a random 5 - word integer ( for example ) might only occupy 4 or 
    fewer words , but we 're generating uniformly distributed values so the 
    probability of that happening should be at most 1 in 2 ^ 64 .
   fewer words, but we're generating uniformly distributed values so the
   probability of that happening should be at most 1 in 2^64. -}
randNwords :: Int -> StdGen -> (Integer, StdGen)
randNwords n gen = randomR (lb,ub) gen
    where lb = 2^(64*(n-1))
          ub = 2^(64*n) - 1

  Generate a random  ( just here for consistency )
randBool :: StdGen -> (Bool, StdGen)
randBool gen = randomR (False, True) gen


  Hedgehog generators
seedA :: H.Seed
seedA = H.Seed 42 43

seedB :: H.Seed
seedB = H.Seed 44 45

genSample :: H.Seed -> G.Gen a -> a
genSample seed gen = Prelude.maybe (Prelude.error "Couldn't create a sample") T.treeValue $ G.evalGen (H.Size 1) seed gen

  Given a list [ n_1 , n_2 , ... ] create a list [ m_1 ,  , ... ] where m_i is an n_i - word random integer
makeSizedIntegers :: StdGen -> [Int] -> ([Integer], StdGen)
makeSizedIntegers g [] = ([], g)
makeSizedIntegers g (n:ns) =
    let (m,g1) = randNwords n g
        (ms,g2) = makeSizedIntegers g1 ns
    in (m:ms,g2)

-- Create a bytestring whose memory usage is n.  Since we measure memory usage
  in 64 - bit words we have to create a bytestring containing 8*n bytes .
makeSizedByteString :: H.Seed -> Int -> ByteString
makeSizedByteString seed n = genSample seed (G.bytes (R.singleton (8*n)))

-- FIXME: this is terrible
makeSizedByteStrings :: H.Seed -> [Int] -> [ByteString]
makeSizedByteStrings seed l = map (makeSizedByteString seed) l


  TODO : do n't use Hedgehog 's ' sample ' below : it silently resizes the generator
  to size 30 , so listOfSizedByteStrings and listOfByteStrings are biased
-- towards low byte values.

-- Create a list containing m bytestrings of length n (also terrible)
listOfSizedByteStrings :: Int -> Int -> [ByteString]
listOfSizedByteStrings m n = unsafePerformIO . G.sample $ G.list (R.singleton m) (G.bytes (R.singleton n))

-- Create a list containing m bytestrings of random lengths
listOfByteStrings :: Int -> [ByteString]
listOfByteStrings m = unsafePerformIO . G.sample $ G.list (R.singleton m) (G.bytes (R.linear 0 10000))


---------------- Strings (Hedgehog) ----------------

  This makes a Text string containing n unicode characters .   We use the unicode 
  generator since that mostly produces 4 bytes per character , which is the worst 
  case . If we were to use the ascii generator that would give us two bytes per 
  character .   See Note [ Choosing the inputs for costing benchmarks ] in Strings.hs .
 generator since that mostly produces 4 bytes per character, which is the worst
 case. If we were to use the ascii generator that would give us two bytes per
 character.  See Note [Choosing the inputs for costing benchmarks] in Strings.hs. -}
makeSizedTextString :: H.Seed -> Int -> Text
makeSizedTextString seed n = genSample seed (G.text (R.singleton (2*n)) G.unicode)

makeSizedTextStrings :: H.Seed -> [Integer] -> [Text]
makeSizedTextStrings seed sizes = fmap (makeSizedTextString seed . fromInteger) sizes


  | Generate a valid UTF-8 bytestring with memory usage approximately n for 
    benchmarking decodeUtf8 .   We use the ' unicode ' generator beacuse that gives 
    the worst - case behaviour : see Note [ Choosing the inputs for costing 
    benchmarks ] in Strings.hs ) .
   benchmarking decodeUtf8.  We use the 'unicode' generator beacuse that gives
   the worst-case behaviour: see Note [Choosing the inputs for costing
   benchmarks] in Strings.hs).-}
makeSizedUtf8ByteString :: H.Seed -> Int -> ByteString
makeSizedUtf8ByteString seed n = genSample seed (G.utf8 (R.singleton (2*n)) G.unicode)

makeSizedUtf8ByteStrings :: H.Seed -> [Integer] -> [ByteString]
makeSizedUtf8ByteStrings seed sizes = (makeSizedUtf8ByteString seed . fromInteger) <$> sizes


---------------- Data (QuickCheck) ----------------

  Create a large Integer .    does n't provide this by default ( and 
  arbitrary : :   behaves oddly if you resize it to large values like 
  2 ^ 62 ( " negative size " error ) , 2 ^ 63 ( never terminates ) , or 2 ^ 64 or larger 
  ( always returns 0 ) .   We need really big integers , so we generate a ( signed 
  64 - bit ) Int and list of unsigned Word64 values and make them into an Integer 
  using bit operations . This is designed so that arbitraryBigInteger n returns 
  integers with memoryUsage equal to n. 
 
 arbitrary :: Gen Integer behaves oddly if you resize it to large values like
 2^62 ("negative size" error ), 2^63 (never terminates), or 2^64 or larger
 (always returns 0).  We need really big integers, so we generate a (signed
 64-bit) Int and list of unsigned Word64 values and make them into an Integer
 using bit operations. This is designed so that arbitraryBigInteger n returns
 integers with memoryUsage equal to n.
 -}
genBigInteger :: Int -> Gen Integer
genBigInteger n = do
  body :: [Word64] <- vectorOf (n-1) arbitrary
  first :: Int64 <- arbitrary
  pure $ foldl' go (fromIntegral first) body
      where go :: Integer -> Word64 -> Integer
            go acc w = acc `shiftL` 64 + fromIntegral w

{- | Generate an arbitrary integer of size n (words) -}
genI :: Int -> Gen Data
genI n = do
  I <$> genBigInteger n

{- | Generate an arbitrary bytestring of size n (words) -}
genB :: Int -> Gen Data
genB n = do
  let size = 4*n
  B <$> resize size arbitrary

  | Generate an arbitrary Data object of depth at most n containing btyrestrings 
    with memory usage bmem and integers with memory usage imem .   The ` memUsage ` 
    instance for Data can only return a single number , and we could have for 
    example a deep tree containing small bytsetrings / integers which has exactly 
    the same memory usage as an object consisting of a single B constructor with 
    a very large bytestring .   The costing function for certain operations 
    ( specifically ' equalsData ' ) would assign the same cost to processing both 
    objects even though they have very different structures .   This generator 
    allows us to explore a large space of objects to see how the costs of 
    operations vary with the size .   For size 5 we generally get trees of depth 3 
    and up to 4000 nodes in total ( although skewed towards smaller numbers of 
    nodes ) ; for size 10 we get trees of depth 4 and up to about 20000 nodes ; for 
    size 20 , depth 5 and up to about one million nodes . 

   with memory usage bmem and integers with memory usage imem.  The `memUsage`
   instance for Data can only return a single number, and we could have for
   example a deep tree containing small bytsetrings/integers which has exactly
   the same memory usage as an object consisting of a single B constructor with
   a very large bytestring.  The costing function for certain operations
   (specifically 'equalsData') would assign the same cost to processing both
   objects even though they have very different structures.  This generator
   allows us to explore a large space of objects to see how the costs of
   operations vary with the size.  For size 5 we generally get trees of depth 3
   and up to 4000 nodes in total (although skewed towards smaller numbers of
   nodes); for size 10 we get trees of depth 4 and up to about 20000 nodes; for
   size 20, depth 5 and up to about one million nodes.
-}
genBoundedData :: Int -> Int -> Int -> Gen Data
genBoundedData imem bmem size = genD size
    where genD n =
              if n <= 1
              then
                  Test.QuickCheck.oneof [ genI imem, genB bmem]
              else
                  Test.QuickCheck.oneof [ Constr <$> choose (1,100) <*> resize 5 (listOf (genD (n `div` 2)))
  Constr is unilkely to have very many arguments .
                                        , List <$> listOf' (genD (n `div` 2))
                                        , Map <$> (listOf' $ (,) <$> genD (n `div` 2) <*> genD (n `div` 2))
                                        ]
                  where listOf' g = frequency [ (800, resize   10 (listOf g))
                                              , (200, resize  100 (listOf g))
                                              , (2,   resize 1000 (listOf g))
                                              ]
                  -- We probably will get large lists occasionally in practice,
                  -- but if we generate them too often we get enormous objects.
                  -- listOf' attempts to give us very occasional large lists.

  | Given a list [ ( n1 , ( i1 , b1 , s1 ) ) , ( n2 , ( i2 , b2 , s2 ) ) , ... ] , return a list 
  containing n1 samples generated by     followed by n2 
  samples from genBoundedData i2 b2 s2 , and so on .  
 containing n1 samples generated by genBoundedData i1 b1 s1 followed by n2
 samples from genBoundedData i2 b2 s2, and so on.  -}
genDataSample :: [(Int, (Int, Int, Int))] -> [Data]
genDataSample l =
    unsafePerformIO $ concat <$> mapM gen1 l
        where gen1 (count, (imem, bmem, size)) = replicateM count . generate $ genBoundedData imem bmem size


-- A list of parameters for genDataSample
dataParams :: [(Int, (Int, Int, Int))]
dataParams = [ (10, ( 10,  10,  1))
             , (10, (  5,  10,  2))
             , (10, (100, 100,  2))
             , (10, (  1,   1,  3))
             , (10, (  2,   4,  4))
             , (30, ( 10,   7,  5))
             , (20, (  1,   1, 10))
             , (20, ( 30,  30, 10))
             , (10, (  4,   6, 15))
             , (10, (  5,   5, 20))
             , (10, (  1,  10, 25))
  150 entries in all

-- We want a list of random data, but for benchmarking purposes we also want to
-- be able to filter out sublists for various constructors.  To do this we
-- generate a large number of samples and always take some sublist for the
-- inputs to the benchmarks.
dataSample :: [Data]
dataSample = genDataSample (take 500 $ cycle dataParams)


  A list of data for EqualsData , which is difficult to cost . We want some very
-- small objects in here to give us an idea of what the intercept of regression
-- line should be (but note tht the minimum ExMemoryUsage of a Data object is
  currently 4 , since each node costs that much ) .   We also exclude really large
-- objects.
dataSampleForEq :: [Data]
dataSampleForEq = take 400 . filter (\d -> memoryUsage d < 1000000) . genDataSample . take 1000 $ cycle ((20, (1,1,1)):dataParams)

 editorconfig-checker-disable-file
 |  Generators for various types of data for benchmarking built-in functions 
-------------- Creating things of a given size ----------------
 Create a bytestring whose memory usage is n.  Since we measure memory usage
 FIXME: this is terrible
 towards low byte values.
 Create a list containing m bytestrings of length n (also terrible)
 Create a list containing m bytestrings of random lengths
-------------- Strings (Hedgehog) ----------------
-------------- Data (QuickCheck) ----------------
 | Generate an arbitrary integer of size n (words) 
 | Generate an arbitrary bytestring of size n (words) 
 We probably will get large lists occasionally in practice,
 but if we generate them too often we get enormous objects.
 listOf' attempts to give us very occasional large lists.
 A list of parameters for genDataSample
 We want a list of random data, but for benchmarking purposes we also want to
 be able to filter out sublists for various constructors.  To do this we
 generate a large number of samples and always take some sublist for the
 inputs to the benchmarks.
 small objects in here to give us an idea of what the intercept of regression
 line should be (but note tht the minimum ExMemoryUsage of a Data object is
 objects.

module Generators where

import PlutusCore.Data
import PlutusCore.Evaluation.Machine.ExMemory (ExMemoryUsage (..))

import Control.Monad
import Data.Bits
import Data.ByteString (ByteString)
import Data.Int (Int64)
import Data.List (foldl')
import Data.Text (Text)
import Data.Word (Word64)

import Hedgehog qualified as H
import Hedgehog.Internal.Gen qualified as G
import Hedgehog.Internal.Range qualified as R
import Hedgehog.Internal.Tree qualified as T
import System.IO.Unsafe (unsafePerformIO)
import System.Random (StdGen, randomR)
import Test.QuickCheck
import Test.QuickCheck.Instances.ByteString ()


  TODO : we 're using Hedgehog for some things , QuickCheck for others , and 
    System . Random for others .   We should rationalise this .     is 
    more future - proof since it can produce random instances of a wide variety of 
    types .   On the other hand , you have to be careful with Hedgehog Ranges since 
    they do n't necessarily do what you might expect .   It might be a bit tricky to 
    use Hedgehog everywhere because we 'd need a lot of monadic code to take care 
    of generator states .
   System.Random for others.  We should rationalise this.  Pehaps Hedgehog is
   more future-proof since it can produce random instances of a wide variety of
   types.  On the other hand, you have to be careful with Hedgehog Ranges since
   they don't necessarily do what you might expect.  It might be a bit tricky to
   use Hedgehog everywhere because we'd need a lot of monadic code to take care
   of generator states. -}



  Generate a random n - word ( ie , 64n - bit ) integer
  In principle a random 5 - word integer ( for example ) might only occupy 4 or 
    fewer words , but we 're generating uniformly distributed values so the 
    probability of that happening should be at most 1 in 2 ^ 64 .
   fewer words, but we're generating uniformly distributed values so the
   probability of that happening should be at most 1 in 2^64. -}
randNwords :: Int -> StdGen -> (Integer, StdGen)
randNwords n gen = randomR (lb,ub) gen
    where lb = 2^(64*(n-1))
          ub = 2^(64*n) - 1

  Generate a random  ( just here for consistency )
randBool :: StdGen -> (Bool, StdGen)
randBool gen = randomR (False, True) gen


  Hedgehog generators
seedA :: H.Seed
seedA = H.Seed 42 43

seedB :: H.Seed
seedB = H.Seed 44 45

genSample :: H.Seed -> G.Gen a -> a
genSample seed gen = Prelude.maybe (Prelude.error "Couldn't create a sample") T.treeValue $ G.evalGen (H.Size 1) seed gen

  Given a list [ n_1 , n_2 , ... ] create a list [ m_1 ,  , ... ] where m_i is an n_i - word random integer
makeSizedIntegers :: StdGen -> [Int] -> ([Integer], StdGen)
makeSizedIntegers g [] = ([], g)
makeSizedIntegers g (n:ns) =
    let (m,g1) = randNwords n g
        (ms,g2) = makeSizedIntegers g1 ns
    in (m:ms,g2)

  in 64 - bit words we have to create a bytestring containing 8*n bytes .
makeSizedByteString :: H.Seed -> Int -> ByteString
makeSizedByteString seed n = genSample seed (G.bytes (R.singleton (8*n)))

makeSizedByteStrings :: H.Seed -> [Int] -> [ByteString]
makeSizedByteStrings seed l = map (makeSizedByteString seed) l


  TODO : do n't use Hedgehog 's ' sample ' below : it silently resizes the generator
  to size 30 , so listOfSizedByteStrings and listOfByteStrings are biased

listOfSizedByteStrings :: Int -> Int -> [ByteString]
listOfSizedByteStrings m n = unsafePerformIO . G.sample $ G.list (R.singleton m) (G.bytes (R.singleton n))

listOfByteStrings :: Int -> [ByteString]
listOfByteStrings m = unsafePerformIO . G.sample $ G.list (R.singleton m) (G.bytes (R.linear 0 10000))



  This makes a Text string containing n unicode characters .   We use the unicode 
  generator since that mostly produces 4 bytes per character , which is the worst 
  case . If we were to use the ascii generator that would give us two bytes per 
  character .   See Note [ Choosing the inputs for costing benchmarks ] in Strings.hs .
 generator since that mostly produces 4 bytes per character, which is the worst
 case. If we were to use the ascii generator that would give us two bytes per
 character.  See Note [Choosing the inputs for costing benchmarks] in Strings.hs. -}
makeSizedTextString :: H.Seed -> Int -> Text
makeSizedTextString seed n = genSample seed (G.text (R.singleton (2*n)) G.unicode)

makeSizedTextStrings :: H.Seed -> [Integer] -> [Text]
makeSizedTextStrings seed sizes = fmap (makeSizedTextString seed . fromInteger) sizes


  | Generate a valid UTF-8 bytestring with memory usage approximately n for 
    benchmarking decodeUtf8 .   We use the ' unicode ' generator beacuse that gives 
    the worst - case behaviour : see Note [ Choosing the inputs for costing 
    benchmarks ] in Strings.hs ) .
   benchmarking decodeUtf8.  We use the 'unicode' generator beacuse that gives
   the worst-case behaviour: see Note [Choosing the inputs for costing
   benchmarks] in Strings.hs).-}
makeSizedUtf8ByteString :: H.Seed -> Int -> ByteString
makeSizedUtf8ByteString seed n = genSample seed (G.utf8 (R.singleton (2*n)) G.unicode)

makeSizedUtf8ByteStrings :: H.Seed -> [Integer] -> [ByteString]
makeSizedUtf8ByteStrings seed sizes = (makeSizedUtf8ByteString seed . fromInteger) <$> sizes



  Create a large Integer .    does n't provide this by default ( and 
  arbitrary : :   behaves oddly if you resize it to large values like 
  2 ^ 62 ( " negative size " error ) , 2 ^ 63 ( never terminates ) , or 2 ^ 64 or larger 
  ( always returns 0 ) .   We need really big integers , so we generate a ( signed 
  64 - bit ) Int and list of unsigned Word64 values and make them into an Integer 
  using bit operations . This is designed so that arbitraryBigInteger n returns 
  integers with memoryUsage equal to n. 
 
 arbitrary :: Gen Integer behaves oddly if you resize it to large values like
 2^62 ("negative size" error ), 2^63 (never terminates), or 2^64 or larger
 (always returns 0).  We need really big integers, so we generate a (signed
 64-bit) Int and list of unsigned Word64 values and make them into an Integer
 using bit operations. This is designed so that arbitraryBigInteger n returns
 integers with memoryUsage equal to n.
 -}
genBigInteger :: Int -> Gen Integer
genBigInteger n = do
  body :: [Word64] <- vectorOf (n-1) arbitrary
  first :: Int64 <- arbitrary
  pure $ foldl' go (fromIntegral first) body
      where go :: Integer -> Word64 -> Integer
            go acc w = acc `shiftL` 64 + fromIntegral w

genI :: Int -> Gen Data
genI n = do
  I <$> genBigInteger n

genB :: Int -> Gen Data
genB n = do
  let size = 4*n
  B <$> resize size arbitrary

  | Generate an arbitrary Data object of depth at most n containing btyrestrings 
    with memory usage bmem and integers with memory usage imem .   The ` memUsage ` 
    instance for Data can only return a single number , and we could have for 
    example a deep tree containing small bytsetrings / integers which has exactly 
    the same memory usage as an object consisting of a single B constructor with 
    a very large bytestring .   The costing function for certain operations 
    ( specifically ' equalsData ' ) would assign the same cost to processing both 
    objects even though they have very different structures .   This generator 
    allows us to explore a large space of objects to see how the costs of 
    operations vary with the size .   For size 5 we generally get trees of depth 3 
    and up to 4000 nodes in total ( although skewed towards smaller numbers of 
    nodes ) ; for size 10 we get trees of depth 4 and up to about 20000 nodes ; for 
    size 20 , depth 5 and up to about one million nodes . 

   with memory usage bmem and integers with memory usage imem.  The `memUsage`
   instance for Data can only return a single number, and we could have for
   example a deep tree containing small bytsetrings/integers which has exactly
   the same memory usage as an object consisting of a single B constructor with
   a very large bytestring.  The costing function for certain operations
   (specifically 'equalsData') would assign the same cost to processing both
   objects even though they have very different structures.  This generator
   allows us to explore a large space of objects to see how the costs of
   operations vary with the size.  For size 5 we generally get trees of depth 3
   and up to 4000 nodes in total (although skewed towards smaller numbers of
   nodes); for size 10 we get trees of depth 4 and up to about 20000 nodes; for
   size 20, depth 5 and up to about one million nodes.
-}
genBoundedData :: Int -> Int -> Int -> Gen Data
genBoundedData imem bmem size = genD size
    where genD n =
              if n <= 1
              then
                  Test.QuickCheck.oneof [ genI imem, genB bmem]
              else
                  Test.QuickCheck.oneof [ Constr <$> choose (1,100) <*> resize 5 (listOf (genD (n `div` 2)))
  Constr is unilkely to have very many arguments .
                                        , List <$> listOf' (genD (n `div` 2))
                                        , Map <$> (listOf' $ (,) <$> genD (n `div` 2) <*> genD (n `div` 2))
                                        ]
                  where listOf' g = frequency [ (800, resize   10 (listOf g))
                                              , (200, resize  100 (listOf g))
                                              , (2,   resize 1000 (listOf g))
                                              ]

  | Given a list [ ( n1 , ( i1 , b1 , s1 ) ) , ( n2 , ( i2 , b2 , s2 ) ) , ... ] , return a list 
  containing n1 samples generated by     followed by n2 
  samples from genBoundedData i2 b2 s2 , and so on .  
 containing n1 samples generated by genBoundedData i1 b1 s1 followed by n2
 samples from genBoundedData i2 b2 s2, and so on.  -}
genDataSample :: [(Int, (Int, Int, Int))] -> [Data]
genDataSample l =
    unsafePerformIO $ concat <$> mapM gen1 l
        where gen1 (count, (imem, bmem, size)) = replicateM count . generate $ genBoundedData imem bmem size


dataParams :: [(Int, (Int, Int, Int))]
dataParams = [ (10, ( 10,  10,  1))
             , (10, (  5,  10,  2))
             , (10, (100, 100,  2))
             , (10, (  1,   1,  3))
             , (10, (  2,   4,  4))
             , (30, ( 10,   7,  5))
             , (20, (  1,   1, 10))
             , (20, ( 30,  30, 10))
             , (10, (  4,   6, 15))
             , (10, (  5,   5, 20))
             , (10, (  1,  10, 25))
  150 entries in all

dataSample :: [Data]
dataSample = genDataSample (take 500 $ cycle dataParams)


  A list of data for EqualsData , which is difficult to cost . We want some very
  currently 4 , since each node costs that much ) .   We also exclude really large
dataSampleForEq :: [Data]
dataSampleForEq = take 400 . filter (\d -> memoryUsage d < 1000000) . genDataSample . take 1000 $ cycle ((20, (1,1,1)):dataParams)
