speed up shrinkRawBytes

Before:

  Test.Database.LSMTree.Generators
    RawBytes
      Shrinking satisfies invariant:    OK (0.01s)
        +++ OK, passed 100 tests (1% no shrinks).

        number of shrinks (99 in total):
        64% 100 <= n < 1000
        30% 10 <= n < 100
         6% 1 <= n < 10
    lists of key/op pairs
      Shrinking satisfies invariant: OK (4.51s)
        +++ OK, passed 100 tests (4% no shrinks).

        number of shrinks (96 in total):
        45% 1000 <= n < 10000
        44% 10000 <= n < 100000
        10% 100 <= n < 1000
         1% 10 <= n < 100

After:

  Test.Database.LSMTree.Generators
    RawBytes
      Shrinking satisfies invariant:   OK
        +++ OK, passed 100 tests (4% no shrinks).

        number of shrinks (96 in total):
        52% 10 <= n < 100
        45% 100 <= n < 1000
         3% 1 <= n < 10
    lists of key/op pairs
      Shrinking satisfies invariant:   OK (1.57s)
        +++ OK, passed 100 tests (6% no shrinks).

        number of shrinks (94 in total):
        56% 1000 <= n < 10000
        21% 100 <= n < 1000
        19% 10000 <= n < 100000
         3% 10 <= n < 100

Author: mheinzel

src-extras/Database/LSMTree/Extras/Generators.hs

@@ -37,14 +37,16 @@ module Database.LSMTree.Extras.Generators (
   , isKeyForIndexCompact
   , KeyForIndexCompact (..)
   , BiasedKeyForIndexCompact (..)
+    -- * 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 (sort)
+import           Data.List (nub, sort)
 import qualified Data.Primitive.ByteArray as BA
 import qualified Data.Vector.Primitive as VP
 import           Data.Word
@@ -57,7 +59,8 @@ 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.RawBytes as RB
+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,
@@ -454,8 +457,33 @@ packRawBytesPinnedOrUnpinned True  = \ws ->
       return mba
 
 shrinkRawBytes :: RawBytes -> [RawBytes]
-shrinkRawBytes (RawBytes pvec) = [ RawBytes (VP.fromList ws)
-                                 | ws <- QC.shrink (VP.toList pvec) ]
+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
@@ -465,10 +493,14 @@ genSlice (RawBytes pvec) = do
 
 shrinkSlice :: RawBytes -> [RawBytes]
 shrinkSlice (RawBytes pvec) =
-    [ RawBytes (VP.slice m n pvec)
-    | n <- QC.shrink (VP.length pvec)
-    , m <- QC.shrink (VP.length pvec - n)
+    [ 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
 

test/Test/Database/LSMTree/Generators.hs

@@ -25,7 +25,8 @@ import           Database.LSMTree.Internal.Serialise
 import qualified Test.QuickCheck as QC
 import           Test.QuickCheck (Property)
 import           Test.Tasty (TestTree, localOption, testGroup)
-import           Test.Tasty.QuickCheck (QuickCheckMaxSize (..), testProperty)
+import           Test.Tasty.QuickCheck (QuickCheckMaxSize (..), testProperty,
+                     (===))
 import           Test.Util.Arbitrary
 
 tests :: TestTree
@@ -62,6 +63,10 @@ tests = testGroup "Test.Database.LSMTree.Generators" [
     , testGroup "lists of key/op pairs" $
         prop_arbitraryAndShrinkPreserveInvariant labelTestKOps $
           deepseqInvariant
+    , testGroup "helpers"
+        [ testProperty "prop_shrinkVec" $ \vec ->
+            shrinkVec (QC.shrink @Int) vec === map VP.fromList (QC.shrink (VP.toList vec))
+        ]
     ]
 
 prop_packRawBytesPinnedOrUnpinned :: Bool -> [Word8] -> Bool