lsm-tree-1.0.0.0: src-core/Database/LSMTree/Internal/RawPage.hs
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# OPTIONS_HADDOCK not-home #-}
module Database.LSMTree.Internal.RawPage (
RawPage (..),
emptyRawPage,
makeRawPage,
unsafeMakeRawPage,
rawPageRawBytes,
rawPageNumKeys,
rawPageNumBlobs,
rawPageLookup,
RawPageLookup(..),
rawPageOverflowPages,
rawPageFindKey,
rawPageIndex,
RawPageIndex(..),
getRawPageIndexKey,
-- * Test and debug
rawPageKeyOffsets,
rawPageValueOffsets,
rawPageValueOffsets1,
rawPageHasBlobSpanAt,
rawPageBlobSpanIndex,
rawPageOpAt,
rawPageKeys,
rawPageValues,
) where
import Control.DeepSeq (NFData (rnf))
import Control.Exception (assert)
import Data.Bits (complement, popCount, unsafeShiftL, unsafeShiftR,
(.&.))
import Data.Primitive.ByteArray (ByteArray (..), byteArrayFromList,
copyByteArray, fillByteArray, indexByteArray,
isByteArrayPinned, newAlignedPinnedByteArray, runByteArray,
sizeofByteArray)
import qualified Data.Vector as V
import qualified Data.Vector.Primitive as VP
import Data.Word (Word16, Word32, Word64, Word8)
import Database.LSMTree.Internal.BitMath
import Database.LSMTree.Internal.BlobRef (BlobSpan (..))
import Database.LSMTree.Internal.Entry (Entry (..))
import Database.LSMTree.Internal.RawBytes (RawBytes (..))
import qualified Database.LSMTree.Internal.RawBytes as RB
import Database.LSMTree.Internal.Serialise (SerialisedKey (..),
SerialisedValue (..))
import Database.LSMTree.Internal.Vector
import qualified GHC.List as List
-------------------------------------------------------------------------------
-- RawPage type
-------------------------------------------------------------------------------
data RawPage = RawPage
!Int -- ^ offset in Word16s.
!ByteArray
deriving stock (Show)
emptyRawPage :: RawPage
emptyRawPage = flip makeRawPage 0 $ byteArrayFromList
[ 0, 0, 0, 0
, 8, 0, 0, 0
, 10 :: Word8
]
invariant :: RawPage -> Bool
invariant (RawPage off ba) = and
[ off >= 0 -- offset is positive
, mod8 offsetInBytes == 0 -- 64 bit/8 byte alignment
, (sizeofByteArray ba - offsetInBytes) >= 4096 -- there is always 4096 bytes
, isByteArrayPinned ba -- bytearray should be pinned
]
where
offsetInBytes = mul2 off
instance NFData RawPage where
rnf (RawPage _ _) = ()
-- | This instance assumes pages are 4096 bytes in size
instance Eq RawPage where
RawPage off1 ba1 == RawPage off2 ba2 = v1 == v2
where
v1, v2 :: VP.Vector Word16
v1 = mkPrimVector off1 2048 ba1
v2 = mkPrimVector off2 2048 ba2
-- | Create 'RawPage'.
--
-- This function may copy data to satisfy internal 'RawPage' invariants.
-- Use 'unsafeMakeRawPage' if you don't want copy.
makeRawPage ::
ByteArray -- ^ bytearray, must contain 4096 bytes (after offset)
-> Int -- ^ offset in bytes, must be 8 byte aligned.
-> RawPage
makeRawPage ba off
| invariant page = page
| otherwise = RawPage 0 $ runByteArray $ do
mba <- newAlignedPinnedByteArray 4096 8
fillByteArray mba 0 4096 0
copyByteArray mba 0 ba off (clamp 0 4096 (sizeofByteArray ba - off))
pure mba
where
page = RawPage (div2 off) ba
clamp l u x = max l (min u x)
unsafeMakeRawPage ::
ByteArray -- ^ bytearray, must be pinned and contain 4096 bytes (after offset)
-> Int -- ^ offset in bytes, must be 8 byte aligned.
-> RawPage
unsafeMakeRawPage ba off = assert (invariant page) page
where
page = RawPage (div2 off) ba
rawPageRawBytes :: RawPage -> RawBytes
rawPageRawBytes (RawPage off ba) =
RB.fromByteArray (mul2 off) 4096 ba
-------------------------------------------------------------------------------
-- Lookup function
-------------------------------------------------------------------------------
data RawPageLookup entry =
-- | The key is not present on the page.
LookupEntryNotPresent
-- | The key is present and corresponds to a normal entry that fits
-- fully within the page.
| LookupEntry !entry
-- | The key is present and corresponds to an entry where the value
-- may have overflowed onto subsequent pages. In this case the entry
-- contains the /prefix/ of the value (that did fit on the page). The
-- length of the suffix is returned separately.
| LookupEntryOverflow !entry !Word32
deriving stock (Eq, Functor, Show)
-- |
-- __Time:__ \( \mathcal{O}\left( \log_2 n \right) \)
-- where \( n \) is the number of entries in the 'RawPage'.
--
-- Return the 'Entry' corresponding to the supplied 'SerialisedKey' if it exists
-- within the 'RawPage'.
rawPageLookup ::
RawPage
-> SerialisedKey
-> RawPageLookup (Entry SerialisedValue BlobSpan)
rawPageLookup !page !key
| dirNumKeys == 1 = lookup1
| otherwise = case bisectPageToKey (fromIntegral dirNumKeys) page key of
KeyWasFoundAt i -> LookupEntry $ rawPageEntryAt page i
_ -> LookupEntryNotPresent
where
!dirNumKeys = rawPageNumKeys page
lookup1
| key == rawPageKeyAt page 0
= let !entry = rawPageEntry1 page
!suffix = rawPageSingleValueSuffix page
in if suffix > 0
then LookupEntryOverflow entry suffix
else LookupEntry entry
| otherwise
= LookupEntryNotPresent
-- |
-- __Time:__ \( \mathcal{O}\left( \log_2 n \right) \)
-- where \( n \) is the number of entries in the 'RawPage'.
--
-- Return the least entry number in the 'RawPage' (if it exists)
-- which is greater than or equal to the supplied 'SerialisedKey'.
--
-- The following law always holds \( \forall \mathtt{key} \mathtt{page} \):
--
-- * maybe True (key <=) (getRawPageIndexKey . rawPageIndex page . id =<< rawPageFindKey page key)
--
-- * maybe True (key > ) (getRawPageIndexKey . rawPageIndex page . pred =<< rawPageFindKey page key)
--
-- * maybe (maximum (rawPageKeys page) < key) (rawPageFindKey page key)
rawPageFindKey ::
RawPage
-> SerialisedKey
-> Maybe Word16 -- ^ entry number of first entry greater or equal to the key
rawPageFindKey !page !key
| dirNumKeys == 1 = lookup1
| otherwise = case bisectPageToKey (fromIntegral dirNumKeys) page key of
KeyNotFoundExceedsPage -> Nothing
KeyNotFoundNearestIsAt i -> Just $ fromIntegral i
KeyWasFoundAt i -> Just $ fromIntegral i
where
!dirNumKeys = rawPageNumKeys page
lookup1
| key <= rawPageKeyAt page 0 = Just 0
| otherwise = Nothing
-- | This data-type facilitates the code reuse of 'bisectPageToKey' between the
-- and 'rawPageLookup' and 'rawPageFindKey'.
data BinarySearchResult
= KeyNotFoundExceedsPage
| KeyNotFoundNearestIsAt {-# UNPACK #-} !Int
| KeyWasFoundAt {-# UNPACK #-} !Int
-- | Binary search procedure shared between 'rawPageLookup' and 'rawPageFindKey'.
{-# INLINE bisectPageToKey #-}
bisectPageToKey ::
Int
-> RawPage
-> SerialisedKey
-> BinarySearchResult
bisectPageToKey !numKeys !page !key = go 0 numKeys
where
-- Binary search for within the range [i, j)
-- for the index k corresponding to key.
--
-- If k > key, search [ k + 1, j )
-- If k < key, search [ i, j )
go :: Int -> Int -> BinarySearchResult
go !i !j
| i >= j =
if i == numKeys
then KeyNotFoundExceedsPage
else KeyNotFoundNearestIsAt j
| otherwise =
let !k = i + div2 (j - i)
in case key `compare` rawPageKeyAt page k of
GT -> go (k + 1) j
EQ -> KeyWasFoundAt k
LT -> go i k
data RawPageIndex entry =
IndexNotPresent
-- | The index is present and corresponds to a normal entry that fits
-- fully within the page (but might be the only entry on the page).
| IndexEntry !SerialisedKey !entry
-- | The index is present and corresponds to an entry where the value
-- has overflowed onto subsequent pages. In this case only prefix entry
-- and the length of the suffix are returned.
-- The caller can copy the full serialised pages themselves.
| IndexEntryOverflow !SerialisedKey !entry !Word32
deriving stock (Eq, Functor, Show)
-- |
-- __Time:__ \( \mathcal{O}\left( 1 \right) \)
--
-- Conveniently access the 'SerialisedKey' of a 'RawPageIndex'.
{-# INLINE getRawPageIndexKey #-}
getRawPageIndexKey :: RawPageIndex e -> Maybe SerialisedKey
getRawPageIndexKey = \case
IndexEntry k _ -> Just k
IndexEntryOverflow k _ _ -> Just k
IndexNotPresent -> Nothing
{-# INLINE rawPageIndex #-}
rawPageIndex ::
RawPage
-> Word16
-> RawPageIndex (Entry SerialisedValue BlobSpan)
rawPageIndex !page !ix
| ix >= dirNumKeys =
IndexNotPresent
| dirNumKeys > 1 =
let ix' = fromIntegral ix
in IndexEntry (rawPageKeyAt page ix') (rawPageEntryAt page ix')
| otherwise =
let key = rawPageKeyAt page 0
entry = rawPageEntry1 page
!suffix = rawPageSingleValueSuffix page
in if suffix <= 0
then IndexEntry key entry
else IndexEntryOverflow key entry suffix
where
!dirNumKeys = rawPageNumKeys page
-- | for non-single key page case
rawPageEntryAt :: RawPage -> Int -> Entry SerialisedValue BlobSpan
rawPageEntryAt page i =
assert (i < fromIntegral (rawPageNumKeys page)) $
assert (rawPageNumKeys page > 1) $
case rawPageOpAt page i of
0 -> if rawPageHasBlobSpanAt page i == 0
then Insert (rawPageValueAt page i)
else InsertWithBlob (rawPageValueAt page i)
(rawPageBlobSpanIndex page
(rawPageCalculateBlobIndex page i))
1 -> Upsert (rawPageValueAt page i)
_ -> Delete
-- | Single key page case
rawPageEntry1 :: RawPage -> Entry SerialisedValue BlobSpan
rawPageEntry1 page =
case rawPageOpAt page 0 of
0 -> if rawPageHasBlobSpanAt page 0 == 0
then Insert (rawPageSingleValuePrefix page)
else InsertWithBlob (rawPageSingleValuePrefix page)
(rawPageBlobSpanIndex page
(rawPageCalculateBlobIndex page 0))
1 -> Upsert (rawPageSingleValuePrefix page)
_ -> Delete
-- | Calculate the number of overflow pages that are expected to follow this
-- page.
--
-- This will be non-zero when the page contains a single key\/op entry that is
-- itself too large to fit within the page.
--
rawPageOverflowPages :: RawPage -> Int
rawPageOverflowPages page
| rawPageNumKeys page == 1
, let (_, end) = rawPageValueOffsets1 page
= fromIntegral (ceilDivPageSize end - 1) -- don't count the first page
| otherwise = 0
-------------------------------------------------------------------------------
-- Accessors
-------------------------------------------------------------------------------
-- Note: indexByteArray's offset is given in elements of type a rather than in bytes.
rawPageNumKeys :: RawPage -> Word16
rawPageNumKeys (RawPage off ba) = indexByteArray ba off
rawPageNumBlobs :: RawPage -> Word16
rawPageNumBlobs (RawPage off ba) = indexByteArray ba (off + 1)
rawPageKeysOffset :: RawPage -> Word16
rawPageKeysOffset (RawPage off ba) = indexByteArray ba (off + 2)
type KeyOffset = Word16
type ValueOffset = Word16
rawPageKeyOffsets :: RawPage -> VP.Vector KeyOffset
rawPageKeyOffsets page@(RawPage off ba) =
mkPrimVector
(off + fromIntegral (div2 dirOffset))
(fromIntegral dirNumKeys + 1)
ba
where
!dirNumKeys = rawPageNumKeys page
!dirOffset = rawPageKeysOffset page
-- | for non-single key page case
rawPageValueOffsets :: RawPage -> VP.Vector ValueOffset
rawPageValueOffsets page@(RawPage off ba) =
assert (dirNumKeys /= 1) $
mkPrimVector
(off + fromIntegral (div2 dirOffset) + fromIntegral dirNumKeys)
(fromIntegral dirNumKeys + 1)
ba
where
!dirNumKeys = rawPageNumKeys page
!dirOffset = rawPageKeysOffset page
-- | single key page case
{-# INLINE rawPageValueOffsets1 #-}
rawPageValueOffsets1 :: RawPage -> (Word16, Word32)
rawPageValueOffsets1 page@(RawPage off ba) =
assert (rawPageNumKeys page == 1) $
( indexByteArray ba (off + fromIntegral (div2 dirOffset) + 1)
, indexByteArray ba (div2 (off + fromIntegral (div2 dirOffset)) + 1)
)
where
!dirOffset = rawPageKeysOffset page
rawPageHasBlobSpanAt :: RawPage -> Int -> Word64
rawPageHasBlobSpanAt _page@(RawPage off ba) i = do
let j = div64 i
let k = mod64 i
let word = indexByteArray ba (div4 off + 1 + j)
unsafeShiftR word k .&. 1
rawPageOpAt :: RawPage -> Int -> Word64
rawPageOpAt page@(RawPage off ba) i = do
let j = div32 i
let k = mod32 i
let word = indexByteArray ba (div4 off + 1 + ceilDiv64 (fromIntegral dirNumKeys) + j)
unsafeShiftR word (mul2 k) .&. 3
where
!dirNumKeys = rawPageNumKeys page
rawPageKeys :: RawPage -> V.Vector SerialisedKey
rawPageKeys page@(RawPage off ba) = do
let offs = rawPageKeyOffsets page
V.fromList
[ SerialisedKey (RB.fromByteArray (mul2 off + start) (end - start) ba)
| i <- [ 0 .. fromIntegral dirNumKeys - 1 ] :: [Int]
, let start = fromIntegral (VP.unsafeIndex offs i) :: Int
, let end = fromIntegral (VP.unsafeIndex offs (i + 1)) :: Int
]
where
!dirNumKeys = rawPageNumKeys page
rawPageKeyAt :: RawPage -> Int -> SerialisedKey
rawPageKeyAt page@(RawPage off ba) i = do
SerialisedKey (RB.fromByteArray (mul2 off + start) (end - start) ba)
where
offs = rawPageKeyOffsets page
start = fromIntegral (VP.unsafeIndex offs i) :: Int
end = fromIntegral (VP.unsafeIndex offs (i + 1)) :: Int
-- | Non-single page case
rawPageValues :: RawPage -> V.Vector SerialisedValue
rawPageValues page@(RawPage off ba) =
let offs = rawPageValueOffsets page in
V.fromList
[ SerialisedValue $ RB.fromByteArray (mul2 off + start) (end - start) ba
| i <- [ 0 .. fromIntegral dirNumKeys - 1 ] :: [Int]
, let start = fromIntegral (VP.unsafeIndex offs i) :: Int
, let end = fromIntegral (VP.unsafeIndex offs (i + 1)) :: Int
]
where
!dirNumKeys = rawPageNumKeys page
-- | Non-single page case
rawPageValueAt :: RawPage -> Int -> SerialisedValue
rawPageValueAt page@(RawPage off ba) i =
SerialisedValue (RB.fromByteArray (mul2 off + start) (end - start) ba)
where
offs = rawPageValueOffsets page
start = fromIntegral (VP.unsafeIndex offs i) :: Int
end = fromIntegral (VP.unsafeIndex offs (i + 1)) :: Int
-- | Single key page case
rawPageSingleValuePrefix :: RawPage -> SerialisedValue
rawPageSingleValuePrefix page@(RawPage off ba) =
SerialisedValue $
RB.fromByteArray
(mul2 off + fromIntegral start)
(fromIntegral prefix_end - fromIntegral start)
ba
where
(start, end) = rawPageValueOffsets1 page
prefix_end = min 4096 end
-- | Single key page case
rawPageSingleValueSuffix :: RawPage -> Word32
rawPageSingleValueSuffix page
| end > 4096 = end - 4096
| otherwise = 0
where
(_, end) = rawPageValueOffsets1 page
-- we could create unboxed array.
{-# INLINE rawPageBlobSpanIndex #-}
rawPageBlobSpanIndex :: RawPage
-> Int -- ^ blobspan index. Calculate with 'rawPageCalculateBlobIndex'
-> BlobSpan
rawPageBlobSpanIndex page@(RawPage off ba) i = BlobSpan
( indexByteArray ba (off1 + i) )
( indexByteArray ba (off2 + i) )
where
!dirNumKeys = rawPageNumKeys page
!dirNumBlobs = rawPageNumBlobs page
-- offset to start of blobspan arr
off1 = div4 off + 1 + ceilDiv64 (fromIntegral dirNumKeys) + ceilDiv64 (fromIntegral (mul2 dirNumKeys))
off2 = mul2 (off1 + fromIntegral dirNumBlobs)
rawPageCalculateBlobIndex ::
RawPage
-> Int -- ^ key index
-> Int -- ^ blobspan index
rawPageCalculateBlobIndex (RawPage off ba) i = do
let j = unsafeShiftR i 6 -- `div` 64
let k = i .&. 63 -- `mod` 64
-- generic sum isn't too great
let s = List.foldl' (+) 0 [ popCount (indexByteArray ba (div4 off + 1 + jj) :: Word64) | jj <- [0 .. j-1 ] ]
let word = indexByteArray ba (div4 off + 1 + j) :: Word64
s + popCount (word .&. complement (unsafeShiftL 0xffffffffffffffff k))