{-# LANGUAGE FlexibleInstances #-}

module DataFrame.IO.Utils.RandomAccess where

import Control.Monad.IO.Class (MonadIO (..))
import Data.ByteString (ByteString)
import Data.ByteString.Internal (ByteString (PS))
import qualified Data.Vector.Storable as VS
import Data.Word (Word8)
import DataFrame.IO.Parquet.Seeking (
    FileBufferedOrSeekable,
    fGet,
    fSeek,
    readLastBytes,
 )
import Foreign (castForeignPtr)
import System.IO (
    SeekMode (AbsoluteSeek),
 )

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

data Range = Range {Range -> Integer
offset :: !Integer, Range -> Int
length :: !Int} deriving (Range -> Range -> Bool
(Range -> Range -> Bool) -> (Range -> Range -> Bool) -> Eq Range
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Range -> Range -> Bool
== :: Range -> Range -> Bool
$c/= :: Range -> Range -> Bool
/= :: Range -> Range -> Bool
Eq, Int -> Range -> ShowS
[Range] -> ShowS
Range -> String
(Int -> Range -> ShowS)
-> (Range -> String) -> ([Range] -> ShowS) -> Show Range
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Range -> ShowS
showsPrec :: Int -> Range -> ShowS
$cshow :: Range -> String
show :: Range -> String
$cshowList :: [Range] -> ShowS
showList :: [Range] -> ShowS
Show)

class (Monad m) => RandomAccess m where
    readBytes :: Range -> m ByteString
    readRanges :: [Range] -> m [ByteString]
    readRanges = (Range -> m ByteString) -> [Range] -> m [ByteString]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Range -> m ByteString
forall (m :: * -> *). RandomAccess m => Range -> m ByteString
readBytes
    readSuffix :: Int -> m ByteString

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

instance Functor (ReaderIO r) where
    fmap :: forall a b. (a -> b) -> ReaderIO r a -> ReaderIO r b
fmap a -> b
f (ReaderIO r -> IO a
run) = (r -> IO b) -> ReaderIO r b
forall r a. (r -> IO a) -> ReaderIO r a
ReaderIO ((r -> IO b) -> ReaderIO r b) -> (r -> IO b) -> ReaderIO r b
forall a b. (a -> b) -> a -> b
$ (a -> b) -> IO a -> IO b
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f (IO a -> IO b) -> (r -> IO a) -> r -> IO b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. r -> IO a
run

instance Applicative (ReaderIO r) where
    pure :: forall a. a -> ReaderIO r a
pure a
a = (r -> IO a) -> ReaderIO r a
forall r a. (r -> IO a) -> ReaderIO r a
ReaderIO ((r -> IO a) -> ReaderIO r a) -> (r -> IO a) -> ReaderIO r a
forall a b. (a -> b) -> a -> b
$ \r
_ -> a -> IO a
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a
    (ReaderIO r -> IO (a -> b)
fg) <*> :: forall a b. ReaderIO r (a -> b) -> ReaderIO r a -> ReaderIO r b
<*> (ReaderIO r -> IO a
fa) = (r -> IO b) -> ReaderIO r b
forall r a. (r -> IO a) -> ReaderIO r a
ReaderIO ((r -> IO b) -> ReaderIO r b) -> (r -> IO b) -> ReaderIO r b
forall a b. (a -> b) -> a -> b
$ \r
r -> do
        a
a <- r -> IO a
fa r
r
        a -> b
g <- r -> IO (a -> b)
fg r
r
        b -> IO b
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> b
g a
a)

instance Monad (ReaderIO r) where
    return :: forall a. a -> ReaderIO r a
return = a -> ReaderIO r a
forall a. a -> ReaderIO r a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    (ReaderIO r -> IO a
ma) >>= :: forall a b. ReaderIO r a -> (a -> ReaderIO r b) -> ReaderIO r b
>>= a -> ReaderIO r b
f = (r -> IO b) -> ReaderIO r b
forall r a. (r -> IO a) -> ReaderIO r a
ReaderIO ((r -> IO b) -> ReaderIO r b) -> (r -> IO b) -> ReaderIO r b
forall a b. (a -> b) -> a -> b
$ \r
r -> do
        a
a <- r -> IO a
ma r
r
        ReaderIO r b -> r -> IO b
forall r a. ReaderIO r a -> r -> IO a
runReaderIO (a -> ReaderIO r b
f a
a) r
r

instance MonadIO (ReaderIO r) where
    liftIO :: forall a. IO a -> ReaderIO r a
liftIO IO a
io = (r -> IO a) -> ReaderIO r a
forall r a. (r -> IO a) -> ReaderIO r a
ReaderIO ((r -> IO a) -> ReaderIO r a) -> (r -> IO a) -> ReaderIO r a
forall a b. (a -> b) -> a -> b
$ IO a -> r -> IO a
forall a b. a -> b -> a
const IO a
io

type LocalFile = ReaderIO FileBufferedOrSeekable

instance RandomAccess LocalFile where
    readBytes :: Range -> LocalFile ByteString
readBytes (Range Integer
offset' Int
length') = (FileBufferedOrSeekable -> IO ByteString) -> LocalFile ByteString
forall r a. (r -> IO a) -> ReaderIO r a
ReaderIO ((FileBufferedOrSeekable -> IO ByteString) -> LocalFile ByteString)
-> (FileBufferedOrSeekable -> IO ByteString)
-> LocalFile ByteString
forall a b. (a -> b) -> a -> b
$ \FileBufferedOrSeekable
handle -> do
        FileBufferedOrSeekable -> SeekMode -> Integer -> IO ()
fSeek FileBufferedOrSeekable
handle SeekMode
AbsoluteSeek Integer
offset'
        FileBufferedOrSeekable -> Int -> IO ByteString
fGet FileBufferedOrSeekable
handle Int
length'
    readSuffix :: Int -> LocalFile ByteString
readSuffix Int
n = (FileBufferedOrSeekable -> IO ByteString) -> LocalFile ByteString
forall r a. (r -> IO a) -> ReaderIO r a
ReaderIO (Integer -> FileBufferedOrSeekable -> IO ByteString
readLastBytes (Integer -> FileBufferedOrSeekable -> IO ByteString)
-> Integer -> FileBufferedOrSeekable -> IO ByteString
forall a b. (a -> b) -> a -> b
$ Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n)

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 -> MMappedFile ByteString
readBytes (Range Integer
offset' Int
length') =
        (Vector Word8 -> IO ByteString) -> MMappedFile ByteString
forall r a. (r -> IO a) -> ReaderIO r a
ReaderIO ((Vector Word8 -> IO ByteString) -> MMappedFile ByteString)
-> (Vector Word8 -> IO ByteString) -> MMappedFile ByteString
forall a b. (a -> b) -> a -> b
$
            ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> IO ByteString)
-> (Vector Word8 -> ByteString) -> Vector Word8 -> IO ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Word8 -> ByteString
unsafeToByteString (Vector Word8 -> ByteString)
-> (Vector Word8 -> Vector Word8) -> Vector Word8 -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> Vector Word8 -> Vector Word8
forall a. Storable a => Int -> Int -> Vector a -> Vector a
VS.slice (Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
offset') Int
length'
    readSuffix :: Int -> MMappedFile ByteString
readSuffix Int
n =
        (Vector Word8 -> IO ByteString) -> MMappedFile ByteString
forall r a. (r -> IO a) -> ReaderIO r a
ReaderIO ((Vector Word8 -> IO ByteString) -> MMappedFile ByteString)
-> (Vector Word8 -> IO ByteString) -> MMappedFile ByteString
forall a b. (a -> b) -> a -> b
$ \Vector Word8
v ->
            let len :: Int
len = Vector Word8 -> Int
forall a. Storable a => Vector a -> Int
VS.length Vector Word8
v
                n' :: Int
n' = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
n Int
len
                start :: Int
start = Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n'
             in ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> IO ByteString)
-> (Vector Word8 -> ByteString) -> Vector Word8 -> IO ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Word8 -> ByteString
unsafeToByteString (Vector Word8 -> IO ByteString) -> Vector Word8 -> IO ByteString
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Vector Word8 -> Vector Word8
forall a. Storable a => Int -> Int -> Vector a -> Vector a
VS.slice Int
start Int
n' Vector Word8
v

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