rrb-vector-0.2.2.0: test/Arbitrary.hs

{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE LambdaCase #-}

module Arbitrary where

#if !(MIN_VERSION_base(4,18,0))
import Control.Applicative (liftA2)
#endif
import Control.Monad.ST
import Data.Bool (bool)
import Debug.Trace

import Test.Tasty.QuickCheck

import qualified Data.RRBVector.Internal as V
import qualified Data.RRBVector.Internal.Buffer as Buffer
import Data.RRBVector.Internal.Debug

arbitraryVectorOf :: Gen a -> Gen (V.Vector a)
arbitraryVectorOf gen = sized $ \n -> do
    xs <- vectorOf n gen
    if n <= V.blockSize
        then pure $ V.fromList xs
        else nodes V.Leaf xs >>= \case
            [tree] -> pure $ V.Root (V.treeSize 0 tree) 0 tree
            ts -> iterateNodes V.blockShift ts
  where
    nodes f trees = do
        sizes <- infiniteListOf (arbitrary >>= bool (pure V.blockSize) (choose (1, V.blockSize - 1)))
        pure $ runST $ do
            buffer <- Buffer.new V.blockSize
            let loop _ [] = do
                    result <- Buffer.get buffer
                    pure [f result]
                loop (sz : sizes') (t : ts) = do
                    size <- Buffer.size buffer
                    if size == sz then do
                        result <- Buffer.get buffer
                        Buffer.push buffer t
                        rest <- loop sizes' ts
                        pure (f result : rest)
                    else do
                        Buffer.push buffer t
                        loop (sz : sizes') ts
            loop sizes trees
    {-# INLINE nodes #-}

    iterateNodes sh trees = do
        ts <- nodes (V.computeSizes sh) trees
        case ts of
            [tree] -> pure $ V.Root (V.treeSize sh tree) sh tree
            trees' -> iterateNodes (V.up sh) trees'

instance (Arbitrary a) => Arbitrary (V.Vector a) where
    arbitrary = arbitrary1
    shrink = shrink1

-- TODO: improve instance
instance Arbitrary1 V.Vector where
    liftArbitrary = arbitraryVectorOf

-- A custom 'Testable' instance to use 'showTree'.
instance {-# OVERLAPPING #-} (Arbitrary a, Show a, Testable prop) => Testable (V.Vector a -> prop) where
    property = propertyForAllShrinkShow arbitrary shrink (pure . showTree)

    propertyForAllShrinkShow gen shr shw f =
        propertyForAllShrinkShow
            (liftA2 (,) gen arbitrary)
            (liftShrink2 shr shrink)
            (\(x, y) -> shw x ++ [showTree y])
            (uncurry f)