module Data.Store.Core
(
Poke(..), PokeException(..), pokeException
, Peek(..), PeekException(..), peekException, tooManyBytes
, Offset
, unsafeEncodeWith
, decodeWith
, decodeExWith, decodeExPortionWith
, decodeIOWith, decodeIOPortionWith
, decodeIOWithFromPtr, decodeIOPortionWithFromPtr
, pokeStorable, peekStorable, peekStorableTy
, pokeFromForeignPtr, peekToPlainForeignPtr, pokeFromPtr
, pokeFromByteArray, peekToByteArray
) where
import Control.Applicative
import Control.Exception (Exception(..), throwIO, try)
import Control.Monad (when)
import qualified Control.Monad.Fail as Fail
import Control.Monad.IO.Class (MonadIO(..))
import Control.Monad.Primitive (PrimMonad (..))
import Data.ByteString (ByteString)
import qualified Data.ByteString.Internal as BS
import Data.Monoid ((<>))
import Data.Primitive.ByteArray
import qualified Data.Text as T
import Data.Typeable
import Data.Word
import Foreign.ForeignPtr (ForeignPtr, withForeignPtr, castForeignPtr)
import Foreign.Ptr (plusPtr, minusPtr, castPtr)
import Foreign.Storable as Storable
import GHC.Prim (unsafeCoerce#, RealWorld, copyByteArrayToAddr#, copyAddrToByteArray#)
import GHC.Ptr (Ptr(..))
import GHC.Types (IO(..), Int(..))
import Prelude
import System.IO.Unsafe (unsafePerformIO)
type Offset = Int
newtype Poke a = Poke
{ runPoke :: forall byte. Ptr byte -> Offset -> IO (Offset, a)
}
deriving Functor
instance Applicative Poke where
pure x = Poke $ \_ptr offset -> pure (offset, x)
Poke f <*> Poke g = Poke $ \ptr offset1 -> do
(offset2, f') <- f ptr offset1
(offset3, g') <- g ptr offset2
return (offset3, f' g')
Poke f *> Poke g = Poke $ \ptr offset1 -> do
(offset2, _) <- f ptr offset1
g ptr offset2
instance Monad Poke where
return = pure
(>>) = (*>)
Poke x >>= f = Poke $ \ptr offset1 -> do
(offset2, x') <- x ptr offset1
runPoke (f x') ptr offset2
fail = Fail.fail
instance Fail.MonadFail Poke where
fail = pokeException . T.pack
instance MonadIO Poke where
liftIO f = Poke $ \_ offset -> (offset, ) <$> f
data PokeException = PokeException
{ pokeExByteIndex :: Offset
, pokeExMessage :: T.Text
}
deriving (Eq, Show, Typeable)
instance Exception PokeException where
#if MIN_VERSION_base(4,8,0)
displayException (PokeException offset msg) =
"Exception while poking, at byte index " ++
show offset ++
" : " ++
T.unpack msg
#endif
pokeException :: T.Text -> Poke a
pokeException msg = Poke $ \_ off -> throwIO (PokeException off msg)
newtype Peek a = Peek
{ runPeek :: forall byte. Ptr byte -> Ptr byte -> IO (Ptr byte, a)
}
deriving Functor
instance Applicative Peek where
pure x = Peek (\_ ptr -> return (ptr, x))
Peek f <*> Peek g = Peek $ \end ptr1 -> do
(ptr2, f') <- f end ptr1
(ptr3, g') <- g end ptr2
return (ptr3, f' g')
Peek f *> Peek g = Peek $ \end ptr1 -> do
(ptr2, _) <- f end ptr1
g end ptr2
instance Monad Peek where
return = pure
(>>) = (*>)
Peek x >>= f = Peek $ \end ptr1 -> do
(ptr2, x') <- x end ptr1
runPeek (f x') end ptr2
fail = Fail.fail
instance Fail.MonadFail Peek where
fail = peekException . T.pack
instance PrimMonad Peek where
type PrimState Peek = RealWorld
primitive action = Peek $ \_ ptr -> do
x <- primitive (unsafeCoerce# action)
return (ptr, x)
instance MonadIO Peek where
liftIO f = Peek $ \_ ptr -> (ptr, ) <$> f
data PeekException = PeekException
{ peekExBytesFromEnd :: Offset
, peekExMessage :: T.Text
} deriving (Eq, Show, Typeable)
instance Exception PeekException where
#if MIN_VERSION_base(4,8,0)
displayException (PeekException offset msg) =
"Exception while peeking, " ++
show offset ++
" bytes from end: " ++
T.unpack msg
#endif
peekException :: T.Text -> Peek a
peekException msg = Peek $ \end ptr -> throwIO (PeekException (end `minusPtr` ptr) msg)
tooManyBytes :: Int -> Int -> String -> IO void
tooManyBytes needed remaining ty =
throwIO $ PeekException remaining $ T.pack $
"Attempted to read too many bytes for " ++
ty ++
". Needed " ++
show needed ++ ", but only " ++
show remaining ++ " remain."
negativeBytes :: Int -> Int -> String -> IO void
negativeBytes needed remaining ty =
throwIO $ PeekException remaining $ T.pack $
"Attempted to read negative number of bytes for " ++
ty ++
". Tried to read " ++
show needed ++ ". This probably means that we're trying to read invalid data."
unsafeEncodeWith :: Poke () -> Int -> ByteString
unsafeEncodeWith f l = BS.unsafeCreate l $ \p -> do
(o, ()) <- runPoke f p 0
checkOffset o l
checkOffset :: Int -> Int -> IO ()
checkOffset o l
| o > l = throwIO $ PokeException o $ T.pack $
"encode overshot end of " ++
show l ++
" byte long buffer"
| o < l = throwIO $ PokeException o $ T.pack $
"encode undershot end of " <>
show l <>
" byte long buffer"
| otherwise = return ()
decodeWith :: Peek a -> ByteString -> Either PeekException a
decodeWith mypeek = unsafePerformIO . try . decodeIOWith mypeek
decodeExWith :: Peek a -> ByteString -> a
decodeExWith f = unsafePerformIO . decodeIOWith f
decodeExPortionWith :: Peek a -> ByteString -> (Offset, a)
decodeExPortionWith f = unsafePerformIO . decodeIOPortionWith f
decodeIOWith :: Peek a -> ByteString -> IO a
decodeIOWith mypeek (BS.PS x s len) =
withForeignPtr x $ \ptr0 ->
let ptr = ptr0 `plusPtr` s
in decodeIOWithFromPtr mypeek ptr len
decodeIOPortionWith :: Peek a -> ByteString -> IO (Offset, a)
decodeIOPortionWith mypeek (BS.PS x s len) =
withForeignPtr x $ \ptr0 ->
let ptr = ptr0 `plusPtr` s
in decodeIOPortionWithFromPtr mypeek ptr len
decodeIOWithFromPtr :: Peek a -> Ptr Word8 -> Int -> IO a
decodeIOWithFromPtr mypeek ptr len = do
(offset, x) <- decodeIOPortionWithFromPtr mypeek ptr len
if len /= offset
then throwIO $ PeekException (len offset) "Didn't consume all input."
else return x
decodeIOPortionWithFromPtr :: Peek a -> Ptr Word8 -> Int -> IO (Offset, a)
decodeIOPortionWithFromPtr mypeek ptr len =
let end = ptr `plusPtr` len
remaining = end `minusPtr` ptr
in do
(ptr2, x') <- runPeek mypeek end ptr
if len > remaining
then throwIO $ PeekException (end `minusPtr` ptr2) "Overshot end of buffer"
else return (ptr2 `minusPtr` ptr, x')
pokeStorable :: Storable a => a -> Poke ()
pokeStorable x = Poke $ \ptr offset -> do
y <- pokeByteOff ptr offset x
let !newOffset = offset + sizeOf x
return (newOffset, y)
peekStorable :: forall a. (Storable a, Typeable a) => Peek a
peekStorable = peekStorableTy (show (typeRep (Proxy :: Proxy a)))
peekStorableTy :: forall a. Storable a => String -> Peek a
peekStorableTy ty = Peek $ \end ptr ->
let ptr' = ptr `plusPtr` needed
needed = sizeOf (undefined :: a)
remaining = end `minusPtr` ptr
in do
when (needed > remaining) $
tooManyBytes needed remaining ty
x <- Storable.peek (castPtr ptr)
return (ptr', x)
pokeFromForeignPtr :: ForeignPtr a -> Int -> Int -> Poke ()
pokeFromForeignPtr sourceFp sourceOffset len =
Poke $ \targetPtr targetOffset -> do
withForeignPtr sourceFp $ \sourcePtr ->
BS.memcpy (targetPtr `plusPtr` targetOffset)
(sourcePtr `plusPtr` sourceOffset)
len
let !newOffset = targetOffset + len
return (newOffset, ())
peekToPlainForeignPtr :: String -> Int -> Peek (ForeignPtr a)
peekToPlainForeignPtr ty len =
Peek $ \end sourcePtr -> do
let ptr2 = sourcePtr `plusPtr` len
remaining = end `minusPtr` sourcePtr
when (len > remaining) $
tooManyBytes len remaining ty
when (len < 0) $
negativeBytes len remaining ty
fp <- BS.mallocByteString len
withForeignPtr fp $ \targetPtr ->
BS.memcpy targetPtr (castPtr sourcePtr) len
return (ptr2, castForeignPtr fp)
pokeFromPtr :: Ptr a -> Int -> Int -> Poke ()
pokeFromPtr sourcePtr sourceOffset len =
Poke $ \targetPtr targetOffset -> do
BS.memcpy (targetPtr `plusPtr` targetOffset)
(sourcePtr `plusPtr` sourceOffset)
len
let !newOffset = targetOffset + len
return (newOffset, ())
pokeFromByteArray :: ByteArray# -> Int -> Int -> Poke ()
pokeFromByteArray sourceArr sourceOffset len =
Poke $ \targetPtr targetOffset -> do
let target = targetPtr `plusPtr` targetOffset
copyByteArrayToAddr sourceArr sourceOffset target len
let !newOffset = targetOffset + len
return (newOffset, ())
peekToByteArray :: String -> Int -> Peek ByteArray
peekToByteArray ty len =
Peek $ \end sourcePtr -> do
let ptr2 = sourcePtr `plusPtr` len
remaining = end `minusPtr` sourcePtr
when (len > remaining) $
tooManyBytes len remaining ty
when (len < 0) $
negativeBytes len remaining ty
marr <- newByteArray len
copyAddrToByteArray sourcePtr marr 0 len
x <- unsafeFreezeByteArray marr
return (ptr2, x)
copyByteArrayToAddr :: ByteArray# -> Int -> Ptr a -> Int -> IO ()
copyByteArrayToAddr arr (I# offset) (Ptr addr) (I# len) =
IO (\s -> (# copyByteArrayToAddr# arr offset addr len s, () #))
copyAddrToByteArray :: Ptr a -> MutableByteArray (PrimState IO) -> Int -> Int -> IO ()
copyAddrToByteArray (Ptr addr) (MutableByteArray arr) (I# offset) (I# len) =
IO (\s -> (# copyAddrToByteArray# addr arr offset len s, () #))