{-# 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.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.Unsafe as BSU
import qualified Data.Text as T
import qualified Data.Text.Encoding as TE
import qualified Data.Text.Internal.StrictBuilder as TB
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 -> 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.
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 (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
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
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
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
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, 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 = forall a. Program a => a -> String
programName prog
prog
      (SouffleM IO a
result) = do
        Maybe (Handle prog)
maybeHandle <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (String -> IO (Maybe (ForeignPtr Souffle))
Internal.init String
progName) forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
          Maybe (ForeignPtr Souffle)
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
          Just ForeignPtr Souffle
souffleHandle -> do
            MVar BufData
bufData <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
              ForeignPtr ByteBuf
ptr <- forall a. Ptr a -> IO (ForeignPtr a)
newForeignPtr_ forall a. Ptr a
nullPtr
              forall a. a -> IO (MVar a)
newMVar forall a b. (a -> b) -> a -> b
$ ForeignPtr ByteBuf -> Int -> BufData
BufData ForeignPtr ByteBuf
ptr Int
0
            forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ 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 (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
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
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
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) = 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 <- forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets forall a b. Ptr a -> Ptr b
castPtr
  forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Storable a => Ptr a -> a -> IO ()
S.poke Ptr a
ptr a
a
  forall s (m :: * -> *). MonadState s m => s -> m ()
put forall a b. (a -> b) -> a -> b
$ Ptr a
ptr 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 <- forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets forall a b. Ptr a -> Ptr b
castPtr
  a
a <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Storable a => Ptr a -> IO a
S.peek Ptr a
ptr
  forall s (m :: * -> *). MonadState s m => s -> m ()
put forall a b. (a -> b) -> a -> b
$ Ptr a
ptr forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
ramDomainSize
  forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a
{-# INLINABLE readAsBytes #-}

instance MonadPush CMarshalFast where
  pushInt32 :: Int32 -> CMarshalFast ()
pushInt32 = forall a. (Storable a, Marshal a) => a -> CMarshalFast ()
writeAsBytes
  {-# INLINABLE pushInt32 #-}
  pushUInt32 :: Word32 -> CMarshalFast ()
pushUInt32 = forall a. (Storable a, Marshal a) => a -> CMarshalFast ()
writeAsBytes
  {-# INLINABLE pushUInt32 #-}
  pushFloat :: Float -> CMarshalFast ()
pushFloat = forall a. (Storable a, Marshal a) => a -> CMarshalFast ()
writeAsBytes
  {-# INLINABLE pushFloat #-}
  pushString :: String -> CMarshalFast ()
pushString String
str = forall (m :: * -> *). MonadPush m => Text -> m ()
pushText forall a b. (a -> b) -> a -> b
$ String -> Text
T.pack String
str
  {-# INLINABLE pushString #-}
  pushText :: Text -> CMarshalFast ()
pushText Text
_ =
    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 = forall a. (Storable a, Marshal a) => CMarshalFast a
readAsBytes
  {-# INLINABLE popInt32 #-}
  popUInt32 :: CMarshalFast Word32
popUInt32 = forall a. (Storable a, Marshal a) => CMarshalFast a
readAsBytes
  {-# INLINABLE popUInt32 #-}
  popFloat :: CMarshalFast Float
popFloat = forall a. (Storable a, Marshal a) => CMarshalFast a
readAsBytes
  {-# INLINABLE popFloat #-}
  popString :: CMarshalFast String
popString = Text -> String
T.unpack forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *). MonadPop m => m Text
popText
  {-# INLINABLE popString #-}
  popText :: CMarshalFast Text
popText = do
    Word32
byteCount <- forall (m :: * -> *). MonadPop m => m Word32
popUInt32
    if Word32
byteCount forall a. Eq a => a -> a -> Bool
== Word32
0
      then forall (f :: * -> *) a. Applicative f => a -> f a
pure Text
T.empty
      else do
        Ptr CChar
ptr <- forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets forall a b. Ptr a -> Ptr b
castPtr
        ByteString
bs <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ CStringLen -> IO ByteString
BSU.unsafePackCStringLen (Ptr CChar
ptr, forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
byteCount)
        forall s (m :: * -> *). MonadState s m => s -> m ()
put forall a b. (a -> b) -> a -> b
$ Ptr CChar
ptr forall a b. Ptr a -> Int -> Ptr b
`plusPtr` 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, before the bytearray is overwritten.
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$! StrictBuilder -> Text
TB.toText forall a b. (a -> b) -> a -> b
$ ByteString -> StrictBuilder
TB.unsafeFromByteString ByteString
bs
  {-# INLINABLE popText #-}


type MarshalState :: Type
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).
type CMarshalSlow :: Type -> Type
newtype CMarshalSlow a = CMarshalSlow (StateT MarshalState IO a)
  deriving (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
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
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
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 forall a. Ord a => a -> a -> Bool
> Int
byteCount
    then forall (f :: * -> *) a. Applicative f => a -> f a
pure BufData
bufData
    else forall a b c. (a -> b -> c) -> b -> a -> c
flip ForeignPtr ByteBuf -> Int -> BufData
BufData Int
byteCount forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *). MonadIO m => Int -> m (ForeignPtr ByteBuf)
allocateBuf Int
byteCount
  let ptr :: Ptr ByteBuf
ptr = forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr (BufData -> ForeignPtr ByteBuf
bufPtr BufData
bufData')
  forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT StateT MarshalState IO a
m 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 <- forall s (m :: * -> *). MonadState s m => m s
get
  let totalByteCount :: Int
totalByteCount = BufData -> Int
bufSize BufData
bufData
  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
byteCount forall a. Num a => a -> a -> a
+ Int
offset forall a. Ord a => a -> a -> Bool
> Int
totalByteCount) 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 <- 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 = forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr ByteBuf
newBuf
        bufData' :: BufData
bufData' = ForeignPtr ByteBuf -> Int -> BufData
BufData ForeignPtr ByteBuf
newBuf Int
newTotalByteCount
    forall s (m :: * -> *). MonadState s m => s -> m ()
put forall a b. (a -> b) -> a -> b
$ BufData -> Ptr ByteBuf -> Int -> MarshalState
MarshalState BufData
bufData' (Ptr ByteBuf
newPtr 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 = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$
  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 = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$
  forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr ByteBuf
src forall a b. (a -> b) -> a -> b
$ \Ptr ByteBuf
srcPtr ->
  forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr ByteBuf
dst forall a b. (a -> b) -> a -> b
$ \Ptr ByteBuf
dstPtr ->
    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 forall a. Num a => a -> a -> a
+ Int
offset forall a. Ord a => a -> a -> Bool
> Int
totalByteCount = Int -> Int
go (Int
totalByteCount forall a. Num a => a -> a -> a
* Int
2)
    | Bool
otherwise = Int
totalByteCount
{-# INLINABLE getNewTotalByteCount #-}

incrementPtr :: ByteCount -> CMarshalSlow ()
incrementPtr :: Int -> CMarshalSlow ()
incrementPtr Int
byteCount =
  forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify 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 forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
byteCount) (Int
offset forall a. Num a => a -> a -> a
+ Int
byteCount)
{-# INLINABLE incrementPtr #-}

instance MonadPush CMarshalSlow where
  pushInt32 :: Int32 -> CMarshalSlow ()
pushInt32 = forall a. (Storable a, Marshal a) => a -> CMarshalSlow ()
writeAsBytesSlow
  {-# INLINABLE pushInt32 #-}
  pushUInt32 :: Word32 -> CMarshalSlow ()
pushUInt32 = forall a. (Storable a, Marshal a) => a -> CMarshalSlow ()
writeAsBytesSlow
  {-# INLINABLE pushUInt32 #-}
  pushFloat :: Float -> CMarshalSlow ()
pushFloat = forall a. (Storable a, Marshal a) => a -> CMarshalSlow ()
writeAsBytesSlow
  {-# INLINABLE pushFloat #-}
  pushString :: String -> CMarshalSlow ()
pushString String
str = forall (m :: * -> *). MonadPush m => Text -> m ()
pushText forall a b. (a -> b) -> a -> b
$ String -> Text
T.pack String
str
  {-# INLINABLE pushString #-}
  pushText :: Text -> CMarshalSlow ()
pushText Text
txt = do
    let bs :: ByteString
bs = Text -> ByteString
TE.encodeUtf8 Text
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 forall a. Num a => a -> a -> a
+ Int
len)
    forall (m :: * -> *). MonadPush m => Word32 -> m ()
pushUInt32 (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len)
    if Int
len forall a. Eq a => a -> a -> Bool
== Int
0
      then forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      else do
        Ptr CChar
ptr <- forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (forall a b. Ptr a -> Ptr b
castPtr forall b c a. (b -> c) -> (a -> b) -> a -> c
. MarshalState -> Ptr ByteBuf
_ptr)
        forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. ByteString -> (Ptr CChar -> IO a) -> IO a
BSU.unsafeUseAsCString ByteString
bs forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> b -> a -> c
flip (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 <- forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (forall a b. Ptr a -> Ptr b
castPtr forall b c a. (b -> c) -> (a -> b) -> a -> c
. MarshalState -> Ptr ByteBuf
_ptr)
  forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Storable a => Ptr a -> a -> IO ()
S.poke Ptr a
ptr a
a
  Int -> CMarshalSlow ()
incrementPtr Int
ramDomainSize
{-# INLINABLE writeAsBytesSlow #-}


type Collect :: (Type -> Type) -> Constraint
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 = 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 forall a. Eq a => a -> a -> Bool
== t
0 = forall (f :: * -> *) a. Applicative f => a -> f a
pure [a]
acc
      | Bool
otherwise = do
        !a
x <- forall a (m :: * -> *). (Marshal a, MonadPop m) => m a
pop
        t -> [a] -> f [a]
go (t
count forall a. Num a => a -> a -> a
- t
1) (a
xforall 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 <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ 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' = 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 forall a. Eq a => a -> a -> Bool
== Int
objCount' = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> m (Vector a)
V.unsafeFreeze MVector RealWorld a
vec
        | Bool
otherwise = do
          !a
obj <- forall a (m :: * -> *). (Marshal a, MonadPop m) => m a
pop
          forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
MV.write MVector RealWorld a
vec Int
idx a
obj
          MVector RealWorld a -> Int -> CMarshalFast (Vector a)
collect' MVector RealWorld a
vec (Int
idx 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 <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> m (a i e)
A.newArray_ (Int
0, Int
objCount' forall a. Num a => a -> a -> a
- Int
1)
    forall a.
Marshal a =>
IOArray Int a -> Int -> CMarshalFast (Array Int a)
collect' IOArray Int a
ma Int
0
    where
      objCount' :: Int
objCount' = 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 forall a. Eq a => a -> a -> Bool
== Int
objCount' = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ 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 <- forall a (m :: * -> *). (Marshal a, MonadPop m) => m a
pop
          forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ 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
          forall a.
Marshal a =>
IOArray Int a -> Int -> CMarshalFast (Array Int a)
collect' IOArray Int a
array (Int
idx 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 :: 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 :: forall prog. Handler SouffleM prog -> SouffleM ()
run (Handle ForeignPtr Souffle
prog MVar BufData
_) = forall a. IO a -> SouffleM a
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 =
    forall a. IO a -> SouffleM a
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
_) =
    forall a. IO a -> SouffleM a
SouffleM 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 = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
    let relationName :: String
relationName = forall a. Fact a => Proxy a -> String
factName (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
    forall (f :: * -> *) a.
(Foldable f, Marshal a, Submit a) =>
MVar BufData -> Ptr Relation -> f a -> IO ()
writeBytes MVar BufData
bufVar Ptr Relation
relation (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 = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
    let relationName :: String
relationName = forall a. Fact a => Proxy a -> String
factName (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
    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
_) = forall a. IO a -> SouffleM a
SouffleM forall a b. (a -> b) -> a -> b
$ do
    let relationName :: String
relationName = forall a. Fact a => Proxy a -> String
factName (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 <- forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Souffle
prog forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> b -> a -> c
flip Ptr Souffle -> Ptr Relation -> IO (Ptr ByteBuf)
Internal.popFacts Ptr Relation
relation
    forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a. CMarshalFast a -> Ptr ByteBuf -> IO a
runMarshalFastM Ptr ByteBuf
buf forall a b. (a -> b) -> a -> b
$ forall (c :: * -> *) a.
(Collect c, Marshal a) =>
Word32 -> CMarshalFast (c a)
collect forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< 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 = forall a. IO a -> SouffleM a
SouffleM forall a b. (a -> b) -> a -> b
$ do
    let relationName :: String
relationName = forall a. Fact a => Proxy a -> String
factName (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 forall a. Submit a => Proxy a -> ByteSize
estimateNumBytes (forall {k} (t :: k). Proxy t
Proxy @a) of
      Exact Int
numBytes -> do
        forall a b. MVar a -> (a -> IO (a, b)) -> IO b
modifyMVarMasked MVar BufData
bufVar forall a b. (a -> b) -> a -> b
$ \BufData
bufData -> do
          BufData
bufData' <- if BufData -> Int
bufSize BufData
bufData forall a. Ord a => a -> a -> Bool
> Int
numBytes
            then forall (f :: * -> *) a. Applicative f => a -> f a
pure BufData
bufData
            else forall a b c. (a -> b -> c) -> b -> a -> c
flip ForeignPtr ByteBuf -> Int -> BufData
BufData Int
numBytes forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *). MonadIO m => Int -> m (ForeignPtr ByteBuf)
allocateBuf Int
numBytes
          Bool
found <- forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr (BufData -> ForeignPtr ByteBuf
bufPtr BufData
bufData') forall a b. (a -> b) -> a -> b
$ \Ptr ByteBuf
ptr -> do
            forall a. CMarshalFast a -> Ptr ByteBuf -> IO a
runMarshalFastM (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
          forall (f :: * -> *) a. Applicative f => a -> f a
pure (BufData
bufData', Bool
found)
      Estimated Int
numBytes -> forall a b. MVar a -> (a -> IO (a, b)) -> IO b
modifyMVarMasked MVar BufData
bufVar forall a b. (a -> b) -> a -> b
$ \BufData
bufData ->
        forall a. BufData -> Int -> CMarshalSlow a -> IO a
runMarshalSlowM BufData
bufData Int
numBytes forall a b. (a -> b) -> a -> b
$ do
          forall a (m :: * -> *). (Marshal a, MonadPush m) => a -> m ()
push a
fact
          BufData
bufData' <- forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets MarshalState -> BufData
_buf
          forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr (BufData -> ForeignPtr ByteBuf
bufPtr BufData
bufData') 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
            forall (f :: * -> *) a. Applicative f => a -> f a
pure (BufData
bufData', Bool
found)
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ if Bool
found then forall a. a -> Maybe a
Just a
fact else 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
_) = forall a. IO a -> SouffleM a
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
_) = forall a. IO a -> SouffleM a
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 Int
s1 <> :: ByteSize -> ByteSize -> ByteSize
<> Exact Int
s2 = Int -> ByteSize
Exact (Int
s1 forall a. Num a => a -> a -> a
+ Int
s2)
  Exact Int
s1 <> Estimated Int
s2 = Int -> ByteSize
Estimated (Int
s1 forall a. Num a => a -> a -> a
+ Int
s2)
  Estimated Int
s1 <> Exact Int
s2 = Int -> ByteSize
Estimated (Int
s1 forall a. Num a => a -> a -> a
+ Int
s2)
  Estimated Int
s1 <> Estimated Int
s2 = Int -> ByteSize
Estimated (Int
s1 forall a. Num a => a -> a -> a
+ Int
s2)
  {-# INLINABLE (<>) #-}

type ToByteSize :: k -> Constraint
class ToByteSize a where
  toByteSize :: Proxy a -> ByteSize

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

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

instance ToByteSize Float where
  toByteSize :: Proxy Float -> ByteSize
toByteSize = forall a b. a -> b -> a
const 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 = forall a b. a -> b -> a
const 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 = forall a b. a -> b -> a
const 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 = forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ Int -> ByteSize
Estimated Int
36
  {-# INLINABLE toByteSize #-}

instance ToByteSize '[] where
  toByteSize :: Proxy '[] -> ByteSize
toByteSize = forall a b. a -> b -> a
const 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 =
    forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ forall k (a :: k). ToByteSize a => Proxy a -> ByteSize
toByteSize (forall {k} (t :: k). Proxy t
Proxy @a) forall a. Semigroup a => a -> a -> a
<> forall k (a :: k). ToByteSize a => Proxy a -> ByteSize
toByteSize (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 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 :: forall a. Submit a => Proxy a -> ByteSize
estimateNumBytes Proxy a
_ = forall k (a :: k). ToByteSize a => Proxy a -> ByteSize
toByteSize (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 forall a. Submit a => Proxy a -> ByteSize
estimateNumBytes (forall {k} (t :: k). Proxy t
Proxy @a) of
  Exact Int
numBytes -> forall a. MVar a -> (a -> IO a) -> IO ()
modifyMVarMasked_ MVar BufData
bufVar forall a b. (a -> b) -> a -> b
$ \BufData
bufData -> do
    let totalByteCount :: Int
totalByteCount = Int
numBytes forall a. Num a => a -> a -> a
* Int
objCount
    BufData
bufData' <- if BufData -> Int
bufSize BufData
bufData forall a. Ord a => a -> a -> Bool
> Int
totalByteCount
      then forall (f :: * -> *) a. Applicative f => a -> f a
pure BufData
bufData
      else forall a b c. (a -> b -> c) -> b -> a -> c
flip ForeignPtr ByteBuf -> Int -> BufData
BufData Int
totalByteCount forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *). MonadIO m => Int -> m (ForeignPtr ByteBuf)
allocateBuf Int
totalByteCount
    forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr (BufData -> ForeignPtr ByteBuf
bufPtr BufData
bufData') forall a b. (a -> b) -> a -> b
$ \Ptr ByteBuf
ptr -> do
      forall a. CMarshalFast a -> Ptr ByteBuf -> IO a
runMarshalFastM (forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ 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 (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
objCount)
    forall (f :: * -> *) a. Applicative f => a -> f a
pure BufData
bufData'

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