RandomAccess.hs

{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE OverloadedRecordDot #-}

module DataFrame.IO.Utils.RandomAccess (
    uncurry3,
    Range (..),
    RandomAccess (..),
    ReaderIO (runReaderIO),
    LocalFile,
    MMappedFile,
    ChunkBuffer (..),
    unsafeToByteString,
    WritableBinaryHandle,
    openWritableBinaryFile,
    withWritableBinaryFile,
    atomicallyWriteFile,
    MemoryBuffer (..),
    ensureCapacity,
    mallocBuffer,
    writeByteString,
    appendTextArraySlice,
    writeWord8,
    writeWord32LE,
    writeWord64LE,
    writeInteger64,
    writeWord32At,
    writeWord64At,
    writeInteger64At,
    writeFloatLE,
    writeDoubleLE,
    bufferResidency,
    bufferToByteString,
    unsafeBufferView,
    flushBufferToBuffer,
    resetPosition,
    flushBufferToFile,
    writeByteStringToFile,
) where

import Control.Exception (bracket, bracketOnError, finally, throw)
import Control.Monad (when)
import Control.Monad.IO.Class (MonadIO (..))
import Control.Monad.ST (RealWorld, ST, stToIO)
import Data.Bits (shiftR)
import Data.ByteString.Internal (ByteString (BS, PS), create)
import qualified Data.ByteString.Unsafe as BU
import Data.Int (Int64)
import Data.Primitive.ByteArray (
    MutableByteArray (..),
    copyMutableByteArray,
    getSizeofMutableByteArray,
    mutableByteArrayContents,
    newPinnedByteArray,
    withMutableByteArrayContents,
    writeByteArray,
 )
import Data.STRef (STRef, newSTRef, readSTRef, writeSTRef)
import qualified Data.Text.Array as TA
import qualified Data.Vector.Storable as VS
import Data.Word (Word32, Word64, Word8)
import DataFrame.IO.Parquet.Seeking (
    FileBufferedOrSeekable,
    fGet,
    fSeek,
    readLastBytes,
 )
import Foreign (castForeignPtr, castPtr, copyBytes, plusPtr)
import GHC.Float (castDoubleToWord64, castFloatToWord32)
import GHC.ForeignPtr (ForeignPtr (..), ForeignPtrContents (PlainPtr))
import GHC.Ptr (Ptr (..))
import System.Directory (copyPermissions, doesFileExist, removeFile, renameFile)
import System.FilePath (takeDirectory)
import System.IO (
    BufferMode (NoBuffering),
    Handle,
    IOMode (WriteMode),
    SeekMode (AbsoluteSeek),
    hClose,
    hPutBuf,
    hSetBinaryMode,
    hSetBuffering,
    openBinaryFile,
    openBinaryTempFileWithDefaultPermissions,
 )

uncurry3 :: (a -> b -> c -> d) -> (a, b, c) -> d
uncurry3 f (a, b, c) = f a b c

data Range = Range {offset :: !Integer, length :: !Int} deriving (Eq, Show)

class (Monad m) => RandomAccess m where
    readBytes :: Range -> m ByteString
    readRanges :: [Range] -> m [ByteString]
    readRanges = mapM readBytes
    readSuffix :: Int -> m ByteString

newtype ReaderIO r a = ReaderIO {runReaderIO :: r -> IO a}

instance Functor (ReaderIO r) where
    fmap f (ReaderIO run) = ReaderIO $ fmap f . run

instance Applicative (ReaderIO r) where
    pure a = ReaderIO $ \_ -> pure a
    (ReaderIO fg) <*> (ReaderIO fa) = ReaderIO $ \r -> do
        a <- fa r
        g <- fg r
        pure (g a)

instance Monad (ReaderIO r) where
    return = pure
    (ReaderIO ma) >>= f = ReaderIO $ \r -> do
        a <- ma r
        runReaderIO (f a) r

instance MonadIO (ReaderIO r) where
    liftIO io = ReaderIO $ const io

type LocalFile = ReaderIO FileBufferedOrSeekable

instance RandomAccess LocalFile where
    readBytes (Range offset' length') = ReaderIO $ \handle -> do
        fSeek handle AbsoluteSeek offset'
        fGet handle length'
    readSuffix n = ReaderIO (readLastBytes $ fromIntegral n)

-- | Bytes of one column chunk, already read; @base@ is their file offset.
data ChunkBuffer = ChunkBuffer !Integer !ByteString

instance RandomAccess (ReaderIO ChunkBuffer) where
    readBytes (Range offset' length') = ReaderIO $ \(ChunkBuffer base bs) ->
        pure (BU.unsafeTake length' (BU.unsafeDrop (fromInteger (offset' - base)) bs))
    readSuffix _ = ReaderIO $ \_ -> error "ChunkBuffer: no file footer"

type MMappedFile = ReaderIO (VS.Vector Word8)

-- The instance exists but we don't have the means to mmap the file currently
instance RandomAccess MMappedFile where
    readBytes (Range offset' length') =
        ReaderIO $
            pure . unsafeToByteString . VS.slice (fromInteger offset') length'
    readSuffix n =
        ReaderIO $ \v ->
            let len = VS.length v
                n' = min n len
                start = len - n'
             in pure . unsafeToByteString $ VS.slice start n' v

unsafeToByteString :: VS.Vector Word8 -> ByteString
unsafeToByteString v = PS (castForeignPtr ptr) offset' len
  where
    (ptr, offset', len) = VS.unsafeToForeignPtr v

-- Writer buffers are pinned MutableByteArrays so hPutBuf can take a Ptr.

newtype WritableBinaryHandle = WritableBinaryHandle {unHandle :: Handle}

openWritableBinaryFile :: FilePath -> IO WritableBinaryHandle
openWritableBinaryFile filepath = do
    h <- openBinaryFile filepath WriteMode
    hSetBinaryMode h True
    hSetBuffering h NoBuffering
    pure . WritableBinaryHandle $ h

atomicallyWriteFile ::
    FilePath ->
    (FilePath -> IO a) ->
    IO a
atomicallyWriteFile path action =
    bracketOnError
        openAction
        removeFile
        ( \tmpFile -> do
            result <- action tmpFile
            renameFile tmpFile path
            pure result
        )
  where
    openAction =
        bracketOnError
            ( openBinaryTempFileWithDefaultPermissions
                (takeDirectory path)
                "dataframe-parquet.incomplete"
            )
            (\(tmpFile, h) -> hClose h `finally` removeFile tmpFile)
            ( \(tmpFile, h) -> do
                hClose h
                destinationExists <- doesFileExist path
                when destinationExists (copyPermissions path tmpFile)
                pure tmpFile
            )

withWritableBinaryFile ::
    FilePath ->
    (WritableBinaryHandle -> IO a) ->
    IO a
withWritableBinaryFile filepath =
    bracket
        (openWritableBinaryFile filepath)
        (hClose . unHandle)

data MemoryBuffer s = MemoryBuffer
    { arrayRef :: !(STRef s (MutableByteArray s))
    , positionRef :: !(STRef s Int)
    }

mallocBuffer :: Int -> ST s (MemoryBuffer s)
mallocBuffer capacity
    | capacity < 0 = throw (userError "mallocBuffer: negative capacity")
    | otherwise = do
        array <- newPinnedByteArray capacity
        MemoryBuffer <$> newSTRef array <*> newSTRef 0

-- Pinned arrays cannot grow in place, so allocate a new one and copy.
ensureCapacity :: MemoryBuffer s -> Int -> ST s (MutableByteArray s)
ensureCapacity buffer needed = do
    array <- readSTRef buffer.arrayRef
    maxSize <- getSizeofMutableByteArray array
    if needed <= maxSize
        then pure array
        else do
            position <- readSTRef buffer.positionRef
            grown <- newPinnedByteArray (needed + (needed `div` 2))
            copyMutableByteArray grown 0 array 0 position
            writeSTRef buffer.arrayRef grown
            pure grown
{-# INLINE ensureCapacity #-}

writeWord8 :: MemoryBuffer s -> Word8 -> ST s ()
writeWord8 buffer b = do
    position <- readSTRef buffer.positionRef
    array <- ensureCapacity buffer (position + 1)
    writeByteArray array position b
    writeSTRef buffer.positionRef (position + 1)
{-# INLINE writeWord8 #-}

-- Reading the ByteString's pointer needs IO, so this stays out of ST.
writeByteString :: MemoryBuffer RealWorld -> ByteString -> IO ()
writeByteString buffer bs =
    BU.unsafeUseAsCStringLen bs $ \(source, len) -> do
        position <- stToIO (readSTRef buffer.positionRef)
        array <- stToIO (ensureCapacity buffer (position + len))
        withMutableByteArrayContents array $ \dst ->
            copyBytes (dst `plusPtr` position) (castPtr source) len
        stToIO (writeSTRef buffer.positionRef (position + len))
{-# INLINE writeByteString #-}

writeWord32LE :: MemoryBuffer s -> Word32 -> ST s ()
writeWord32LE buffer w = do
    position <- readSTRef buffer.positionRef
    writeWord32At buffer position w
    writeSTRef buffer.positionRef (position + 4)
{-# INLINE writeWord32LE #-}

writeWord64LE :: MemoryBuffer s -> Word64 -> ST s ()
writeWord64LE buffer w = do
    position <- readSTRef buffer.positionRef
    writeWord64At buffer position w
    writeSTRef buffer.positionRef (position + 8)
{-# INLINE writeWord64LE #-}

writeWord32At :: MemoryBuffer s -> Int -> Word32 -> ST s ()
writeWord32At buffer position w = do
    array <- ensureCapacity buffer (position + 4)
    writeByteArray array position (fromIntegral w :: Word8)
    writeByteArray array (position + 1) (fromIntegral (w `shiftR` 8) :: Word8)
    writeByteArray array (position + 2) (fromIntegral (w `shiftR` 16) :: Word8)
    writeByteArray array (position + 3) (fromIntegral (w `shiftR` 24) :: Word8)
{-# INLINE writeWord32At #-}

writeWord64At :: MemoryBuffer s -> Int -> Word64 -> ST s ()
writeWord64At buffer position w = do
    array <- ensureCapacity buffer (position + 8)
    writeByteArray array position (fromIntegral w :: Word8)
    writeByteArray array (position + 1) (fromIntegral (w `shiftR` 8) :: Word8)
    writeByteArray array (position + 2) (fromIntegral (w `shiftR` 16) :: Word8)
    writeByteArray array (position + 3) (fromIntegral (w `shiftR` 24) :: Word8)
    writeByteArray array (position + 4) (fromIntegral (w `shiftR` 32) :: Word8)
    writeByteArray array (position + 5) (fromIntegral (w `shiftR` 40) :: Word8)
    writeByteArray array (position + 6) (fromIntegral (w `shiftR` 48) :: Word8)
    writeByteArray array (position + 7) (fromIntegral (w `shiftR` 56) :: Word8)
{-# INLINE writeWord64At #-}

writeInteger64 :: MemoryBuffer s -> Integer -> ST s ()
writeInteger64 buffer value = do
    position <- readSTRef buffer.positionRef
    newPosition <- writeInteger64At buffer position value
    writeSTRef buffer.positionRef newPosition
{-# INLINE writeInteger64 #-}

writeInteger64At :: MemoryBuffer s -> Int -> Integer -> ST s Int
writeInteger64At buffer position value
    | value < toInteger (minBound :: Int64) = outOfRange
    | value > toInteger (maxBound :: Int64) = outOfRange
    | otherwise = do
        writeWord64At buffer position (fromIntegral value)
        pure (position + 8)
  where
    outOfRange =
        throw (userError "writeParquet: Integer value is outside the INT64 range")
{-# INLINE writeInteger64At #-}

writeFloatLE :: MemoryBuffer s -> Float -> ST s ()
writeFloatLE buffer = writeWord32LE buffer . castFloatToWord32
{-# INLINE writeFloatLE #-}

writeDoubleLE :: MemoryBuffer s -> Double -> ST s ()
writeDoubleLE buffer = writeWord64LE buffer . castDoubleToWord64
{-# INLINE writeDoubleLE #-}

flushBufferToBuffer :: MemoryBuffer s -> MemoryBuffer s -> ST s ()
flushBufferToBuffer source destination
    | source.arrayRef == destination.arrayRef = pure ()
    | otherwise = do
        sourceArray <- readSTRef source.arrayRef
        sourcePosition <- readSTRef source.positionRef
        destinationPosition <- readSTRef destination.positionRef
        destinationArray <-
            ensureCapacity destination (destinationPosition + sourcePosition)
        copyMutableByteArray
            destinationArray
            destinationPosition
            sourceArray
            0
            sourcePosition
        writeSTRef destination.positionRef (destinationPosition + sourcePosition)
        writeSTRef source.positionRef 0
{-# INLINE flushBufferToBuffer #-}

bufferToByteString :: MemoryBuffer RealWorld -> IO ByteString
bufferToByteString buffer = do
    array <- stToIO (readSTRef buffer.arrayRef)
    position <- stToIO (readSTRef buffer.positionRef)
    create position $ \dst ->
        withMutableByteArrayContents array $ \src ->
            copyBytes dst (castPtr src) position

-- | The buffer's contents as a ByteString, without copying.
unsafeBufferView :: MemoryBuffer RealWorld -> IO ByteString
unsafeBufferView buffer = do
    array@(MutableByteArray array#) <- stToIO (readSTRef buffer.arrayRef)
    position <- stToIO (readSTRef buffer.positionRef)
    let Ptr address# = mutableByteArrayContents array
    pure (BS (ForeignPtr address# (PlainPtr array#)) position)

bufferResidency :: MemoryBuffer s -> ST s Int
bufferResidency buffer = readSTRef buffer.positionRef
{-# INLINE bufferResidency #-}

resetPosition :: MemoryBuffer s -> ST s ()
resetPosition buffer = writeSTRef buffer.positionRef 0
{-# INLINE resetPosition #-}

-- Flush in 256 KiB chunks; larger blocks gave no write speed in dd tests.
flushBufferToFile :: WritableBinaryHandle -> MemoryBuffer RealWorld -> IO ()
flushBufferToFile (WritableBinaryHandle h) buffer = do
    array <- stToIO (readSTRef buffer.arrayRef)
    position <- stToIO (readSTRef buffer.positionRef)
    withMutableByteArrayContents array $ \ptr -> do
        let chunkSize = 262144
            go offset
                | offset >= position = pure ()
                | otherwise = do
                    let n = min chunkSize (position - offset)
                    hPutBuf h (ptr `plusPtr` offset) n
                    go (offset + n)
        go 0
    stToIO (resetPosition buffer)

writeByteStringToFile :: WritableBinaryHandle -> ByteString -> IO ()
writeByteStringToFile (WritableBinaryHandle h) bs =
    BU.unsafeUseAsCStringLen bs $ \(source, len) -> do
        let chunkSize = 262144
            go offset
                | offset >= len = pure ()
                | otherwise = do
                    let n = min chunkSize (len - offset)
                    hPutBuf h (source `plusPtr` offset) n
                    go (offset + n)
        go 0

appendTextArraySlice :: MemoryBuffer s -> TA.Array -> Int -> Int -> ST s ()
appendTextArraySlice buffer source offset count
    | count < 0 = throw (userError "appendTextArraySlice: negative length")
    | otherwise = do
        position <- readSTRef buffer.positionRef
        array <- ensureCapacity buffer (position + count)
        withMutableByteArrayContents array $ \destination ->
            TA.copyToPointer source offset (destination `plusPtr` position) count
        writeSTRef buffer.positionRef (position + count)
{-# INLINE appendTextArraySlice #-}