{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE RecordWildCards #-}
module Database.Franz.Contents
  ( Contents
  , Database.Franz.Contents.indexNames
  , Item(..)
  , toList
  , last
  , length
  , index
  , lookupIndex
  -- * Internal
  , getResponse
  , readContents
  ) where

import Data.Serialize hiding (getInt64le)
import qualified Data.ByteString.Char8 as B
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as U
import Database.Franz.Internal
import Database.Franz.Protocol
import Database.Franz.Reader
import Data.Int
import Prelude hiding (length, last)

data Item = Item
  { Item -> Int
seqNo :: !Int
  , Item -> Vector Int64
indices :: !(U.Vector Int64)
  , Item -> ByteString
payload :: !B.ByteString
  } deriving (Int -> Item -> ShowS
[Item] -> ShowS
Item -> String
(Int -> Item -> ShowS)
-> (Item -> String) -> ([Item] -> ShowS) -> Show Item
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Item] -> ShowS
$cshowList :: [Item] -> ShowS
show :: Item -> String
$cshow :: Item -> String
showsPrec :: Int -> Item -> ShowS
$cshowsPrec :: Int -> Item -> ShowS
Show, Item -> Item -> Bool
(Item -> Item -> Bool) -> (Item -> Item -> Bool) -> Eq Item
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Item -> Item -> Bool
$c/= :: Item -> Item -> Bool
== :: Item -> Item -> Bool
$c== :: Item -> Item -> Bool
Eq)

data Contents = Contents
  { Contents -> Vector ByteString
indexNames :: !(V.Vector IndexName)
  , Contents -> ByteString
payloads :: !B.ByteString
  , Contents -> IndexVec
indicess :: !IndexVec
  , Contents -> Int
length :: !Int
  , Contents -> Int
payloadOffset :: !Int
  , Contents -> Int
seqnoOffset :: !Int
  }

toList :: Contents -> [Item]
toList :: Contents -> [Item]
toList Contents
contents = [Contents -> Int -> Item
unsafeIndex Contents
contents Int
i | Int
i <- [Int
0..Contents -> Int
length Contents
contents Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]

-- A vector containing file offsets and extra indices
type IndexVec = V.Vector (Int, U.Vector Int64)

getIndexVec :: V.Vector IndexName -> Int -> Get IndexVec
getIndexVec :: Vector ByteString -> Int -> Get IndexVec
getIndexVec Vector ByteString
names Int
len = Int -> Get (Int, Vector Int64) -> Get IndexVec
forall (m :: * -> *) a. Monad m => Int -> m a -> m (Vector a)
V.replicateM Int
len
  (Get (Int, Vector Int64) -> Get IndexVec)
-> Get (Int, Vector Int64) -> Get IndexVec
forall a b. (a -> b) -> a -> b
$ (,) (Int -> Vector Int64 -> (Int, Vector Int64))
-> Get Int -> Get (Vector Int64 -> (Int, Vector Int64))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get Int
forall a. Num a => Get a
getInt64le Get (Vector Int64 -> (Int, Vector Int64))
-> Get (Vector Int64) -> Get (Int, Vector Int64)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Vector Int64 -> Vector Int64
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
U.convert (Vector Int64 -> Vector Int64)
-> Get (Vector Int64) -> Get (Vector Int64)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` (ByteString -> Get Int64)
-> Vector ByteString -> Get (Vector Int64)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (Get Int64 -> ByteString -> Get Int64
forall a b. a -> b -> a
const Get Int64
forall a. Num a => Get a
getInt64le) Vector ByteString
names

getResponse :: Get Contents
getResponse :: Get Contents
getResponse = do
  PayloadHeader Int
seqnoOffset Int
s1 Int
payloadOffset Vector ByteString
indexNames <- Get PayloadHeader
forall t. Serialize t => Get t
get
  let length :: Int
length = Int
s1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
seqnoOffset
  if Int
length Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0
    then Contents -> Get Contents
forall (f :: * -> *) a. Applicative f => a -> f a
pure Contents :: Vector ByteString
-> ByteString -> IndexVec -> Int -> Int -> Int -> Contents
Contents{ payloads :: ByteString
payloads = ByteString
B.empty, indicess :: IndexVec
indicess = IndexVec
forall a. Vector a
V.empty, Int
Vector ByteString
length :: Int
indexNames :: Vector ByteString
payloadOffset :: Int
seqnoOffset :: Int
seqnoOffset :: Int
payloadOffset :: Int
length :: Int
indexNames :: Vector ByteString
..}
    else do
      IndexVec
indicess <- Vector ByteString -> Int -> Get IndexVec
getIndexVec Vector ByteString
indexNames Int
length
      ByteString
payloads <- Int -> Get ByteString
getByteString (Int -> Get ByteString) -> Int -> Get ByteString
forall a b. (a -> b) -> a -> b
$ (Int, Vector Int64) -> Int
forall a b. (a, b) -> a
fst (IndexVec -> (Int, Vector Int64)
forall a. Vector a -> a
V.unsafeLast IndexVec
indicess) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
payloadOffset
      Contents -> Get Contents
forall (f :: * -> *) a. Applicative f => a -> f a
pure Contents :: Vector ByteString
-> ByteString -> IndexVec -> Int -> Int -> Int -> Contents
Contents{Int
ByteString
IndexVec
Vector ByteString
payloads :: ByteString
indicess :: IndexVec
length :: Int
indexNames :: Vector ByteString
payloadOffset :: Int
seqnoOffset :: Int
seqnoOffset :: Int
payloadOffset :: Int
indicess :: IndexVec
payloads :: ByteString
length :: Int
indexNames :: Vector ByteString
..}

last :: Contents -> Maybe Item
last :: Contents -> Maybe Item
last Contents
contents
  | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 = Item -> Maybe Item
forall a. a -> Maybe a
Just (Item -> Maybe Item) -> Item -> Maybe Item
forall a b. (a -> b) -> a -> b
$ Contents -> Int -> Item
unsafeIndex Contents
contents Int
i
  | Bool
otherwise = Maybe Item
forall a. Maybe a
Nothing
  where
    i :: Int
i = Contents -> Int
length Contents
contents Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1

index :: Contents -> Int -> Maybe Item
index :: Contents -> Int -> Maybe Item
index Contents
contents Int
i
  | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Contents -> Int
length Contents
contents Bool -> Bool -> Bool
|| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = Maybe Item
forall a. Maybe a
Nothing
  | Bool
otherwise = Item -> Maybe Item
forall a. a -> Maybe a
Just (Item -> Maybe Item) -> Item -> Maybe Item
forall a b. (a -> b) -> a -> b
$ Contents -> Int -> Item
unsafeIndex Contents
contents Int
i

unsafeIndex :: Contents -> Int -> Item
unsafeIndex :: Contents -> Int -> Item
unsafeIndex Contents{Int
ByteString
IndexVec
Vector ByteString
seqnoOffset :: Int
payloadOffset :: Int
length :: Int
indicess :: IndexVec
payloads :: ByteString
indexNames :: Vector ByteString
seqnoOffset :: Contents -> Int
payloadOffset :: Contents -> Int
indicess :: Contents -> IndexVec
payloads :: Contents -> ByteString
length :: Contents -> Int
indexNames :: Contents -> Vector ByteString
..} Int
i = Item :: Int -> Vector Int64 -> ByteString -> Item
Item{Int
ByteString
Vector Int64
payload :: ByteString
seqNo :: Int
indices :: Vector Int64
payload :: ByteString
indices :: Vector Int64
seqNo :: Int
..}
  where
    ofs0 :: Int
ofs0 = Int
-> ((Int, Vector Int64) -> Int) -> Maybe (Int, Vector Int64) -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
payloadOffset (Int, Vector Int64) -> Int
forall a b. (a, b) -> a
fst (Maybe (Int, Vector Int64) -> Int)
-> Maybe (Int, Vector Int64) -> Int
forall a b. (a -> b) -> a -> b
$ IndexVec
indicess IndexVec -> Int -> Maybe (Int, Vector Int64)
forall a. Vector a -> Int -> Maybe a
V.!? (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
    (Int
ofs1, Vector Int64
indices) = IndexVec
indicess IndexVec -> Int -> (Int, Vector Int64)
forall a. Vector a -> Int -> a
V.! Int
i
    seqNo :: Int
seqNo = Int
seqnoOffset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
    payload :: ByteString
payload = Int -> ByteString -> ByteString
B.take (Int
ofs1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
payloadOffset) (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ Int -> ByteString -> ByteString
B.drop (Int
ofs0 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
payloadOffset) ByteString
payloads

lookupIndex :: Contents -> IndexName -> Maybe (Item -> Int64)
lookupIndex :: Contents -> ByteString -> Maybe (Item -> Int64)
lookupIndex Contents{Vector ByteString
indexNames :: Vector ByteString
indexNames :: Contents -> Vector ByteString
indexNames} ByteString
name
  = (\Int
j Item{Vector Int64
indices :: Vector Int64
indices :: Item -> Vector Int64
indices} -> Vector Int64
indices Vector Int64 -> Int -> Int64
forall a. Unbox a => Vector a -> Int -> a
U.! Int
j) (Int -> Item -> Int64) -> Maybe Int -> Maybe (Item -> Int64)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ByteString -> Vector ByteString -> Maybe Int
forall a. Eq a => a -> Vector a -> Maybe Int
V.elemIndex ByteString
name Vector ByteString
indexNames

readContents :: Stream -> QueryResult -> IO Contents
readContents :: Stream -> QueryResult -> IO Contents
readContents Stream{Vector ByteString
indexNames :: Stream -> Vector ByteString
indexNames :: Vector ByteString
indexNames, Handle
payloadHandle :: Stream -> Handle
payloadHandle :: Handle
payloadHandle, Handle
indexHandle :: Stream -> Handle
indexHandle :: Handle
indexHandle} ((Int
seqnoOffset, Int
payloadOffset), (Int
s1, Int
p1)) = do
  let length :: Int
length = Int
s1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
seqnoOffset
  -- byte offset + number of indices
  let indexSize :: Int
indexSize = Int
8 Int -> Int -> Int
forall a. Num a => a -> a -> a
* (Vector ByteString -> Int
forall a. Vector a -> Int
V.length Vector ByteString
indexNames Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  ByteString
indexBS <- Handle -> Int -> Int -> IO ByteString
hGetRange Handle
indexHandle (Int
indexSize Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
length) (Int
indexSize Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int -> Int
forall a. Enum a => a -> a
succ Int
seqnoOffset)
  ByteString
payloads <- Handle -> Int -> Int -> IO ByteString
hGetRange Handle
payloadHandle (Int
p1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
payloadOffset) Int
payloadOffset
  let indicess :: IndexVec
indicess = (String -> IndexVec)
-> (IndexVec -> IndexVec) -> Either String IndexVec -> IndexVec
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either String -> IndexVec
forall a. HasCallStack => String -> a
error IndexVec -> IndexVec
forall a. a -> a
id (Either String IndexVec -> IndexVec)
-> Either String IndexVec -> IndexVec
forall a b. (a -> b) -> a -> b
$ Get IndexVec -> ByteString -> Either String IndexVec
forall a. Get a -> ByteString -> Either String a
runGet (Vector ByteString -> Int -> Get IndexVec
getIndexVec Vector ByteString
indexNames Int
length) ByteString
indexBS
  Contents -> IO Contents
forall (f :: * -> *) a. Applicative f => a -> f a
pure Contents :: Vector ByteString
-> ByteString -> IndexVec -> Int -> Int -> Int -> Contents
Contents{Int
ByteString
IndexVec
Vector ByteString
indicess :: IndexVec
payloads :: ByteString
length :: Int
payloadOffset :: Int
seqnoOffset :: Int
indexNames :: Vector ByteString
seqnoOffset :: Int
payloadOffset :: Int
indicess :: IndexVec
payloads :: ByteString
length :: Int
indexNames :: Vector ByteString
..}