-- | Contains functions for constructing and working with foreign simdjson instances.

module Data.Hermes.SIMDJSON.Wrapper
  ( getDocumentInfo
  , mkSIMDParser
  , mkSIMDDocument
  , mkSIMDPaddedStr
  , withInputBuffer
  )
  where

import           Control.Monad.IO.Class (liftIO)
import           Data.ByteString (ByteString)
import qualified Data.ByteString.Unsafe as Unsafe
import           Data.Maybe (fromMaybe)
import           Data.Text (Text)
import qualified Data.Text as T
import           UnliftIO.Exception (bracket, mask_)
import           UnliftIO.Foreign
  ( ForeignPtr
  , alloca
  , allocaBytes
  , finalizeForeignPtr
  , newForeignPtr
  , peek
  , peekCString
  , peekCStringLen
  , withForeignPtr
  )

import           Data.Hermes.SIMDJSON.Bindings
  ( currentLocationImpl
  , deleteDocumentImpl
  , deleteInputImpl
  , makeDocumentImpl
  , makeInputImpl
  , parserDestroy
  , parserInit
  , toDebugStringImpl
  )
import           Data.Hermes.SIMDJSON.Types
  ( Document
  , InputBuffer(..)
  , PaddedString
  , SIMDDocument
  , SIMDErrorCode(..)
  , SIMDParser
  )

mkSIMDParser :: Maybe Int -> IO (ForeignPtr SIMDParser)
mkSIMDParser :: Maybe Int -> IO (ForeignPtr SIMDParser)
mkSIMDParser Maybe Int
mCap = IO (ForeignPtr SIMDParser) -> IO (ForeignPtr SIMDParser)
forall (m :: * -> *) a. MonadUnliftIO m => m a -> m a
mask_ (IO (ForeignPtr SIMDParser) -> IO (ForeignPtr SIMDParser))
-> IO (ForeignPtr SIMDParser) -> IO (ForeignPtr SIMDParser)
forall a b. (a -> b) -> a -> b
$ do
  let maxCap :: Int
maxCap = Int
4000000000; -- 4GB
  Ptr SIMDParser
ptr <- CSize -> IO (Ptr SIMDParser)
parserInit (CSize -> IO (Ptr SIMDParser))
-> (Int -> CSize) -> Int -> IO (Ptr SIMDParser)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> CSize
forall a. Enum a => Int -> a
toEnum (Int -> IO (Ptr SIMDParser)) -> Int -> IO (Ptr SIMDParser)
forall a b. (a -> b) -> a -> b
$ Int -> Maybe Int -> Int
forall a. a -> Maybe a -> a
fromMaybe Int
maxCap Maybe Int
mCap
  FinalizerPtr SIMDParser
-> Ptr SIMDParser -> IO (ForeignPtr SIMDParser)
forall (m :: * -> *) a.
MonadIO m =>
FinalizerPtr a -> Ptr a -> m (ForeignPtr a)
newForeignPtr FinalizerPtr SIMDParser
parserDestroy Ptr SIMDParser
ptr

mkSIMDDocument :: IO (ForeignPtr SIMDDocument)
mkSIMDDocument :: IO (ForeignPtr SIMDDocument)
mkSIMDDocument = IO (ForeignPtr SIMDDocument) -> IO (ForeignPtr SIMDDocument)
forall (m :: * -> *) a. MonadUnliftIO m => m a -> m a
mask_ (IO (ForeignPtr SIMDDocument) -> IO (ForeignPtr SIMDDocument))
-> IO (ForeignPtr SIMDDocument) -> IO (ForeignPtr SIMDDocument)
forall a b. (a -> b) -> a -> b
$ do
  Ptr SIMDDocument
ptr <- IO (Ptr SIMDDocument)
makeDocumentImpl
  FinalizerPtr SIMDDocument
-> Ptr SIMDDocument -> IO (ForeignPtr SIMDDocument)
forall (m :: * -> *) a.
MonadIO m =>
FinalizerPtr a -> Ptr a -> m (ForeignPtr a)
newForeignPtr FinalizerPtr SIMDDocument
deleteDocumentImpl Ptr SIMDDocument
ptr

mkSIMDPaddedStr :: ByteString -> IO (ForeignPtr PaddedString)
mkSIMDPaddedStr :: ByteString -> IO (ForeignPtr PaddedString)
mkSIMDPaddedStr ByteString
input = IO (ForeignPtr PaddedString) -> IO (ForeignPtr PaddedString)
forall (m :: * -> *) a. MonadUnliftIO m => m a -> m a
mask_ (IO (ForeignPtr PaddedString) -> IO (ForeignPtr PaddedString))
-> IO (ForeignPtr PaddedString) -> IO (ForeignPtr PaddedString)
forall a b. (a -> b) -> a -> b
$
  ByteString
-> (CStringLen -> IO (ForeignPtr PaddedString))
-> IO (ForeignPtr PaddedString)
forall a. ByteString -> (CStringLen -> IO a) -> IO a
Unsafe.unsafeUseAsCStringLen ByteString
input ((CStringLen -> IO (ForeignPtr PaddedString))
 -> IO (ForeignPtr PaddedString))
-> (CStringLen -> IO (ForeignPtr PaddedString))
-> IO (ForeignPtr PaddedString)
forall a b. (a -> b) -> a -> b
$ \(Ptr CChar
cstr, Int
len) -> do
    Ptr PaddedString
ptr <- Ptr CChar -> CSize -> IO (Ptr PaddedString)
makeInputImpl Ptr CChar
cstr (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len)
    FinalizerPtr PaddedString
-> Ptr PaddedString -> IO (ForeignPtr PaddedString)
forall (m :: * -> *) a.
MonadIO m =>
FinalizerPtr a -> Ptr a -> m (ForeignPtr a)
newForeignPtr FinalizerPtr PaddedString
deleteInputImpl Ptr PaddedString
ptr

-- | Construct a simdjson:padded_string from a Haskell `ByteString`, and pass
-- it to a monadic action. The instance lifetime is managed by the `bracket` function.
withInputBuffer :: ByteString -> (InputBuffer -> IO a) -> IO a
withInputBuffer :: ByteString -> (InputBuffer -> IO a) -> IO a
withInputBuffer ByteString
bs InputBuffer -> IO a
f =
  IO (ForeignPtr PaddedString)
-> (ForeignPtr PaddedString -> IO ())
-> (ForeignPtr PaddedString -> IO a)
-> IO a
forall (m :: * -> *) a b c.
MonadUnliftIO m =>
m a -> (a -> m b) -> (a -> m c) -> m c
bracket IO (ForeignPtr PaddedString)
acquire ForeignPtr PaddedString -> IO ()
forall a. ForeignPtr a -> IO ()
release ((ForeignPtr PaddedString -> IO a) -> IO a)
-> (ForeignPtr PaddedString -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \ForeignPtr PaddedString
fPtr -> ForeignPtr PaddedString -> (Ptr PaddedString -> IO a) -> IO a
forall (m :: * -> *) a b.
MonadUnliftIO m =>
ForeignPtr a -> (Ptr a -> m b) -> m b
withForeignPtr ForeignPtr PaddedString
fPtr ((Ptr PaddedString -> IO a) -> IO a)
-> (Ptr PaddedString -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ InputBuffer -> IO a
f (InputBuffer -> IO a)
-> (Ptr PaddedString -> InputBuffer) -> Ptr PaddedString -> IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ptr PaddedString -> InputBuffer
InputBuffer
  where
    acquire :: IO (ForeignPtr PaddedString)
acquire = IO (ForeignPtr PaddedString) -> IO (ForeignPtr PaddedString)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (ForeignPtr PaddedString) -> IO (ForeignPtr PaddedString))
-> IO (ForeignPtr PaddedString) -> IO (ForeignPtr PaddedString)
forall a b. (a -> b) -> a -> b
$ ByteString -> IO (ForeignPtr PaddedString)
mkSIMDPaddedStr ByteString
bs
    release :: ForeignPtr a -> IO ()
release = IO () -> IO ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> IO ())
-> (ForeignPtr a -> IO ()) -> ForeignPtr a -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ForeignPtr a -> IO ()
forall (m :: * -> *) a. MonadIO m => ForeignPtr a -> m ()
finalizeForeignPtr

-- | Read the document location and debug string. If the iterator is out of bounds
-- then we abort reading from the iterator buffers to prevent reading garbage.
getDocumentInfo :: Document -> IO (Text, Text)
getDocumentInfo :: Document -> IO (Text, Text)
getDocumentInfo Document
docPtr = (Ptr (Ptr CChar) -> IO (Text, Text)) -> IO (Text, Text)
forall (m :: * -> *) a b.
(MonadUnliftIO m, Storable a) =>
(Ptr a -> m b) -> m b
alloca ((Ptr (Ptr CChar) -> IO (Text, Text)) -> IO (Text, Text))
-> (Ptr (Ptr CChar) -> IO (Text, Text)) -> IO (Text, Text)
forall a b. (a -> b) -> a -> b
$ \Ptr (Ptr CChar)
locStrPtr -> (Ptr CSize -> IO (Text, Text)) -> IO (Text, Text)
forall (m :: * -> *) a b.
(MonadUnliftIO m, Storable a) =>
(Ptr a -> m b) -> m b
alloca ((Ptr CSize -> IO (Text, Text)) -> IO (Text, Text))
-> (Ptr CSize -> IO (Text, Text)) -> IO (Text, Text)
forall a b. (a -> b) -> a -> b
$ \Ptr CSize
lenPtr -> do
  CInt
err <- IO CInt -> IO CInt
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO CInt -> IO CInt) -> IO CInt -> IO CInt
forall a b. (a -> b) -> a -> b
$ Document -> Ptr (Ptr CChar) -> IO CInt
currentLocationImpl Document
docPtr Ptr (Ptr CChar)
locStrPtr
  let errCode :: SIMDErrorCode
errCode = Int -> SIMDErrorCode
forall a. Enum a => Int -> a
toEnum (Int -> SIMDErrorCode) -> Int -> SIMDErrorCode
forall a b. (a -> b) -> a -> b
$ CInt -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral CInt
err
  if SIMDErrorCode
errCode SIMDErrorCode -> SIMDErrorCode -> Bool
forall a. Eq a => a -> a -> Bool
== SIMDErrorCode
OUT_OF_BOUNDS
    then (Text, Text) -> IO (Text, Text)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Text
"out of bounds", Text
"")
    else Int -> (Ptr CChar -> IO (Text, Text)) -> IO (Text, Text)
forall (m :: * -> *) a b.
MonadUnliftIO m =>
Int -> (Ptr a -> m b) -> m b
allocaBytes Int
128 ((Ptr CChar -> IO (Text, Text)) -> IO (Text, Text))
-> (Ptr CChar -> IO (Text, Text)) -> IO (Text, Text)
forall a b. (a -> b) -> a -> b
$ \Ptr CChar
dbStrPtr -> do
      Text
locStr <- (String -> Text) -> IO String -> IO Text
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap String -> Text
T.pack (IO String -> IO Text) -> IO String -> IO Text
forall a b. (a -> b) -> a -> b
$ Ptr CChar -> IO String
forall (m :: * -> *). MonadIO m => Ptr CChar -> m String
peekCString (Ptr CChar -> IO String) -> IO (Ptr CChar) -> IO String
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Ptr (Ptr CChar) -> IO (Ptr CChar)
forall a. Storable a => Ptr a -> IO a
peek Ptr (Ptr CChar)
locStrPtr
      Document -> Ptr CChar -> Ptr CSize -> IO ()
toDebugStringImpl Document
docPtr Ptr CChar
dbStrPtr Ptr CSize
lenPtr
      Int
len <- (CSize -> Int) -> IO CSize -> IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap CSize -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (IO CSize -> IO Int) -> IO CSize -> IO Int
forall a b. (a -> b) -> a -> b
$ Ptr CSize -> IO CSize
forall a. Storable a => Ptr a -> IO a
peek Ptr CSize
lenPtr
      Text
debugStr <- (String -> Text) -> IO String -> IO Text
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap String -> Text
T.pack (IO String -> IO Text) -> IO String -> IO Text
forall a b. (a -> b) -> a -> b
$ CStringLen -> IO String
forall (m :: * -> *). MonadIO m => CStringLen -> m String
peekCStringLen (Ptr CChar
dbStrPtr, Int
len)
      (Text, Text) -> IO (Text, Text)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Text
locStr, Text
debugStr)