{-# OPTIONS_HADDOCK not-home #-}

module Database.LSMTree.Internal.ChecksumHandle
  (
    -- * Checksum handles
    -- $checksum-handles
    ChecksumHandle (..),
    makeHandle,
    readChecksum,
    dropCache,
    closeHandle,
    writeToHandle,
    -- * Specialised writers
    writeRawPage,
    writeRawOverflowPages,
    writeBlob,
    copyBlob,
    writeFilter,
    writeIndexHeader,
    writeIndexChunk,
    writeIndexFinal,
  ) where

import           Control.Monad.Class.MonadSTM (MonadSTM (..))
import           Control.Monad.Class.MonadThrow (MonadThrow)
import           Control.Monad.Primitive
import           Data.BloomFilter (Bloom)
import qualified Data.ByteString.Lazy as BSL
import           Data.Primitive.PrimVar
import           Data.Word (Word64)
import           Database.LSMTree.Internal.BlobRef (BlobSpan (..), RawBlobRef)
import qualified Database.LSMTree.Internal.BlobRef as BlobRef
import           Database.LSMTree.Internal.BloomFilter (bloomFilterToLBS)
import           Database.LSMTree.Internal.Chunk (Chunk)
import qualified Database.LSMTree.Internal.Chunk as Chunk (toByteString)
import           Database.LSMTree.Internal.CRC32C (CRC32C)
import qualified Database.LSMTree.Internal.CRC32C as CRC
import           Database.LSMTree.Internal.Entry
import           Database.LSMTree.Internal.Index (Index, IndexType)
import qualified Database.LSMTree.Internal.Index as Index (finalLBS, headerLBS)
import           Database.LSMTree.Internal.Paths (ForBlob (..), ForFilter (..),
                     ForIndex (..), ForKOps (..))
import qualified Database.LSMTree.Internal.RawBytes as RB
import           Database.LSMTree.Internal.RawOverflowPage (RawOverflowPage)
import qualified Database.LSMTree.Internal.RawOverflowPage as RawOverflowPage
import           Database.LSMTree.Internal.RawPage (RawPage)
import qualified Database.LSMTree.Internal.RawPage as RawPage
import           Database.LSMTree.Internal.Serialise
import qualified System.FS.API as FS
import           System.FS.API
import qualified System.FS.BlockIO.API as FS
import           System.FS.BlockIO.API (HasBlockIO)

{-------------------------------------------------------------------------------
  ChecksumHandle
-------------------------------------------------------------------------------}

{- $checksum-handles
  A handle ('ChecksumHandle') that maintains a running CRC32 checksum.
-}

-- | Tracks the checksum of a (write mode) file handle.
data ChecksumHandle s h = ChecksumHandle !(FS.Handle h) !(PrimVar s CRC32C)

{-# SPECIALISE makeHandle ::
     HasFS IO h
  -> FS.FsPath
  -> IO (ChecksumHandle RealWorld h) #-}
makeHandle ::
     (MonadSTM m, PrimMonad m)
  => HasFS m h
  -> FS.FsPath
  -> m (ChecksumHandle (PrimState m) h)
makeHandle :: forall (m :: * -> *) h.
(MonadSTM m, PrimMonad m) =>
HasFS m h -> FsPath -> m (ChecksumHandle (PrimState m) h)
makeHandle HasFS m h
fs FsPath
path =
    Handle h
-> PrimVar (PrimState m) CRC32C -> ChecksumHandle (PrimState m) h
forall s h. Handle h -> PrimVar s CRC32C -> ChecksumHandle s h
ChecksumHandle
      (Handle h
 -> PrimVar (PrimState m) CRC32C -> ChecksumHandle (PrimState m) h)
-> m (Handle h)
-> m (PrimVar (PrimState m) CRC32C
      -> ChecksumHandle (PrimState m) h)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> HasFS m h -> HasCallStack => FsPath -> OpenMode -> m (Handle h)
forall (m :: * -> *) h.
HasFS m h -> HasCallStack => FsPath -> OpenMode -> m (Handle h)
FS.hOpen HasFS m h
fs FsPath
path (AllowExisting -> OpenMode
FS.WriteMode AllowExisting
FS.MustBeNew)
      m (PrimVar (PrimState m) CRC32C -> ChecksumHandle (PrimState m) h)
-> m (PrimVar (PrimState m) CRC32C)
-> m (ChecksumHandle (PrimState m) h)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> CRC32C -> m (PrimVar (PrimState m) CRC32C)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
a -> m (PrimVar (PrimState m) a)
newPrimVar CRC32C
CRC.initialCRC32C

{-# SPECIALISE readChecksum ::
     ChecksumHandle RealWorld h
  -> IO CRC32C #-}
readChecksum ::
     PrimMonad m
  => ChecksumHandle (PrimState m) h
  -> m CRC32C
readChecksum :: forall (m :: * -> *) h.
PrimMonad m =>
ChecksumHandle (PrimState m) h -> m CRC32C
readChecksum (ChecksumHandle Handle h
_h PrimVar (PrimState m) CRC32C
checksum) = PrimVar (PrimState m) CRC32C -> m CRC32C
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimVar (PrimState m) a -> m a
readPrimVar PrimVar (PrimState m) CRC32C
checksum

dropCache :: HasBlockIO m h -> ChecksumHandle (PrimState m) h -> m ()
dropCache :: forall (m :: * -> *) h.
HasBlockIO m h -> ChecksumHandle (PrimState m) h -> m ()
dropCache HasBlockIO m h
hbio (ChecksumHandle Handle h
h PrimVar (PrimState m) CRC32C
_) = HasBlockIO m h -> Handle h -> m ()
forall (m :: * -> *) h. HasBlockIO m h -> Handle h -> m ()
FS.hDropCacheAll HasBlockIO m h
hbio Handle h
h

closeHandle :: HasFS m h -> ChecksumHandle (PrimState m) h -> m ()
closeHandle :: forall (m :: * -> *) h.
HasFS m h -> ChecksumHandle (PrimState m) h -> m ()
closeHandle HasFS m h
fs (ChecksumHandle Handle h
h PrimVar (PrimState m) CRC32C
_checksum) = HasFS m h -> HasCallStack => Handle h -> m ()
forall (m :: * -> *) h.
HasFS m h -> HasCallStack => Handle h -> m ()
FS.hClose HasFS m h
fs Handle h
h

{-# SPECIALISE writeToHandle ::
     HasFS IO h
  -> ChecksumHandle RealWorld h
  -> BSL.ByteString
  -> IO () #-}
writeToHandle ::
     (MonadSTM m, PrimMonad m)
  => HasFS m h
  -> ChecksumHandle (PrimState m) h
  -> BSL.ByteString
  -> m ()
writeToHandle :: forall (m :: * -> *) h.
(MonadSTM m, PrimMonad m) =>
HasFS m h -> ChecksumHandle (PrimState m) h -> ByteString -> m ()
writeToHandle HasFS m h
fs (ChecksumHandle Handle h
h PrimVar (PrimState m) CRC32C
checksum) ByteString
lbs = do
    CRC32C
crc <- PrimVar (PrimState m) CRC32C -> m CRC32C
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimVar (PrimState m) a -> m a
readPrimVar PrimVar (PrimState m) CRC32C
checksum
    (Word64
_, CRC32C
crc') <- HasFS m h -> Handle h -> ByteString -> CRC32C -> m (Word64, CRC32C)
forall (m :: * -> *) h.
Monad m =>
HasFS m h -> Handle h -> ByteString -> CRC32C -> m (Word64, CRC32C)
CRC.hPutAllChunksCRC32C HasFS m h
fs Handle h
h ByteString
lbs CRC32C
crc
    PrimVar (PrimState m) CRC32C -> CRC32C -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimVar (PrimState m) a -> a -> m ()
writePrimVar PrimVar (PrimState m) CRC32C
checksum CRC32C
crc'

{-------------------------------------------------------------------------------
  Specialised Writers for ChecksumHandle
-------------------------------------------------------------------------------}

{-# SPECIALISE writeRawPage ::
     HasFS IO h
  -> ForKOps (ChecksumHandle RealWorld h)
  -> RawPage
  -> IO () #-}
writeRawPage ::
     (MonadSTM m, PrimMonad m)
  => HasFS m h
  -> ForKOps (ChecksumHandle (PrimState m) h)
  -> RawPage
  -> m ()
writeRawPage :: forall (m :: * -> *) h.
(MonadSTM m, PrimMonad m) =>
HasFS m h
-> ForKOps (ChecksumHandle (PrimState m) h) -> RawPage -> m ()
writeRawPage HasFS m h
hfs ForKOps (ChecksumHandle (PrimState m) h)
kOpsHandle =
    HasFS m h -> ChecksumHandle (PrimState m) h -> ByteString -> m ()
forall (m :: * -> *) h.
(MonadSTM m, PrimMonad m) =>
HasFS m h -> ChecksumHandle (PrimState m) h -> ByteString -> m ()
writeToHandle HasFS m h
hfs (ForKOps (ChecksumHandle (PrimState m) h)
-> ChecksumHandle (PrimState m) h
forall a. ForKOps a -> a
unForKOps ForKOps (ChecksumHandle (PrimState m) h)
kOpsHandle)
  (ByteString -> m ()) -> (RawPage -> ByteString) -> RawPage -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
BSL.fromStrict
  (ByteString -> ByteString)
-> (RawPage -> ByteString) -> RawPage -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HasCallStack => RawBytes -> ByteString
RawBytes -> ByteString
RB.unsafePinnedToByteString -- 'RawPage' is guaranteed to be pinned
  (RawBytes -> ByteString)
-> (RawPage -> RawBytes) -> RawPage -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RawPage -> RawBytes
RawPage.rawPageRawBytes

{-# SPECIALISE writeRawOverflowPages ::
     HasFS IO h
  -> ForKOps (ChecksumHandle RealWorld h)
  -> [RawOverflowPage]
  -> IO () #-}
writeRawOverflowPages ::
     (MonadSTM m, PrimMonad m)
  => HasFS m h
  -> ForKOps (ChecksumHandle (PrimState m) h)
  -> [RawOverflowPage]
  -> m ()
writeRawOverflowPages :: forall (m :: * -> *) h.
(MonadSTM m, PrimMonad m) =>
HasFS m h
-> ForKOps (ChecksumHandle (PrimState m) h)
-> [RawOverflowPage]
-> m ()
writeRawOverflowPages HasFS m h
hfs ForKOps (ChecksumHandle (PrimState m) h)
kOpsHandle =
    HasFS m h -> ChecksumHandle (PrimState m) h -> ByteString -> m ()
forall (m :: * -> *) h.
(MonadSTM m, PrimMonad m) =>
HasFS m h -> ChecksumHandle (PrimState m) h -> ByteString -> m ()
writeToHandle HasFS m h
hfs (ForKOps (ChecksumHandle (PrimState m) h)
-> ChecksumHandle (PrimState m) h
forall a. ForKOps a -> a
unForKOps ForKOps (ChecksumHandle (PrimState m) h)
kOpsHandle)
  (ByteString -> m ())
-> ([RawOverflowPage] -> ByteString) -> [RawOverflowPage] -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [ByteString] -> ByteString
BSL.fromChunks
  ([ByteString] -> ByteString)
-> ([RawOverflowPage] -> [ByteString])
-> [RawOverflowPage]
-> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (RawOverflowPage -> ByteString)
-> [RawOverflowPage] -> [ByteString]
forall a b. (a -> b) -> [a] -> [b]
map (RawOverflowPage -> ByteString
RawOverflowPage.rawOverflowPageToByteString)

{-# SPECIALISE writeBlob ::
     HasFS IO h
  -> PrimVar RealWorld Word64
  -> ForBlob (ChecksumHandle RealWorld h)
  -> SerialisedBlob
  -> IO BlobSpan #-}
writeBlob ::
     (MonadSTM m, PrimMonad m)
  => HasFS m h
  -> PrimVar (PrimState m) Word64
  -> ForBlob (ChecksumHandle (PrimState m) h)
  -> SerialisedBlob
  -> m BlobSpan
writeBlob :: forall (m :: * -> *) h.
(MonadSTM m, PrimMonad m) =>
HasFS m h
-> PrimVar (PrimState m) Word64
-> ForBlob (ChecksumHandle (PrimState m) h)
-> SerialisedBlob
-> m BlobSpan
writeBlob HasFS m h
hfs PrimVar (PrimState m) Word64
blobOffset ForBlob (ChecksumHandle (PrimState m) h)
blobHandle SerialisedBlob
blob = do
    -- NOTE: This is different from BlobFile.writeBlob. This is because BlobFile
    --  internalises a regular Handle, rather than a ChecksumHandle. These two
    --  functions cannot be easily unified, because BlobFile.writeBlob permits
    --  writing blobs to arbitrary positions in the blob file, whereas, by the
    --  very nature of CRC32 checksums, ChecksumHandle.writeBlob only supports
    --  sequential writes.
    let size :: Word64
size = SerialisedBlob -> Word64
sizeofBlob64 SerialisedBlob
blob
    Word64
offset <- PrimVar (PrimState m) Word64 -> m Word64
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimVar (PrimState m) a -> m a
readPrimVar PrimVar (PrimState m) Word64
blobOffset
    PrimVar (PrimState m) Word64 -> (Word64 -> Word64) -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimVar (PrimState m) a -> (a -> a) -> m ()
modifyPrimVar PrimVar (PrimState m) Word64
blobOffset (Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+Word64
size)
    let SerialisedBlob RawBytes
rb = SerialisedBlob
blob
    let lbs :: ByteString
lbs = ByteString -> ByteString
BSL.fromStrict (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ RawBytes -> ByteString
RB.toByteString RawBytes
rb
    HasFS m h -> ChecksumHandle (PrimState m) h -> ByteString -> m ()
forall (m :: * -> *) h.
(MonadSTM m, PrimMonad m) =>
HasFS m h -> ChecksumHandle (PrimState m) h -> ByteString -> m ()
writeToHandle HasFS m h
hfs (ForBlob (ChecksumHandle (PrimState m) h)
-> ChecksumHandle (PrimState m) h
forall a. ForBlob a -> a
unForBlob ForBlob (ChecksumHandle (PrimState m) h)
blobHandle) ByteString
lbs
    BlobSpan -> m BlobSpan
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Word64 -> Word32 -> BlobSpan
BlobSpan Word64
offset (Word64 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
size))

{-# SPECIALISE copyBlob ::
     HasFS IO h
  -> PrimVar RealWorld Word64
  -> ForBlob (ChecksumHandle RealWorld h)
  -> RawBlobRef IO h
  -> IO BlobSpan #-}
copyBlob ::
     (MonadSTM m, MonadThrow m, PrimMonad m)
  => HasFS m h
  -> PrimVar (PrimState m) Word64
  -> ForBlob (ChecksumHandle (PrimState m) h)
  -> RawBlobRef m h
  -> m BlobSpan
copyBlob :: forall (m :: * -> *) h.
(MonadSTM m, MonadThrow m, PrimMonad m) =>
HasFS m h
-> PrimVar (PrimState m) Word64
-> ForBlob (ChecksumHandle (PrimState m) h)
-> RawBlobRef m h
-> m BlobSpan
copyBlob HasFS m h
hfs PrimVar (PrimState m) Word64
blobOffset ForBlob (ChecksumHandle (PrimState m) h)
blobHandle RawBlobRef m h
blobref = do
    SerialisedBlob
blob <- HasFS m h -> RawBlobRef m h -> m SerialisedBlob
forall (m :: * -> *) h.
(MonadThrow m, PrimMonad m) =>
HasFS m h -> RawBlobRef m h -> m SerialisedBlob
BlobRef.readRawBlobRef HasFS m h
hfs RawBlobRef m h
blobref
    HasFS m h
-> PrimVar (PrimState m) Word64
-> ForBlob (ChecksumHandle (PrimState m) h)
-> SerialisedBlob
-> m BlobSpan
forall (m :: * -> *) h.
(MonadSTM m, PrimMonad m) =>
HasFS m h
-> PrimVar (PrimState m) Word64
-> ForBlob (ChecksumHandle (PrimState m) h)
-> SerialisedBlob
-> m BlobSpan
writeBlob HasFS m h
hfs PrimVar (PrimState m) Word64
blobOffset ForBlob (ChecksumHandle (PrimState m) h)
blobHandle SerialisedBlob
blob

{-# SPECIALISE writeFilter ::
     HasFS IO h
  -> ForFilter (ChecksumHandle RealWorld h)
  -> Bloom SerialisedKey
  -> IO () #-}
writeFilter ::
     (MonadSTM m, PrimMonad m)
  => HasFS m h
  -> ForFilter (ChecksumHandle (PrimState m) h)
  -> Bloom SerialisedKey
  -> m ()
writeFilter :: forall (m :: * -> *) h.
(MonadSTM m, PrimMonad m) =>
HasFS m h
-> ForFilter (ChecksumHandle (PrimState m) h)
-> Bloom SerialisedKey
-> m ()
writeFilter HasFS m h
hfs ForFilter (ChecksumHandle (PrimState m) h)
filterHandle Bloom SerialisedKey
bf =
    HasFS m h -> ChecksumHandle (PrimState m) h -> ByteString -> m ()
forall (m :: * -> *) h.
(MonadSTM m, PrimMonad m) =>
HasFS m h -> ChecksumHandle (PrimState m) h -> ByteString -> m ()
writeToHandle HasFS m h
hfs (ForFilter (ChecksumHandle (PrimState m) h)
-> ChecksumHandle (PrimState m) h
forall a. ForFilter a -> a
unForFilter ForFilter (ChecksumHandle (PrimState m) h)
filterHandle) (Bloom SerialisedKey -> ByteString
forall a. Bloom a -> ByteString
bloomFilterToLBS Bloom SerialisedKey
bf)

{-# SPECIALISE writeIndexHeader ::
     HasFS IO h
  -> ForIndex (ChecksumHandle RealWorld h)
  -> IndexType
  -> IO () #-}
writeIndexHeader ::
     (MonadSTM m, PrimMonad m)
  => HasFS m h
  -> ForIndex (ChecksumHandle (PrimState m) h)
  -> IndexType
  -> m ()
writeIndexHeader :: forall (m :: * -> *) h.
(MonadSTM m, PrimMonad m) =>
HasFS m h
-> ForIndex (ChecksumHandle (PrimState m) h) -> IndexType -> m ()
writeIndexHeader HasFS m h
hfs ForIndex (ChecksumHandle (PrimState m) h)
indexHandle IndexType
indexType =
    HasFS m h -> ChecksumHandle (PrimState m) h -> ByteString -> m ()
forall (m :: * -> *) h.
(MonadSTM m, PrimMonad m) =>
HasFS m h -> ChecksumHandle (PrimState m) h -> ByteString -> m ()
writeToHandle HasFS m h
hfs (ForIndex (ChecksumHandle (PrimState m) h)
-> ChecksumHandle (PrimState m) h
forall a. ForIndex a -> a
unForIndex ForIndex (ChecksumHandle (PrimState m) h)
indexHandle) (ByteString -> m ()) -> ByteString -> m ()
forall a b. (a -> b) -> a -> b
$
      IndexType -> ByteString
Index.headerLBS IndexType
indexType

{-# SPECIALISE writeIndexChunk ::
     HasFS IO h
  -> ForIndex (ChecksumHandle RealWorld h)
  -> Chunk
  -> IO () #-}
writeIndexChunk ::
     (MonadSTM m, PrimMonad m)
  => HasFS m h
  -> ForIndex (ChecksumHandle (PrimState m) h)
  -> Chunk
  -> m ()
writeIndexChunk :: forall (m :: * -> *) h.
(MonadSTM m, PrimMonad m) =>
HasFS m h
-> ForIndex (ChecksumHandle (PrimState m) h) -> Chunk -> m ()
writeIndexChunk HasFS m h
hfs ForIndex (ChecksumHandle (PrimState m) h)
indexHandle Chunk
chunk =
    HasFS m h -> ChecksumHandle (PrimState m) h -> ByteString -> m ()
forall (m :: * -> *) h.
(MonadSTM m, PrimMonad m) =>
HasFS m h -> ChecksumHandle (PrimState m) h -> ByteString -> m ()
writeToHandle HasFS m h
hfs (ForIndex (ChecksumHandle (PrimState m) h)
-> ChecksumHandle (PrimState m) h
forall a. ForIndex a -> a
unForIndex ForIndex (ChecksumHandle (PrimState m) h)
indexHandle) (ByteString -> m ()) -> ByteString -> m ()
forall a b. (a -> b) -> a -> b
$
      ByteString -> ByteString
BSL.fromStrict (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ Chunk -> ByteString
Chunk.toByteString Chunk
chunk

{-# SPECIALISE writeIndexFinal ::
     HasFS IO h
  -> ForIndex (ChecksumHandle RealWorld h)
  -> NumEntries
  -> Index
  -> IO () #-}
writeIndexFinal ::
     (MonadSTM m, PrimMonad m)
  => HasFS m h
  -> ForIndex (ChecksumHandle (PrimState m) h)
  -> NumEntries
  -> Index
  -> m ()
writeIndexFinal :: forall (m :: * -> *) h.
(MonadSTM m, PrimMonad m) =>
HasFS m h
-> ForIndex (ChecksumHandle (PrimState m) h)
-> NumEntries
-> Index
-> m ()
writeIndexFinal HasFS m h
hfs ForIndex (ChecksumHandle (PrimState m) h)
indexHandle NumEntries
numEntries Index
index =
    HasFS m h -> ChecksumHandle (PrimState m) h -> ByteString -> m ()
forall (m :: * -> *) h.
(MonadSTM m, PrimMonad m) =>
HasFS m h -> ChecksumHandle (PrimState m) h -> ByteString -> m ()
writeToHandle HasFS m h
hfs (ForIndex (ChecksumHandle (PrimState m) h)
-> ChecksumHandle (PrimState m) h
forall a. ForIndex a -> a
unForIndex ForIndex (ChecksumHandle (PrimState m) h)
indexHandle) (ByteString -> m ()) -> ByteString -> m ()
forall a b. (a -> b) -> a -> b
$
      NumEntries -> Index -> ByteString
Index.finalLBS NumEntries
numEntries Index
index