{-# LANGUAGE CPP #-}
{-# OPTIONS_HADDOCK not-home #-}
{- HLINT ignore "Avoid restricted alias" -}

{-|
    Incremental construction functionality for the general-purpose fence pointer
    index.
-}
module Database.LSMTree.Internal.Index.OrdinaryAcc
(
    IndexOrdinaryAcc (IndexOrdinaryAcc),
    new,
    newWithDefaults,
    appendSingle,
    appendMulti,
    unsafeEnd
)
where

import           Prelude hiding (take)

import           Control.Exception (assert)
import           Control.Monad.ST.Strict (ST)
import           Data.Maybe (maybeToList)
import qualified Data.Vector.Primitive as Primitive (Vector, length)
import           Data.Word (Word16, Word32, Word8)
import           Database.LSMTree.Internal.Chunk (Baler, Chunk, createBaler,
                     feedBaler, unsafeEndBaler)
import           Database.LSMTree.Internal.Index.Ordinary
                     (IndexOrdinary (IndexOrdinary))
import           Database.LSMTree.Internal.Serialise
                     (SerialisedKey (SerialisedKey'))
import           Database.LSMTree.Internal.Unsliced (Unsliced, makeUnslicedKey)
import           Database.LSMTree.Internal.Vector (byteVectorFromPrim)
import           Database.LSMTree.Internal.Vector.Growing (GrowingVector)
import qualified Database.LSMTree.Internal.Vector.Growing as Growing (append,
                     freeze, new)
#ifdef NO_IGNORE_ASSERTS
import           Database.LSMTree.Internal.Unsliced (fromUnslicedKey)
import qualified Database.LSMTree.Internal.Vector.Growing as Growing
                     (readMaybeLast)
#endif

{-|
    A general-purpose fence pointer index under incremental construction.

    A value @IndexOrdinaryAcc unslicedLastKeys baler@ denotes a partially
    constructed index that assigns keys to pages according to @unslicedLastKeys@
    and uses @baler@ for incremental output of the serialised key list.
-}
data IndexOrdinaryAcc s = IndexOrdinaryAcc
                              !(GrowingVector s (Unsliced SerialisedKey))
                              !(Baler s)

-- | Creates a new, initially empty, index.
new :: Int                       -- ^ Initial size of the key buffer
    -> Int                       -- ^ Minimum chunk size in bytes
    -> ST s (IndexOrdinaryAcc s) -- ^ Construction of the index
new :: forall s. Int -> Int -> ST s (IndexOrdinaryAcc s)
new Int
initialKeyBufferSize Int
minChunkSize = GrowingVector s (Unsliced SerialisedKey)
-> Baler s -> IndexOrdinaryAcc s
forall s.
GrowingVector s (Unsliced SerialisedKey)
-> Baler s -> IndexOrdinaryAcc s
IndexOrdinaryAcc                 (GrowingVector s (Unsliced SerialisedKey)
 -> Baler s -> IndexOrdinaryAcc s)
-> ST s (GrowingVector s (Unsliced SerialisedKey))
-> ST s (Baler s -> IndexOrdinaryAcc s)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
                                        Int -> ST s (GrowingVector s (Unsliced SerialisedKey))
forall s a. Int -> ST s (GrowingVector s a)
Growing.new Int
initialKeyBufferSize ST s (Baler s -> IndexOrdinaryAcc s)
-> ST s (Baler s) -> ST s (IndexOrdinaryAcc s)
forall a b. ST s (a -> b) -> ST s a -> ST s b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*>
                                        Int -> ST s (Baler s)
forall s. Int -> ST s (Baler s)
createBaler Int
minChunkSize

{-|
    For a specification of this operation, see the documentation of [its
    type-agnostic version]('Database.LSMTree.Internal.Index.newWithDefaults').
-}
newWithDefaults :: ST s (IndexOrdinaryAcc s)
newWithDefaults :: forall s. ST s (IndexOrdinaryAcc s)
newWithDefaults = Int -> Int -> ST s (IndexOrdinaryAcc s)
forall s. Int -> Int -> ST s (IndexOrdinaryAcc s)
new Int
1024 Int
4096

-- | Yields the serialisation of an element of a key list.
keyListElem :: SerialisedKey -> [Primitive.Vector Word8]
keyListElem :: SerialisedKey -> [Vector Word8]
keyListElem (SerialisedKey' Vector Word8
keyBytes) = [Vector Word8
keySizeBytes, Vector Word8
keyBytes] where

    keySize :: Int
    !keySize :: Int
keySize = Vector Word8 -> Int
forall a. Prim a => Vector a -> Int
Primitive.length Vector Word8
keyBytes

    keySizeAsWord16 :: Word16
    !keySizeAsWord16 :: Word16
keySizeAsWord16 = Bool -> Word16 -> Word16
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
keySize Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Word16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word16
forall a. Bounded a => a
maxBound :: Word16)) (Word16 -> Word16) -> Word16 -> Word16
forall a b. (a -> b) -> a -> b
$
                       Int -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
keySize

    keySizeBytes :: Primitive.Vector Word8
    !keySizeBytes :: Vector Word8
keySizeBytes = Word16 -> Vector Word8
forall a. Prim a => a -> Vector Word8
byteVectorFromPrim Word16
keySizeAsWord16

{-|
    For a specification of this operation, see the documentation of [its
    type-agnostic version]('Database.LSMTree.Internal.Index.appendSingle').
-}
appendSingle :: (SerialisedKey, SerialisedKey)
             -> IndexOrdinaryAcc s
             -> ST s (Maybe Chunk)
appendSingle :: forall s.
(SerialisedKey, SerialisedKey)
-> IndexOrdinaryAcc s -> ST s (Maybe Chunk)
appendSingle (SerialisedKey
firstKey, SerialisedKey
lastKey) (IndexOrdinaryAcc GrowingVector s (Unsliced SerialisedKey)
unslicedLastKeys Baler s
baler)
    = Bool -> ST s (Maybe Chunk) -> ST s (Maybe Chunk)
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (SerialisedKey
firstKey SerialisedKey -> SerialisedKey -> Bool
forall a. Ord a => a -> a -> Bool
<= SerialisedKey
lastKey) (ST s (Maybe Chunk) -> ST s (Maybe Chunk))
-> ST s (Maybe Chunk) -> ST s (Maybe Chunk)
forall a b. (a -> b) -> a -> b
$
      do
#ifdef NO_IGNORE_ASSERTS
          maybeLastUnslicedLastKey <- Growing.readMaybeLast unslicedLastKeys
          assert
              (all (< firstKey) (fromUnslicedKey <$> maybeLastUnslicedLastKey))
              (return ())
#endif
          GrowingVector s (Unsliced SerialisedKey)
-> Int -> Unsliced SerialisedKey -> ST s ()
forall s a. GrowingVector s a -> Int -> a -> ST s ()
Growing.append GrowingVector s (Unsliced SerialisedKey)
unslicedLastKeys Int
1 (SerialisedKey -> Unsliced SerialisedKey
makeUnslicedKey SerialisedKey
lastKey)
          [Vector Word8] -> Baler s -> ST s (Maybe Chunk)
forall s. [Vector Word8] -> Baler s -> ST s (Maybe Chunk)
feedBaler (SerialisedKey -> [Vector Word8]
keyListElem SerialisedKey
lastKey) Baler s
baler

{-|
    For a specification of this operation, see the documentation of [its
    type-agnostic version]('Database.LSMTree.Internal.Index.appendMulti').
-}
appendMulti :: (SerialisedKey, Word32)
            -> IndexOrdinaryAcc s
            -> ST s [Chunk]
appendMulti :: forall s.
(SerialisedKey, Word32) -> IndexOrdinaryAcc s -> ST s [Chunk]
appendMulti (SerialisedKey
key, Word32
overflowPageCount) (IndexOrdinaryAcc GrowingVector s (Unsliced SerialisedKey)
unslicedLastKeys Baler s
baler)
    = do
#ifdef NO_IGNORE_ASSERTS
          maybeLastUnslicedLastKey <- Growing.readMaybeLast unslicedLastKeys
          assert (all (< key) (fromUnslicedKey <$> maybeLastUnslicedLastKey))
                 (return ())
#endif
          GrowingVector s (Unsliced SerialisedKey)
-> Int -> Unsliced SerialisedKey -> ST s ()
forall s a. GrowingVector s a -> Int -> a -> ST s ()
Growing.append GrowingVector s (Unsliced SerialisedKey)
unslicedLastKeys Int
pageCount (SerialisedKey -> Unsliced SerialisedKey
makeUnslicedKey SerialisedKey
key)
          Maybe Chunk -> [Chunk]
forall a. Maybe a -> [a]
maybeToList (Maybe Chunk -> [Chunk]) -> ST s (Maybe Chunk) -> ST s [Chunk]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Vector Word8] -> Baler s -> ST s (Maybe Chunk)
forall s. [Vector Word8] -> Baler s -> ST s (Maybe Chunk)
feedBaler [Vector Word8]
keyListElems Baler s
baler
    where

    pageCount :: Int
    !pageCount :: Int
pageCount = Int -> Int
forall a. Enum a => a -> a
succ (Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
overflowPageCount)

    keyListElems :: [Primitive.Vector Word8]
    keyListElems :: [Vector Word8]
keyListElems = [[Vector Word8]] -> [Vector Word8]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (Int -> [Vector Word8] -> [[Vector Word8]]
forall a. Int -> a -> [a]
replicate Int
pageCount (SerialisedKey -> [Vector Word8]
keyListElem SerialisedKey
key))

{-|
    For a specification of this operation, see the documentation of [its
    type-agnostic version]('Database.LSMTree.Internal.Index.unsafeEnd').
-}
unsafeEnd :: IndexOrdinaryAcc s -> ST s (Maybe Chunk, IndexOrdinary)
unsafeEnd :: forall s. IndexOrdinaryAcc s -> ST s (Maybe Chunk, IndexOrdinary)
unsafeEnd (IndexOrdinaryAcc GrowingVector s (Unsliced SerialisedKey)
unslicedLastKeys Baler s
baler) = do
    Vector (Unsliced SerialisedKey)
frozenUnslicedLastKeys <- GrowingVector s (Unsliced SerialisedKey)
-> ST s (Vector (Unsliced SerialisedKey))
forall s a. GrowingVector s a -> ST s (Vector a)
Growing.freeze GrowingVector s (Unsliced SerialisedKey)
unslicedLastKeys
    Maybe Chunk
remnant <- Baler s -> ST s (Maybe Chunk)
forall s. Baler s -> ST s (Maybe Chunk)
unsafeEndBaler Baler s
baler
    (Maybe Chunk, IndexOrdinary) -> ST s (Maybe Chunk, IndexOrdinary)
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Chunk
remnant, Vector (Unsliced SerialisedKey) -> IndexOrdinary
IndexOrdinary Vector (Unsliced SerialisedKey)
frozenUnslicedLastKeys)