module Data.Store.Impl where
import Control.Applicative
import Control.Exception (Exception(..), throwIO)
import Control.Exception (try)
import Control.Monad
import qualified Control.Monad.Fail as Fail
import Control.Monad.IO.Class (MonadIO(..))
import Control.Monad.Primitive (PrimMonad (..))
import qualified Data.ByteString as BS
import qualified Data.ByteString.Internal as BS
import Data.Monoid ((<>))
import Data.Primitive.ByteArray
import Data.Proxy
import qualified Data.Text as T
import Data.Typeable
import Data.Word
import Foreign.ForeignPtr (ForeignPtr, withForeignPtr, castForeignPtr)
import Foreign.Ptr (Ptr, plusPtr, minusPtr, castPtr)
import Foreign.Storable (pokeByteOff, Storable, sizeOf)
import qualified Foreign.Storable as Storable
import GHC.Generics
import GHC.Prim ( unsafeCoerce#, RealWorld )
import GHC.Prim (copyByteArrayToAddr#, copyAddrToByteArray#)
import GHC.Ptr (Ptr(..))
import GHC.TypeLits
import GHC.Types (IO(..), Int(..))
import Prelude
import System.IO.Unsafe (unsafePerformIO)
class Store a where
size :: Size a
poke :: a -> Poke ()
peek :: Peek a
default size :: (Generic a, GStoreSize (Rep a)) => Size a
size = genericSize
default poke :: (Generic a, GStorePoke (Rep a)) => a -> Poke ()
poke = genericPoke
default peek :: (Generic a , GStorePeek (Rep a)) => Peek a
peek = genericPeek
encode :: Store a => a -> BS.ByteString
encode x = BS.unsafeCreate l $ \p -> do
(o, ()) <- runPoke (poke x) p 0
checkOffset o l
where
l = getSize x
checkOffset :: Int -> Int -> IO ()
checkOffset o l
| o > l = throwIO $ PokeException o $
"encode overshot end of " <>
T.pack (show l) <>
" byte long buffer"
| o < l = throwIO $ PokeException o $
"encode undershot end of " <>
T.pack (show l) <>
" byte long buffer"
| otherwise = return ()
decode :: Store a => BS.ByteString -> Either PeekException a
decode = unsafePerformIO . try . decodeIO
decodeWith :: Peek a -> BS.ByteString -> Either PeekException a
decodeWith mypeek = unsafePerformIO . try . decodeIOWith mypeek
decodeEx :: Store a => BS.ByteString -> a
decodeEx = unsafePerformIO . decodeIO
decodeExWith :: Peek a -> BS.ByteString -> a
decodeExWith f = unsafePerformIO . decodeIOWith f
decodeExPortionWith :: Peek a -> BS.ByteString -> (Offset, a)
decodeExPortionWith f = unsafePerformIO . decodeIOPortionWith f
decodeIO :: Store a => BS.ByteString -> IO a
decodeIO = decodeIOWith peek
decodeIOWith :: Peek a -> BS.ByteString -> IO a
decodeIOWith mypeek bs = do
(offset, x) <- decodeIOPortionWith mypeek bs
let remaining = BS.length bs offset
if remaining > 0
then throwIO $ PeekException remaining "Didn't consume all input."
else return x
decodeIOPortionWith :: Peek a -> BS.ByteString -> IO (Offset, a)
decodeIOPortionWith mypeek (BS.PS x s len) =
withForeignPtr x $ \ptr0 ->
let ptr = ptr0 `plusPtr` s
end = ptr `plusPtr` len
in do
(ptr2, x') <- runPeek mypeek end ptr
if ptr2 > end
then throwIO $ PeekException (end `minusPtr` ptr2) "Overshot end of buffer"
else return (ptr2 `minusPtr` ptr, x')
genericSize :: (Generic a, GStoreSize (Rep a)) => Size a
genericSize = contramapSize from gsize
genericPoke :: (Generic a, GStorePoke (Rep a)) => a -> Poke ()
genericPoke = gpoke . from
genericPeek :: (Generic a , GStorePeek (Rep a)) => Peek a
genericPeek = to <$> gpeek
type family SumArity (a :: * -> *) :: Nat where
SumArity (C1 c a) = 1
SumArity (x :+: y) = SumArity x + SumArity y
class GStoreSize f where gsize :: Size (f a)
class GStorePoke f where gpoke :: f a -> Poke ()
class GStorePeek f where gpeek :: Peek (f a)
instance GStoreSize f => GStoreSize (M1 i c f) where
gsize = contramapSize unM1 gsize
instance GStorePoke f => GStorePoke (M1 i c f) where
gpoke = gpoke . unM1
instance GStorePeek f => GStorePeek (M1 i c f) where
gpeek = fmap M1 gpeek
instance Store a => GStoreSize (K1 i a) where
gsize = contramapSize unK1 size
instance Store a => GStorePoke (K1 i a) where
gpoke = poke . unK1
instance Store a => GStorePeek (K1 i a) where
gpeek = fmap K1 peek
instance GStoreSize U1 where
gsize = ConstSize 0
instance GStorePoke U1 where
gpoke _ = return ()
instance GStorePeek U1 where
gpeek = return U1
instance GStoreSize V1 where
gsize = ConstSize 0
instance GStorePoke V1 where
gpoke x = case x of {}
instance GStorePeek V1 where
gpeek = undefined
instance (GStoreSize a, GStoreSize b) => GStoreSize (a :*: b) where
gsize = combineSize' (\(x :*: _) -> x) (\(_ :*: y) -> y) gsize gsize
instance (GStorePoke a, GStorePoke b) => GStorePoke (a :*: b) where
gpoke (a :*: b) = gpoke a >> gpoke b
instance (GStorePeek a, GStorePeek b) => GStorePeek (a :*: b) where
gpeek = (:*:) <$> gpeek <*> gpeek
instance (SumArity (a :+: b) <= 255, GStoreSizeSum 0 (a :+: b))
=> GStoreSize (a :+: b) where
gsize = VarSize $ \x -> sizeOf (undefined :: Word8) + gsizeSum x (Proxy :: Proxy 0)
instance (SumArity (a :+: b) <= 255, GStorePokeSum 0 (a :+: b))
=> GStorePoke (a :+: b) where
gpoke x = gpokeSum x (Proxy :: Proxy 0)
instance (SumArity (a :+: b) <= 255, GStorePeekSum 0 (a :+: b))
=> GStorePeek (a :+: b) where
gpeek = do
tag <- peekStorable
gpeekSum tag (Proxy :: Proxy 0)
class KnownNat n => GStoreSizeSum (n :: Nat) (f :: * -> *) where gsizeSum :: f a -> Proxy n -> Int
class KnownNat n => GStorePokeSum (n :: Nat) (f :: * -> *) where gpokeSum :: f p -> Proxy n -> Poke ()
class KnownNat n => GStorePeekSum (n :: Nat) (f :: * -> *) where gpeekSum :: Word8 -> Proxy n -> Peek (f p)
instance (GStoreSizeSum n a, GStoreSizeSum (n + SumArity a) b, KnownNat n)
=> GStoreSizeSum n (a :+: b) where
gsizeSum (L1 l) _ = gsizeSum l (Proxy :: Proxy n)
gsizeSum (R1 r) _ = gsizeSum r (Proxy :: Proxy (n + SumArity a))
instance (GStorePokeSum n a, GStorePokeSum (n + SumArity a) b, KnownNat n)
=> GStorePokeSum n (a :+: b) where
gpokeSum (L1 l) _ = gpokeSum l (Proxy :: Proxy n)
gpokeSum (R1 r) _ = gpokeSum r (Proxy :: Proxy (n + SumArity a))
instance (GStorePeekSum n a, GStorePeekSum (n + SumArity a) b, KnownNat n)
=> GStorePeekSum n (a :+: b) where
gpeekSum tag proxyL
| tag < sizeL = L1 <$> gpeekSum tag proxyL
| otherwise = R1 <$> gpeekSum tag (Proxy :: Proxy (n + SumArity a))
where
sizeL = fromInteger (natVal (Proxy :: Proxy (n + SumArity a)))
instance (GStoreSize a, KnownNat n) => GStoreSizeSum n (C1 c a) where
gsizeSum x _ = getSizeWith gsize x
instance (GStorePoke a, KnownNat n) => GStorePokeSum n (C1 c a) where
gpokeSum x _ = do
pokeStorable (fromInteger (natVal (Proxy :: Proxy n)) :: Word8)
gpoke x
instance (GStorePeek a, KnownNat n) => GStorePeekSum n (C1 c a) where
gpeekSum tag _
| tag == cur = gpeek
| tag > cur = peekException "Sum tag invalid"
| otherwise = peekException "Error in implementation of Store Generics"
where
cur = fromInteger (natVal (Proxy :: Proxy n))
sizeStorable :: forall a. (Storable a, Typeable a) => Size a
sizeStorable = sizeStorableTy (show (typeRep (Proxy :: Proxy a)))
sizeStorableTy :: forall a. Storable a => String -> Size a
sizeStorableTy ty = ConstSize (sizeOf (error msg :: a))
where
msg = "In Data.Store.storableSize: " ++ ty ++ "'s sizeOf evaluated its argument."
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 (ptr' > end) $
tooManyBytes needed remaining ty
x <- Storable.peek (castPtr ptr)
return (ptr', x)
type Total = Int
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."
data Size a
= VarSize (a -> Int)
| ConstSize !Int
deriving Typeable
getSize :: Store a => a -> Int
getSize = getSizeWith size
getSizeWith :: Size a -> a -> Int
getSizeWith (VarSize f) x = f x
getSizeWith (ConstSize n) _ = n
contramapSize :: (a -> b) -> Size b -> Size a
contramapSize f (VarSize g) = VarSize (g . f)
contramapSize _ (ConstSize n) = ConstSize n
combineSize :: forall a b c. (Store a, Store b) => (c -> a) -> (c -> b) -> Size c
combineSize toA toB = combineSize' toA toB size size
combineSize' :: forall a b c. (c -> a) -> (c -> b) -> Size a -> Size b -> Size c
combineSize' toA toB sizeA sizeB =
case (sizeA, sizeB) of
(VarSize f, VarSize g) -> VarSize (\x -> f (toA x) + g (toB x))
(VarSize f, ConstSize m) -> VarSize (\x -> f (toA x) + m)
(ConstSize n, VarSize g) -> VarSize (\x -> n + g (toB x))
(ConstSize n, ConstSize m) -> ConstSize (n + m)
scaleSize :: Int -> Size a -> Size a
scaleSize s (ConstSize n) = ConstSize (s * n)
scaleSize s (VarSize f) = VarSize ((s *) . f)
addSize :: Int -> Size a -> Size a
addSize x (ConstSize n) = ConstSize (x + n)
addSize x (VarSize f) = VarSize ((x +) . f)
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
when (ptr2 > end) $
tooManyBytes len (end `minusPtr` sourcePtr) 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
when (ptr2 > end) $
tooManyBytes len (end `minusPtr` sourcePtr) 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, () #))