{-# 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 :: Type
type ByteCount = Int
type ByteBuf :: Type
type ByteBuf = Internal.ByteBuf

type BufData :: Type
data BufData
  = BufData
  { BufData -> ForeignPtr ByteBuf
bufPtr :: {-# UNPACK #-} !(ForeignPtr ByteBuf)
  , BufData -> ByteCount
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.
type Handle :: Type -> Type
data Handle prog
  = Handle {-# UNPACK #-} !(ForeignPtr Internal.Souffle)
           {-# UNPACK #-} !(MVar BufData)
type role Handle nominal

-- | A monad for executing Souffle-related actions in.
type SouffleM :: Type -> Type
newtype SouffleM a = SouffleM (IO a)
  deriving (a -> SouffleM b -> SouffleM a
(a -> b) -> SouffleM a -> SouffleM b
(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
<$ :: a -> SouffleM b -> SouffleM a
$c<$ :: forall a b. a -> SouffleM b -> SouffleM a
fmap :: (a -> b) -> SouffleM a -> SouffleM b
$cfmap :: forall a b. (a -> b) -> SouffleM a -> SouffleM b
Functor, Functor SouffleM
a -> SouffleM a
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
SouffleM a -> SouffleM b -> SouffleM b
SouffleM a -> SouffleM b -> SouffleM a
SouffleM (a -> b) -> SouffleM a -> SouffleM b
(a -> b -> c) -> SouffleM a -> SouffleM b -> SouffleM c
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
<* :: SouffleM a -> SouffleM b -> SouffleM a
$c<* :: forall a b. SouffleM a -> SouffleM b -> SouffleM a
*> :: SouffleM a -> SouffleM b -> SouffleM b
$c*> :: forall a b. SouffleM a -> SouffleM b -> SouffleM b
liftA2 :: (a -> b -> c) -> SouffleM a -> SouffleM b -> SouffleM c
$cliftA2 :: forall a b c.
(a -> b -> c) -> SouffleM a -> SouffleM b -> SouffleM c
<*> :: SouffleM (a -> b) -> SouffleM a -> SouffleM b
$c<*> :: forall a b. SouffleM (a -> b) -> SouffleM a -> SouffleM b
pure :: a -> SouffleM a
$cpure :: forall a. a -> SouffleM a
$cp1Applicative :: Functor SouffleM
Applicative, Applicative SouffleM
a -> SouffleM a
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
SouffleM a -> (a -> SouffleM b) -> SouffleM b
SouffleM a -> SouffleM b -> SouffleM b
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 :: a -> SouffleM a
$creturn :: forall a. a -> SouffleM a
>> :: SouffleM a -> SouffleM b -> SouffleM b
$c>> :: forall a b. SouffleM a -> SouffleM b -> SouffleM b
>>= :: SouffleM a -> (a -> SouffleM b) -> SouffleM b
$c>>= :: forall a b. SouffleM a -> (a -> SouffleM b) -> SouffleM b
$cp1Monad :: Applicative SouffleM
Monad, Monad SouffleM
Monad SouffleM
-> (forall a. IO a -> SouffleM a) -> MonadIO SouffleM
IO a -> SouffleM a
forall a. IO a -> SouffleM a
forall (m :: * -> *).
Monad m -> (forall a. IO a -> m a) -> MonadIO m
liftIO :: IO a -> SouffleM a
$cliftIO :: forall a. IO a -> SouffleM a
$cp1MonadIO :: Monad SouffleM
MonadIO) via IO
  deriving (b -> SouffleM a -> SouffleM a
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 :: 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
$cp1Monoid :: forall a. Monoid a => Semigroup (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 :: 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 -> ByteCount -> BufData
BufData ForeignPtr ByteBuf
ptr ByteCount
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 prog. 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++).
type CMarshalFast :: Type -> Type
newtype CMarshalFast a = CMarshalFast (StateT (Ptr ByteBuf) IO a)
  deriving (a -> CMarshalFast b -> CMarshalFast a
(a -> b) -> CMarshalFast a -> CMarshalFast b
(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
<$ :: a -> CMarshalFast b -> CMarshalFast a
$c<$ :: forall a b. a -> CMarshalFast b -> CMarshalFast a
fmap :: (a -> b) -> CMarshalFast a -> CMarshalFast b
$cfmap :: forall a b. (a -> b) -> CMarshalFast a -> CMarshalFast b
Functor, Functor CMarshalFast
a -> CMarshalFast a
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
CMarshalFast a -> CMarshalFast b -> CMarshalFast b
CMarshalFast a -> CMarshalFast b -> CMarshalFast a
CMarshalFast (a -> b) -> CMarshalFast a -> CMarshalFast b
(a -> b -> c) -> CMarshalFast a -> CMarshalFast b -> CMarshalFast c
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
<* :: CMarshalFast a -> CMarshalFast b -> CMarshalFast a
$c<* :: forall a b. CMarshalFast a -> CMarshalFast b -> CMarshalFast a
*> :: CMarshalFast a -> CMarshalFast b -> CMarshalFast b
$c*> :: forall a b. CMarshalFast a -> CMarshalFast b -> CMarshalFast b
liftA2 :: (a -> b -> c) -> CMarshalFast a -> CMarshalFast b -> CMarshalFast c
$cliftA2 :: forall a b c.
(a -> b -> c) -> CMarshalFast a -> CMarshalFast b -> CMarshalFast c
<*> :: CMarshalFast (a -> b) -> CMarshalFast a -> CMarshalFast b
$c<*> :: forall a b.
CMarshalFast (a -> b) -> CMarshalFast a -> CMarshalFast b
pure :: a -> CMarshalFast a
$cpure :: forall a. a -> CMarshalFast a
$cp1Applicative :: Functor CMarshalFast
Applicative, Applicative CMarshalFast
a -> CMarshalFast a
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
CMarshalFast a -> (a -> CMarshalFast b) -> CMarshalFast b
CMarshalFast a -> CMarshalFast b -> CMarshalFast b
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 :: a -> CMarshalFast a
$creturn :: forall a. a -> CMarshalFast a
>> :: CMarshalFast a -> CMarshalFast b -> CMarshalFast b
$c>> :: forall a b. CMarshalFast a -> CMarshalFast b -> CMarshalFast b
>>= :: CMarshalFast a -> (a -> CMarshalFast b) -> CMarshalFast b
$c>>= :: forall a b.
CMarshalFast a -> (a -> CMarshalFast b) -> CMarshalFast b
$cp1Monad :: Applicative CMarshalFast
Monad, Monad CMarshalFast
Monad CMarshalFast
-> (forall a. IO a -> CMarshalFast a) -> MonadIO CMarshalFast
IO a -> CMarshalFast a
forall a. IO a -> CMarshalFast a
forall (m :: * -> *).
Monad m -> (forall a. IO a -> m a) -> MonadIO m
liftIO :: IO a -> CMarshalFast a
$cliftIO :: forall a. IO a -> CMarshalFast a
$cp1MonadIO :: Monad CMarshalFast
MonadIO, MonadState (Ptr ByteBuf))
  via (StateT (Ptr ByteBuf) IO)

runMarshalFastM :: CMarshalFast a -> Ptr ByteBuf -> IO a
runMarshalFastM :: 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 :: ByteCount
ramDomainSize = ByteCount
4

writeAsBytes :: (S.Storable a, Marshal a) => a -> CMarshalFast ()
writeAsBytes :: 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 -> ByteCount -> Ptr ByteBuf
forall a b. Ptr a -> ByteCount -> Ptr b
`plusPtr` ByteCount
ramDomainSize
{-# INLINABLE writeAsBytes #-}

readAsBytes :: (S.Storable a, Marshal a) => CMarshalFast a
readAsBytes :: 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 -> ByteCount -> Ptr ByteBuf
forall a b. Ptr a -> ByteCount -> Ptr b
`plusPtr` ByteCount
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 -> ByteCount
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 -> ByteCount -> Ptr ByteBuf
forall a b. Ptr a -> ByteCount -> Ptr b
`plusPtr` Word32 -> ByteCount
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 #-}


type MarshalState :: Type
data MarshalState
  = MarshalState
  { MarshalState -> BufData
_buf :: {-# UNPACK #-} !BufData
  , MarshalState -> Ptr ByteBuf
_ptr :: {-# UNPACK #-} !(Ptr ByteBuf)
  , MarshalState -> ByteCount
_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).
type CMarshalSlow :: Type -> Type
newtype CMarshalSlow a = CMarshalSlow (StateT MarshalState IO a)
  deriving (a -> CMarshalSlow b -> CMarshalSlow a
(a -> b) -> CMarshalSlow a -> CMarshalSlow b
(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
<$ :: a -> CMarshalSlow b -> CMarshalSlow a
$c<$ :: forall a b. a -> CMarshalSlow b -> CMarshalSlow a
fmap :: (a -> b) -> CMarshalSlow a -> CMarshalSlow b
$cfmap :: forall a b. (a -> b) -> CMarshalSlow a -> CMarshalSlow b
Functor, Functor CMarshalSlow
a -> CMarshalSlow a
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
CMarshalSlow a -> CMarshalSlow b -> CMarshalSlow b
CMarshalSlow a -> CMarshalSlow b -> CMarshalSlow a
CMarshalSlow (a -> b) -> CMarshalSlow a -> CMarshalSlow b
(a -> b -> c) -> CMarshalSlow a -> CMarshalSlow b -> CMarshalSlow c
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
<* :: CMarshalSlow a -> CMarshalSlow b -> CMarshalSlow a
$c<* :: forall a b. CMarshalSlow a -> CMarshalSlow b -> CMarshalSlow a
*> :: CMarshalSlow a -> CMarshalSlow b -> CMarshalSlow b
$c*> :: forall a b. CMarshalSlow a -> CMarshalSlow b -> CMarshalSlow b
liftA2 :: (a -> b -> c) -> CMarshalSlow a -> CMarshalSlow b -> CMarshalSlow c
$cliftA2 :: forall a b c.
(a -> b -> c) -> CMarshalSlow a -> CMarshalSlow b -> CMarshalSlow c
<*> :: CMarshalSlow (a -> b) -> CMarshalSlow a -> CMarshalSlow b
$c<*> :: forall a b.
CMarshalSlow (a -> b) -> CMarshalSlow a -> CMarshalSlow b
pure :: a -> CMarshalSlow a
$cpure :: forall a. a -> CMarshalSlow a
$cp1Applicative :: Functor CMarshalSlow
Applicative, Applicative CMarshalSlow
a -> CMarshalSlow a
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
CMarshalSlow a -> (a -> CMarshalSlow b) -> CMarshalSlow b
CMarshalSlow a -> CMarshalSlow b -> CMarshalSlow b
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 :: a -> CMarshalSlow a
$creturn :: forall a. a -> CMarshalSlow a
>> :: CMarshalSlow a -> CMarshalSlow b -> CMarshalSlow b
$c>> :: forall a b. CMarshalSlow a -> CMarshalSlow b -> CMarshalSlow b
>>= :: CMarshalSlow a -> (a -> CMarshalSlow b) -> CMarshalSlow b
$c>>= :: forall a b.
CMarshalSlow a -> (a -> CMarshalSlow b) -> CMarshalSlow b
$cp1Monad :: Applicative CMarshalSlow
Monad, Monad CMarshalSlow
Monad CMarshalSlow
-> (forall a. IO a -> CMarshalSlow a) -> MonadIO CMarshalSlow
IO a -> CMarshalSlow a
forall a. IO a -> CMarshalSlow a
forall (m :: * -> *).
Monad m -> (forall a. IO a -> m a) -> MonadIO m
liftIO :: IO a -> CMarshalSlow a
$cliftIO :: forall a. IO a -> CMarshalSlow a
$cp1MonadIO :: Monad CMarshalSlow
MonadIO, MonadState MarshalState)
  via (StateT MarshalState IO)

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

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

allocateBuf :: MonadIO m => ByteCount -> m (ForeignPtr ByteBuf)
allocateBuf :: ByteCount -> m (ForeignPtr ByteBuf)
allocateBuf ByteCount
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
$
  ByteCount -> IO (ForeignPtr ByteBuf)
forall a. ByteCount -> IO (ForeignPtr a)
mallocForeignPtrBytes ByteCount
byteCount
{-# INLINABLE allocateBuf #-}

copyBuf :: ForeignPtr ByteBuf -> ForeignPtr ByteBuf -> Int -> CMarshalSlow ()
copyBuf :: ForeignPtr ByteBuf
-> ForeignPtr ByteBuf -> ByteCount -> CMarshalSlow ()
copyBuf ForeignPtr ByteBuf
dst ForeignPtr ByteBuf
src ByteCount
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 -> ByteCount -> IO ()
forall a. Ptr a -> Ptr a -> ByteCount -> IO ()
copyBytes Ptr ByteBuf
dstPtr Ptr ByteBuf
srcPtr ByteCount
byteCount
{-# INLINABLE copyBuf #-}

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

incrementPtr :: ByteCount -> CMarshalSlow ()
incrementPtr :: ByteCount -> CMarshalSlow ()
incrementPtr ByteCount
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 ByteCount
offset) ->
    BufData -> Ptr ByteBuf -> ByteCount -> MarshalState
MarshalState BufData
buf (Ptr ByteBuf
ptr Ptr ByteBuf -> ByteCount -> Ptr ByteBuf
forall a b. Ptr a -> ByteCount -> Ptr b
`plusPtr` ByteCount
byteCount) (ByteCount
offset ByteCount -> ByteCount -> ByteCount
forall a. Num a => a -> a -> a
+ ByteCount
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 :: ByteCount
len = ByteString -> ByteCount
BS.length ByteString
bs
    ByteCount -> CMarshalSlow ()
resizeBufWhenNeeded (ByteCount
ramDomainSize ByteCount -> ByteCount -> ByteCount
forall a. Num a => a -> a -> a
+ ByteCount
len)
    Word32 -> CMarshalSlow ()
forall (m :: * -> *). MonadPush m => Word32 -> m ()
pushUInt32 (ByteCount -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral ByteCount
len)
    if ByteCount
len ByteCount -> ByteCount -> Bool
forall a. Eq a => a -> a -> Bool
== ByteCount
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 -> ByteCount -> IO ())
-> ByteCount -> Ptr CChar -> IO ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip (Ptr CChar -> Ptr CChar -> ByteCount -> IO ()
forall a. Ptr a -> Ptr a -> ByteCount -> IO ()
copyBytes Ptr CChar
ptr) ByteCount
len
        ByteCount -> CMarshalSlow ()
incrementPtr ByteCount
len
  {-# INLINABLE pushText #-}

writeAsBytesSlow :: (S.Storable a, Marshal a) => a -> CMarshalSlow ()
writeAsBytesSlow :: a -> CMarshalSlow ()
writeAsBytesSlow a
a = do
  ByteCount -> CMarshalSlow ()
resizeBufWhenNeeded ByteCount
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
  ByteCount -> CMarshalSlow ()
incrementPtr ByteCount
ramDomainSize
{-# INLINABLE writeAsBytesSlow #-}


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

instance Collect [] where
  collect :: 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 :: 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
$ ByteCount -> IO (MVector (PrimState IO) a)
forall (m :: * -> *) a.
PrimMonad m =>
ByteCount -> m (MVector (PrimState m) a)
MV.unsafeNew ByteCount
objCount'
    MVector RealWorld a -> ByteCount -> CMarshalFast (Vector a)
collect' MVector RealWorld a
vm ByteCount
0
    where
      objCount' :: ByteCount
objCount' = Word32 -> ByteCount
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
objCount
      collect' :: MVector RealWorld a -> ByteCount -> CMarshalFast (Vector a)
collect' MVector RealWorld a
vec ByteCount
idx
        | ByteCount
idx ByteCount -> ByteCount -> Bool
forall a. Eq a => a -> a -> Bool
== ByteCount
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 -> ByteCount -> a -> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> ByteCount -> a -> m ()
MV.write MVector RealWorld a
MVector (PrimState IO) a
vec ByteCount
idx a
obj
          MVector RealWorld a -> ByteCount -> CMarshalFast (Vector a)
collect' MVector RealWorld a
vec (ByteCount
idx ByteCount -> ByteCount -> ByteCount
forall a. Num a => a -> a -> a
+ ByteCount
1)
  {-# INLINABLE collect #-}

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

-- | A helper typeclass constraint, needed to serialize Datalog facts from
--   Haskell to C++.
type Submit :: Type -> Constraint
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 :: Handler SouffleM prog -> SouffleM ()
run (Handle prog _) = 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 :: Handler SouffleM prog -> Word64 -> SouffleM ()
setNumThreads (Handle prog _) 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 :: Handler SouffleM prog -> SouffleM Word64
getNumThreads (Handle prog _) =
    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 :: 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 :: 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 :: 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 :: 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 (Proxy a
forall k (t :: k). Proxy t
Proxy @a) of
      Exact ByteCount
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 -> ByteCount
bufSize BufData
bufData ByteCount -> ByteCount -> Bool
forall a. Ord a => a -> a -> Bool
> ByteCount
numBytes
            then BufData -> IO BufData
forall (f :: * -> *) a. Applicative f => a -> f a
pure BufData
bufData
            else (ForeignPtr ByteBuf -> ByteCount -> BufData)
-> ByteCount -> ForeignPtr ByteBuf -> BufData
forall a b c. (a -> b -> c) -> b -> a -> c
flip ForeignPtr ByteBuf -> ByteCount -> BufData
BufData ByteCount
numBytes (ForeignPtr ByteBuf -> BufData)
-> IO (ForeignPtr ByteBuf) -> IO BufData
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ByteCount -> IO (ForeignPtr ByteBuf)
forall (m :: * -> *).
MonadIO m =>
ByteCount -> m (ForeignPtr ByteBuf)
allocateBuf ByteCount
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 ByteCount
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
-> ByteCount -> CMarshalSlow (BufData, Bool) -> IO (BufData, Bool)
forall a. BufData -> ByteCount -> CMarshalSlow a -> IO a
runMarshalSlowM BufData
bufData ByteCount
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 :: Handler SouffleM prog -> String -> SouffleM ()
loadFiles (Handle prog _) = 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 :: Handler SouffleM prog -> String -> SouffleM ()
writeFiles (Handle prog _) = 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 #-}


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

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

type ToByteSize :: k -> Constraint
class ToByteSize a 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
$ ByteCount -> ByteSize
Exact ByteCount
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
$ ByteCount -> ByteSize
Exact ByteCount
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
$ ByteCount -> ByteSize
Exact ByteCount
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
$ ByteCount -> ByteSize
Estimated ByteCount
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
$ ByteCount -> ByteSize
Estimated ByteCount
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
$ ByteCount -> ByteSize
Estimated ByteCount
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
$ ByteCount -> ByteSize
Estimated ByteCount
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
$ ByteCount -> ByteSize
Exact ByteCount
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 (Proxy a
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 (Proxy as
forall k (t :: k). Proxy t
Proxy @as)
  {-# INLINABLE toByteSize #-}

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

type DoGetFields :: Type -> [Type]
type family DoGetFields a 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 (++) :: [Type] -> [Type] -> [Type]
type family a ++ b where
  '[] ++ b = b
  (a ': as) ++ bs = a ': as ++ bs

estimateNumBytes :: forall a. Submit a => Proxy a -> ByteSize
estimateNumBytes :: Proxy a -> ByteSize
estimateNumBytes Proxy a
_ = Proxy (GetFields (Rep a)) -> ByteSize
forall k (a :: k). ToByteSize a => Proxy a -> ByteSize
toByteSize (Proxy (GetFields (Rep a))
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 :: 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 (Proxy a
forall k (t :: k). Proxy t
Proxy @a) of
  Exact ByteCount
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 :: ByteCount
totalByteCount = ByteCount
numBytes ByteCount -> ByteCount -> ByteCount
forall a. Num a => a -> a -> a
* ByteCount
objCount
    BufData
bufData' <- if BufData -> ByteCount
bufSize BufData
bufData ByteCount -> ByteCount -> Bool
forall a. Ord a => a -> a -> Bool
> ByteCount
totalByteCount
      then BufData -> IO BufData
forall (f :: * -> *) a. Applicative f => a -> f a
pure BufData
bufData
      else (ForeignPtr ByteBuf -> ByteCount -> BufData)
-> ByteCount -> ForeignPtr ByteBuf -> BufData
forall a b c. (a -> b -> c) -> b -> a -> c
flip ForeignPtr ByteBuf -> ByteCount -> BufData
BufData ByteCount
totalByteCount (ForeignPtr ByteBuf -> BufData)
-> IO (ForeignPtr ByteBuf) -> IO BufData
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ByteCount -> IO (ForeignPtr ByteBuf)
forall (m :: * -> *).
MonadIO m =>
ByteCount -> m (ForeignPtr ByteBuf)
allocateBuf ByteCount
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 (ByteCount -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral ByteCount
objCount)
    BufData -> IO BufData
forall (f :: * -> *) a. Applicative f => a -> f a
pure BufData
bufData'

  Estimated ByteCount
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 -> ByteCount -> CMarshalSlow BufData -> IO BufData
forall a. BufData -> ByteCount -> CMarshalSlow a -> IO a
runMarshalSlowM BufData
bufData (ByteCount
numBytes ByteCount -> ByteCount -> ByteCount
forall a. Num a => a -> a -> a
* ByteCount
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 (ByteCount -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral ByteCount
objCount)
        BufData -> IO BufData
forall (f :: * -> *) a. Applicative f => a -> f a
pure BufData
bufData'
  where objCount :: ByteCount
objCount = f a -> ByteCount
forall (t :: * -> *) a. Foldable t => t a -> ByteCount
length f a
fa
{-# INLINABLE writeBytes #-}