{-# LANGUAGE TypeApplications #-}

module Cardano.Wallet.Primitive.Types.TokenQuantity.Gen
    ( chooseTokenQuantity
    , genTokenQuantity
    , genTokenQuantityPositive
    , genTokenQuantityFullRange
    , shrinkTokenQuantity
    , shrinkTokenQuantityPositive
    , shrinkTokenQuantityFullRange
    , genTokenQuantityPartition
    ) where

import Prelude

import Cardano.Wallet.Primitive.Types.TokenQuantity
    ( TokenQuantity (..) )
import Control.Monad
    ( replicateM )
import Data.Coerce
    ( coerce )
import Data.List.NonEmpty
    ( NonEmpty )
import Data.Word
    ( Word64 )
import Test.QuickCheck
    ( Gen, choose, frequency, shrink, sized )
import Test.QuickCheck.Extra
    ( chooseNatural )

import qualified Cardano.Wallet.Primitive.Types.TokenQuantity as TokenQuantity
import qualified Data.List.NonEmpty as NE

--------------------------------------------------------------------------------
-- Choosing token quantities from a range.
--------------------------------------------------------------------------------

chooseTokenQuantity :: (TokenQuantity, TokenQuantity) -> Gen TokenQuantity
chooseTokenQuantity :: (TokenQuantity, TokenQuantity) -> Gen TokenQuantity
chooseTokenQuantity = ((Natural, Natural) -> Gen Natural)
-> (TokenQuantity, TokenQuantity) -> Gen TokenQuantity
coerce (Natural, Natural) -> Gen Natural
chooseNatural

--------------------------------------------------------------------------------
-- Token quantities chosen according to the size parameter.
--------------------------------------------------------------------------------

genTokenQuantity :: Gen TokenQuantity
genTokenQuantity :: Gen TokenQuantity
genTokenQuantity = (Int -> Gen TokenQuantity) -> Gen TokenQuantity
forall a. (Int -> Gen a) -> Gen a
sized ((Int -> Gen TokenQuantity) -> Gen TokenQuantity)
-> (Int -> Gen TokenQuantity) -> Gen TokenQuantity
forall a b. (a -> b) -> a -> b
$ \Int
n -> Int -> TokenQuantity
quantityFromInt (Int -> TokenQuantity) -> Gen Int -> Gen TokenQuantity
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int, Int) -> Gen Int
forall a. Random a => (a, a) -> Gen a
choose (Int
0, Int
n)

shrinkTokenQuantity :: TokenQuantity -> [TokenQuantity]
shrinkTokenQuantity :: TokenQuantity -> [TokenQuantity]
shrinkTokenQuantity = (Integer -> TokenQuantity) -> [Integer] -> [TokenQuantity]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Integer -> TokenQuantity
quantityFromInteger ([Integer] -> [TokenQuantity])
-> (TokenQuantity -> [Integer]) -> TokenQuantity -> [TokenQuantity]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> [Integer]
forall a. Arbitrary a => a -> [a]
shrink (Integer -> [Integer])
-> (TokenQuantity -> Integer) -> TokenQuantity -> [Integer]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TokenQuantity -> Integer
quantityToInteger

--------------------------------------------------------------------------------
-- Token quantities chosen according to the size parameter, but strictly
-- positive.
--------------------------------------------------------------------------------

genTokenQuantityPositive :: Gen TokenQuantity
genTokenQuantityPositive :: Gen TokenQuantity
genTokenQuantityPositive = (Int -> Gen TokenQuantity) -> Gen TokenQuantity
forall a. (Int -> Gen a) -> Gen a
sized ((Int -> Gen TokenQuantity) -> Gen TokenQuantity)
-> (Int -> Gen TokenQuantity) -> Gen TokenQuantity
forall a b. (a -> b) -> a -> b
$ \Int
n -> Int -> TokenQuantity
quantityFromInt (Int -> TokenQuantity) -> Gen Int -> Gen TokenQuantity
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int, Int) -> Gen Int
forall a. Random a => (a, a) -> Gen a
choose (Int
1, Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 Int
n)

shrinkTokenQuantityPositive :: TokenQuantity -> [TokenQuantity]
shrinkTokenQuantityPositive :: TokenQuantity -> [TokenQuantity]
shrinkTokenQuantityPositive
    = (Integer -> TokenQuantity) -> [Integer] -> [TokenQuantity]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Integer -> TokenQuantity
quantityFromInteger
    ([Integer] -> [TokenQuantity])
-> (TokenQuantity -> [Integer]) -> TokenQuantity -> [TokenQuantity]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer -> Bool) -> [Integer] -> [Integer]
forall a. (a -> Bool) -> [a] -> [a]
filter (Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
0)
    ([Integer] -> [Integer])
-> (TokenQuantity -> [Integer]) -> TokenQuantity -> [Integer]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> [Integer]
forall a. Arbitrary a => a -> [a]
shrink
    (Integer -> [Integer])
-> (TokenQuantity -> Integer) -> TokenQuantity -> [Integer]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TokenQuantity -> Integer
quantityToInteger

--------------------------------------------------------------------------------
-- Token quantities chosen from the full range available.
--------------------------------------------------------------------------------

-- | Generates token quantities across the full range of what may be encoded
--   within a single on-chain token bundle.
--
-- This generator has a slight bias towards the limits of the range, but
-- otherwise generates values uniformly across the whole range.
--
-- This can be useful when testing roundtrip conversions between different
-- types.
--
genTokenQuantityFullRange :: Gen TokenQuantity
genTokenQuantityFullRange :: Gen TokenQuantity
genTokenQuantityFullRange = [(Int, Gen TokenQuantity)] -> Gen TokenQuantity
forall a. [(Int, Gen a)] -> Gen a
frequency
    [ ( Int
1, TokenQuantity -> Gen TokenQuantity
forall (f :: * -> *) a. Applicative f => a -> f a
pure TokenQuantity
minTokenQuantity )
    , ( Int
1, TokenQuantity -> Gen TokenQuantity
forall (f :: * -> *) a. Applicative f => a -> f a
pure TokenQuantity
maxTokenQuantity )
    , ( Int
8
      , Integer -> TokenQuantity
quantityFromInteger (Integer -> TokenQuantity) -> Gen Integer -> Gen TokenQuantity
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
        (Integer, Integer) -> Gen Integer
forall a. Random a => (a, a) -> Gen a
choose (Integer
1, TokenQuantity -> Integer
quantityToInteger TokenQuantity
maxTokenQuantity Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1)
      )
    ]
  where
    minTokenQuantity :: TokenQuantity
    minTokenQuantity :: TokenQuantity
minTokenQuantity = Natural -> TokenQuantity
TokenQuantity Natural
0
    maxTokenQuantity :: TokenQuantity
    maxTokenQuantity :: TokenQuantity
maxTokenQuantity = Natural -> TokenQuantity
TokenQuantity (Natural -> TokenQuantity) -> Natural -> TokenQuantity
forall a b. (a -> b) -> a -> b
$ Word64 -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64 -> Natural) -> Word64 -> Natural
forall a b. (a -> b) -> a -> b
$ Bounded Word64 => Word64
forall a. Bounded a => a
maxBound @Word64

shrinkTokenQuantityFullRange :: TokenQuantity -> [TokenQuantity]
shrinkTokenQuantityFullRange :: TokenQuantity -> [TokenQuantity]
shrinkTokenQuantityFullRange =
    -- Given that we may have a large value, we limit the number of results
    -- returned in order to avoid processing long lists of shrunken values.
    Int -> [TokenQuantity] -> [TokenQuantity]
forall a. Int -> [a] -> [a]
take Int
8 ([TokenQuantity] -> [TokenQuantity])
-> (TokenQuantity -> [TokenQuantity])
-> TokenQuantity
-> [TokenQuantity]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TokenQuantity -> [TokenQuantity]
shrinkTokenQuantity

--------------------------------------------------------------------------------
-- Partitioning token quantities
--------------------------------------------------------------------------------

-- | Partitions a token quantity randomly into a given number of parts.
--
-- Satisfies the following properties:
--
-- prop> forAll (genTokenQuantityPartition q i) $ (==       q) . fold
-- prop> forAll (genTokenQuantityPartition q i) $ (== max 1 i) . length
--
genTokenQuantityPartition
    :: TokenQuantity -> Int -> Gen (NonEmpty TokenQuantity)
genTokenQuantityPartition :: TokenQuantity -> Int -> Gen (NonEmpty TokenQuantity)
genTokenQuantityPartition TokenQuantity
c Int
i =
    TokenQuantity -> NonEmpty TokenQuantity -> NonEmpty TokenQuantity
TokenQuantity.partitionDefault TokenQuantity
c (NonEmpty TokenQuantity -> NonEmpty TokenQuantity)
-> Gen (NonEmpty TokenQuantity) -> Gen (NonEmpty TokenQuantity)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen (NonEmpty TokenQuantity)
genWeights
  where
    genWeights :: Gen (NonEmpty TokenQuantity)
    genWeights :: Gen (NonEmpty TokenQuantity)
genWeights = [TokenQuantity] -> NonEmpty TokenQuantity
forall a. [a] -> NonEmpty a
NE.fromList ([TokenQuantity] -> NonEmpty TokenQuantity)
-> Gen [TokenQuantity] -> Gen (NonEmpty TokenQuantity)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Gen TokenQuantity -> Gen [TokenQuantity]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 Int
i)
        ((TokenQuantity, TokenQuantity) -> Gen TokenQuantity
chooseTokenQuantity (Natural -> TokenQuantity
TokenQuantity Natural
1, TokenQuantity -> TokenQuantity -> TokenQuantity
forall a. Ord a => a -> a -> a
max (Natural -> TokenQuantity
TokenQuantity Natural
1) TokenQuantity
c))

--------------------------------------------------------------------------------
-- Internal functions
--------------------------------------------------------------------------------

quantityToInteger :: TokenQuantity -> Integer
quantityToInteger :: TokenQuantity -> Integer
quantityToInteger (TokenQuantity Natural
q) = Natural -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Natural
q

quantityFromInt :: Int -> TokenQuantity
quantityFromInt :: Int -> TokenQuantity
quantityFromInt Int
i
    | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = [Char] -> TokenQuantity
forall a. HasCallStack => [Char] -> a
error ([Char] -> TokenQuantity) -> [Char] -> TokenQuantity
forall a b. (a -> b) -> a -> b
$ [Char]
"Unable to convert integer to token quantity: " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Int -> [Char]
forall a. Show a => a -> [Char]
show Int
i
    | Bool
otherwise = Natural -> TokenQuantity
TokenQuantity (Natural -> TokenQuantity) -> Natural -> TokenQuantity
forall a b. (a -> b) -> a -> b
$ Int -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i

quantityFromInteger :: Integer -> TokenQuantity
quantityFromInteger :: Integer -> TokenQuantity
quantityFromInteger Integer
i
    | Integer
i Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0 = [Char] -> TokenQuantity
forall a. HasCallStack => [Char] -> a
error ([Char] -> TokenQuantity) -> [Char] -> TokenQuantity
forall a b. (a -> b) -> a -> b
$ [Char]
"Unable to convert integer to token quantity: " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Integer -> [Char]
forall a. Show a => a -> [Char]
show Integer
i
    | Bool
otherwise = Natural -> TokenQuantity
TokenQuantity (Natural -> TokenQuantity) -> Natural -> TokenQuantity
forall a b. (a -> b) -> a -> b
$ Integer -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
i