{-# OPTIONS_GHC -fno-warn-orphans #-}

module Data.GenValidity.ByteString where

import qualified Data.ByteString as SB
import qualified Data.ByteString.Lazy as LB
import qualified Data.ByteString.Lazy.Internal as LB
import qualified Data.ByteString.Short as Short
import Data.GenValidity
import Data.Validity.ByteString ()
import Data.Word (Word8)
import System.Random as Random
import Test.QuickCheck
import Test.QuickCheck.Gen
import Test.QuickCheck.Random

-- |
--
-- > genValid = SB.pack <$> genValid
-- > shrinkValid = fmap SB.pack . shrinkValid . SB.unpack
instance GenValid SB.ByteString where
  genValid :: Gen ByteString
genValid = Gen Word8 -> Gen ByteString
genStrictByteStringBy Gen Word8
forall a. GenValid a => Gen a
genValid
  shrinkValid :: ByteString -> [ByteString]
shrinkValid = ([Word8] -> ByteString) -> [[Word8]] -> [ByteString]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [Word8] -> ByteString
SB.pack ([[Word8]] -> [ByteString])
-> (ByteString -> [[Word8]]) -> ByteString -> [ByteString]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Word8] -> [[Word8]]
forall a. GenValid a => a -> [a]
shrinkValid ([Word8] -> [[Word8]])
-> (ByteString -> [Word8]) -> ByteString -> [[Word8]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> [Word8]
SB.unpack

genStrictByteStringBy :: Gen Word8 -> Gen SB.ByteString
genStrictByteStringBy :: Gen Word8 -> Gen ByteString
genStrictByteStringBy (MkGen QCGen -> Int -> Word8
word8Func) = do
  Int
len <- Gen Int
genListLength
  (QCGen -> Int -> ByteString) -> Gen ByteString
forall a. (QCGen -> Int -> a) -> Gen a
MkGen ((QCGen -> Int -> ByteString) -> Gen ByteString)
-> (QCGen -> Int -> ByteString) -> Gen ByteString
forall a b. (a -> b) -> a -> b
$ \QCGen
qcgen Int
size ->
    let go :: QCGen -> Maybe (Word8, QCGen)
        go :: QCGen -> Maybe (Word8, QCGen)
go QCGen
qcg =
          let (QCGen
qc1, QCGen
qc2) = QCGen -> (QCGen, QCGen)
forall g. RandomGen g => g -> (g, g)
Random.split QCGen
qcg
           in (Word8, QCGen) -> Maybe (Word8, QCGen)
forall a. a -> Maybe a
Just (QCGen -> Int -> Word8
word8Func QCGen
qc1 Int
size, QCGen
qc2)
     in (ByteString, Maybe QCGen) -> ByteString
forall a b. (a, b) -> a
fst ((ByteString, Maybe QCGen) -> ByteString)
-> (ByteString, Maybe QCGen) -> ByteString
forall a b. (a -> b) -> a -> b
$ Int
-> (QCGen -> Maybe (Word8, QCGen))
-> QCGen
-> (ByteString, Maybe QCGen)
forall a.
Int -> (a -> Maybe (Word8, a)) -> a -> (ByteString, Maybe a)
SB.unfoldrN Int
len QCGen -> Maybe (Word8, QCGen)
go QCGen
qcgen

instance GenValid LB.ByteString where
  genValid :: Gen ByteString
genValid = Gen ByteString -> Gen ByteString
genLazyByteStringByStrictByteString Gen ByteString
forall a. GenValid a => Gen a
genValid
  shrinkValid :: ByteString -> [ByteString]
shrinkValid = ([Word8] -> ByteString) -> [[Word8]] -> [ByteString]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [Word8] -> ByteString
LB.pack ([[Word8]] -> [ByteString])
-> (ByteString -> [[Word8]]) -> ByteString -> [ByteString]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Word8] -> [[Word8]]
forall a. GenValid a => a -> [a]
shrinkValid ([Word8] -> [[Word8]])
-> (ByteString -> [Word8]) -> ByteString -> [[Word8]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> [Word8]
LB.unpack

genLazyByteStringBy :: Gen Word8 -> Gen LB.ByteString
genLazyByteStringBy :: Gen Word8 -> Gen ByteString
genLazyByteStringBy Gen Word8
gen = Gen ByteString -> Gen ByteString
genLazyByteStringByStrictByteString (Gen Word8 -> Gen ByteString
genStrictByteStringBy Gen Word8
gen)

genLazyByteStringByStrictByteString :: Gen SB.ByteString -> Gen LB.ByteString
genLazyByteStringByStrictByteString :: Gen ByteString -> Gen ByteString
genLazyByteStringByStrictByteString Gen ByteString
gen =
  (Int -> Gen ByteString) -> Gen ByteString
forall a. (Int -> Gen a) -> Gen a
sized ((Int -> Gen ByteString) -> Gen ByteString)
-> (Int -> Gen ByteString) -> Gen ByteString
forall a b. (a -> b) -> a -> b
$ \Int
s -> do
    [Int]
ss <- Int -> Gen [Int]
arbPartition Int
s
    [Int] -> Gen ByteString
go [Int]
ss
  where
    go :: [Int] -> Gen ByteString
go [] = ByteString -> Gen ByteString
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
LB.Empty
    go (Int
s : [Int]
ss) = ByteString -> ByteString -> ByteString
LB.Chunk (ByteString -> ByteString -> ByteString)
-> Gen ByteString -> Gen (ByteString -> ByteString)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Gen ByteString -> Gen ByteString
forall a. Int -> Gen a -> Gen a
resize Int
s Gen ByteString
gen Gen (ByteString -> ByteString) -> Gen ByteString -> Gen ByteString
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Int] -> Gen ByteString
go [Int]
ss

instance GenValid Short.ShortByteString where
  genValid :: Gen ShortByteString
genValid = [Word8] -> ShortByteString
Short.pack ([Word8] -> ShortByteString) -> Gen [Word8] -> Gen ShortByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen [Word8]
forall a. GenValid a => Gen a
genValid
  shrinkValid :: ShortByteString -> [ShortByteString]
shrinkValid = ([Word8] -> ShortByteString) -> [[Word8]] -> [ShortByteString]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [Word8] -> ShortByteString
Short.pack ([[Word8]] -> [ShortByteString])
-> (ShortByteString -> [[Word8]])
-> ShortByteString
-> [ShortByteString]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Word8] -> [[Word8]]
forall a. GenValid a => a -> [a]
shrinkValid ([Word8] -> [[Word8]])
-> (ShortByteString -> [Word8]) -> ShortByteString -> [[Word8]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShortByteString -> [Word8]
Short.unpack