Add a custom TestKey type for compact index tests

Outside of direct tests on the compact index, clashes are arguably not so
important: we can assume that the compact index handles clashes correctly
because we already test this extensively.

This also fixes #580, restoring the generation of clashing keys.

Author: jorisdral

test/Test/Database/LSMTree/Internal/Index/Compact.hs

@@ -27,7 +27,10 @@ import qualified Data.Vector.Unboxed as VU
 import qualified Data.Vector.Unboxed.Base as VU
 import           Data.Word
 import           Database.LSMTree.Extras
-import           Database.LSMTree.Extras.Generators as Gen
+import           Database.LSMTree.Extras.Generators (BiasedKeyForIndexCompact,
+                     ChunkSize (..), LogicalPageSummaries,
+                     LogicalPageSummary (..), Pages (..), genRawBytes,
+                     isKeyForIndexCompact, labelPages, toAppends)
 import           Database.LSMTree.Extras.Index (Append (..), appendToCompact)
 import           Database.LSMTree.Internal.BitMath
 import           Database.LSMTree.Internal.Chunk as Chunk (toByteString)
@@ -36,6 +39,8 @@ import           Database.LSMTree.Internal.Index.Compact
 import           Database.LSMTree.Internal.Index.CompactAcc
 import           Database.LSMTree.Internal.Page (PageNo (PageNo), PageSpan,
                      multiPage, singlePage)
+import           Database.LSMTree.Internal.RawBytes (RawBytes (..))
+import qualified Database.LSMTree.Internal.RawBytes as RB
 import           Database.LSMTree.Internal.Serialise
 import           Numeric (showHex)
 import           Prelude hiding (max, min, pi)
@@ -51,7 +56,9 @@ import           Text.Printf (printf)
 
 tests :: TestTree
 tests = testGroup "Test.Database.LSMTree.Internal.Index.Compact" [
-    testProperty "prop_distribution @BiasedKeyForIndexCompact" $
+    testGroup "TestKey" $
+      prop_arbitraryAndShrinkPreserveInvariant @TestKey noTags isTestKey
+  , testProperty "prop_distribution @BiasedKeyForIndexCompact" $
       prop_distribution @BiasedKeyForIndexCompact
   , testProperty "prop_searchMinMaxKeysAfterConstruction" $
       prop_searchMinMaxKeysAfterConstruction @BiasedKeyForIndexCompact 100
@@ -127,6 +134,57 @@ tests = testGroup "Test.Database.LSMTree.Internal.Index.Compact" [
       ]
   ]
 
+{-------------------------------------------------------------------------------
+  Test key
+-------------------------------------------------------------------------------}
+
+-- | Key type for compact index tests
+--
+-- Tests outside this module don't have to worry about generating clashing keys.
+-- We can assume that the compact index handles clashes correctly, because we
+-- test this extensively in this module already.
+newtype TestKey = TestKey RawBytes
+  deriving stock (Show, Eq, Ord)
+  deriving newtype SerialiseKey
+
+-- | Generate keys with a non-neglible probability of clashes. This generates
+-- sliced keys too.
+--
+-- Note: recall that keys /clash/ only if their primary bits (first 8 bytes)
+-- match. It does not matter whether the other bytes do not match.
+instance Arbitrary TestKey where
+  arbitrary = do
+      -- Generate primary bits from a relatively small distribution. This
+      -- ensures that we get clashes between keys with a non-negligible
+      -- probability.
+      primBits <- do
+        lastPrefixByte <- QC.getSmall <$> arbitrary
+        pure $ RB.pack ([0,0,0,0,0,0,0] <> [lastPrefixByte])
+      -- The rest of the bits after the primary bits can be anything
+      restBits <- genRawBytes
+      -- The compact index should store keys without retaining unused memory.
+      -- Therefore, we generate slices of keys too.
+      prefix <- elements [RB.pack [], RB.pack [0]]
+      suffix <- elements [RB.pack [], RB.pack [0]]
+      -- Combine the bytes and make sure to take out only the slice we need.
+      let bytes = prefix <> primBits <> restBits <> suffix
+          n = RB.size primBits + RB.size restBits
+          bytes' = RB.take n $ RB.drop (RB.size prefix) bytes
+      pure $ TestKey bytes'
+
+  -- Shrink keys extensively: most failures will occur in small counterexamples,
+  -- so we don't have to limit the number of shrinks as much.
+  shrink (TestKey bytes) = [
+        TestKey bytes'
+      | let RawBytes vec = bytes
+      , vec' <- VP.fromList <$> shrink (VP.toList vec)
+      , let bytes' = RawBytes vec'
+      , isKeyForIndexCompact bytes'
+      ]
+
+isTestKey :: TestKey -> Bool
+isTestKey (TestKey bytes) = isKeyForIndexCompact bytes
+
 {-------------------------------------------------------------------------------
   Properties
 -------------------------------------------------------------------------------}