Make topBits64 safe for raw bytes of size less than 8

When the size of the raw bytes is less than 8, then a fallback method is used
that is likely slower. When the size of the raw bytes is 8 or higher, then an
additional integer comparison is made, but this has very little impact on
performance, which I verified using micro-benchmarks. The upside to making the
function safe for any input raw bytes is that the API and test generators become
simpler, because there is one fewer constraint to satisfy: the minimum size of 8
bytes for serialised keys.

Author: jorisdral

src/Database/LSMTree/Internal/RawBytes.hs

@@ -51,7 +51,7 @@ module Database.LSMTree.Internal.RawBytes (
   ) where
 
 import           Control.DeepSeq (NFData)
-import           Control.Exception (assert)
+import           Data.Bits (Bits (..))
 import           Data.BloomFilter.Hash (Hashable (..), hashByteArray)
 import qualified Data.ByteString as BS
 import qualified Data.ByteString.Builder as BB
@@ -71,6 +71,9 @@ import           GHC.Stack
 import           GHC.Word
 import           Text.Printf (printf)
 
+-- $setup
+-- >>> import Numeric
+
 {- Note: [Export structure]
    ~~~~~~~~~~~~~~~~~~~~~~~
    Since RawBytes are very similar to Primitive Vectors, the code is sectioned
@@ -172,14 +175,29 @@ drop = coerce VP.drop
 --
 -- The /top/ corresponds to the most significant bit (big-endian).
 --
--- PRECONDITION: The byte-size of the raw bytes should be at least 8 bytes.
+-- If the number of bits is smaller than @64@, then any missing bits default to
+-- @0@s.
+--
+-- >>> showHex (topBits64 (pack [1,0,0,0,0,0,0,0])) ""
+-- "100000000000000"
+--
+-- >>> showHex (topBits64 (pack [1,0,0])) ""
+-- "100000000000000"
 --
 -- TODO: optimisation ideas: use unsafe shift/byteswap primops, look at GHC
 -- core, find other opportunities for using primops.
 --
 topBits64 :: RawBytes -> Word64
-topBits64 rb@(RawBytes (VP.Vector (I# off#) _size (ByteArray k#))) =
-    assert (size rb >= 8) $ toWord64 (indexWord8ArrayAsWord64# k# off#)
+topBits64 rb@(RawBytes v@(VP.Vector (I# off#) _size (ByteArray k#)))
+  | n >= 8
+  = toWord64 (indexWord8ArrayAsWord64# k# off#)
+  | otherwise
+  = VP.foldl' f 0 v `unsafeShiftL` ((8 - n) * 8)
+  where
+    !n = size rb
+
+    f :: Word64 -> Word8 -> Word64
+    f acc w = acc `unsafeShiftL` 8 + fromIntegral w
 
 #if (MIN_VERSION_GLASGOW_HASKELL(9, 4, 0, 0))
 toWord64 :: Word64# -> Word64

test/Test/Database/LSMTree/Internal/RawBytes.hs

@@ -1,8 +1,13 @@
+{-# LANGUAGE OverloadedLists #-}
+
 module Test.Database.LSMTree.Internal.RawBytes (tests) where
 
+import           Data.Bits (Bits (shiftL))
+import qualified Data.List as List
+import qualified Data.Vector.Primitive as VP
 import           Database.LSMTree.Extras.Generators ()
-import           Database.LSMTree.Internal.RawBytes (RawBytes)
-import qualified Database.LSMTree.Internal.RawBytes as RB (size)
+import           Database.LSMTree.Internal.RawBytes (RawBytes (RawBytes))
+import qualified Database.LSMTree.Internal.RawBytes as RB
 import           Test.QuickCheck (Property, classify, collect, mapSize,
                      withDiscardRatio, withMaxSuccess, (.||.), (===), (==>))
 import           Test.Tasty (TestTree, testGroup)
@@ -26,7 +31,9 @@ tests = testGroup "Test.Database.LSMTree.Internal.RawBytes" $
                 testProperty "Transitivity"  prop_ordTransitivity,
                 testProperty "Reflexivity"   prop_ordReflexivity,
                 testProperty "Antisymmetry"  prop_ordAntisymmetry
-            ]
+            ],
+            testProperty "prop_topBits64" prop_topBits64,
+            testProperty "prop_topBits64_default0s" prop_topBits64_default0s
         ]
 
 -- * Utilities
@@ -92,3 +99,22 @@ prop_ordAntisymmetry = mapSize (const 4)    $
     untunedProp block1 block2
         = withFirstBlockSizeInfo block1 $
           block1 <= block2 && block2 <= block1 ==> block1 === block2
+
+{-------------------------------------------------------------------------------
+  Accessors
+-------------------------------------------------------------------------------}
+
+-- | Compare 'topBits64' against a model
+prop_topBits64 :: RawBytes -> Property
+prop_topBits64 x@(RawBytes v) =
+    expected === RB.topBits64 x
+  where
+    expected =
+      let ws = take 8 (VP.toList v ++ repeat 0)
+      in  List.foldl' (\acc w -> acc `shiftL` 8 + fromIntegral w) 0 ws
+
+-- | If @x@ has fewer than 8 bytes, then all missing bits in the result default
+-- to 0s.
+prop_topBits64_default0s :: RawBytes -> Property
+prop_topBits64_default0s x =
+    RB.topBits64 x === RB.topBits64 (x <> mconcat (replicate 8 [0]))