lsm-tree-1.0.0.0: src-extras/Database/LSMTree/Extras/Generators.hs
{-# OPTIONS_GHC -Wno-orphans #-}
module Database.LSMTree.Extras.Generators (
-- * WithSerialised
WithSerialised (..)
-- * A (logical\/true) page
-- ** A true page
, TruePageSummary (..)
, flattenLogicalPageSummary
-- ** A logical page
, LogicalPageSummary (..)
, shrinkLogicalPageSummary
, toAppend
-- * Sequences of (logical\/true) pages
, Pages (..)
-- ** Sequences of true pages
, TruePageSummaries
, flattenLogicalPageSummaries
-- ** Sequences of logical pages
, LogicalPageSummaries
, toAppends
, labelPages
, shrinkPages
, genPages
, mkPages
, pagesInvariant
-- * Chunking size
, ChunkSize (..)
, chunkSizeInvariant
-- * Serialised keys\/values\/blobs
, genRawBytes
, genRawBytesN
, genRawBytesSized
, packRawBytesPinnedOrUnpinned
, LargeRawBytes (..)
, BiasedKey (..)
-- * helpers
, shrinkVec
) where
import Control.DeepSeq (NFData)
import Control.Exception (assert)
import Data.Coerce (coerce)
import Data.Containers.ListUtils (nubOrd)
import Data.Function ((&))
import Data.List (nub, sort)
import qualified Data.Primitive.ByteArray as BA
import qualified Data.Vector.Primitive as VP
import Data.Word
import qualified Database.LSMTree as Full
import Database.LSMTree.Extras
import Database.LSMTree.Extras.Index (Append (..))
import Database.LSMTree.Extras.Orphans ()
import Database.LSMTree.Internal.BlobRef (BlobSpan (..))
import Database.LSMTree.Internal.Entry (Entry (..), NumEntries (..))
import qualified Database.LSMTree.Internal.Merge as Merge
import Database.LSMTree.Internal.Page (PageNo (..))
import Database.LSMTree.Internal.Range (Range (..))
import Database.LSMTree.Internal.RawBytes (RawBytes (RawBytes))
import qualified Database.LSMTree.Internal.RawBytes as RB
import Database.LSMTree.Internal.Serialise
import qualified Database.LSMTree.Internal.Serialise.Class as S.Class
import Database.LSMTree.Internal.Unsliced (Unsliced, fromUnslicedKey,
makeUnslicedKey)
import Database.LSMTree.Internal.Vector (mkPrimVector)
import GHC.Generics (Generic)
import qualified Test.QuickCheck as QC
import Test.QuickCheck (Arbitrary (..), Arbitrary1 (..),
Arbitrary2 (..), Gen, Property, elements, frequency)
import Test.QuickCheck.Gen (genDouble)
import Test.QuickCheck.Instances ()
{-------------------------------------------------------------------------------
Common LSMTree types
-------------------------------------------------------------------------------}
instance (Arbitrary v, Arbitrary b) => Arbitrary (Full.Update v b) where
arbitrary = QC.arbitrary2
shrink = QC.shrink2
instance Arbitrary2 Full.Update where
liftArbitrary2 genVal genBlob = frequency
[ (10, Full.Insert <$> genVal <*> liftArbitrary genBlob)
, (5, Full.Upsert <$> genVal)
, (1, pure Full.Delete)
]
liftShrink2 shrinkVal shrinkBlob = \case
Full.Insert v blob ->
Full.Delete
: map (uncurry Full.Insert)
(liftShrink2 shrinkVal (liftShrink shrinkBlob) (v, blob))
Full.Upsert v -> Full.Insert v Nothing : map Full.Upsert (shrinkVal v)
Full.Delete -> []
instance (Arbitrary k, Ord k) => Arbitrary (Range k) where
arbitrary = do
key1 <- arbitrary
key2 <- arbitrary `QC.suchThat` (/= key1)
(lb, ub) <- frequency
[ (1, pure (key1, key1)) -- lb == ub
, (1, pure (max key1 key2, min key1 key2)) -- lb > ub
, (8, pure (min key1 key2, max key1 key2)) -- lb < ub
]
elements
[ FromToExcluding lb ub
, FromToIncluding lb ub
]
shrink (FromToExcluding f t) =
uncurry FromToExcluding <$> shrink (f, t)
shrink (FromToIncluding f t) =
uncurry FromToIncluding <$> shrink (f, t)
{-------------------------------------------------------------------------------
Entry
-------------------------------------------------------------------------------}
instance (Arbitrary v, Arbitrary b) => Arbitrary (Entry v b) where
arbitrary = QC.arbitrary2
shrink = QC.shrink2
instance Arbitrary2 Entry where
liftArbitrary2 genVal genBlob = frequency
[ (5, Insert <$> genVal)
, (1, InsertWithBlob <$> genVal <*> genBlob)
, (1, Upsert <$> genVal)
, (1, pure Delete)
]
liftShrink2 shrinkVal shrinkBlob = \case
Insert v -> Delete : (Insert <$> shrinkVal v)
InsertWithBlob v b -> [Delete, Insert v]
++ fmap (uncurry InsertWithBlob)
(liftShrink2 shrinkVal shrinkBlob (v, b))
Upsert v -> Delete : Insert v : (Upsert <$> shrinkVal v)
Delete -> []
{-------------------------------------------------------------------------------
WithSerialised
-------------------------------------------------------------------------------}
-- | Cache serialised keys
--
-- Also useful for failing tests that have keys as inputs, because the printed
-- 'WithSerialised' values will show both keys and their serialised form.
data WithSerialised k = WithSerialised k SerialisedKey
deriving stock Show
instance Eq k => Eq (WithSerialised k) where
WithSerialised k1 _ == WithSerialised k2 _ = k1 == k2
instance Ord k => Ord (WithSerialised k) where
WithSerialised k1 _ `compare` WithSerialised k2 _ = k1 `compare` k2
instance (Arbitrary k, SerialiseKey k) => Arbitrary (WithSerialised k) where
arbitrary = do
x <- arbitrary
pure $ WithSerialised x (serialiseKey x)
shrink (WithSerialised k _) = [WithSerialised k' (serialiseKey k') | k' <- shrink k]
instance SerialiseKey k => SerialiseKey (WithSerialised k) where
serialiseKey (WithSerialised _ (SerialisedKey bytes)) = bytes
deserialiseKey bytes = WithSerialised (S.Class.deserialiseKey bytes) (SerialisedKey bytes)
{-------------------------------------------------------------------------------
Other number newtypes
-------------------------------------------------------------------------------}
instance Arbitrary PageNo where
arbitrary = coerce (arbitrary @(QC.NonNegative Int))
shrink = coerce (shrink @(QC.NonNegative Int))
instance Arbitrary NumEntries where
arbitrary = coerce (arbitrary @(QC.NonNegative Int))
shrink = coerce (shrink @(QC.NonNegative Int))
{-------------------------------------------------------------------------------
True page
-------------------------------------------------------------------------------}
-- | A summary of min/max information for keys on a /true/ page.
--
-- A true page corresponds directly to a disk page. See 'LogicalPageSummary' for
-- contrast.
data TruePageSummary k = TruePageSummary { tpsMinKey :: k, tpsMaxKey :: k }
flattenLogicalPageSummary :: LogicalPageSummary k -> [TruePageSummary k]
flattenLogicalPageSummary = \case
OnePageOneKey k -> [TruePageSummary k k]
OnePageManyKeys k1 k2 -> [TruePageSummary k1 k2]
MultiPageOneKey k n -> replicate (fromIntegral n+1) (TruePageSummary k k)
{-------------------------------------------------------------------------------
Logical page
-------------------------------------------------------------------------------}
-- | A summary of min/max information for keys on a /logical/ page.
--
-- A key\/operation pair can fit onto a single page, or the operation is so
-- large that its bytes flow over into subsequent pages. A logical page makes
-- this overflow explicit. Making these cases explicit in the representation
-- makes generating and shrinking test cases easier.
data LogicalPageSummary k =
OnePageOneKey k
| OnePageManyKeys k k
| MultiPageOneKey k Word32 -- ^ number of overflow pages
deriving stock (Show, Generic, Functor)
deriving anyclass NFData
toAppend :: LogicalPageSummary SerialisedKey -> Append
toAppend (OnePageOneKey k) = AppendSinglePage k k
toAppend (OnePageManyKeys k1 k2) = AppendSinglePage k1 k2
toAppend (MultiPageOneKey k n) = AppendMultiPage k n
shrinkLogicalPageSummary :: Arbitrary k => LogicalPageSummary k -> [LogicalPageSummary k]
shrinkLogicalPageSummary = \case
OnePageOneKey k -> OnePageOneKey <$> shrink k
OnePageManyKeys k1 k2 -> [
OnePageManyKeys k1' k2'
| (k1', k2') <- shrink (k1, k2)
]
MultiPageOneKey k n -> [
MultiPageOneKey k' n'
| (k', n') <- shrink (k, n)
]
{-------------------------------------------------------------------------------
Sequences of (logical\/true) pages
-------------------------------------------------------------------------------}
-- | Sequences of (logical\/true) pages
--
-- INVARIANT: The sequence consists of multiple pages in sorted order (keys are
-- sorted within a page and across pages).
newtype Pages fp k = Pages { getPages :: [fp k] }
deriving stock (Show, Generic, Functor)
deriving anyclass NFData
class TrueNumberOfPages fp where
trueNumberOfPages :: Pages fp k -> Int
instance TrueNumberOfPages LogicalPageSummary where
trueNumberOfPages :: LogicalPageSummaries k -> Int
trueNumberOfPages = length . getPages . flattenLogicalPageSummaries
instance TrueNumberOfPages TruePageSummary where
trueNumberOfPages :: TruePageSummaries k -> Int
trueNumberOfPages = length . getPages
{-------------------------------------------------------------------------------
Sequences of true pages
-------------------------------------------------------------------------------}
type TruePageSummaries k = Pages TruePageSummary k
flattenLogicalPageSummaries :: LogicalPageSummaries k -> TruePageSummaries k
flattenLogicalPageSummaries (Pages ps) = Pages (concatMap flattenLogicalPageSummary ps)
{-------------------------------------------------------------------------------
Sequences of logical pages
-------------------------------------------------------------------------------}
type LogicalPageSummaries k = Pages LogicalPageSummary k
toAppends :: SerialiseKey k => LogicalPageSummaries k -> [Append]
toAppends (Pages ps) = fmap (toAppend . fmap serialiseKey) ps
--
-- Labelling
--
labelPages :: LogicalPageSummaries k -> (Property -> Property)
labelPages ps =
QC.tabulate "# True pages" [showPowersOf10 nTruePages]
. QC.tabulate "# Logical pages" [showPowersOf10 nLogicalPages]
. QC.tabulate "# OnePageOneKey logical pages" [showPowersOf10 n1]
. QC.tabulate "# OnePageManyKeys logical pages" [showPowersOf10 n2]
. QC.tabulate "# MultiPageOneKey logical pages" [showPowersOf10 n3]
where
nLogicalPages = length $ getPages ps
nTruePages = trueNumberOfPages ps
(n1,n2,n3) = counts (getPages ps)
counts :: [LogicalPageSummary k] -> (Int, Int, Int)
counts [] = (0, 0, 0)
counts (lp:lps) = let (x, y, z) = counts lps
in case lp of
OnePageOneKey{} -> (x+1, y, z)
OnePageManyKeys{} -> (x, y+1, z)
MultiPageOneKey{} -> (x, y, z+1)
--
-- Generation and shrinking
--
instance (Arbitrary k, Ord k)
=> Arbitrary (LogicalPageSummaries k) where
arbitrary = genPages 0.03 (QC.choose (0, 16)) 0.01
shrink = shrinkPages
shrinkPages ::
(Arbitrary k, Ord k)
=> LogicalPageSummaries k
-> [LogicalPageSummaries k]
shrinkPages (Pages ps) = [
Pages ps'
| ps' <- QC.shrinkList shrinkLogicalPageSummary ps, pagesInvariant (Pages ps')
]
genPages ::
(Arbitrary k, Ord k)
=> Double -- ^ Probability of a value being larger-than-page
-> Gen Word32 -- ^ Number of overflow pages for a larger-than-page value
-> Double -- ^ Probability of generating a page with only one key and value,
-- which does /not/ span multiple pages.
-> Gen (LogicalPageSummaries k)
genPages p genN p' = do
ks <- arbitrary
mkPages p genN p' ks
mkPages ::
forall k. Ord k
=> Double -- ^ Probability of a value being larger-than-page
-> Gen Word32 -- ^ Number of overflow pages for a larger-than-page value
-> Double -- ^ Probability of generating a page with only one key and value,
-- which does /not/ span multiple pages.
-> [k]
-> Gen (LogicalPageSummaries k)
mkPages p genN p' =
fmap Pages . go . nubOrd . sort
where
go :: [k] -> Gen [LogicalPageSummary k]
go [] = pure []
go [k] = do
b <- largerThanPage
if b then pure . MultiPageOneKey k <$> genN
else pure [OnePageOneKey k]
-- the min and max key are allowed to be the same
go (k1:k2:ks) = do
b <- largerThanPage
b' <- onePageOneKey
if b then (:) <$> (MultiPageOneKey k1 <$> genN) <*> go (k2 : ks)
else if b' then (OnePageOneKey k1 :) <$> go (k2 : ks)
else (OnePageManyKeys k1 k2 :) <$> go ks
largerThanPage :: Gen Bool
largerThanPage = genDouble >>= \x -> pure (x < p)
onePageOneKey :: Gen Bool
onePageOneKey = genDouble >>= \x -> pure (x < p')
pagesInvariant :: Ord k => LogicalPageSummaries k -> Bool
pagesInvariant (Pages ps0) =
sort ks == ks
&& nubOrd ks == ks
where
ks = flatten ps0
flatten :: Eq k => [LogicalPageSummary k] -> [k]
flatten [] = []
-- the min and max key are allowed to be the same
flatten (p:ps) = case p of
OnePageOneKey k -> k : flatten ps
OnePageManyKeys k1 k2 -> k1 : k2 : flatten ps
MultiPageOneKey k _ -> k : flatten ps
{-------------------------------------------------------------------------------
Chunking size
-------------------------------------------------------------------------------}
newtype ChunkSize = ChunkSize Int
deriving stock Show
deriving newtype Num
instance Arbitrary ChunkSize where
arbitrary = ChunkSize <$> QC.chooseInt (chunkSizeLB, chunkSizeUB)
shrink (ChunkSize csize) = [
ChunkSize csize'
| csize' <- shrink csize
, chunkSizeInvariant (ChunkSize csize')
]
chunkSizeLB, chunkSizeUB :: Int
chunkSizeLB = 1
chunkSizeUB = 20
chunkSizeInvariant :: ChunkSize -> Bool
chunkSizeInvariant (ChunkSize csize) = chunkSizeLB <= csize && csize <= chunkSizeUB
{-------------------------------------------------------------------------------
Serialised keys/values/blobs
-------------------------------------------------------------------------------}
instance Arbitrary RawBytes where
arbitrary = do
QC.NonNegative (QC.Small prefixLength) <- arbitrary
QC.NonNegative (QC.Small payloadLength) <- arbitrary
QC.NonNegative (QC.Small suffixLength) <- arbitrary
base <- genRawBytesN (prefixLength + payloadLength + suffixLength)
pure (base & RB.drop prefixLength & RB.take payloadLength)
shrink rb = shrinkSlice rb ++ shrinkRawBytes rb
genRawBytesN :: Int -> Gen RawBytes
genRawBytesN n =
packRawBytesPinnedOrUnpinned <$> arbitrary <*> QC.vectorOf n arbitrary
genRawBytes :: Gen RawBytes
genRawBytes =
packRawBytesPinnedOrUnpinned <$> arbitrary <*> QC.listOf arbitrary
genRawBytesSized :: Int -> Gen RawBytes
genRawBytesSized n = QC.resize n genRawBytes
packRawBytesPinnedOrUnpinned :: Bool -> [Word8] -> RawBytes
packRawBytesPinnedOrUnpinned False = RB.pack
packRawBytesPinnedOrUnpinned True = \ws ->
let len = length ws in
RB.RawBytes $ mkPrimVector 0 len $ BA.runByteArray $ do
mba <- BA.newPinnedByteArray len
sequence_ [ BA.writeByteArray mba i w | (i, w) <- zip [0..] ws ]
pure mba
shrinkRawBytes :: RawBytes -> [RawBytes]
shrinkRawBytes (RawBytes pvec) =
[ RawBytes pvec'
| pvec' <- shrinkVec shrinkByte pvec
]
where
-- no need to try harder shrinking individual bytes
shrinkByte b = nub (takeWhile (< b) [0, b `div` 2])
-- | Based on QuickCheck's 'shrinkList' (behaves identically, see tests).
shrinkVec :: VP.Prim a => (a -> [a]) -> VP.Vector a -> [VP.Vector a]
shrinkVec shr vec =
concat [ removeBlockOf k | k <- takeWhile (> 0) (iterate (`div` 2) len) ]
++ shrinkOne
where
len = VP.length vec
shrinkOne =
[ vec VP.// [(i, x')]
| i <- [0 .. len-1]
, let x = vec VP.! i
, x' <- shr x
]
removeBlockOf k =
[ VP.take i vec VP.++ VP.drop (i + k) vec
| i <- [0, k .. len - k]
]
genSlice :: RawBytes -> Gen RawBytes
genSlice (RawBytes pvec) = do
n <- QC.chooseInt (0, VP.length pvec)
m <- QC.chooseInt (0, VP.length pvec - n)
pure $ RawBytes (VP.slice m n pvec)
shrinkSlice :: RawBytes -> [RawBytes]
shrinkSlice (RawBytes pvec) =
[ RawBytes (VP.take len' pvec)
| len' <- QC.shrink len
] ++
[ RawBytes (VP.drop (len - len') pvec)
| len' <- QC.shrink len
]
where
len = VP.length pvec
deriving newtype instance Arbitrary SerialisedKey
instance Arbitrary SerialisedValue where
-- good mix of sizes, including larger than two pages, also some slices
arbitrary = SerialisedValue <$> frequency
[ (16, arbitrary)
, ( 4, genRawBytesN =<< QC.chooseInt ( 100, 1000))
, ( 2, genRawBytesN =<< QC.chooseInt (1000, 4000))
, ( 1, genRawBytesN =<< QC.chooseInt (4000, 10000))
, ( 1, genSlice =<< genRawBytesN =<< QC.chooseInt (0, 10000))
]
shrink (SerialisedValue rb)
| RB.size rb > 64 = coerce (shrink (LargeRawBytes rb))
| otherwise = coerce (shrink rb)
deriving newtype instance Arbitrary SerialisedBlob
newtype LargeRawBytes = LargeRawBytes RawBytes
deriving stock Show
deriving newtype NFData
instance Arbitrary LargeRawBytes where
arbitrary = genRawBytesSized (4096*3) >>= fmap LargeRawBytes . genSlice
shrink (LargeRawBytes rb) =
map LargeRawBytes (shrinkSlice rb)
-- After shrinking length, don't shrink content using normal list shrink
-- as that's too slow. We try zeroing out long suffixes of the bytes
-- (since for large raw bytes in page format, the interesting information
-- is at the start and the suffix is just the value.
++ [ LargeRawBytes (RawBytes pvec')
| let (RawBytes pvec) = rb
, n <- QC.shrink (VP.length pvec)
, assert (n >= 0) True -- negative values would make pvec' longer
, let pvec' = VP.take n pvec VP.++ VP.replicate (VP.length pvec - n) 0
, assert (VP.length pvec' == VP.length pvec) $
pvec' /= pvec
]
deriving newtype instance SerialiseValue LargeRawBytes
-- we try to make collisions and close keys more likely (very crudely)
arbitraryBiasedKey :: (RawBytes -> k) -> Gen RawBytes -> Gen k
arbitraryBiasedKey fromRB genUnbiased = fromRB <$> frequency
[ (6, genUnbiased)
, (1, do
lastByte <- QC.sized $ skewedWithMax . fromIntegral
pure (RB.pack ([1,3,3,7,0,1,7] <> [lastByte]))
)
]
where
-- generates a value in range from 0 to ub, but skewed towards low end
skewedWithMax ub0 = do
ub1 <- QC.chooseBoundedIntegral (0, ub0)
ub2 <- QC.chooseBoundedIntegral (0, ub1)
QC.chooseBoundedIntegral (0, ub2)
newtype BiasedKey = BiasedKey { getBiasedKey :: RawBytes }
deriving stock (Eq, Ord, Show)
deriving newtype NFData
instance Arbitrary BiasedKey where
arbitrary = arbitraryBiasedKey BiasedKey arbitrary
shrink (BiasedKey rb) = [BiasedKey rb' | rb' <- shrink rb]
deriving newtype instance SerialiseKey BiasedKey
{-------------------------------------------------------------------------------
Unsliced
-------------------------------------------------------------------------------}
instance Arbitrary (Unsliced SerialisedKey) where
arbitrary = makeUnslicedKey <$> arbitrary
shrink = fmap makeUnslicedKey . shrink . fromUnslicedKey
{-------------------------------------------------------------------------------
BlobRef
-------------------------------------------------------------------------------}
instance Arbitrary BlobSpan where
arbitrary = BlobSpan <$> arbitrary <*> arbitrary
shrink (BlobSpan x y) = [ BlobSpan x' y' | (x', y') <- shrink (x, y) ]
{-------------------------------------------------------------------------------
Merge
-------------------------------------------------------------------------------}
instance Arbitrary Merge.MergeType where
arbitrary = QC.elements
[Merge.MergeTypeMidLevel, Merge.MergeTypeLastLevel, Merge.MergeTypeUnion]
shrink Merge.MergeTypeMidLevel = []
shrink Merge.MergeTypeLastLevel = [Merge.MergeTypeMidLevel]
shrink Merge.MergeTypeUnion = [Merge.MergeTypeLastLevel]
instance Arbitrary Merge.LevelMergeType where
arbitrary = QC.elements [Merge.MergeMidLevel, Merge.MergeLastLevel]
shrink Merge.MergeMidLevel = []
shrink Merge.MergeLastLevel = [Merge.MergeMidLevel]
instance Arbitrary Merge.TreeMergeType where
arbitrary = QC.elements [Merge.MergeLevel, Merge.MergeUnion]
shrink Merge.MergeLevel = []
shrink Merge.MergeUnion = [Merge.MergeLevel]