{-# LANGUAGE DerivingStrategies         #-}
{-# LANGUAGE GeneralisedNewtypeDeriving #-}
{-# OPTIONS_HADDOCK not-home #-}

-- | The in-memory LSM level 0.
--
module Database.LSMTree.Internal.WriteBuffer (
    WriteBuffer,
    empty,
    numEntries,
    fromMap,
    toMap,
    fromList,
    toList,
    addEntry,
    null,
    lookups,
    lookup,
    rangeLookups,
) where

import           Control.DeepSeq (NFData (..))
import           Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import qualified Data.Vector as V
import           Database.LSMTree.Internal.BlobRef (BlobSpan)
import           Database.LSMTree.Internal.Entry
import qualified Database.LSMTree.Internal.Map.Range as Map.R
import           Database.LSMTree.Internal.Range (Range (..))
import           Database.LSMTree.Internal.Serialise
import qualified Database.LSMTree.Internal.Vector as V
import           Prelude hiding (lookup, null)

{-------------------------------------------------------------------------------
  Writebuffer type
-------------------------------------------------------------------------------}

newtype WriteBuffer =
  WB { WriteBuffer -> Map SerialisedKey (Entry SerialisedValue BlobSpan)
unWB :: Map SerialisedKey (Entry SerialisedValue BlobSpan) }
  deriving stock (WriteBuffer -> WriteBuffer -> Bool
(WriteBuffer -> WriteBuffer -> Bool)
-> (WriteBuffer -> WriteBuffer -> Bool) -> Eq WriteBuffer
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: WriteBuffer -> WriteBuffer -> Bool
== :: WriteBuffer -> WriteBuffer -> Bool
$c/= :: WriteBuffer -> WriteBuffer -> Bool
/= :: WriteBuffer -> WriteBuffer -> Bool
Eq, Int -> WriteBuffer -> ShowS
[WriteBuffer] -> ShowS
WriteBuffer -> String
(Int -> WriteBuffer -> ShowS)
-> (WriteBuffer -> String)
-> ([WriteBuffer] -> ShowS)
-> Show WriteBuffer
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> WriteBuffer -> ShowS
showsPrec :: Int -> WriteBuffer -> ShowS
$cshow :: WriteBuffer -> String
show :: WriteBuffer -> String
$cshowList :: [WriteBuffer] -> ShowS
showList :: [WriteBuffer] -> ShowS
Show)
  deriving newtype WriteBuffer -> ()
(WriteBuffer -> ()) -> NFData WriteBuffer
forall a. (a -> ()) -> NFData a
$crnf :: WriteBuffer -> ()
rnf :: WriteBuffer -> ()
NFData

empty :: WriteBuffer
empty :: WriteBuffer
empty = Map SerialisedKey (Entry SerialisedValue BlobSpan) -> WriteBuffer
WB Map SerialisedKey (Entry SerialisedValue BlobSpan)
forall k a. Map k a
Map.empty

-- | \( O(1) \)
numEntries :: WriteBuffer -> NumEntries
numEntries :: WriteBuffer -> NumEntries
numEntries (WB Map SerialisedKey (Entry SerialisedValue BlobSpan)
m) = Int -> NumEntries
NumEntries (Map SerialisedKey (Entry SerialisedValue BlobSpan) -> Int
forall k a. Map k a -> Int
Map.size Map SerialisedKey (Entry SerialisedValue BlobSpan)
m)

-- | \( O(1)) \)
fromMap ::
     Map SerialisedKey (Entry SerialisedValue BlobSpan)
  -> WriteBuffer
fromMap :: Map SerialisedKey (Entry SerialisedValue BlobSpan) -> WriteBuffer
fromMap Map SerialisedKey (Entry SerialisedValue BlobSpan)
m = Map SerialisedKey (Entry SerialisedValue BlobSpan) -> WriteBuffer
WB Map SerialisedKey (Entry SerialisedValue BlobSpan)
m

-- | \( O(1) \)
toMap :: WriteBuffer -> Map SerialisedKey (Entry SerialisedValue BlobSpan)
toMap :: WriteBuffer -> Map SerialisedKey (Entry SerialisedValue BlobSpan)
toMap = WriteBuffer -> Map SerialisedKey (Entry SerialisedValue BlobSpan)
unWB

-- | \( O(n \log n) \)
fromList ::
     (SerialisedValue -> SerialisedValue -> SerialisedValue) -- ^ merge function
  -> [(SerialisedKey, Entry SerialisedValue BlobSpan)]
  -> WriteBuffer
fromList :: (SerialisedValue -> SerialisedValue -> SerialisedValue)
-> [(SerialisedKey, Entry SerialisedValue BlobSpan)] -> WriteBuffer
fromList SerialisedValue -> SerialisedValue -> SerialisedValue
f [(SerialisedKey, Entry SerialisedValue BlobSpan)]
es = Map SerialisedKey (Entry SerialisedValue BlobSpan) -> WriteBuffer
WB (Map SerialisedKey (Entry SerialisedValue BlobSpan) -> WriteBuffer)
-> Map SerialisedKey (Entry SerialisedValue BlobSpan)
-> WriteBuffer
forall a b. (a -> b) -> a -> b
$ (Entry SerialisedValue BlobSpan
 -> Entry SerialisedValue BlobSpan
 -> Entry SerialisedValue BlobSpan)
-> [(SerialisedKey, Entry SerialisedValue BlobSpan)]
-> Map SerialisedKey (Entry SerialisedValue BlobSpan)
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith ((SerialisedValue -> SerialisedValue -> SerialisedValue)
-> Entry SerialisedValue BlobSpan
-> Entry SerialisedValue BlobSpan
-> Entry SerialisedValue BlobSpan
forall v b. (v -> v -> v) -> Entry v b -> Entry v b -> Entry v b
combine SerialisedValue -> SerialisedValue -> SerialisedValue
f) [(SerialisedKey, Entry SerialisedValue BlobSpan)]
es

-- | \( O(n) \)
toList :: WriteBuffer -> [(SerialisedKey, Entry SerialisedValue BlobSpan)]
toList :: WriteBuffer -> [(SerialisedKey, Entry SerialisedValue BlobSpan)]
toList (WB Map SerialisedKey (Entry SerialisedValue BlobSpan)
m) = Map SerialisedKey (Entry SerialisedValue BlobSpan)
-> [(SerialisedKey, Entry SerialisedValue BlobSpan)]
forall k a. Map k a -> [(k, a)]
Map.assocs Map SerialisedKey (Entry SerialisedValue BlobSpan)
m

{-------------------------------------------------------------------------------
  Updates
-------------------------------------------------------------------------------}

addEntry ::
     (SerialisedValue -> SerialisedValue -> SerialisedValue) -- ^ merge function
  -> SerialisedKey
  -> Entry SerialisedValue BlobSpan
  -> WriteBuffer
  -> WriteBuffer
addEntry :: (SerialisedValue -> SerialisedValue -> SerialisedValue)
-> SerialisedKey
-> Entry SerialisedValue BlobSpan
-> WriteBuffer
-> WriteBuffer
addEntry SerialisedValue -> SerialisedValue -> SerialisedValue
f SerialisedKey
k Entry SerialisedValue BlobSpan
e (WB Map SerialisedKey (Entry SerialisedValue BlobSpan)
wb) =
    Map SerialisedKey (Entry SerialisedValue BlobSpan) -> WriteBuffer
WB ((Entry SerialisedValue BlobSpan
 -> Entry SerialisedValue BlobSpan
 -> Entry SerialisedValue BlobSpan)
-> SerialisedKey
-> Entry SerialisedValue BlobSpan
-> Map SerialisedKey (Entry SerialisedValue BlobSpan)
-> Map SerialisedKey (Entry SerialisedValue BlobSpan)
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith ((SerialisedValue -> SerialisedValue -> SerialisedValue)
-> Entry SerialisedValue BlobSpan
-> Entry SerialisedValue BlobSpan
-> Entry SerialisedValue BlobSpan
forall v b. (v -> v -> v) -> Entry v b -> Entry v b -> Entry v b
combine SerialisedValue -> SerialisedValue -> SerialisedValue
f) SerialisedKey
k Entry SerialisedValue BlobSpan
e Map SerialisedKey (Entry SerialisedValue BlobSpan)
wb)

{-------------------------------------------------------------------------------
  Querying
-------------------------------------------------------------------------------}

null :: WriteBuffer -> Bool
null :: WriteBuffer -> Bool
null (WB Map SerialisedKey (Entry SerialisedValue BlobSpan)
m) = Map SerialisedKey (Entry SerialisedValue BlobSpan) -> Bool
forall k a. Map k a -> Bool
Map.null Map SerialisedKey (Entry SerialisedValue BlobSpan)
m

-- We return an 'Entry' with serialised values, so it can be properly combined
-- with the lookups in other runs. Deserialisation only occurs afterwards.
--
-- Note: the entry may be 'Delete'.
--
lookups ::
     WriteBuffer
  -> V.Vector SerialisedKey
  -> V.Vector (Maybe (Entry SerialisedValue BlobSpan))
lookups :: WriteBuffer
-> Vector SerialisedKey
-> Vector (Maybe (Entry SerialisedValue BlobSpan))
lookups (WB !Map SerialisedKey (Entry SerialisedValue BlobSpan)
m) !Vector SerialisedKey
ks = (SerialisedKey -> Maybe (Entry SerialisedValue BlobSpan))
-> Vector SerialisedKey
-> Vector (Maybe (Entry SerialisedValue BlobSpan))
forall a b. (a -> b) -> Vector a -> Vector b
V.mapStrict (SerialisedKey
-> Map SerialisedKey (Entry SerialisedValue BlobSpan)
-> Maybe (Entry SerialisedValue BlobSpan)
forall k a. Ord k => k -> Map k a -> Maybe a
`Map.lookup` Map SerialisedKey (Entry SerialisedValue BlobSpan)
m) Vector SerialisedKey
ks

lookup ::
     WriteBuffer
  -> SerialisedKey
  -> Maybe (Entry SerialisedValue BlobSpan)
lookup :: WriteBuffer
-> SerialisedKey -> Maybe (Entry SerialisedValue BlobSpan)
lookup (WB !Map SerialisedKey (Entry SerialisedValue BlobSpan)
m) !SerialisedKey
k = SerialisedKey
-> Map SerialisedKey (Entry SerialisedValue BlobSpan)
-> Maybe (Entry SerialisedValue BlobSpan)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup SerialisedKey
k Map SerialisedKey (Entry SerialisedValue BlobSpan)
m

{-------------------------------------------------------------------------------
  RangeQueries
-------------------------------------------------------------------------------}

-- | We return 'Entry' instead of either @QueryResult@,
-- so we can properly combine lookup results.
--
-- Note: 'Delete's are not filtered out.
--
rangeLookups ::
     WriteBuffer
  -> Range SerialisedKey
  -> [(SerialisedKey, Entry SerialisedValue BlobSpan)]
rangeLookups :: WriteBuffer
-> Range SerialisedKey
-> [(SerialisedKey, Entry SerialisedValue BlobSpan)]
rangeLookups (WB Map SerialisedKey (Entry SerialisedValue BlobSpan)
m) Range SerialisedKey
r =
    [ (SerialisedKey
k, Entry SerialisedValue BlobSpan
e)
    | let (Bound SerialisedKey
lb, Bound SerialisedKey
ub) = Range SerialisedKey -> (Bound SerialisedKey, Bound SerialisedKey)
forall k. Range k -> (Bound k, Bound k)
convertRange Range SerialisedKey
r
    , (SerialisedKey
k, Entry SerialisedValue BlobSpan
e) <- Bound SerialisedKey
-> Bound SerialisedKey
-> Map SerialisedKey (Entry SerialisedValue BlobSpan)
-> [(SerialisedKey, Entry SerialisedValue BlobSpan)]
forall k v. Ord k => Bound k -> Bound k -> Map k v -> [(k, v)]
Map.R.rangeLookup Bound SerialisedKey
lb Bound SerialisedKey
ub Map SerialisedKey (Entry SerialisedValue BlobSpan)
m
    ]

convertRange :: Range k -> (Map.R.Bound k, Map.R.Bound k)
convertRange :: forall k. Range k -> (Bound k, Bound k)
convertRange (FromToExcluding k
lb k
ub) =
    ( k -> Clusive -> Bound k
forall k. k -> Clusive -> Bound k
Map.R.Bound k
lb Clusive
Map.R.Inclusive
    , k -> Clusive -> Bound k
forall k. k -> Clusive -> Bound k
Map.R.Bound k
ub Clusive
Map.R.Exclusive )
convertRange (FromToIncluding k
lb k
ub) =
    ( k -> Clusive -> Bound k
forall k. k -> Clusive -> Bound k
Map.R.Bound k
lb Clusive
Map.R.Inclusive
    , k -> Clusive -> Bound k
forall k. k -> Clusive -> Bound k
Map.R.Bound k
ub Clusive
Map.R.Inclusive )