{-# LANGUAGE ScopedTypeVariables #-}

module Hedgehog.Classes.Storable (storableLaws) where

import Hedgehog
import Hedgehog.Classes.Common
import Hedgehog.Internal.Gen (sample)

import qualified Data.List as List
import qualified Hedgehog.Gen as Gen
import qualified Hedgehog.Range as Range

import Foreign.Marshal.Alloc
import Foreign.Marshal.Array
import GHC.Ptr (Ptr(..), nullPtr, plusPtr, minusPtr, alignPtr)
import Foreign.Storable (Storable(..))
import System.IO.Unsafe (unsafePerformIO)

-- | Tests the following 'Storable' laws:
--
-- [__Set-Get__]: @'pokeElemOff' ptr ix a '>>' 'peekElemOff' ptr ix@ ≡ @'pure' a@
-- [__Get-Set__]: @'peekElemOff' ptr ix '>>=' 'pokeElemOff' ptr ix@ ≡ @'pure' ()@ (Putting back what you got out has no effect)
-- [__List Conversion Roundtrips__]: Mallocing a list and then reconstructing it gives you the same list
-- [__PeekElemOff/Peek__]: @'peekElemOff' a i@ ≡ @'peek' ('plusPtr' a (i '*' 'sizeOf' 'undefined'))@
-- [__PokeElemOff/Poke__]: @'pokeElemOff' a i x@ ≡ @'poke' ('plusPtr' a (i '*' 'sizeOf' 'undefined')) x@
-- [__PeekByteOff/Peek__]: @'peekByteOff' a i@ ≡ @'peek' ('plusPtr' a i)@
-- [__PokeByteOff/Peek__]: @'pokeByteOff' a i x@ ≡ @'poke' ('plusPtr' a i) x@
storableLaws :: (Eq a, Show a, Storable a) => Gen a -> Laws
storableLaws :: forall a. (Eq a, Show a, Storable a) => Gen a -> Laws
storableLaws Gen a
gen = String -> [(String, Property)] -> Laws
Laws String
"Storable"
  [ (String
"Set-Get (you get back what you put in)", Gen a -> Property
forall a. (Eq a, Show a, Storable a) => Gen a -> Property
storableSetGet Gen a
gen)
  , (String
"Get-Set (putting back what you got out has no effect)", Gen a -> Property
forall a. (Eq a, Show a, Storable a) => Gen a -> Property
storableGetSet Gen a
gen)
  , (String
"List Conversion Roundtrips", Gen a -> Property
forall a. (Eq a, Show a, Storable a) => Gen a -> Property
storableList Gen a
gen)
  , (String
"peekElemOff a i ≡ peek (plusPtr a (i * sizeOf undefined))", Gen a -> Property
forall a. (Eq a, Show a, Storable a) => Gen a -> Property
storablePeekElem Gen a
gen)
  , (String
"pokeElemOff a i x ≡ poke (plusPtr a (i * sizeOf undefined)) x ≡ id ", Gen a -> Property
forall a. (Eq a, Show a, Storable a) => Gen a -> Property
storablePokeElem Gen a
gen)
  , (String
"peekByteOff a i ≡ peek (plusPtr a i)", Gen a -> Property
forall a. (Eq a, Show a, Storable a) => Gen a -> Property
storablePeekByte Gen a
gen)
  , (String
"pokeByteOff a i x ≡ poke (plusPtr a i) x ≡ id ", Gen a -> Property
forall a. (Eq a, Show a, Storable a) => Gen a -> Property
storablePokeByte Gen a
gen)
  ]

genArray :: forall a. (Storable a) => Gen a -> Int -> IO (Ptr a)
genArray :: forall a. Storable a => Gen a -> Int -> IO (Ptr a)
genArray Gen a
gen Int
len = do
  let go :: Int -> [a] -> f [a]
go Int
ix [a]
xs = if Int
ix Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len
        then [a] -> f [a]
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [a]
xs
        else do
          x <- Gen a -> f a
forall (m :: * -> *) a. MonadIO m => Gen a -> m a
sample Gen a
gen
          go (ix + 1) (x : xs)
  as <- Int -> [a] -> IO [a]
forall {f :: * -> *}. MonadIO f => Int -> [a] -> f [a]
go Int
0 []
  newArray as

storablePeekElem :: forall a. (Eq a, Show a, Storable a) => Gen a -> Property
storablePeekElem :: forall a. (Eq a, Show a, Storable a) => Gen a -> Property
storablePeekElem Gen a
gen = HasCallStack => PropertyT IO () -> Property
PropertyT IO () -> Property
property (PropertyT IO () -> Property) -> PropertyT IO () -> Property
forall a b. (a -> b) -> a -> b
$ do
  as <- Gen [a] -> PropertyT IO [a]
forall (m :: * -> *) a.
(Monad m, Show a, HasCallStack) =>
Gen a -> PropertyT m a
forAll (Gen [a] -> PropertyT IO [a]) -> Gen [a] -> PropertyT IO [a]
forall a b. (a -> b) -> a -> b
$ Gen a -> Gen [a]
forall a. Gen a -> Gen [a]
genSmallNonEmptyList Gen a
gen
  let len = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
List.length [a]
as
  ix <- forAll $ Gen.int (Range.linear 0 (len - 1))
  unsafePerformIO $ do
    addr <- genArray gen len
    x <- peekElemOff addr ix
    y <- peek (addr `plusPtr` (ix * sizeOf (undefined :: a)))
    free addr
    pure (x === y)

storablePokeElem :: forall a. (Eq a, Show a, Storable a) => Gen a -> Property
storablePokeElem :: forall a. (Eq a, Show a, Storable a) => Gen a -> Property
storablePokeElem Gen a
gen = HasCallStack => PropertyT IO () -> Property
PropertyT IO () -> Property
property (PropertyT IO () -> Property) -> PropertyT IO () -> Property
forall a b. (a -> b) -> a -> b
$ do
  as <- Gen [a] -> PropertyT IO [a]
forall (m :: * -> *) a.
(Monad m, Show a, HasCallStack) =>
Gen a -> PropertyT m a
forAll (Gen [a] -> PropertyT IO [a]) -> Gen [a] -> PropertyT IO [a]
forall a b. (a -> b) -> a -> b
$ Gen a -> Gen [a]
forall a. Gen a -> Gen [a]
genSmallNonEmptyList Gen a
gen
  x <- forAll gen
  let len = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
List.length [a]
as
  ix <- forAll $ Gen.int (Range.linear 0 (len - 1))
  unsafePerformIO $ do
    addr <- genArray gen len
    pokeElemOff addr ix x
    u <- peekElemOff addr ix
    poke (addr `plusPtr` (ix * sizeOf x)) x
    v <- peekElemOff addr ix
    free addr
    pure (u === v)

storablePeekByte :: forall a. (Eq a, Show a, Storable a) => Gen a -> Property
storablePeekByte :: forall a. (Eq a, Show a, Storable a) => Gen a -> Property
storablePeekByte Gen a
gen = HasCallStack => PropertyT IO () -> Property
PropertyT IO () -> Property
property (PropertyT IO () -> Property) -> PropertyT IO () -> Property
forall a b. (a -> b) -> a -> b
$ do
  as <- Gen [a] -> PropertyT IO [a]
forall (m :: * -> *) a.
(Monad m, Show a, HasCallStack) =>
Gen a -> PropertyT m a
forAll (Gen [a] -> PropertyT IO [a]) -> Gen [a] -> PropertyT IO [a]
forall a b. (a -> b) -> a -> b
$ Gen a -> Gen [a]
forall a. Gen a -> Gen [a]
genSmallNonEmptyList Gen a
gen
  let len = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
List.length [a]
as
  ix <- forAll $ Gen.int (Range.linear 0 (len - 1))
  let off = Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
* (Ptr (ZonkAny 2)
forall a. Ptr a
nullPtr Ptr (ZonkAny 2) -> Int -> Ptr (ZonkAny 1)
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` a -> Int
forall a. Storable a => a -> Int
sizeOf ([a] -> a
forall a. HasCallStack => [a] -> a
head [a]
as)) Ptr (ZonkAny 1) -> Int -> Ptr (ZonkAny 1)
forall a. Ptr a -> Int -> Ptr a
`alignPtr` a -> Int
forall a. Storable a => a -> Int
alignment ([a] -> a
forall a. HasCallStack => [a] -> a
head [a]
as) Ptr (ZonkAny 1) -> Ptr (ZonkAny 0) -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr (ZonkAny 0)
forall a. Ptr a
nullPtr
  unsafePerformIO $ do
    addr <- genArray gen len
    x :: a <- peekByteOff addr off
    y :: a <- peek (addr `plusPtr` off)
    free addr
    pure (x === y)

storablePokeByte :: forall a. (Eq a, Show a, Storable a) => Gen a -> Property
storablePokeByte :: forall a. (Eq a, Show a, Storable a) => Gen a -> Property
storablePokeByte Gen a
gen = HasCallStack => PropertyT IO () -> Property
PropertyT IO () -> Property
property (PropertyT IO () -> Property) -> PropertyT IO () -> Property
forall a b. (a -> b) -> a -> b
$ do
  as <- Gen [a] -> PropertyT IO [a]
forall (m :: * -> *) a.
(Monad m, Show a, HasCallStack) =>
Gen a -> PropertyT m a
forAll (Gen [a] -> PropertyT IO [a]) -> Gen [a] -> PropertyT IO [a]
forall a b. (a -> b) -> a -> b
$ Gen a -> Gen [a]
forall a. Gen a -> Gen [a]
genSmallNonEmptyList Gen a
gen
  x <- forAll gen
  let len = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
List.length [a]
as
  off <- forAll $ Gen.int (Range.linear 0 (len - 1))
  unsafePerformIO $ do
    addr <- genArray gen len
    pokeByteOff addr off x
    u :: a <- peekByteOff addr off
    poke (addr `plusPtr` off) x
    v :: a <- peekByteOff addr off
    free addr
    pure (u === v)

storableSetGet :: forall a. (Eq a, Show a, Storable a) => Gen a -> Property
storableSetGet :: forall a. (Eq a, Show a, Storable a) => Gen a -> Property
storableSetGet Gen a
gen = HasCallStack => PropertyT IO () -> Property
PropertyT IO () -> Property
property (PropertyT IO () -> Property) -> PropertyT IO () -> Property
forall a b. (a -> b) -> a -> b
$ do
  a <- Gen a -> PropertyT IO a
forall (m :: * -> *) a.
(Monad m, Show a, HasCallStack) =>
Gen a -> PropertyT m a
forAll Gen a
gen
  len <- forAll $ Gen.int (Range.linear 1 20)
  ix <- forAll $ Gen.int (Range.linear 0 (len - 1))
  unsafePerformIO $ do
    ptr <- genArray gen len
    pokeElemOff ptr ix a
    a' <- peekElemOff ptr ix
    free ptr
    pure (a === a')

storableGetSet :: forall a. (Eq a, Show a, Storable a) => Gen a -> Property
storableGetSet :: forall a. (Eq a, Show a, Storable a) => Gen a -> Property
storableGetSet Gen a
gen = HasCallStack => PropertyT IO () -> Property
PropertyT IO () -> Property
property (PropertyT IO () -> Property) -> PropertyT IO () -> Property
forall a b. (a -> b) -> a -> b
$ do
  as <- Gen [a] -> PropertyT IO [a]
forall (m :: * -> *) a.
(Monad m, Show a, HasCallStack) =>
Gen a -> PropertyT m a
forAll (Gen [a] -> PropertyT IO [a]) -> Gen [a] -> PropertyT IO [a]
forall a b. (a -> b) -> a -> b
$ Gen a -> Gen [a]
forall a. Gen a -> Gen [a]
genSmallNonEmptyList Gen a
gen
  let len = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
List.length [a]
as
  ix <- forAll $ Gen.int (Range.linear 0 (len - 1))
  unsafePerformIO $ do
    ptrA <- newArray as
    ptrB <- genArray gen len
    copyArray ptrB ptrA len
    a <- peekElemOff ptrA ix
    pokeElemOff ptrA ix a
    res <- arrayEq ptrA ptrB len
    free ptrA
    free ptrB
    pure (res === True)

storableList :: forall a. (Eq a, Show a, Storable a) => Gen a -> Property
storableList :: forall a. (Eq a, Show a, Storable a) => Gen a -> Property
storableList Gen a
gen = HasCallStack => PropertyT IO () -> Property
PropertyT IO () -> Property
property (PropertyT IO () -> Property) -> PropertyT IO () -> Property
forall a b. (a -> b) -> a -> b
$ do
  as <- Gen [a] -> PropertyT IO [a]
forall (m :: * -> *) a.
(Monad m, Show a, HasCallStack) =>
Gen a -> PropertyT m a
forAll (Gen [a] -> PropertyT IO [a]) -> Gen [a] -> PropertyT IO [a]
forall a b. (a -> b) -> a -> b
$ Gen a -> Gen [a]
forall a. Gen a -> Gen [a]
genSmallNonEmptyList Gen a
gen 
  unsafePerformIO $ do
    let len = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
List.length [a]
as
    ptr <- newArray as
    let rebuild :: Int -> IO [a]
        rebuild Int
ix = if Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
len
          then (:) (a -> [a] -> [a]) -> IO a -> IO ([a] -> [a])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ptr a -> Int -> IO a
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr a
ptr Int
ix IO ([a] -> [a]) -> IO [a] -> IO [a]
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> IO [a]
rebuild (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
          else [a] -> IO [a]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
    asNew <- rebuild 0
    free ptr
    pure (as === asNew)

arrayEq :: forall a. (Eq a, Storable a) => Ptr a -> Ptr a -> Int -> IO Bool
arrayEq :: forall a. (Eq a, Storable a) => Ptr a -> Ptr a -> Int -> IO Bool
arrayEq Ptr a
ptrA Ptr a
ptrB Int
len = Int -> IO Bool
go Int
0 where
  go :: Int -> IO Bool
go Int
i = if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
len
    then do
      a <- Ptr a -> Int -> IO a
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr a
ptrA Int
i
      b <- peekElemOff ptrB i
      if a == b
        then go (i + 1)
        else pure False
    else Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True