{-# OPTIONS_GHC -Wno-redundant-constraints #-}
{-# LANGUAGE FlexibleInstances, FlexibleContexts, TypeFamilies, DerivingVia #-}
{-# LANGUAGE BangPatterns, RoleAnnotations, MultiParamTypeClasses #-}
{-# LANGUAGE InstanceSigs, DataKinds, TypeApplications, TypeOperators #-}
{-# LANGUAGE ConstraintKinds, PolyKinds, UndecidableInstances #-}

-- | This module provides an implementation for the typeclasses defined in
--   "Language.Souffle.Class".
--   It makes use of the low level Souffle C++ API to offer a much more
--   performant alternative implementation to the implementation in
--   "Language.Souffle.Interpreted".
--
--   This module is mainly intended to be used after the prototyping phase is
--   over since the iteration cycle is slower due to the additional
--   C++ compilation times.
module Language.Souffle.Compiled
  ( Program(..)
  , ProgramOptions(..)
  , Fact(..)
  , FactOptions(..)
  , Marshal(..)
  , Direction(..)
  , ContainsInputFact
  , ContainsOutputFact
  , Submit
  , Handle
  , SouffleM
  , MonadSouffle(..)
  , MonadSouffleFileIO(..)
  , runSouffle
  ) where

import Prelude hiding ( init )
import Control.Monad.Except
import Control.Monad.State.Strict
import Data.Foldable ( traverse_ )
import Data.Functor.Identity
import Data.Proxy
import Data.Kind
import qualified Data.Array as A
import qualified Data.Array.IO as A
import qualified Data.Array.Unsafe as A
import qualified Data.ByteString as BS
import qualified Data.ByteString.Short as BSS
import qualified Data.ByteString.Unsafe as BSU
import qualified Data.Text as T
import qualified Data.Text.Short as TS
import qualified Data.Text.Short.Unsafe as TSU
import qualified Data.Text.Lazy as TL
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV
import Data.Int
import Data.Word
import Foreign.ForeignPtr
import Foreign.ForeignPtr.Unsafe
import Foreign (copyBytes)
import Foreign.Ptr
import qualified Foreign.Storable as S
import GHC.Generics
import Language.Souffle.Class
import qualified Language.Souffle.Internal as Internal
import Language.Souffle.Marshal
import Control.Concurrent


type ByteCount = Int
type ByteBuf = Internal.ByteBuf

data BufData
  = BufData
  { BufData -> ForeignPtr ByteBuf
bufPtr :: {-# UNPACK #-} !(ForeignPtr ByteBuf)
  , BufData -> Int
bufSize :: {-# UNPACK #-} !ByteCount
  }

-- | A datatype representing a handle to a datalog program.
--   The type parameter is used for keeping track of which program
--   type the handle belongs to for additional type safety.
data Handle prog
  = Handle {-# UNPACK #-} !(ForeignPtr Internal.Souffle)
           {-# UNPACK #-} !(MVar BufData)
type role Handle nominal

-- | A monad for executing Souffle-related actions in.
newtype SouffleM a = SouffleM (IO a)
  deriving ((forall a b. (a -> b) -> SouffleM a -> SouffleM b)
-> (forall a b. a -> SouffleM b -> SouffleM a) -> Functor SouffleM
forall a b. a -> SouffleM b -> SouffleM a
forall a b. (a -> b) -> SouffleM a -> SouffleM b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> SouffleM b -> SouffleM a
$c<$ :: forall a b. a -> SouffleM b -> SouffleM a
fmap :: forall a b. (a -> b) -> SouffleM a -> SouffleM b
$cfmap :: forall a b. (a -> b) -> SouffleM a -> SouffleM b
Functor, Functor SouffleM
Functor SouffleM
-> (forall a. a -> SouffleM a)
-> (forall a b. SouffleM (a -> b) -> SouffleM a -> SouffleM b)
-> (forall a b c.
    (a -> b -> c) -> SouffleM a -> SouffleM b -> SouffleM c)
-> (forall a b. SouffleM a -> SouffleM b -> SouffleM b)
-> (forall a b. SouffleM a -> SouffleM b -> SouffleM a)
-> Applicative SouffleM
forall a. a -> SouffleM a
forall a b. SouffleM a -> SouffleM b -> SouffleM a
forall a b. SouffleM a -> SouffleM b -> SouffleM b
forall a b. SouffleM (a -> b) -> SouffleM a -> SouffleM b
forall a b c.
(a -> b -> c) -> SouffleM a -> SouffleM b -> SouffleM c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
<* :: forall a b. SouffleM a -> SouffleM b -> SouffleM a
$c<* :: forall a b. SouffleM a -> SouffleM b -> SouffleM a
*> :: forall a b. SouffleM a -> SouffleM b -> SouffleM b
$c*> :: forall a b. SouffleM a -> SouffleM b -> SouffleM b
liftA2 :: forall a b c.
(a -> b -> c) -> SouffleM a -> SouffleM b -> SouffleM c
$cliftA2 :: forall a b c.
(a -> b -> c) -> SouffleM a -> SouffleM b -> SouffleM c
<*> :: forall a b. SouffleM (a -> b) -> SouffleM a -> SouffleM b
$c<*> :: forall a b. SouffleM (a -> b) -> SouffleM a -> SouffleM b
pure :: forall a. a -> SouffleM a
$cpure :: forall a. a -> SouffleM a
Applicative, Applicative SouffleM
Applicative SouffleM
-> (forall a b. SouffleM a -> (a -> SouffleM b) -> SouffleM b)
-> (forall a b. SouffleM a -> SouffleM b -> SouffleM b)
-> (forall a. a -> SouffleM a)
-> Monad SouffleM
forall a. a -> SouffleM a
forall a b. SouffleM a -> SouffleM b -> SouffleM b
forall a b. SouffleM a -> (a -> SouffleM b) -> SouffleM b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
return :: forall a. a -> SouffleM a
$creturn :: forall a. a -> SouffleM a
>> :: forall a b. SouffleM a -> SouffleM b -> SouffleM b
$c>> :: forall a b. SouffleM a -> SouffleM b -> SouffleM b
>>= :: forall a b. SouffleM a -> (a -> SouffleM b) -> SouffleM b
$c>>= :: forall a b. SouffleM a -> (a -> SouffleM b) -> SouffleM b
Monad, Monad SouffleM
Monad SouffleM
-> (forall a. IO a -> SouffleM a) -> MonadIO SouffleM
forall a. IO a -> SouffleM a
forall (m :: * -> *).
Monad m -> (forall a. IO a -> m a) -> MonadIO m
liftIO :: forall a. IO a -> SouffleM a
$cliftIO :: forall a. IO a -> SouffleM a
MonadIO) via IO
  deriving (NonEmpty (SouffleM a) -> SouffleM a
SouffleM a -> SouffleM a -> SouffleM a
(SouffleM a -> SouffleM a -> SouffleM a)
-> (NonEmpty (SouffleM a) -> SouffleM a)
-> (forall b. Integral b => b -> SouffleM a -> SouffleM a)
-> Semigroup (SouffleM a)
forall b. Integral b => b -> SouffleM a -> SouffleM a
forall a. Semigroup a => NonEmpty (SouffleM a) -> SouffleM a
forall a. Semigroup a => SouffleM a -> SouffleM a -> SouffleM a
forall a b.
(Semigroup a, Integral b) =>
b -> SouffleM a -> SouffleM a
forall a.
(a -> a -> a)
-> (NonEmpty a -> a)
-> (forall b. Integral b => b -> a -> a)
-> Semigroup a
stimes :: forall b. Integral b => b -> SouffleM a -> SouffleM a
$cstimes :: forall a b.
(Semigroup a, Integral b) =>
b -> SouffleM a -> SouffleM a
sconcat :: NonEmpty (SouffleM a) -> SouffleM a
$csconcat :: forall a. Semigroup a => NonEmpty (SouffleM a) -> SouffleM a
<> :: SouffleM a -> SouffleM a -> SouffleM a
$c<> :: forall a. Semigroup a => SouffleM a -> SouffleM a -> SouffleM a
Semigroup, Semigroup (SouffleM a)
SouffleM a
Semigroup (SouffleM a)
-> SouffleM a
-> (SouffleM a -> SouffleM a -> SouffleM a)
-> ([SouffleM a] -> SouffleM a)
-> Monoid (SouffleM a)
[SouffleM a] -> SouffleM a
SouffleM a -> SouffleM a -> SouffleM a
forall a.
Semigroup a -> a -> (a -> a -> a) -> ([a] -> a) -> Monoid a
forall {a}. Monoid a => Semigroup (SouffleM a)
forall a. Monoid a => SouffleM a
forall a. Monoid a => [SouffleM a] -> SouffleM a
forall a. Monoid a => SouffleM a -> SouffleM a -> SouffleM a
mconcat :: [SouffleM a] -> SouffleM a
$cmconcat :: forall a. Monoid a => [SouffleM a] -> SouffleM a
mappend :: SouffleM a -> SouffleM a -> SouffleM a
$cmappend :: forall a. Monoid a => SouffleM a -> SouffleM a -> SouffleM a
mempty :: SouffleM a
$cmempty :: forall a. Monoid a => SouffleM a
Monoid) via (IO a)

{- | Initializes and runs a Souffle program.

     The 2nd argument is passed in a handle after initialization of the
     Souffle program. The handle will contain 'Nothing' if it failed to
     load the Souffle C++ program. In the successful case it will contain
     a handle that can be used for performing Souffle related actions
     using the other functions in this module.
-}
runSouffle :: forall prog a. Program prog
           => prog -> (Maybe (Handle prog) -> SouffleM a) -> IO a
runSouffle :: forall prog a.
Program prog =>
prog -> (Maybe (Handle prog) -> SouffleM a) -> IO a
runSouffle prog
prog Maybe (Handle prog) -> SouffleM a
action =
  let progName :: String
progName = prog -> String
forall a. Program a => a -> String
programName prog
prog
      (SouffleM IO a
result) = do
        Maybe (Handle prog)
maybeHandle <- IO (Maybe (ForeignPtr Souffle))
-> SouffleM (Maybe (ForeignPtr Souffle))
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (String -> IO (Maybe (ForeignPtr Souffle))
Internal.init String
progName) SouffleM (Maybe (ForeignPtr Souffle))
-> (Maybe (ForeignPtr Souffle) -> SouffleM (Maybe (Handle prog)))
-> SouffleM (Maybe (Handle prog))
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
          Maybe (ForeignPtr Souffle)
Nothing -> Maybe (Handle prog) -> SouffleM (Maybe (Handle prog))
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Handle prog)
forall a. Maybe a
Nothing
          Just ForeignPtr Souffle
souffleHandle -> do
            MVar BufData
bufData <- IO (MVar BufData) -> SouffleM (MVar BufData)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (MVar BufData) -> SouffleM (MVar BufData))
-> IO (MVar BufData) -> SouffleM (MVar BufData)
forall a b. (a -> b) -> a -> b
$ do
              ForeignPtr ByteBuf
ptr <- Ptr ByteBuf -> IO (ForeignPtr ByteBuf)
forall a. Ptr a -> IO (ForeignPtr a)
newForeignPtr_ Ptr ByteBuf
forall a. Ptr a
nullPtr
              BufData -> IO (MVar BufData)
forall a. a -> IO (MVar a)
newMVar (BufData -> IO (MVar BufData)) -> BufData -> IO (MVar BufData)
forall a b. (a -> b) -> a -> b
$ ForeignPtr ByteBuf -> Int -> BufData
BufData ForeignPtr ByteBuf
ptr Int
0
            Maybe (Handle prog) -> SouffleM (Maybe (Handle prog))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (Handle prog) -> SouffleM (Maybe (Handle prog)))
-> Maybe (Handle prog) -> SouffleM (Maybe (Handle prog))
forall a b. (a -> b) -> a -> b
$ Handle prog -> Maybe (Handle prog)
forall a. a -> Maybe a
Just (Handle prog -> Maybe (Handle prog))
-> Handle prog -> Maybe (Handle prog)
forall a b. (a -> b) -> a -> b
$ ForeignPtr Souffle -> MVar BufData -> Handle prog
forall {k} (prog :: k).
ForeignPtr Souffle -> MVar BufData -> Handle prog
Handle ForeignPtr Souffle
souffleHandle MVar BufData
bufData
        Maybe (Handle prog) -> SouffleM a
action Maybe (Handle prog)
maybeHandle
   in IO a
result
{-# INLINABLE runSouffle #-}

-- | A monad used solely for marshalling and unmarshalling
--   between Haskell and Souffle Datalog. This fast variant is used when the
--   marshalling from Haskell to C++ and the exact size of a datastructure
--   is statically known (read: data type contains no string-like types),
--   or when marshalling from C++ to Haskell (pointer is then managed by C++).
newtype CMarshalFast a = CMarshalFast (StateT (Ptr ByteBuf) IO a)
  deriving ((forall a b. (a -> b) -> CMarshalFast a -> CMarshalFast b)
-> (forall a b. a -> CMarshalFast b -> CMarshalFast a)
-> Functor CMarshalFast
forall a b. a -> CMarshalFast b -> CMarshalFast a
forall a b. (a -> b) -> CMarshalFast a -> CMarshalFast b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> CMarshalFast b -> CMarshalFast a
$c<$ :: forall a b. a -> CMarshalFast b -> CMarshalFast a
fmap :: forall a b. (a -> b) -> CMarshalFast a -> CMarshalFast b
$cfmap :: forall a b. (a -> b) -> CMarshalFast a -> CMarshalFast b
Functor, Functor CMarshalFast
Functor CMarshalFast
-> (forall a. a -> CMarshalFast a)
-> (forall a b.
    CMarshalFast (a -> b) -> CMarshalFast a -> CMarshalFast b)
-> (forall a b c.
    (a -> b -> c)
    -> CMarshalFast a -> CMarshalFast b -> CMarshalFast c)
-> (forall a b. CMarshalFast a -> CMarshalFast b -> CMarshalFast b)
-> (forall a b. CMarshalFast a -> CMarshalFast b -> CMarshalFast a)
-> Applicative CMarshalFast
forall a. a -> CMarshalFast a
forall a b. CMarshalFast a -> CMarshalFast b -> CMarshalFast a
forall a b. CMarshalFast a -> CMarshalFast b -> CMarshalFast b
forall a b.
CMarshalFast (a -> b) -> CMarshalFast a -> CMarshalFast b
forall a b c.
(a -> b -> c) -> CMarshalFast a -> CMarshalFast b -> CMarshalFast c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
<* :: forall a b. CMarshalFast a -> CMarshalFast b -> CMarshalFast a
$c<* :: forall a b. CMarshalFast a -> CMarshalFast b -> CMarshalFast a
*> :: forall a b. CMarshalFast a -> CMarshalFast b -> CMarshalFast b
$c*> :: forall a b. CMarshalFast a -> CMarshalFast b -> CMarshalFast b
liftA2 :: forall a b c.
(a -> b -> c) -> CMarshalFast a -> CMarshalFast b -> CMarshalFast c
$cliftA2 :: forall a b c.
(a -> b -> c) -> CMarshalFast a -> CMarshalFast b -> CMarshalFast c
<*> :: forall a b.
CMarshalFast (a -> b) -> CMarshalFast a -> CMarshalFast b
$c<*> :: forall a b.
CMarshalFast (a -> b) -> CMarshalFast a -> CMarshalFast b
pure :: forall a. a -> CMarshalFast a
$cpure :: forall a. a -> CMarshalFast a
Applicative, Applicative CMarshalFast
Applicative CMarshalFast
-> (forall a b.
    CMarshalFast a -> (a -> CMarshalFast b) -> CMarshalFast b)
-> (forall a b. CMarshalFast a -> CMarshalFast b -> CMarshalFast b)
-> (forall a. a -> CMarshalFast a)
-> Monad CMarshalFast
forall a. a -> CMarshalFast a
forall a b. CMarshalFast a -> CMarshalFast b -> CMarshalFast b
forall a b.
CMarshalFast a -> (a -> CMarshalFast b) -> CMarshalFast b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
return :: forall a. a -> CMarshalFast a
$creturn :: forall a. a -> CMarshalFast a
>> :: forall a b. CMarshalFast a -> CMarshalFast b -> CMarshalFast b
$c>> :: forall a b. CMarshalFast a -> CMarshalFast b -> CMarshalFast b
>>= :: forall a b.
CMarshalFast a -> (a -> CMarshalFast b) -> CMarshalFast b
$c>>= :: forall a b.
CMarshalFast a -> (a -> CMarshalFast b) -> CMarshalFast b
Monad, Monad CMarshalFast
Monad CMarshalFast
-> (forall a. IO a -> CMarshalFast a) -> MonadIO CMarshalFast
forall a. IO a -> CMarshalFast a
forall (m :: * -> *).
Monad m -> (forall a. IO a -> m a) -> MonadIO m
liftIO :: forall a. IO a -> CMarshalFast a
$cliftIO :: forall a. IO a -> CMarshalFast a
MonadIO, MonadState (Ptr ByteBuf))
  via (StateT (Ptr ByteBuf) IO)

runMarshalFastM :: CMarshalFast a -> Ptr ByteBuf -> IO a
runMarshalFastM :: forall a. CMarshalFast a -> Ptr ByteBuf -> IO a
runMarshalFastM (CMarshalFast StateT (Ptr ByteBuf) IO a
m) = StateT (Ptr ByteBuf) IO a -> Ptr ByteBuf -> IO a
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT StateT (Ptr ByteBuf) IO a
m
{-# INLINABLE runMarshalFastM #-}

-- NOTE: assumes Souffle is compiled with 32-bit RAM domain.
ramDomainSize :: Int
ramDomainSize :: Int
ramDomainSize = Int
4

writeAsBytes :: (S.Storable a, Marshal a) => a -> CMarshalFast ()
writeAsBytes :: forall a. (Storable a, Marshal a) => a -> CMarshalFast ()
writeAsBytes a
a = do
  Ptr a
ptr <- (Ptr ByteBuf -> Ptr a) -> CMarshalFast (Ptr a)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets Ptr ByteBuf -> Ptr a
forall a b. Ptr a -> Ptr b
castPtr
  IO () -> CMarshalFast ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> CMarshalFast ()) -> IO () -> CMarshalFast ()
forall a b. (a -> b) -> a -> b
$ Ptr a -> a -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
S.poke Ptr a
ptr a
a
  Ptr ByteBuf -> CMarshalFast ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Ptr ByteBuf -> CMarshalFast ()) -> Ptr ByteBuf -> CMarshalFast ()
forall a b. (a -> b) -> a -> b
$ Ptr a
ptr Ptr a -> Int -> Ptr ByteBuf
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
ramDomainSize
{-# INLINABLE writeAsBytes #-}

readAsBytes :: (S.Storable a, Marshal a) => CMarshalFast a
readAsBytes :: forall a. (Storable a, Marshal a) => CMarshalFast a
readAsBytes = do
  Ptr a
ptr <- (Ptr ByteBuf -> Ptr a) -> CMarshalFast (Ptr a)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets Ptr ByteBuf -> Ptr a
forall a b. Ptr a -> Ptr b
castPtr
  a
a <- IO a -> CMarshalFast a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO a -> CMarshalFast a) -> IO a -> CMarshalFast a
forall a b. (a -> b) -> a -> b
$ Ptr a -> IO a
forall a. Storable a => Ptr a -> IO a
S.peek Ptr a
ptr
  Ptr ByteBuf -> CMarshalFast ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Ptr ByteBuf -> CMarshalFast ()) -> Ptr ByteBuf -> CMarshalFast ()
forall a b. (a -> b) -> a -> b
$ Ptr a
ptr Ptr a -> Int -> Ptr ByteBuf
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
ramDomainSize
  a -> CMarshalFast a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a
{-# INLINABLE readAsBytes #-}

instance MonadPush CMarshalFast where
  pushInt32 :: Int32 -> CMarshalFast ()
pushInt32 = Int32 -> CMarshalFast ()
forall a. (Storable a, Marshal a) => a -> CMarshalFast ()
writeAsBytes
  {-# INLINABLE pushInt32 #-}
  pushUInt32 :: Word32 -> CMarshalFast ()
pushUInt32 = Word32 -> CMarshalFast ()
forall a. (Storable a, Marshal a) => a -> CMarshalFast ()
writeAsBytes
  {-# INLINABLE pushUInt32 #-}
  pushFloat :: Float -> CMarshalFast ()
pushFloat = Float -> CMarshalFast ()
forall a. (Storable a, Marshal a) => a -> CMarshalFast ()
writeAsBytes
  {-# INLINABLE pushFloat #-}
  pushString :: String -> CMarshalFast ()
pushString String
str = ShortText -> CMarshalFast ()
forall (m :: * -> *). MonadPush m => ShortText -> m ()
pushText (ShortText -> CMarshalFast ()) -> ShortText -> CMarshalFast ()
forall a b. (a -> b) -> a -> b
$ String -> ShortText
TS.pack String
str
  {-# INLINABLE pushString #-}
  pushTextUtf16 :: Text -> CMarshalFast ()
pushTextUtf16 Text
str = ShortText -> CMarshalFast ()
forall (m :: * -> *). MonadPush m => ShortText -> m ()
pushText (ShortText -> CMarshalFast ()) -> ShortText -> CMarshalFast ()
forall a b. (a -> b) -> a -> b
$ Text -> ShortText
TS.fromText Text
str
  {-# INLINABLE pushTextUtf16 #-}
  pushText :: ShortText -> CMarshalFast ()
pushText ShortText
_ =
    String -> CMarshalFast ()
forall a. HasCallStack => String -> a
error String
"Fast marshalling does not support serializing string-like values."
  {-# INLINABLE pushText #-}

instance MonadPop CMarshalFast where
  popInt32 :: CMarshalFast Int32
popInt32 = CMarshalFast Int32
forall a. (Storable a, Marshal a) => CMarshalFast a
readAsBytes
  {-# INLINABLE popInt32 #-}
  popUInt32 :: CMarshalFast Word32
popUInt32 = CMarshalFast Word32
forall a. (Storable a, Marshal a) => CMarshalFast a
readAsBytes
  {-# INLINABLE popUInt32 #-}
  popFloat :: CMarshalFast Float
popFloat = CMarshalFast Float
forall a. (Storable a, Marshal a) => CMarshalFast a
readAsBytes
  {-# INLINABLE popFloat #-}
  popString :: CMarshalFast String
popString = ShortText -> String
TS.unpack (ShortText -> String)
-> CMarshalFast ShortText -> CMarshalFast String
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> CMarshalFast ShortText
forall (m :: * -> *). MonadPop m => m ShortText
popText
  {-# INLINABLE popString #-}
  popTextUtf16 :: CMarshalFast Text
popTextUtf16 = ShortText -> Text
TS.toText (ShortText -> Text) -> CMarshalFast ShortText -> CMarshalFast Text
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> CMarshalFast ShortText
forall (m :: * -> *). MonadPop m => m ShortText
popText
  {-# INLINABLE popTextUtf16 #-}
  popText :: CMarshalFast ShortText
popText = do
    Word32
byteCount <- CMarshalFast Word32
forall (m :: * -> *). MonadPop m => m Word32
popUInt32
    if Word32
byteCount Word32 -> Word32 -> Bool
forall a. Eq a => a -> a -> Bool
== Word32
0
      then ShortText -> CMarshalFast ShortText
forall (f :: * -> *) a. Applicative f => a -> f a
pure ShortText
TS.empty
      else do
        Ptr CChar
ptr <- (Ptr ByteBuf -> Ptr CChar) -> CMarshalFast (Ptr CChar)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets Ptr ByteBuf -> Ptr CChar
forall a b. Ptr a -> Ptr b
castPtr
        ByteString
bs <- IO ByteString -> CMarshalFast ByteString
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO ByteString -> CMarshalFast ByteString)
-> IO ByteString -> CMarshalFast ByteString
forall a b. (a -> b) -> a -> b
$ CStringLen -> IO ByteString
BSU.unsafePackCStringLen (Ptr CChar
ptr, Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
byteCount)
        Ptr ByteBuf -> CMarshalFast ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Ptr ByteBuf -> CMarshalFast ()) -> Ptr ByteBuf -> CMarshalFast ()
forall a b. (a -> b) -> a -> b
$ Ptr CChar
ptr Ptr CChar -> Int -> Ptr ByteBuf
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
byteCount
        -- NOTE: $! is needed here to force the text value. A copy needs to
        -- be made (using toShort), before the bytearray is overwritten.
        ShortText -> CMarshalFast ShortText
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ShortText -> CMarshalFast ShortText)
-> ShortText -> CMarshalFast ShortText
forall a b. (a -> b) -> a -> b
$! ShortByteString -> ShortText
TSU.fromShortByteStringUnsafe (ShortByteString -> ShortText) -> ShortByteString -> ShortText
forall a b. (a -> b) -> a -> b
$ ByteString -> ShortByteString
BSS.toShort ByteString
bs
  {-# INLINABLE popText #-}


data MarshalState
  = MarshalState
  { MarshalState -> BufData
_buf :: {-# UNPACK #-} !BufData
  , MarshalState -> Ptr ByteBuf
_ptr :: {-# UNPACK #-} !(Ptr ByteBuf)
  , MarshalState -> Int
_ptrOffset :: {-# UNPACK #-} !Int
  }

-- | A monad used solely for marshalling from Haskell to Souffle Datalog (C++).
--   This slow variant is used when the exact size of a datastructure is *not*
--   statically known (read: data type contains string-like types).
newtype CMarshalSlow a = CMarshalSlow (StateT MarshalState IO a)
  deriving ((forall a b. (a -> b) -> CMarshalSlow a -> CMarshalSlow b)
-> (forall a b. a -> CMarshalSlow b -> CMarshalSlow a)
-> Functor CMarshalSlow
forall a b. a -> CMarshalSlow b -> CMarshalSlow a
forall a b. (a -> b) -> CMarshalSlow a -> CMarshalSlow b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> CMarshalSlow b -> CMarshalSlow a
$c<$ :: forall a b. a -> CMarshalSlow b -> CMarshalSlow a
fmap :: forall a b. (a -> b) -> CMarshalSlow a -> CMarshalSlow b
$cfmap :: forall a b. (a -> b) -> CMarshalSlow a -> CMarshalSlow b
Functor, Functor CMarshalSlow
Functor CMarshalSlow
-> (forall a. a -> CMarshalSlow a)
-> (forall a b.
    CMarshalSlow (a -> b) -> CMarshalSlow a -> CMarshalSlow b)
-> (forall a b c.
    (a -> b -> c)
    -> CMarshalSlow a -> CMarshalSlow b -> CMarshalSlow c)
-> (forall a b. CMarshalSlow a -> CMarshalSlow b -> CMarshalSlow b)
-> (forall a b. CMarshalSlow a -> CMarshalSlow b -> CMarshalSlow a)
-> Applicative CMarshalSlow
forall a. a -> CMarshalSlow a
forall a b. CMarshalSlow a -> CMarshalSlow b -> CMarshalSlow a
forall a b. CMarshalSlow a -> CMarshalSlow b -> CMarshalSlow b
forall a b.
CMarshalSlow (a -> b) -> CMarshalSlow a -> CMarshalSlow b
forall a b c.
(a -> b -> c) -> CMarshalSlow a -> CMarshalSlow b -> CMarshalSlow c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
<* :: forall a b. CMarshalSlow a -> CMarshalSlow b -> CMarshalSlow a
$c<* :: forall a b. CMarshalSlow a -> CMarshalSlow b -> CMarshalSlow a
*> :: forall a b. CMarshalSlow a -> CMarshalSlow b -> CMarshalSlow b
$c*> :: forall a b. CMarshalSlow a -> CMarshalSlow b -> CMarshalSlow b
liftA2 :: forall a b c.
(a -> b -> c) -> CMarshalSlow a -> CMarshalSlow b -> CMarshalSlow c
$cliftA2 :: forall a b c.
(a -> b -> c) -> CMarshalSlow a -> CMarshalSlow b -> CMarshalSlow c
<*> :: forall a b.
CMarshalSlow (a -> b) -> CMarshalSlow a -> CMarshalSlow b
$c<*> :: forall a b.
CMarshalSlow (a -> b) -> CMarshalSlow a -> CMarshalSlow b
pure :: forall a. a -> CMarshalSlow a
$cpure :: forall a. a -> CMarshalSlow a
Applicative, Applicative CMarshalSlow
Applicative CMarshalSlow
-> (forall a b.
    CMarshalSlow a -> (a -> CMarshalSlow b) -> CMarshalSlow b)
-> (forall a b. CMarshalSlow a -> CMarshalSlow b -> CMarshalSlow b)
-> (forall a. a -> CMarshalSlow a)
-> Monad CMarshalSlow
forall a. a -> CMarshalSlow a
forall a b. CMarshalSlow a -> CMarshalSlow b -> CMarshalSlow b
forall a b.
CMarshalSlow a -> (a -> CMarshalSlow b) -> CMarshalSlow b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
return :: forall a. a -> CMarshalSlow a
$creturn :: forall a. a -> CMarshalSlow a
>> :: forall a b. CMarshalSlow a -> CMarshalSlow b -> CMarshalSlow b
$c>> :: forall a b. CMarshalSlow a -> CMarshalSlow b -> CMarshalSlow b
>>= :: forall a b.
CMarshalSlow a -> (a -> CMarshalSlow b) -> CMarshalSlow b
$c>>= :: forall a b.
CMarshalSlow a -> (a -> CMarshalSlow b) -> CMarshalSlow b
Monad, Monad CMarshalSlow
Monad CMarshalSlow
-> (forall a. IO a -> CMarshalSlow a) -> MonadIO CMarshalSlow
forall a. IO a -> CMarshalSlow a
forall (m :: * -> *).
Monad m -> (forall a. IO a -> m a) -> MonadIO m
liftIO :: forall a. IO a -> CMarshalSlow a
$cliftIO :: forall a. IO a -> CMarshalSlow a
MonadIO, MonadState MarshalState)
  via (StateT MarshalState IO)

runMarshalSlowM :: BufData -> Int -> CMarshalSlow a -> IO a
runMarshalSlowM :: forall a. BufData -> Int -> CMarshalSlow a -> IO a
runMarshalSlowM BufData
bufData Int
byteCount (CMarshalSlow StateT MarshalState IO a
m) = do
  BufData
bufData' <- if BufData -> Int
bufSize BufData
bufData Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
byteCount
    then BufData -> IO BufData
forall (f :: * -> *) a. Applicative f => a -> f a
pure BufData
bufData
    else (ForeignPtr ByteBuf -> Int -> BufData)
-> Int -> ForeignPtr ByteBuf -> BufData
forall a b c. (a -> b -> c) -> b -> a -> c
flip ForeignPtr ByteBuf -> Int -> BufData
BufData Int
byteCount (ForeignPtr ByteBuf -> BufData)
-> IO (ForeignPtr ByteBuf) -> IO BufData
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> IO (ForeignPtr ByteBuf)
forall (m :: * -> *). MonadIO m => Int -> m (ForeignPtr ByteBuf)
allocateBuf Int
byteCount
  let ptr :: Ptr ByteBuf
ptr = ForeignPtr ByteBuf -> Ptr ByteBuf
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr (BufData -> ForeignPtr ByteBuf
bufPtr BufData
bufData')
  StateT MarshalState IO a -> MarshalState -> IO a
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT StateT MarshalState IO a
m (MarshalState -> IO a) -> MarshalState -> IO a
forall a b. (a -> b) -> a -> b
$ BufData -> Ptr ByteBuf -> Int -> MarshalState
MarshalState BufData
bufData' Ptr ByteBuf
ptr Int
0
{-# INLINABLE runMarshalSlowM #-}

resizeBufWhenNeeded :: ByteCount -> CMarshalSlow ()
resizeBufWhenNeeded :: Int -> CMarshalSlow ()
resizeBufWhenNeeded Int
byteCount = do
  MarshalState BufData
bufData Ptr ByteBuf
_ Int
offset <- CMarshalSlow MarshalState
forall s (m :: * -> *). MonadState s m => m s
get
  let totalByteCount :: Int
totalByteCount = BufData -> Int
bufSize BufData
bufData
  Bool -> CMarshalSlow () -> CMarshalSlow ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
byteCount Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
offset Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
totalByteCount) (CMarshalSlow () -> CMarshalSlow ())
-> CMarshalSlow () -> CMarshalSlow ()
forall a b. (a -> b) -> a -> b
$ do
    let newTotalByteCount :: Int
newTotalByteCount = Int -> Int -> Int -> Int
getNewTotalByteCount Int
byteCount Int
offset Int
totalByteCount
    ForeignPtr ByteBuf
newBuf <- Int -> CMarshalSlow (ForeignPtr ByteBuf)
forall (m :: * -> *). MonadIO m => Int -> m (ForeignPtr ByteBuf)
allocateBuf Int
newTotalByteCount
    ForeignPtr ByteBuf -> ForeignPtr ByteBuf -> Int -> CMarshalSlow ()
copyBuf ForeignPtr ByteBuf
newBuf (BufData -> ForeignPtr ByteBuf
bufPtr BufData
bufData) Int
totalByteCount
    let newPtr :: Ptr ByteBuf
newPtr = ForeignPtr ByteBuf -> Ptr ByteBuf
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr ByteBuf
newBuf
        bufData' :: BufData
bufData' = ForeignPtr ByteBuf -> Int -> BufData
BufData ForeignPtr ByteBuf
newBuf Int
newTotalByteCount
    MarshalState -> CMarshalSlow ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (MarshalState -> CMarshalSlow ())
-> MarshalState -> CMarshalSlow ()
forall a b. (a -> b) -> a -> b
$ BufData -> Ptr ByteBuf -> Int -> MarshalState
MarshalState BufData
bufData' (Ptr ByteBuf
newPtr Ptr ByteBuf -> Int -> Ptr ByteBuf
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
offset) Int
offset
{-# INLINABLE resizeBufWhenNeeded #-}

allocateBuf :: MonadIO m => ByteCount -> m (ForeignPtr ByteBuf)
allocateBuf :: forall (m :: * -> *). MonadIO m => Int -> m (ForeignPtr ByteBuf)
allocateBuf Int
byteCount = IO (ForeignPtr ByteBuf) -> m (ForeignPtr ByteBuf)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (ForeignPtr ByteBuf) -> m (ForeignPtr ByteBuf))
-> IO (ForeignPtr ByteBuf) -> m (ForeignPtr ByteBuf)
forall a b. (a -> b) -> a -> b
$
  Int -> IO (ForeignPtr ByteBuf)
forall a. Int -> IO (ForeignPtr a)
mallocForeignPtrBytes Int
byteCount
{-# INLINABLE allocateBuf #-}

copyBuf :: ForeignPtr ByteBuf -> ForeignPtr ByteBuf -> Int -> CMarshalSlow ()
copyBuf :: ForeignPtr ByteBuf -> ForeignPtr ByteBuf -> Int -> CMarshalSlow ()
copyBuf ForeignPtr ByteBuf
dst ForeignPtr ByteBuf
src Int
byteCount = IO () -> CMarshalSlow ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> CMarshalSlow ()) -> IO () -> CMarshalSlow ()
forall a b. (a -> b) -> a -> b
$
  ForeignPtr ByteBuf -> (Ptr ByteBuf -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr ByteBuf
src ((Ptr ByteBuf -> IO ()) -> IO ())
-> (Ptr ByteBuf -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr ByteBuf
srcPtr ->
  ForeignPtr ByteBuf -> (Ptr ByteBuf -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr ByteBuf
dst ((Ptr ByteBuf -> IO ()) -> IO ())
-> (Ptr ByteBuf -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr ByteBuf
dstPtr ->
    Ptr ByteBuf -> Ptr ByteBuf -> Int -> IO ()
forall a. Ptr a -> Ptr a -> Int -> IO ()
copyBytes Ptr ByteBuf
dstPtr Ptr ByteBuf
srcPtr Int
byteCount
{-# INLINABLE copyBuf #-}

getNewTotalByteCount :: ByteCount -> Int -> ByteCount -> ByteCount
getNewTotalByteCount :: Int -> Int -> Int -> Int
getNewTotalByteCount Int
byteCount Int
offset = Int -> Int
go where
  go :: Int -> Int
go Int
totalByteCount
    | Int
byteCount Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
offset Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
totalByteCount = Int -> Int
go (Int
totalByteCount Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
2)
    | Bool
otherwise = Int
totalByteCount
{-# INLINABLE getNewTotalByteCount #-}

incrementPtr :: ByteCount -> CMarshalSlow ()
incrementPtr :: Int -> CMarshalSlow ()
incrementPtr Int
byteCount =
  (MarshalState -> MarshalState) -> CMarshalSlow ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((MarshalState -> MarshalState) -> CMarshalSlow ())
-> (MarshalState -> MarshalState) -> CMarshalSlow ()
forall a b. (a -> b) -> a -> b
$ \(MarshalState BufData
buf Ptr ByteBuf
ptr Int
offset) ->
    BufData -> Ptr ByteBuf -> Int -> MarshalState
MarshalState BufData
buf (Ptr ByteBuf
ptr Ptr ByteBuf -> Int -> Ptr ByteBuf
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
byteCount) (Int
offset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
byteCount)
{-# INLINABLE incrementPtr #-}

instance MonadPush CMarshalSlow where
  pushInt32 :: Int32 -> CMarshalSlow ()
pushInt32 = Int32 -> CMarshalSlow ()
forall a. (Storable a, Marshal a) => a -> CMarshalSlow ()
writeAsBytesSlow
  {-# INLINABLE pushInt32 #-}
  pushUInt32 :: Word32 -> CMarshalSlow ()
pushUInt32 = Word32 -> CMarshalSlow ()
forall a. (Storable a, Marshal a) => a -> CMarshalSlow ()
writeAsBytesSlow
  {-# INLINABLE pushUInt32 #-}
  pushFloat :: Float -> CMarshalSlow ()
pushFloat = Float -> CMarshalSlow ()
forall a. (Storable a, Marshal a) => a -> CMarshalSlow ()
writeAsBytesSlow
  {-# INLINABLE pushFloat #-}
  pushString :: String -> CMarshalSlow ()
pushString String
str = ShortText -> CMarshalSlow ()
forall (m :: * -> *). MonadPush m => ShortText -> m ()
pushText (ShortText -> CMarshalSlow ()) -> ShortText -> CMarshalSlow ()
forall a b. (a -> b) -> a -> b
$ String -> ShortText
TS.pack String
str
  {-# INLINABLE pushString #-}
  pushTextUtf16 :: Text -> CMarshalSlow ()
pushTextUtf16 Text
str = ShortText -> CMarshalSlow ()
forall (m :: * -> *). MonadPush m => ShortText -> m ()
pushText (ShortText -> CMarshalSlow ()) -> ShortText -> CMarshalSlow ()
forall a b. (a -> b) -> a -> b
$ Text -> ShortText
TS.fromText Text
str
  {-# INLINABLE pushTextUtf16 #-}
  pushText :: ShortText -> CMarshalSlow ()
pushText ShortText
txt = do
    let bs :: ByteString
bs = ShortText -> ByteString
TS.toByteString ShortText
txt  -- TODO: is it possible to get rid of this copy?
        len :: Int
len = ByteString -> Int
BS.length ByteString
bs
    Int -> CMarshalSlow ()
resizeBufWhenNeeded (Int
ramDomainSize Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len)
    Word32 -> CMarshalSlow ()
forall (m :: * -> *). MonadPush m => Word32 -> m ()
pushUInt32 (Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len)
    if Int
len Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
      then () -> CMarshalSlow ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      else do
        Ptr CChar
ptr <- (MarshalState -> Ptr CChar) -> CMarshalSlow (Ptr CChar)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (Ptr ByteBuf -> Ptr CChar
forall a b. Ptr a -> Ptr b
castPtr (Ptr ByteBuf -> Ptr CChar)
-> (MarshalState -> Ptr ByteBuf) -> MarshalState -> Ptr CChar
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MarshalState -> Ptr ByteBuf
_ptr)
        IO () -> CMarshalSlow ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> CMarshalSlow ()) -> IO () -> CMarshalSlow ()
forall a b. (a -> b) -> a -> b
$ ByteString -> (Ptr CChar -> IO ()) -> IO ()
forall a. ByteString -> (Ptr CChar -> IO a) -> IO a
BSU.unsafeUseAsCString ByteString
bs ((Ptr CChar -> IO ()) -> IO ()) -> (Ptr CChar -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ (Ptr CChar -> Int -> IO ()) -> Int -> Ptr CChar -> IO ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip (Ptr CChar -> Ptr CChar -> Int -> IO ()
forall a. Ptr a -> Ptr a -> Int -> IO ()
copyBytes Ptr CChar
ptr) Int
len
        Int -> CMarshalSlow ()
incrementPtr Int
len
  {-# INLINABLE pushText #-}

writeAsBytesSlow :: (S.Storable a, Marshal a) => a -> CMarshalSlow ()
writeAsBytesSlow :: forall a. (Storable a, Marshal a) => a -> CMarshalSlow ()
writeAsBytesSlow a
a = do
  Int -> CMarshalSlow ()
resizeBufWhenNeeded Int
ramDomainSize
  Ptr a
ptr <- (MarshalState -> Ptr a) -> CMarshalSlow (Ptr a)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (Ptr ByteBuf -> Ptr a
forall a b. Ptr a -> Ptr b
castPtr (Ptr ByteBuf -> Ptr a)
-> (MarshalState -> Ptr ByteBuf) -> MarshalState -> Ptr a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MarshalState -> Ptr ByteBuf
_ptr)
  IO () -> CMarshalSlow ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> CMarshalSlow ()) -> IO () -> CMarshalSlow ()
forall a b. (a -> b) -> a -> b
$ Ptr a -> a -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
S.poke Ptr a
ptr a
a
  Int -> CMarshalSlow ()
incrementPtr Int
ramDomainSize
{-# INLINABLE writeAsBytesSlow #-}


class Collect c where
  collect :: Marshal a => Word32 -> CMarshalFast (c a)

instance Collect [] where
  collect :: forall a. Marshal a => Word32 -> CMarshalFast [a]
collect Word32
objCount = Word32 -> [a] -> CMarshalFast [a]
forall {t} {f :: * -> *} {a}.
(Eq t, Num t, Marshal a, MonadPop f) =>
t -> [a] -> f [a]
go Word32
objCount [] where
    go :: t -> [a] -> f [a]
go t
count [a]
acc
      | t
count t -> t -> Bool
forall a. Eq a => a -> a -> Bool
== t
0 = [a] -> f [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [a]
acc
      | Bool
otherwise = do
        !a
x <- f a
forall a (m :: * -> *). (Marshal a, MonadPop m) => m a
pop
        t -> [a] -> f [a]
go (t
count t -> t -> t
forall a. Num a => a -> a -> a
- t
1) (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
acc)
  {-# INLINABLE collect #-}

instance Collect V.Vector where
  collect :: forall a. Marshal a => Word32 -> CMarshalFast (Vector a)
collect Word32
objCount = do
    MVector RealWorld a
vm <- IO (MVector RealWorld a) -> CMarshalFast (MVector RealWorld a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (MVector RealWorld a) -> CMarshalFast (MVector RealWorld a))
-> IO (MVector RealWorld a) -> CMarshalFast (MVector RealWorld a)
forall a b. (a -> b) -> a -> b
$ Int -> IO (MVector (PrimState IO) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> m (MVector (PrimState m) a)
MV.unsafeNew Int
objCount'
    MVector RealWorld a -> Int -> CMarshalFast (Vector a)
collect' MVector RealWorld a
vm Int
0
    where
      objCount' :: Int
objCount' = Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
objCount
      collect' :: MVector RealWorld a -> Int -> CMarshalFast (Vector a)
collect' MVector RealWorld a
vec Int
idx
        | Int
idx Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
objCount' = IO (Vector a) -> CMarshalFast (Vector a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Vector a) -> CMarshalFast (Vector a))
-> IO (Vector a) -> CMarshalFast (Vector a)
forall a b. (a -> b) -> a -> b
$ MVector (PrimState IO) a -> IO (Vector a)
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> m (Vector a)
V.unsafeFreeze MVector RealWorld a
MVector (PrimState IO) a
vec
        | Bool
otherwise = do
          !a
obj <- CMarshalFast a
forall a (m :: * -> *). (Marshal a, MonadPop m) => m a
pop
          IO () -> CMarshalFast ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> CMarshalFast ()) -> IO () -> CMarshalFast ()
forall a b. (a -> b) -> a -> b
$ MVector (PrimState IO) a -> Int -> a -> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
MV.write MVector RealWorld a
MVector (PrimState IO) a
vec Int
idx a
obj
          MVector RealWorld a -> Int -> CMarshalFast (Vector a)
collect' MVector RealWorld a
vec (Int
idx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  {-# INLINABLE collect #-}

instance Collect (A.Array Int) where
  collect :: forall a. Marshal a => Word32 -> CMarshalFast (Array Int a)
collect Word32
objCount = do
    IOArray Int a
ma <- IO (IOArray Int a) -> CMarshalFast (IOArray Int a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (IOArray Int a) -> CMarshalFast (IOArray Int a))
-> IO (IOArray Int a) -> CMarshalFast (IOArray Int a)
forall a b. (a -> b) -> a -> b
$ (Int, Int) -> IO (IOArray Int a)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> m (a i e)
A.newArray_ (Int
0, Int
objCount' Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
    IOArray Int a -> Int -> CMarshalFast (Array Int a)
forall a.
Marshal a =>
IOArray Int a -> Int -> CMarshalFast (Array Int a)
collect' IOArray Int a
ma Int
0
    where
      objCount' :: Int
objCount' = Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
objCount
      collect' :: Marshal a => A.IOArray Int a -> Int -> CMarshalFast (A.Array Int a)
      collect' :: forall a.
Marshal a =>
IOArray Int a -> Int -> CMarshalFast (Array Int a)
collect' IOArray Int a
array Int
idx
        | Int
idx Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
objCount' = IO (Array Int a) -> CMarshalFast (Array Int a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Array Int a) -> CMarshalFast (Array Int a))
-> IO (Array Int a) -> CMarshalFast (Array Int a)
forall a b. (a -> b) -> a -> b
$ IOArray Int a -> IO (Array Int a)
forall i (a :: * -> * -> *) e (m :: * -> *) (b :: * -> * -> *).
(Ix i, MArray a e m, IArray b e) =>
a i e -> m (b i e)
A.unsafeFreeze IOArray Int a
array
        | Bool
otherwise = do
          !a
obj <- CMarshalFast a
forall a (m :: * -> *). (Marshal a, MonadPop m) => m a
pop
          IO () -> CMarshalFast ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> CMarshalFast ()) -> IO () -> CMarshalFast ()
forall a b. (a -> b) -> a -> b
$ IOArray Int a -> Int -> a -> IO ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
A.writeArray IOArray Int a
array Int
idx a
obj
          IOArray Int a -> Int -> CMarshalFast (Array Int a)
forall a.
Marshal a =>
IOArray Int a -> Int -> CMarshalFast (Array Int a)
collect' IOArray Int a
array (Int
idx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  {-# INLINABLE collect #-}

-- | A helper typeclass constraint, needed to serialize Datalog facts from
--   Haskell to C++.
type Submit a = ToByteSize (GetFields (Rep a))

instance MonadSouffle SouffleM where
  type Handler SouffleM = Handle
  type CollectFacts SouffleM c = Collect c
  type SubmitFacts SouffleM a = Submit a

  run :: forall prog. Handler SouffleM prog -> SouffleM ()
run (Handle ForeignPtr Souffle
prog MVar BufData
_) = IO () -> SouffleM ()
forall a. IO a -> SouffleM a
SouffleM (IO () -> SouffleM ()) -> IO () -> SouffleM ()
forall a b. (a -> b) -> a -> b
$ ForeignPtr Souffle -> IO ()
Internal.run ForeignPtr Souffle
prog
  {-# INLINABLE run #-}

  setNumThreads :: forall prog. Handler SouffleM prog -> Word64 -> SouffleM ()
setNumThreads (Handle ForeignPtr Souffle
prog MVar BufData
_) Word64
numCores =
    IO () -> SouffleM ()
forall a. IO a -> SouffleM a
SouffleM (IO () -> SouffleM ()) -> IO () -> SouffleM ()
forall a b. (a -> b) -> a -> b
$ ForeignPtr Souffle -> Word64 -> IO ()
Internal.setNumThreads ForeignPtr Souffle
prog Word64
numCores
  {-# INLINABLE setNumThreads #-}

  getNumThreads :: forall prog. Handler SouffleM prog -> SouffleM Word64
getNumThreads (Handle ForeignPtr Souffle
prog MVar BufData
_) =
    IO Word64 -> SouffleM Word64
forall a. IO a -> SouffleM a
SouffleM (IO Word64 -> SouffleM Word64) -> IO Word64 -> SouffleM Word64
forall a b. (a -> b) -> a -> b
$ ForeignPtr Souffle -> IO Word64
Internal.getNumThreads ForeignPtr Souffle
prog
  {-# INLINABLE getNumThreads #-}

  addFact :: forall a prog. (Fact a, ContainsInputFact prog a, Submit a)
          => Handle prog -> a -> SouffleM ()
  addFact :: forall a prog.
(Fact a, ContainsInputFact prog a, Submit a) =>
Handle prog -> a -> SouffleM ()
addFact (Handle ForeignPtr Souffle
prog MVar BufData
bufVar) a
fact = IO () -> SouffleM ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> SouffleM ()) -> IO () -> SouffleM ()
forall a b. (a -> b) -> a -> b
$ do
    let relationName :: String
relationName = Proxy a -> String
forall a. Fact a => Proxy a -> String
factName (Proxy a
forall {k} (t :: k). Proxy t
Proxy :: Proxy a)
    Ptr Relation
relation <- ForeignPtr Souffle -> String -> IO (Ptr Relation)
Internal.getRelation ForeignPtr Souffle
prog String
relationName
    MVar BufData -> Ptr Relation -> Identity a -> IO ()
forall (f :: * -> *) a.
(Foldable f, Marshal a, Submit a) =>
MVar BufData -> Ptr Relation -> f a -> IO ()
writeBytes MVar BufData
bufVar Ptr Relation
relation (a -> Identity a
forall a. a -> Identity a
Identity a
fact)
  {-# INLINABLE addFact #-}

  addFacts :: forall t a prog. (Foldable t, Fact a, ContainsInputFact prog a, Submit a)
           => Handle prog -> t a -> SouffleM ()
  addFacts :: forall (t :: * -> *) a prog.
(Foldable t, Fact a, ContainsInputFact prog a, Submit a) =>
Handle prog -> t a -> SouffleM ()
addFacts (Handle ForeignPtr Souffle
prog MVar BufData
bufVar) t a
facts = IO () -> SouffleM ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> SouffleM ()) -> IO () -> SouffleM ()
forall a b. (a -> b) -> a -> b
$ do
    let relationName :: String
relationName = Proxy a -> String
forall a. Fact a => Proxy a -> String
factName (Proxy a
forall {k} (t :: k). Proxy t
Proxy :: Proxy a)
    Ptr Relation
relation <- ForeignPtr Souffle -> String -> IO (Ptr Relation)
Internal.getRelation ForeignPtr Souffle
prog String
relationName
    MVar BufData -> Ptr Relation -> t a -> IO ()
forall (f :: * -> *) a.
(Foldable f, Marshal a, Submit a) =>
MVar BufData -> Ptr Relation -> f a -> IO ()
writeBytes MVar BufData
bufVar Ptr Relation
relation t a
facts
  {-# INLINABLE addFacts #-}

  getFacts :: forall a c prog. (Fact a, ContainsOutputFact prog a, Collect c)
           => Handle prog -> SouffleM (c a)
  getFacts :: forall a (c :: * -> *) prog.
(Fact a, ContainsOutputFact prog a, Collect c) =>
Handle prog -> SouffleM (c a)
getFacts (Handle ForeignPtr Souffle
prog MVar BufData
_) = IO (c a) -> SouffleM (c a)
forall a. IO a -> SouffleM a
SouffleM (IO (c a) -> SouffleM (c a)) -> IO (c a) -> SouffleM (c a)
forall a b. (a -> b) -> a -> b
$ do
    let relationName :: String
relationName = Proxy a -> String
forall a. Fact a => Proxy a -> String
factName (Proxy a
forall {k} (t :: k). Proxy t
Proxy :: Proxy a)
    Ptr Relation
relation <- ForeignPtr Souffle -> String -> IO (Ptr Relation)
Internal.getRelation ForeignPtr Souffle
prog String
relationName
    Ptr ByteBuf
buf <- ForeignPtr Souffle
-> (Ptr Souffle -> IO (Ptr ByteBuf)) -> IO (Ptr ByteBuf)
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Souffle
prog ((Ptr Souffle -> IO (Ptr ByteBuf)) -> IO (Ptr ByteBuf))
-> (Ptr Souffle -> IO (Ptr ByteBuf)) -> IO (Ptr ByteBuf)
forall a b. (a -> b) -> a -> b
$ (Ptr Souffle -> Ptr Relation -> IO (Ptr ByteBuf))
-> Ptr Relation -> Ptr Souffle -> IO (Ptr ByteBuf)
forall a b c. (a -> b -> c) -> b -> a -> c
flip Ptr Souffle -> Ptr Relation -> IO (Ptr ByteBuf)
Internal.popFacts Ptr Relation
relation
    (CMarshalFast (c a) -> Ptr ByteBuf -> IO (c a))
-> Ptr ByteBuf -> CMarshalFast (c a) -> IO (c a)
forall a b c. (a -> b -> c) -> b -> a -> c
flip CMarshalFast (c a) -> Ptr ByteBuf -> IO (c a)
forall a. CMarshalFast a -> Ptr ByteBuf -> IO a
runMarshalFastM Ptr ByteBuf
buf (CMarshalFast (c a) -> IO (c a)) -> CMarshalFast (c a) -> IO (c a)
forall a b. (a -> b) -> a -> b
$ Word32 -> CMarshalFast (c a)
forall (c :: * -> *) a.
(Collect c, Marshal a) =>
Word32 -> CMarshalFast (c a)
collect (Word32 -> CMarshalFast (c a))
-> CMarshalFast Word32 -> CMarshalFast (c a)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< CMarshalFast Word32
forall (m :: * -> *). MonadPop m => m Word32
popUInt32
  {-# INLINABLE getFacts #-}

  findFact :: forall a prog. (Fact a, ContainsOutputFact prog a, Submit a)
           => Handle prog -> a -> SouffleM (Maybe a)
  findFact :: forall a prog.
(Fact a, ContainsOutputFact prog a, Submit a) =>
Handle prog -> a -> SouffleM (Maybe a)
findFact (Handle ForeignPtr Souffle
prog MVar BufData
bufVar) a
fact = IO (Maybe a) -> SouffleM (Maybe a)
forall a. IO a -> SouffleM a
SouffleM (IO (Maybe a) -> SouffleM (Maybe a))
-> IO (Maybe a) -> SouffleM (Maybe a)
forall a b. (a -> b) -> a -> b
$ do
    let relationName :: String
relationName = Proxy a -> String
forall a. Fact a => Proxy a -> String
factName (Proxy a
forall {k} (t :: k). Proxy t
Proxy :: Proxy a)
    Ptr Relation
relation <- ForeignPtr Souffle -> String -> IO (Ptr Relation)
Internal.getRelation ForeignPtr Souffle
prog String
relationName
    Bool
found <- case Proxy a -> ByteSize
forall a. Submit a => Proxy a -> ByteSize
estimateNumBytes (forall {t}. Proxy t
forall {k} (t :: k). Proxy t
Proxy @a) of
      Exact Int
numBytes -> do
        MVar BufData -> (BufData -> IO (BufData, Bool)) -> IO Bool
forall a b. MVar a -> (a -> IO (a, b)) -> IO b
modifyMVarMasked MVar BufData
bufVar ((BufData -> IO (BufData, Bool)) -> IO Bool)
-> (BufData -> IO (BufData, Bool)) -> IO Bool
forall a b. (a -> b) -> a -> b
$ \BufData
bufData -> do
          BufData
bufData' <- if BufData -> Int
bufSize BufData
bufData Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
numBytes
            then BufData -> IO BufData
forall (f :: * -> *) a. Applicative f => a -> f a
pure BufData
bufData
            else (ForeignPtr ByteBuf -> Int -> BufData)
-> Int -> ForeignPtr ByteBuf -> BufData
forall a b c. (a -> b -> c) -> b -> a -> c
flip ForeignPtr ByteBuf -> Int -> BufData
BufData Int
numBytes (ForeignPtr ByteBuf -> BufData)
-> IO (ForeignPtr ByteBuf) -> IO BufData
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> IO (ForeignPtr ByteBuf)
forall (m :: * -> *). MonadIO m => Int -> m (ForeignPtr ByteBuf)
allocateBuf Int
numBytes
          Bool
found <- ForeignPtr ByteBuf -> (Ptr ByteBuf -> IO Bool) -> IO Bool
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr (BufData -> ForeignPtr ByteBuf
bufPtr BufData
bufData') ((Ptr ByteBuf -> IO Bool) -> IO Bool)
-> (Ptr ByteBuf -> IO Bool) -> IO Bool
forall a b. (a -> b) -> a -> b
$ \Ptr ByteBuf
ptr -> do
            CMarshalFast () -> Ptr ByteBuf -> IO ()
forall a. CMarshalFast a -> Ptr ByteBuf -> IO a
runMarshalFastM (a -> CMarshalFast ()
forall a (m :: * -> *). (Marshal a, MonadPush m) => a -> m ()
push a
fact) Ptr ByteBuf
ptr
            Ptr Relation -> Ptr ByteBuf -> IO Bool
Internal.containsFact Ptr Relation
relation Ptr ByteBuf
ptr
          (BufData, Bool) -> IO (BufData, Bool)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (BufData
bufData', Bool
found)
      Estimated Int
numBytes -> MVar BufData -> (BufData -> IO (BufData, Bool)) -> IO Bool
forall a b. MVar a -> (a -> IO (a, b)) -> IO b
modifyMVarMasked MVar BufData
bufVar ((BufData -> IO (BufData, Bool)) -> IO Bool)
-> (BufData -> IO (BufData, Bool)) -> IO Bool
forall a b. (a -> b) -> a -> b
$ \BufData
bufData ->
        BufData
-> Int -> CMarshalSlow (BufData, Bool) -> IO (BufData, Bool)
forall a. BufData -> Int -> CMarshalSlow a -> IO a
runMarshalSlowM BufData
bufData Int
numBytes (CMarshalSlow (BufData, Bool) -> IO (BufData, Bool))
-> CMarshalSlow (BufData, Bool) -> IO (BufData, Bool)
forall a b. (a -> b) -> a -> b
$ do
          a -> CMarshalSlow ()
forall a (m :: * -> *). (Marshal a, MonadPush m) => a -> m ()
push a
fact
          BufData
bufData' <- (MarshalState -> BufData) -> CMarshalSlow BufData
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets MarshalState -> BufData
_buf
          IO (BufData, Bool) -> CMarshalSlow (BufData, Bool)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (BufData, Bool) -> CMarshalSlow (BufData, Bool))
-> IO (BufData, Bool) -> CMarshalSlow (BufData, Bool)
forall a b. (a -> b) -> a -> b
$ ForeignPtr ByteBuf
-> (Ptr ByteBuf -> IO (BufData, Bool)) -> IO (BufData, Bool)
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr (BufData -> ForeignPtr ByteBuf
bufPtr BufData
bufData') ((Ptr ByteBuf -> IO (BufData, Bool)) -> IO (BufData, Bool))
-> (Ptr ByteBuf -> IO (BufData, Bool)) -> IO (BufData, Bool)
forall a b. (a -> b) -> a -> b
$ \Ptr ByteBuf
ptr -> do
            Bool
found <- Ptr Relation -> Ptr ByteBuf -> IO Bool
Internal.containsFact Ptr Relation
relation Ptr ByteBuf
ptr
            (BufData, Bool) -> IO (BufData, Bool)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (BufData
bufData', Bool
found)
    Maybe a -> IO (Maybe a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe a -> IO (Maybe a)) -> Maybe a -> IO (Maybe a)
forall a b. (a -> b) -> a -> b
$ if Bool
found then a -> Maybe a
forall a. a -> Maybe a
Just a
fact else Maybe a
forall a. Maybe a
Nothing
  {-# INLINABLE findFact #-}

instance MonadSouffleFileIO SouffleM where
  loadFiles :: forall prog. Handler SouffleM prog -> String -> SouffleM ()
loadFiles (Handle ForeignPtr Souffle
prog MVar BufData
_) = IO () -> SouffleM ()
forall a. IO a -> SouffleM a
SouffleM (IO () -> SouffleM ())
-> (String -> IO ()) -> String -> SouffleM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ForeignPtr Souffle -> String -> IO ()
Internal.loadAll ForeignPtr Souffle
prog
  {-# INLINABLE loadFiles #-}

  writeFiles :: forall prog. Handler SouffleM prog -> String -> SouffleM ()
writeFiles (Handle ForeignPtr Souffle
prog MVar BufData
_) = IO () -> SouffleM ()
forall a. IO a -> SouffleM a
SouffleM (IO () -> SouffleM ())
-> (String -> IO ()) -> String -> SouffleM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ForeignPtr Souffle -> String -> IO ()
Internal.printAll ForeignPtr Souffle
prog
  {-# INLINABLE writeFiles #-}


data ByteSize
  = Exact {-# UNPACK #-} !ByteCount
  | Estimated {-# UNPACK #-} !ByteCount

instance Semigroup ByteSize where
  Exact Int
s1 <> :: ByteSize -> ByteSize -> ByteSize
<> Exact Int
s2 = Int -> ByteSize
Exact (Int
s1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
s2)
  Exact Int
s1 <> Estimated Int
s2 = Int -> ByteSize
Estimated (Int
s1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
s2)
  Estimated Int
s1 <> Exact Int
s2 = Int -> ByteSize
Estimated (Int
s1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
s2)
  Estimated Int
s1 <> Estimated Int
s2 = Int -> ByteSize
Estimated (Int
s1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
s2)
  {-# INLINABLE (<>) #-}

class ToByteSize (a :: k) where
  toByteSize :: Proxy a -> ByteSize

instance ToByteSize Int32 where
  toByteSize :: Proxy Int32 -> ByteSize
toByteSize = ByteSize -> Proxy Int32 -> ByteSize
forall a b. a -> b -> a
const (ByteSize -> Proxy Int32 -> ByteSize)
-> ByteSize -> Proxy Int32 -> ByteSize
forall a b. (a -> b) -> a -> b
$ Int -> ByteSize
Exact Int
4
  {-# INLINABLE toByteSize #-}

instance ToByteSize Word32 where
  toByteSize :: Proxy Word32 -> ByteSize
toByteSize = ByteSize -> Proxy Word32 -> ByteSize
forall a b. a -> b -> a
const (ByteSize -> Proxy Word32 -> ByteSize)
-> ByteSize -> Proxy Word32 -> ByteSize
forall a b. (a -> b) -> a -> b
$ Int -> ByteSize
Exact Int
4
  {-# INLINABLE toByteSize #-}

instance ToByteSize Float where
  toByteSize :: Proxy Float -> ByteSize
toByteSize = ByteSize -> Proxy Float -> ByteSize
forall a b. a -> b -> a
const (ByteSize -> Proxy Float -> ByteSize)
-> ByteSize -> Proxy Float -> ByteSize
forall a b. (a -> b) -> a -> b
$ Int -> ByteSize
Exact Int
4
  {-# INLINABLE toByteSize #-}

instance ToByteSize String where
  -- 4 for length prefix + 32 for actual string
  toByteSize :: Proxy String -> ByteSize
toByteSize = ByteSize -> Proxy String -> ByteSize
forall a b. a -> b -> a
const (ByteSize -> Proxy String -> ByteSize)
-> ByteSize -> Proxy String -> ByteSize
forall a b. (a -> b) -> a -> b
$ Int -> ByteSize
Estimated Int
36
  {-# INLINABLE toByteSize #-}

instance ToByteSize T.Text where
  -- 4 for length prefix + 32 for actual string
  toByteSize :: Proxy Text -> ByteSize
toByteSize = ByteSize -> Proxy Text -> ByteSize
forall a b. a -> b -> a
const (ByteSize -> Proxy Text -> ByteSize)
-> ByteSize -> Proxy Text -> ByteSize
forall a b. (a -> b) -> a -> b
$ Int -> ByteSize
Estimated Int
36
  {-# INLINABLE toByteSize #-}

instance ToByteSize TL.Text where
  -- 4 for length prefix + 32 for actual string
  toByteSize :: Proxy Text -> ByteSize
toByteSize = ByteSize -> Proxy Text -> ByteSize
forall a b. a -> b -> a
const (ByteSize -> Proxy Text -> ByteSize)
-> ByteSize -> Proxy Text -> ByteSize
forall a b. (a -> b) -> a -> b
$ Int -> ByteSize
Estimated Int
36
  {-# INLINABLE toByteSize #-}

instance ToByteSize TS.ShortText where
  -- 4 for length prefix + 32 for actual string
  toByteSize :: Proxy ShortText -> ByteSize
toByteSize = ByteSize -> Proxy ShortText -> ByteSize
forall a b. a -> b -> a
const (ByteSize -> Proxy ShortText -> ByteSize)
-> ByteSize -> Proxy ShortText -> ByteSize
forall a b. (a -> b) -> a -> b
$ Int -> ByteSize
Estimated Int
36
  {-# INLINABLE toByteSize #-}

instance ToByteSize '[] where
  toByteSize :: Proxy '[] -> ByteSize
toByteSize = ByteSize -> Proxy '[] -> ByteSize
forall a b. a -> b -> a
const (ByteSize -> Proxy '[] -> ByteSize)
-> ByteSize -> Proxy '[] -> ByteSize
forall a b. (a -> b) -> a -> b
$ Int -> ByteSize
Exact Int
0
  {-# INLINABLE toByteSize #-}

instance (ToByteSize a, ToByteSize as) => ToByteSize (a ': as) where
  toByteSize :: Proxy (a : as) -> ByteSize
toByteSize =
    ByteSize -> Proxy (a : as) -> ByteSize
forall a b. a -> b -> a
const (ByteSize -> Proxy (a : as) -> ByteSize)
-> ByteSize -> Proxy (a : as) -> ByteSize
forall a b. (a -> b) -> a -> b
$ Proxy a -> ByteSize
forall k (a :: k). ToByteSize a => Proxy a -> ByteSize
toByteSize (forall {t :: a}. Proxy t
forall {k} (t :: k). Proxy t
Proxy @a) ByteSize -> ByteSize -> ByteSize
forall a. Semigroup a => a -> a -> a
<> Proxy as -> ByteSize
forall k (a :: k). ToByteSize a => Proxy a -> ByteSize
toByteSize (forall {t :: [a]}. Proxy t
forall {k} (t :: k). Proxy t
Proxy @as)
  {-# INLINABLE toByteSize #-}

-- | A helper type family, for getting all directly marshallable fields of a type.
type family GetFields (a :: k) :: [Type] where
  GetFields (K1 _ a) = DoGetFields a
  GetFields (M1 _ _ a) = GetFields a
  GetFields (f :*: g) = GetFields f ++ GetFields g

type family DoGetFields (a :: Type) :: [Type] where
  DoGetFields Int32 = '[Int32]
  DoGetFields Word32 = '[Word32]
  DoGetFields Float = '[Float]
  DoGetFields String = '[String]
  DoGetFields T.Text = '[T.Text]
  DoGetFields TL.Text = '[TL.Text]
  DoGetFields TS.ShortText = '[TS.ShortText]
  DoGetFields a = GetFields (Rep a)

type family a ++ b where
  '[] ++ b = b
  (a ': as) ++ bs = a ': as ++ bs

estimateNumBytes :: forall a. Submit a => Proxy a -> ByteSize
estimateNumBytes :: forall a. Submit a => Proxy a -> ByteSize
estimateNumBytes Proxy a
_ = Proxy (GetFields (Rep a)) -> ByteSize
forall k (a :: k). ToByteSize a => Proxy a -> ByteSize
toByteSize (forall {t :: [*]}. Proxy t
forall {k} (t :: k). Proxy t
Proxy @(GetFields (Rep a)))
{-# INLINABLE estimateNumBytes #-}

writeBytes :: forall f a. (Foldable f, Marshal a, Submit a)
           => MVar BufData -> Ptr Internal.Relation -> f a -> IO ()
writeBytes :: forall (f :: * -> *) a.
(Foldable f, Marshal a, Submit a) =>
MVar BufData -> Ptr Relation -> f a -> IO ()
writeBytes MVar BufData
bufVar Ptr Relation
relation f a
fa = case Proxy a -> ByteSize
forall a. Submit a => Proxy a -> ByteSize
estimateNumBytes (forall {t}. Proxy t
forall {k} (t :: k). Proxy t
Proxy @a) of
  Exact Int
numBytes -> MVar BufData -> (BufData -> IO BufData) -> IO ()
forall a. MVar a -> (a -> IO a) -> IO ()
modifyMVarMasked_ MVar BufData
bufVar ((BufData -> IO BufData) -> IO ())
-> (BufData -> IO BufData) -> IO ()
forall a b. (a -> b) -> a -> b
$ \BufData
bufData -> do
    let totalByteCount :: Int
totalByteCount = Int
numBytes Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
objCount
    BufData
bufData' <- if BufData -> Int
bufSize BufData
bufData Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
totalByteCount
      then BufData -> IO BufData
forall (f :: * -> *) a. Applicative f => a -> f a
pure BufData
bufData
      else (ForeignPtr ByteBuf -> Int -> BufData)
-> Int -> ForeignPtr ByteBuf -> BufData
forall a b c. (a -> b -> c) -> b -> a -> c
flip ForeignPtr ByteBuf -> Int -> BufData
BufData Int
totalByteCount (ForeignPtr ByteBuf -> BufData)
-> IO (ForeignPtr ByteBuf) -> IO BufData
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> IO (ForeignPtr ByteBuf)
forall (m :: * -> *). MonadIO m => Int -> m (ForeignPtr ByteBuf)
allocateBuf Int
totalByteCount
    ForeignPtr ByteBuf -> (Ptr ByteBuf -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr (BufData -> ForeignPtr ByteBuf
bufPtr BufData
bufData') ((Ptr ByteBuf -> IO ()) -> IO ())
-> (Ptr ByteBuf -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr ByteBuf
ptr -> do
      CMarshalFast () -> Ptr ByteBuf -> IO ()
forall a. CMarshalFast a -> Ptr ByteBuf -> IO a
runMarshalFastM ((a -> CMarshalFast ()) -> f a -> CMarshalFast ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ a -> CMarshalFast ()
forall a (m :: * -> *). (Marshal a, MonadPush m) => a -> m ()
push f a
fa) Ptr ByteBuf
ptr
      Ptr Relation -> Ptr ByteBuf -> Word64 -> IO ()
Internal.pushFacts Ptr Relation
relation Ptr ByteBuf
ptr (Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
objCount)
    BufData -> IO BufData
forall (f :: * -> *) a. Applicative f => a -> f a
pure BufData
bufData'

  Estimated Int
numBytes -> MVar BufData -> (BufData -> IO BufData) -> IO ()
forall a. MVar a -> (a -> IO a) -> IO ()
modifyMVarMasked_ MVar BufData
bufVar ((BufData -> IO BufData) -> IO ())
-> (BufData -> IO BufData) -> IO ()
forall a b. (a -> b) -> a -> b
$ \BufData
bufData ->
    BufData -> Int -> CMarshalSlow BufData -> IO BufData
forall a. BufData -> Int -> CMarshalSlow a -> IO a
runMarshalSlowM BufData
bufData (Int
numBytes Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
objCount) (CMarshalSlow BufData -> IO BufData)
-> CMarshalSlow BufData -> IO BufData
forall a b. (a -> b) -> a -> b
$ do
      (a -> CMarshalSlow ()) -> f a -> CMarshalSlow ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ a -> CMarshalSlow ()
forall a (m :: * -> *). (Marshal a, MonadPush m) => a -> m ()
push f a
fa
      BufData
bufData' <- (MarshalState -> BufData) -> CMarshalSlow BufData
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets MarshalState -> BufData
_buf
      IO BufData -> CMarshalSlow BufData
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO BufData -> CMarshalSlow BufData)
-> IO BufData -> CMarshalSlow BufData
forall a b. (a -> b) -> a -> b
$ ForeignPtr ByteBuf -> (Ptr ByteBuf -> IO BufData) -> IO BufData
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr (BufData -> ForeignPtr ByteBuf
bufPtr BufData
bufData') ((Ptr ByteBuf -> IO BufData) -> IO BufData)
-> (Ptr ByteBuf -> IO BufData) -> IO BufData
forall a b. (a -> b) -> a -> b
$ \Ptr ByteBuf
ptr -> do
        Ptr Relation -> Ptr ByteBuf -> Word64 -> IO ()
Internal.pushFacts Ptr Relation
relation Ptr ByteBuf
ptr (Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
objCount)
        BufData -> IO BufData
forall (f :: * -> *) a. Applicative f => a -> f a
pure BufData
bufData'
  where objCount :: Int
objCount = f a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length f a
fa
{-# INLINABLE writeBytes #-}