diff --git a/ChangeLog.md b/ChangeLog.md new file mode 100644 --- /dev/null +++ b/ChangeLog.md @@ -0,0 +1,5 @@ +# ChangeLog + +## 0.1.0.0 + +* First public release diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +The MIT License (MIT) + +Copyright (c) 2016 FP Complete + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,24 @@ +# store + +The 'store' package provides binary serialization of Haskell datatypes. It fills +quite a different niche from packages like 'binary' or 'cereal'. In particular: + +* Its primary goal is speed. Whenever possible, direct machine representations + are used. For numeric types (`Int`, `Double`, `Word32`, etc) and types that + use buffers (`Text`, `ByteString`, `Vector`, etc). This means that much of + serialization uses the equivalent of `memcpy`. + +* By using machine representations, we lose serialization compatibility between + different architectures. Store could in theory be used to describe + machine-independent serialization formats. However, this is not the usecase + it's currently designed for (though utilities might be added for this in the + future!) + +* `Store` will not work at all on architectures which lack unaligned memory + access (for example, older ARM processors). This is not a fundamental + limitation, but we do not currently require ARM support. + +See +[this blog post](https://www.fpcomplete.com/blog/2016/03/efficient-binary-serialization) +which describes the initial motivations and benchmarks that led to the existence +of this package. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/bench/Bench.hs b/bench/Bench.hs new file mode 100644 --- /dev/null +++ b/bench/Bench.hs @@ -0,0 +1,255 @@ +{-# LANGUAGE CPP #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} + +#if !MIN_VERSION_base(4,8,0) +{-# LANGUAGE DeriveDataTypeable #-} +import Control.Applicative ((<$>), (<*>), (*>)) +#endif + +import Control.DeepSeq +import Criterion.Main +import qualified Data.ByteString as BS +import Data.Int +import Data.Store +import Data.Typeable +import qualified Data.Vector as V +import qualified Data.Vector.Storable as SV +import Data.Word +import GHC.Generics + +-- TODO: add packer +#if COMPARISON_BENCH +import qualified Data.Binary as Binary +import qualified Data.Serialize as Cereal +import qualified Data.ByteString.Lazy as BL +import Data.Vector.Serialize () +#endif + +data SomeData = SomeData !Int64 !Word8 !Double + deriving (Eq, Show, Generic) +instance NFData SomeData where + rnf x = x `seq` () +instance Store SomeData +#if COMPARISON_BENCH +instance Cereal.Serialize SomeData +instance Binary.Binary SomeData +#endif + +main :: IO () +main = do +#if SMALL_BENCH + let is = 0::Int + sds = SomeData 1 1 1 + smallprods = (SmallProduct 0 1 2 3) + smallmanualprods = (SmallProductManual 0 1 2 3) + sss = [SS1 1, SS2 2, SS3 3, SS4 4] + ssms = [SSM1 1, SSM2 2, SSM3 3, SSM4 4] + nestedTuples = ((1,2),(3,4)) :: ((Int,Int),(Int,Int)) +#else + let is = V.enumFromTo 1 100 :: V.Vector Int + sds = (\i -> SomeData i (fromIntegral i) (fromIntegral i)) + <$> V.enumFromTo 1 100 + smallprods = (\ i -> SmallProduct i (i+1) (i+2) (i+3)) + <$> V.enumFromTo 1 100 + smallmanualprods = (\ i -> SmallProductManual i (i+1) (i+2) (i+3)) + <$> V.enumFromTo 1 100 + sss = (\i -> case i `mod` 4 of + 0 -> SS1 (fromIntegral i) + 1 -> SS2 (fromIntegral i) + 2 -> SS3 (fromIntegral i) + 3 -> SS4 (fromIntegral i) + _ -> error "This does not compute." + ) <$> V.enumFromTo 1 (100 :: Int) + ssms = (\i -> case i `mod` 4 of + 0 -> SSM1 (fromIntegral i) + 1 -> SSM2 (fromIntegral i) + 2 -> SSM3 (fromIntegral i) + 3 -> SSM4 (fromIntegral i) + _ -> error "This does not compute." + ) <$> V.enumFromTo 1 (100 :: Int) + nestedTuples = (\i -> ((i,i+1),(i+2,i+3))) <$> V.enumFromTo (1::Int) 100 +#endif + defaultMain + [ bgroup "encode" + [ benchEncode is +#if !SMALL_BENCH + , benchEncode' "1kb storable" (SV.fromList ([1..256] :: [Int32])) + , benchEncode' "10kb storable" (SV.fromList ([1..(256 * 10)] :: [Int32])) + , benchEncode' "1kb normal" (V.fromList ([1..256] :: [Int32])) + , benchEncode' "10kb normal" (V.fromList ([1..(256 * 10)] :: [Int32])) +#endif + , benchEncode smallprods + , benchEncode smallmanualprods + , benchEncode sss + , benchEncode ssms + , benchEncode nestedTuples + , benchEncode sds + ] + , bgroup "decode" + [ benchDecode is +#if !SMALL_BENCH + , benchDecode' "1kb storable" (SV.fromList ([1..256] :: [Int32])) + , benchDecode' "10kb storable" (SV.fromList ([1..(256 * 10)] :: [Int32])) + , benchDecode' "1kb normal" (V.fromList ([1..256] :: [Int32])) + , benchDecode' "10kb normal" (V.fromList ([1..(256 * 10)] :: [Int32])) +#endif + , benchDecode smallprods + , benchDecode smallmanualprods + , benchDecode sss + , benchDecode ssms + , benchDecode nestedTuples + , benchDecode sds + ] + ] + +type Ctx a = + ( Store a, Typeable a, NFData a +#if COMPARISON_BENCH + , Binary.Binary a + , Cereal.Serialize a +#endif + ) + +benchEncode :: Ctx a => a -> Benchmark +benchEncode = benchEncode' "" + +benchEncode' :: Ctx a => String -> a -> Benchmark +benchEncode' msg x0 = + env (return x0) $ \x -> + let label = msg ++ " (" ++ show (typeOf x0) ++ ")" + benchStore name = bench name (nf encode x) in +#if COMPARISON_BENCH + bgroup label + [ benchStore "store" + , bench "cereal" (nf Cereal.encode x) + , bench "binary" (nf Binary.encode x) + ] +#else + benchStore label +#endif + +benchDecode :: Ctx a => a -> Benchmark +benchDecode = benchDecode' "" + +-- TODO: comparison bench for decode + +benchDecode' :: forall a. Ctx a => String -> a -> Benchmark +#if COMPARISON_BENCH +benchDecode' prefix x0 = + bgroup label + [ env (return (encode x0)) $ \x -> bench "store" (nf (decodeEx :: BS.ByteString -> a) x) + , env (return (Cereal.encode x0)) $ \x -> bench "cereal" (nf ((ensureRight . Cereal.decode) :: BS.ByteString -> a) x) + , env (return (Binary.encode x0)) $ \x -> bench "binary" (nf (Binary.decode :: BL.ByteString -> a) x) + ] + where + label = prefix ++ " (" ++ show (typeOf x0) ++ ")" + ensureRight (Left x) = error "left!" + ensureRight (Right x) = x +#else +benchDecode' prefix x0 = + env (return (encode x0)) $ \x -> + bench (prefix ++ " (" ++ show (typeOf x0) ++ ")") (nf (decodeEx :: BS.ByteString -> a) x) +#endif + +------------------------------------------------------------------------ +-- Serialized datatypes + +data SmallProduct = SmallProduct Int32 Int32 Int32 Int32 + deriving (Generic, Show, Typeable) + +instance NFData SmallProduct +instance Store SmallProduct + +data SmallProductManual = SmallProductManual Int32 Int32 Int32 Int32 + deriving (Generic, Show, Typeable) + +instance NFData SmallProductManual +instance Store SmallProductManual where + size = ConstSize 16 + peek = SmallProductManual <$> peek <*> peek <*> peek <*> peek + poke (SmallProductManual a b c d) = poke a *> poke b *> poke c *> poke d + +data SmallSum + = SS1 Int8 + | SS2 Int32 + | SS3 Int64 + | SS4 Word32 + deriving (Generic, Show, Typeable) + +instance NFData SmallSum +instance Store SmallSum + +data SmallSumManual + = SSM1 Int8 + | SSM2 Int32 + | SSM3 Int64 + | SSM4 Word32 + deriving (Generic, Show, Typeable) + +instance NFData SmallSumManual +instance Store SmallSumManual where + size = VarSize $ \x -> 1 + case x of + SSM1{} -> 1 + SSM2{} -> 4 + SSM3{} -> 8 + SSM4{} -> 4 + peek = do + tag <- peek + case tag :: Word8 of + 0 -> SSM1 <$> peek + 1 -> SSM2 <$> peek + 2 -> SSM3 <$> peek + 3 -> SSM4 <$> peek + _ -> fail "Invalid tag" + poke (SSM1 x) = poke (0 :: Word8) >> poke x + poke (SSM2 x) = poke (1 :: Word8) >> poke x + poke (SSM3 x) = poke (2 :: Word8) >> poke x + poke (SSM4 x) = poke (3 :: Word8) >> poke x + +-- TODO: add TH generation of the above, and add LargeSum / LargeProduct cases + +#if COMPARISON_BENCH +instance Binary.Binary SmallProduct +instance Binary.Binary SmallSum +instance Cereal.Serialize SmallProduct +instance Cereal.Serialize SmallSum + +instance Binary.Binary SmallProductManual where + get = SmallProductManual <$> Binary.get <*> Binary.get <*> Binary.get <*> Binary.get + put (SmallProductManual a b c d) = Binary.put a *> Binary.put b *> Binary.put c *> Binary.put d + +instance Binary.Binary SmallSumManual where + get = do + tag <- Binary.get + case tag :: Word8 of + 0 -> SSM1 <$> Binary.get + 1 -> SSM2 <$> Binary.get + 2 -> SSM3 <$> Binary.get + 3 -> SSM4 <$> Binary.get + _ -> fail "Invalid tag" + put (SSM1 x) = Binary.put (0 :: Word8) *> Binary.put x + put (SSM2 x) = Binary.put (1 :: Word8) *> Binary.put x + put (SSM3 x) = Binary.put (2 :: Word8) *> Binary.put x + put (SSM4 x) = Binary.put (3 :: Word8) *> Binary.put x + +instance Cereal.Serialize SmallProductManual where + get = SmallProductManual <$> Cereal.get <*> Cereal.get <*> Cereal.get <*> Cereal.get + put (SmallProductManual a b c d) = Cereal.put a *> Cereal.put b *> Cereal.put c *> Cereal.put d + +instance Cereal.Serialize SmallSumManual where + get = do + tag <- Cereal.get + case tag :: Word8 of + 0 -> SSM1 <$> Cereal.get + 1 -> SSM2 <$> Cereal.get + 2 -> SSM3 <$> Cereal.get + 3 -> SSM4 <$> Cereal.get + _ -> fail "Invalid tag" + put (SSM1 x) = Cereal.put (0 :: Word8) *> Cereal.put x + put (SSM2 x) = Cereal.put (1 :: Word8) *> Cereal.put x + put (SSM3 x) = Cereal.put (2 :: Word8) *> Cereal.put x + put (SSM4 x) = Cereal.put (3 :: Word8) *> Cereal.put x +#endif diff --git a/src/Data/Store.hs b/src/Data/Store.hs new file mode 100644 --- /dev/null +++ b/src/Data/Store.hs @@ -0,0 +1,23 @@ +-- | This is the main public API of the store package. The functions +-- exported here are more likely to be stable between versions. +-- +-- Usually you won't need to write your own 'Store' instances, and +-- instead can rely on either using the 'Generic' deriving approach or +-- "Data.Store.TH" for defining 'Store' instances for your datatypes. +-- There are some tradeoffs here - the generics instances do not require +-- @-XTemplateHaskell@, but they do not optimize as well for sum types +-- that only require a constant number of bytes. +module Data.Store + ( + -- * Encoding and decoding strict ByteStrings. + encode, + decode, decodeWith, + decodeEx, decodeExWith, decodeExPortionWith, + decodeIO, decodeIOWith, decodeIOPortionWith + -- * Store class and related types. + , Store(..), Size(..), Poke, Peek + -- ** Exceptions thrown by Peek + , PeekException(..), peekException + ) where + +import Data.Store.Internal diff --git a/src/Data/Store/Impl.hs b/src/Data/Store/Impl.hs new file mode 100644 --- /dev/null +++ b/src/Data/Store/Impl.hs @@ -0,0 +1,644 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DefaultSignatures #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE EmptyCase #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE UnboxedTuples #-} + +-- This module is not exposed. The reason that it is split out from +-- "Data.Store.Internal" is to allow "Data.Store.TH" to refer to these +-- identifiers. "Data.Store.Internal" must be separate from +-- "Data.Store.TH" due to Template Haskell's stage restriction. +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) + +------------------------------------------------------------------------ +-- Store class + +-- TODO: write down more elaborate laws + +-- | The 'Store' typeclass provides efficient serialization and +-- deserialization to raw pointer addresses. +-- +-- The 'peek' and 'poke' methods should be defined such that +-- @ decodeEx (encode x) == x @. +class Store a where + -- | Yields the 'Size' of the buffer, in bytes, required to store + -- the encoded representation of the type. + -- + -- Note that the correctness of this function is crucial for the + -- safety of 'poke', as it does not do any bounds checking. It is + -- the responsibility of the invoker of 'poke' ('encode' and similar + -- functions) to ensure that there's enough space in the output + -- buffer. If 'poke' writes beyond, then arbitrary memory can be + -- overwritten, causing undefined behavior and segmentation faults. + size :: Size a + -- | Serializes a value to bytes. It is the responsibility of the + -- caller to ensure that at least the number of bytes required by + -- 'size' are available. These details are handled by 'encode' and + -- similar utilities. + poke :: a -> Poke () + -- | Serialized a value from bytes, throwing exceptions if it + -- encounters invalid data or runs out of input bytes. + peek :: Peek a + + default size :: (Generic a, GStoreSize (Rep a)) => Size a + size = genericSize + {-# INLINE size #-} + + default poke :: (Generic a, GStorePoke (Rep a)) => a -> Poke () + poke = genericPoke + {-# INLINE poke #-} + + default peek :: (Generic a , GStorePeek (Rep a)) => Peek a + peek = genericPeek + {-# INLINE peek #-} + + +------------------------------------------------------------------------ +-- Utilities for encoding / decoding strict ByteStrings + +-- | Serializes a value to a 'BS.ByteString'. In order to do this, it +-- first allocates a 'BS.ByteString' of the correct size (based on +-- 'size'), and then uses 'poke' to fill it. +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 +{-# INLINE encode #-} + +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 () + +-- | Decodes a value from a 'BS.ByteString'. Returns an exception if +-- there's an error while decoding, or if decoding undershoots / +-- overshoots the end of the buffer. +decode :: Store a => BS.ByteString -> Either PeekException a +decode = unsafePerformIO . try . decodeIO +{-# INLINE decode #-} + +-- | Decodes a value from a 'BS.ByteString', potentially throwing +-- exceptions, and taking a 'Peek' to run. It is an exception to not +-- consume all input. +decodeWith :: Peek a -> BS.ByteString -> Either PeekException a +decodeWith mypeek = unsafePerformIO . try . decodeIOWith mypeek +{-# INLINE decodeWith #-} + +-- | Decodes a value from a 'BS.ByteString', potentially throwing +-- exceptions. It is an exception to not consume all input. +decodeEx :: Store a => BS.ByteString -> a +decodeEx = unsafePerformIO . decodeIO +{-# INLINE decodeEx #-} + +-- | Decodes a value from a 'BS.ByteString', potentially throwing +-- exceptions, and taking a 'Peek' to run. It is an exception to not +-- consume all input. +decodeExWith :: Peek a -> BS.ByteString -> a +decodeExWith f = unsafePerformIO . decodeIOWith f +{-# INLINE decodeExWith #-} + +-- | Similar to 'decodeExWith', but it allows there to be more of the +-- buffer remaining. The 'Offset' of the buffer contents immediately +-- after the decoded value is returned. +decodeExPortionWith :: Peek a -> BS.ByteString -> (Offset, a) +decodeExPortionWith f = unsafePerformIO . decodeIOPortionWith f +{-# INLINE decodeExPortionWith #-} + +-- | Decodes a value from a 'BS.ByteString', potentially throwing +-- exceptions. It is an exception to not consume all input. +decodeIO :: Store a => BS.ByteString -> IO a +decodeIO = decodeIOWith peek +{-# INLINE decodeIO #-} + +-- | Decodes a value from a 'BS.ByteString', potentially throwing +-- exceptions, and taking a 'Peek' to run. It is an exception to not +-- consume all input. +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 +{-# INLINE decodeIOWith #-} + +-- | Similar to 'decodeExPortionWith', but runs in the 'IO' monad. +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') +{-# INLINE decodeIOPortionWith #-} + +------------------------------------------------------------------------ +-- Generics + +genericSize :: (Generic a, GStoreSize (Rep a)) => Size a +genericSize = contramapSize from gsize +{-# INLINE genericSize #-} + +genericPoke :: (Generic a, GStorePoke (Rep a)) => a -> Poke () +genericPoke = gpoke . from +{-# INLINE genericPoke #-} + +genericPeek :: (Generic a , GStorePeek (Rep a)) => Peek a +genericPeek = to <$> gpeek +{-# INLINE genericPeek #-} + +type family SumArity (a :: * -> *) :: Nat where + SumArity (C1 c a) = 1 + SumArity (x :+: y) = SumArity x + SumArity y + +-- This could be just one typeclass, but currently compile times are +-- better with things split up. +-- https://github.com/bos/aeson/pull/335 +-- + +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 + {-# INLINE gsize #-} +instance GStorePoke f => GStorePoke (M1 i c f) where + gpoke = gpoke . unM1 + {-# INLINE gpoke #-} +instance GStorePeek f => GStorePeek (M1 i c f) where + gpeek = fmap M1 gpeek + {-# INLINE gpeek #-} + +instance Store a => GStoreSize (K1 i a) where + gsize = contramapSize unK1 size + {-# INLINE gsize #-} +instance Store a => GStorePoke (K1 i a) where + gpoke = poke . unK1 + {-# INLINE gpoke #-} +instance Store a => GStorePeek (K1 i a) where + gpeek = fmap K1 peek + {-# INLINE gpeek #-} + +instance GStoreSize U1 where + gsize = ConstSize 0 + {-# INLINE gsize #-} +instance GStorePoke U1 where + gpoke _ = return () + {-# INLINE gpoke #-} +instance GStorePeek U1 where + gpeek = return U1 + {-# INLINE gpeek #-} + +instance GStoreSize V1 where + gsize = ConstSize 0 + {-# INLINE gsize #-} +instance GStorePoke V1 where + gpoke x = case x of {} + {-# INLINE gpoke #-} +instance GStorePeek V1 where + gpeek = undefined + {-# INLINE gpeek #-} + +instance (GStoreSize a, GStoreSize b) => GStoreSize (a :*: b) where + gsize = combineSize' (\(x :*: _) -> x) (\(_ :*: y) -> y) gsize gsize + {-# INLINE gsize #-} +instance (GStorePoke a, GStorePoke b) => GStorePoke (a :*: b) where + gpoke (a :*: b) = gpoke a >> gpoke b + {-# INLINE gpoke #-} +instance (GStorePeek a, GStorePeek b) => GStorePeek (a :*: b) where + gpeek = (:*:) <$> gpeek <*> gpeek + {-# INLINE gpeek #-} + +-- The machinery for sum types is why UndecidableInstances is necessary. + +-- FIXME: check that this type level stuff dosen't get turned into +-- costly runtime computation + +instance (SumArity (a :+: b) <= 255, GStoreSizeSum 0 (a :+: b)) + => GStoreSize (a :+: b) where + gsize = VarSize $ \x -> sizeOf (undefined :: Word8) + gsizeSum x (Proxy :: Proxy 0) + {-# INLINE gsize #-} +instance (SumArity (a :+: b) <= 255, GStorePokeSum 0 (a :+: b)) + => GStorePoke (a :+: b) where + gpoke x = gpokeSum x (Proxy :: Proxy 0) + {-# INLINE gpoke #-} +instance (SumArity (a :+: b) <= 255, GStorePeekSum 0 (a :+: b)) + => GStorePeek (a :+: b) where + gpeek = do + tag <- peekStorable + gpeekSum tag (Proxy :: Proxy 0) + {-# INLINE gpeek #-} + +-- Similarly to splitting up the generic class into multiple classes, we +-- also split up the one for sum types. + +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)) + {-# INLINE gsizeSum #-} +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)) + {-# INLINE gpokeSum #-} +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))) + {-# INLINE gpeekSum #-} + +instance (GStoreSize a, KnownNat n) => GStoreSizeSum n (C1 c a) where + gsizeSum x _ = getSizeWith gsize x + {-# INLINE gsizeSum #-} +instance (GStorePoke a, KnownNat n) => GStorePokeSum n (C1 c a) where + gpokeSum x _ = do + pokeStorable (fromInteger (natVal (Proxy :: Proxy n)) :: Word8) + gpoke x + {-# INLINE gpokeSum #-} +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)) + {-# INLINE gpeekSum #-} + +------------------------------------------------------------------------ +-- Utilities for defining 'Store' instances based on 'Storable' + +-- | A 'size' implementation based on an instance of 'Storable' and +-- 'Typeable'. +sizeStorable :: forall a. (Storable a, Typeable a) => Size a +sizeStorable = sizeStorableTy (show (typeRep (Proxy :: Proxy a))) +{-# INLINE sizeStorable #-} + +-- | A 'size' implementation based on an instance of 'Storable'. Use this +-- if the type is not 'Typeable'. +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." +{-# INLINE sizeStorableTy #-} + +-- | A 'poke' implementation based on an instance of 'Storable'. +pokeStorable :: Storable a => a -> Poke () +pokeStorable x = Poke $ \ptr offset -> do + y <- pokeByteOff ptr offset x + let !newOffset = offset + sizeOf x + return (newOffset, y) +{-# INLINE pokeStorable #-} + +-- FIXME: make it the responsibility of the caller to check this. + +-- | A 'peek' implementation based on an instance of 'Storable' and +-- 'Typeable'. +peekStorable :: forall a. (Storable a, Typeable a) => Peek a +peekStorable = peekStorableTy (show (typeRep (Proxy :: Proxy a))) +{-# INLINE peekStorable #-} + +-- | A 'peek' implementation based on an instance of 'Storable'. Use +-- this if the type is not 'Typeable'. +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) +{-# INLINE peekStorableTy #-} + +------------------------------------------------------------------------ +-- Helpful Type Synonyms + +-- | Total byte size of the given Ptr +type Total = Int + +-- | How far into the given Ptr to look +type Offset = Int + +------------------------------------------------------------------------ +-- Poke monad + +newtype Poke a = Poke + { runPoke :: forall byte. Ptr byte -> Offset -> IO (Offset, a) + -- ^ Run the 'Poke' action, with the 'Ptr' to the buffer where + -- data is poked, and the current 'Offset'. The result is the new + -- offset, along with a return value. + } + deriving Functor + +instance Applicative Poke where + pure x = Poke $ \_ptr offset -> pure (offset, x) + {-# INLINE pure #-} + Poke f <*> Poke g = Poke $ \ptr offset1 -> do + (offset2, f') <- f ptr offset1 + (offset3, g') <- g ptr offset2 + return (offset3, f' g') + {-# INLINE (<*>) #-} + Poke f *> Poke g = Poke $ \ptr offset1 -> do + (offset2, _) <- f ptr offset1 + g ptr offset2 + {-# INLINE (*>) #-} + +instance Monad Poke where + return = pure + {-# INLINE return #-} + (>>) = (*>) + {-# INLINE (>>) #-} + Poke x >>= f = Poke $ \ptr offset1 -> do + (offset2, x') <- x ptr offset1 + runPoke (f x') ptr offset2 + {-# INLINE (>>=) #-} + fail = Fail.fail + {-# INLINE fail #-} + +instance Fail.MonadFail Poke where + fail = pokeException . T.pack + {-# INLINE fail #-} + +instance MonadIO Poke where + liftIO f = Poke $ \_ offset -> (offset, ) <$> f + {-# INLINE liftIO #-} + +-- | Exception thrown while running 'poke'. Note that other types of +-- exceptions could also be thrown. Invocations of 'fail' in the 'Poke' +-- monad causes this exception to be thrown. +-- +-- 'PokeException's are not expected to occur in ordinary circumstances, +-- and usually indicate a programming error. +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) + +------------------------------------------------------------------------ +-- Peek monad + +newtype Peek a = Peek + { runPeek :: forall byte. Ptr byte -> Ptr byte -> IO (Ptr byte, a) + -- ^ Run the 'Peek' action, with a 'Ptr' to the end of the buffer + -- where data is poked, and a 'Ptr' to the current position. The + -- result is the 'Ptr', along with a return value. + } + deriving Functor + +instance Applicative Peek where + pure x = Peek (\_ ptr -> return (ptr, x)) + {-# INLINE pure #-} + Peek f <*> Peek g = Peek $ \end ptr1 -> do + (ptr2, f') <- f end ptr1 + (ptr3, g') <- g end ptr2 + return (ptr3, f' g') + {-# INLINE (<*>) #-} + Peek f *> Peek g = Peek $ \end ptr1 -> do + (ptr2, _) <- f end ptr1 + g end ptr2 + {-# INLINE (*>) #-} + +instance Monad Peek where + return = pure + {-# INLINE return #-} + (>>) = (*>) + {-# INLINE (>>) #-} + Peek x >>= f = Peek $ \end ptr1 -> do + (ptr2, x') <- x end ptr1 + runPeek (f x') end ptr2 + {-# INLINE (>>=) #-} + fail = Fail.fail + {-# INLINE fail #-} + +instance Fail.MonadFail Peek where + fail = peekException . T.pack + {-# INLINE fail #-} + +instance PrimMonad Peek where + type PrimState Peek = RealWorld + primitive action = Peek $ \_ ptr -> do + x <- primitive (unsafeCoerce# action) + return (ptr, x) + {-# INLINE primitive #-} + +instance MonadIO Peek where + liftIO f = Peek $ \_ ptr -> (ptr, ) <$> f + {-# INLINE liftIO #-} + +-- | Exception thrown while running 'peek'. Note that other types of +-- exceptions can also be thrown. Invocations of 'fail' in the 'Poke' +-- monad causes this exception to be thrown. +-- +-- 'PeekException' is thrown when the data being decoded is invalid. +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." + +------------------------------------------------------------------------ +-- Size type + +-- | Info about a type's serialized length. Either the length is known +-- independently of the value, or the length depends on the value. +data Size a + = VarSize (a -> Int) + | ConstSize !Int + deriving Typeable + +getSize :: Store a => a -> Int +getSize = getSizeWith size +{-# INLINE getSize #-} + +getSizeWith :: Size a -> a -> Int +getSizeWith (VarSize f) x = f x +getSizeWith (ConstSize n) _ = n +{-# INLINE getSizeWith #-} + +-- TODO: depend on contravariant package? The ConstSize case is a little +-- wonky due to type conversion + +contramapSize :: (a -> b) -> Size b -> Size a +contramapSize f (VarSize g) = VarSize (g . f) +contramapSize _ (ConstSize n) = ConstSize n +{-# INLINE contramapSize #-} + +combineSize :: forall a b c. (Store a, Store b) => (c -> a) -> (c -> b) -> Size c +combineSize toA toB = combineSize' toA toB size size +{-# INLINE combineSize #-} + +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) +{-# INLINE combineSize' #-} + +scaleSize :: Int -> Size a -> Size a +scaleSize s (ConstSize n) = ConstSize (s * n) +scaleSize s (VarSize f) = VarSize ((s *) . f) +{-# INLINE scaleSize #-} + +addSize :: Int -> Size a -> Size a +addSize x (ConstSize n) = ConstSize (x + n) +addSize x (VarSize f) = VarSize ((x +) . f) +{-# INLINE addSize #-} + +------------------------------------------------------------------------ +-- Utilities for implementing 'Store' instances via memcpy + +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, () #)) diff --git a/src/Data/Store/Internal.hs b/src/Data/Store/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Data/Store/Internal.hs @@ -0,0 +1,719 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes#-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE ViewPatterns #-} +{-# LANGUAGE EmptyCase #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} + +-- | Internal API for the store package. The functions here which are +-- not re-exported by "Data.Store" are less likely to have stable APIs. +-- +-- This module also defines most of the included 'Store' instances, for +-- types from the base package and other commonly used packages +-- (bytestring, containers, text, time, etc). +module Data.Store.Internal + ( + -- * Encoding and decoding strict ByteStrings. + encode, + decode, decodeWith, + decodeEx, decodeExWith, decodeExPortionWith + , decodeIO, decodeIOWith, decodeIOPortionWith + -- * Store class and related types. + , Store(..), Poke, Peek, runPeek + -- ** Exceptions thrown by Poke + , PokeException(..), pokeException + -- ** Exceptions thrown by Peek + , PeekException(..), peekException, tooManyBytes + -- ** Size type + , Size(..) + , getSize, getSizeWith + , contramapSize, combineSize, combineSize', scaleSize, addSize + -- ** Store instances in terms of IsSequence + , sizeSequence, pokeSequence, peekSequence + -- ** Store instances in terms of IsSet + , sizeSet, pokeSet, peekSet + -- ** Store instances in terms of IsMap + , sizeMap, pokeMap, peekMap + -- ** Peek utilities + , skip, isolate + -- ** Static Size type + -- + -- This portion of the library is still work-in-progress. + -- 'IsStaticSize' is only supported for strict ByteStrings, in order + -- to support the use case of 'Tagged'. + , IsStaticSize(..), StaticSize(..), toStaticSizeEx, liftStaticSize + ) where + +import Control.Applicative +import Control.DeepSeq (NFData) +import Control.Exception (throwIO) +import Control.Monad (when) +import Control.Monad.IO.Class (liftIO) +import qualified Data.ByteString as BS +import qualified Data.ByteString.Internal as BS +import qualified Data.ByteString.Lazy as LBS +import qualified Data.ByteString.Short.Internal as SBS +import Data.Containers (IsMap, ContainerKey, MapValue, mapFromList, mapToList, IsSet, setFromList) +import Data.Data (Data) +import Data.Fixed (Fixed (..), Pico) +import Data.Foldable (forM_) +import Data.HashMap.Strict (HashMap) +import Data.HashSet (HashSet) +import Data.Hashable (Hashable) +import Data.IntMap (IntMap) +import Data.IntSet (IntSet) +import qualified Data.List.NonEmpty as NE +import Data.Map (Map) +import Data.MonoTraversable +import Data.Monoid +import Data.Orphans () +import Data.Primitive.ByteArray +import Data.Proxy (Proxy(..)) +import Data.Sequence (Seq) +import Data.Sequences (IsSequence, Index, replicateM) +import Data.Set (Set) +import Data.Store.Impl +import Data.Store.TH.Internal +import qualified Data.Text as T +import qualified Data.Text.Array as TA +import qualified Data.Text.Foreign as T +import qualified Data.Text.Internal as T +import qualified Data.Time as Time +import Data.Typeable.Internal (Typeable, Fingerprint) +import qualified Data.Vector as V +import qualified Data.Vector.Mutable as MV +import qualified Data.Vector.Storable as SV +import qualified Data.Vector.Storable.Mutable as MSV +import Data.Void +import Data.Word +import Foreign.Ptr (plusPtr, minusPtr) +import Foreign.Storable (Storable, sizeOf) +import GHC.Generics (Generic) +import qualified GHC.Integer.GMP.Internals as I +import GHC.Real (Ratio(..)) +import GHC.TypeLits +import GHC.Types (Int (I#)) +import Language.Haskell.TH +import Language.Haskell.TH.Instances () +import Language.Haskell.TH.ReifyMany +import Language.Haskell.TH.Syntax +import Prelude +import TH.Derive + +-- Conditional import to avoid warning +#if MIN_VERSION_integer_gmp(1,0,0) +import GHC.Prim (sizeofByteArray#) +#endif + +------------------------------------------------------------------------ +-- Utilities for defining list-like 'Store' instances in terms of 'IsSequence' + +-- | Implement 'size' for an 'IsSequence' of 'Store' instances. +-- +-- Note that many monomorphic containers have more efficient +-- implementations (for example, via memcpy). +sizeSequence :: forall t. (IsSequence t, Store (Element t)) => Size t +sizeSequence = VarSize $ \t -> + case size :: Size (Element t) of + ConstSize n -> n * (olength t) + sizeOf (undefined :: Int) + VarSize f -> ofoldl' (\acc x -> acc + f x) (sizeOf (undefined :: Int)) t +{-# INLINE sizeSequence #-} + +-- | Implement 'poke' for an 'IsSequence' of 'Store' instances. +-- +-- Note that many monomorphic containers have more efficient +-- implementations (for example, via memcpy). +pokeSequence :: (IsSequence t, Store (Element t)) => t -> Poke () +pokeSequence t = + do pokeStorable len + Poke (\ptr offset -> + do offset' <- + ofoldlM (\offset' a -> + do (offset'',_) <- runPoke (poke a) ptr offset' + return offset'') + offset + t + return (offset',())) + where len = olength t +{-# INLINE pokeSequence #-} + +-- | Implement 'peek' for an 'IsSequence' of 'Store' instances. +-- +-- Note that many monomorphic containers have more efficient +-- implementations (for example, via memcpy). +peekSequence :: (IsSequence t, Store (Element t), Index t ~ Int) => Peek t +peekSequence = do + len <- peek + replicateM len peek +{-# INLINE peekSequence #-} + +------------------------------------------------------------------------ +-- Utilities for defining list-like 'Store' instances in terms of 'IsSet' + +-- | Implement 'size' for an 'IsSet' of 'Store' instances. +sizeSet :: forall t. (IsSet t, Store (Element t)) => Size t +sizeSet = VarSize $ \t -> + case size :: Size (Element t) of + ConstSize n -> n * (olength t) + sizeOf (undefined :: Int) + VarSize f -> ofoldl' (\acc x -> acc + f x) (sizeOf (undefined :: Int)) t +{-# INLINE sizeSet #-} + +-- | Implement 'poke' for an 'IsSequence' of 'Store' instances. +pokeSet :: (IsSet t, Store (Element t)) => t -> Poke () +pokeSet t = do + pokeStorable (olength t) + omapM_ poke t +{-# INLINE pokeSet #-} + +-- | Implement 'peek' for an 'IsSequence' of 'Store' instances. +peekSet :: (IsSet t, Store (Element t)) => Peek t +peekSet = do + len <- peek + setFromList <$> replicateM len peek +{-# INLINE peekSet #-} + +------------------------------------------------------------------------ +-- Utilities for defining list-like 'Store' instances in terms of a 'IsMap' + +-- | Implement 'size' for an 'IsMap' of where both 'ContainerKey' and +-- 'MapValue' are 'Store' instances. +sizeMap + :: forall t. (Store (ContainerKey t), Store (MapValue t), IsMap t) + => Size t +sizeMap = VarSize $ \t -> + case (size :: Size (ContainerKey t), size :: Size (MapValue t)) of + (ConstSize nk, ConstSize na) -> (nk + na) * olength t + sizeOf (undefined :: Int) + (szk, sza) -> ofoldl' (\acc (k, a) -> acc + getSizeWith szk k + getSizeWith sza a) + (sizeOf (undefined :: Int)) + (mapToList t) +{-# INLINE sizeMap #-} + +-- | Implement 'poke' for an 'IsMap' of where both 'ContainerKey' and +-- 'MapValue' are 'Store' instances. +pokeMap + :: (Store (ContainerKey t), Store (MapValue t), IsMap t) + => t + -> Poke () +pokeMap t = do + poke (olength t) + ofoldl' (\acc (k, x) -> poke k >> poke x >> acc) + (return ()) + (mapToList t) +{-# INLINE pokeMap #-} + +-- | Implement 'peek' for an 'IsMap' of where both 'ContainerKey' and +-- 'MapValue' are 'Store' instances. +peekMap + :: (Store (ContainerKey t), Store (MapValue t), IsMap t) + => Peek t +peekMap = mapFromList <$> peek +{-# INLINE peekMap #-} + +{- +------------------------------------------------------------------------ +-- Utilities for defining list-like 'Store' instances in terms of Foldable + +-- | Implement 'size' for a 'Foldable' of 'Store' instances. Note that +-- this assumes the extra 'Foldable' structure is discardable - this +-- only serializes the elements. +sizeListLikeFoldable :: forall t a. (Foldable t, Store a) => Size (t a) +sizeListLikeFoldable = VarSize $ \t -> + case size :: Size e of + ConstSize n -> n * length x + sizeOf (undefined :: Int) + VarSize f -> foldl' (\acc x -> acc + f x) (sizeOf (undefined :: Int)) +{-# INLINE sizeSequence #-} + +pokeListLikeFoldable :: forall t a. Foldable t => t a -> Poke () +pokeListLikeFoldable x = do + poke (length x) +-} + +------------------------------------------------------------------------ +-- Utilities for implementing 'Store' instances for list-like mutable things + +-- | Implementation of peek for mutable sequences. The user provides a +-- function for initializing the sequence and a function for mutating an +-- element at a particular index. +peekMutableSequence + :: Store a + => (Int -> IO r) + -> (r -> Int -> a -> IO ()) + -> Peek r +peekMutableSequence new write = do + n <- peek + mut <- liftIO (new n) + forM_ [0..n-1] $ \i -> peek >>= liftIO . write mut i + return mut +{-# INLINE peekMutableSequence #-} + +------------------------------------------------------------------------ +-- Useful combinators + +-- | Skip n bytes forward. +{-# INLINE skip #-} +skip :: Int -> Peek () +skip len = Peek $ \end ptr -> do + let ptr2 = ptr `plusPtr` len + when (ptr2 > end) $ + tooManyBytes len (end `minusPtr` ptr) "skip" + return (ptr2, ()) + +-- | Isolate the input to n bytes, skipping n bytes forward. Fails if @m@ +-- advances the offset beyond the isolated region. +{-# INLINE isolate #-} +isolate :: Int -> Peek a -> Peek a +isolate len m = Peek $ \end ptr -> do + let ptr2 = ptr `plusPtr` len + when (ptr2 > end) $ + tooManyBytes len (end `minusPtr` ptr) "isolate" + (ptr', x) <- runPeek m end ptr + when (ptr' > end) $ + throwIO $ PeekException (ptr' `minusPtr` end) "Overshot end of isolated bytes" + return (ptr2, x) + +------------------------------------------------------------------------ +-- Instances for types based on flat representations + +instance Store a => Store (V.Vector a) where + size = sizeSequence + poke = pokeSequence + peek = V.unsafeFreeze =<< peekMutableSequence MV.new MV.write + {-# INLINE size #-} + {-# INLINE peek #-} + {-# INLINE poke #-} + +instance Storable a => Store (SV.Vector a) where + size = VarSize $ \x -> + sizeOf (undefined :: Int) + + sizeOf (undefined :: a) * SV.length x + poke x = do + let (fptr, len) = SV.unsafeToForeignPtr0 x + poke len + pokeFromForeignPtr fptr 0 (sizeOf (undefined :: a) * len) + peek = do + len <- peek + fp <- peekToPlainForeignPtr "Data.Storable.Vector.Vector" (sizeOf (undefined :: a) * len) + liftIO $ SV.unsafeFreeze (MSV.MVector len fp) + {-# INLINE size #-} + {-# INLINE peek #-} + {-# INLINE poke #-} + +instance Store BS.ByteString where + size = VarSize $ \x -> + sizeOf (undefined :: Int) + + BS.length x + poke x = do + let (sourceFp, sourceOffset, sourceLength) = BS.toForeignPtr x + poke sourceLength + pokeFromForeignPtr sourceFp sourceOffset sourceLength + peek = do + len <- peek + fp <- peekToPlainForeignPtr "Data.ByteString.ByteString" len + return (BS.PS fp 0 len) + {-# INLINE size #-} + {-# INLINE peek #-} + {-# INLINE poke #-} + +instance Store SBS.ShortByteString where + size = VarSize $ \x -> + sizeOf (undefined :: Int) + + SBS.length x + poke x@(SBS.SBS arr) = do + let len = SBS.length x + poke len + pokeFromByteArray arr 0 len + peek = do + len <- peek + ByteArray array <- peekToByteArray "Data.ByteString.Short.ShortByteString" len + return (SBS.SBS array) + {-# INLINE size #-} + {-# INLINE peek #-} + {-# INLINE poke #-} + +instance Store LBS.ByteString where + -- FIXME: faster conversion? Is this ever going to be a problem? + -- + -- I think on 64 bit systems, Int will have 64 bits. On 32 bit + -- systems, we'll never exceed the range of Int by this conversion. + size = VarSize $ \x -> + sizeOf (undefined :: Int) + + fromIntegral (LBS.length x) + -- FIXME: more efficient implementation that avoids the double copy + poke = poke . LBS.toStrict + peek = fmap LBS.fromStrict peek + {-# INLINE size #-} + {-# INLINE peek #-} + {-# INLINE poke #-} + +instance Store T.Text where + size = VarSize $ \x -> + sizeOf (undefined :: Int) + + 2 * (T.lengthWord16 x) + poke x = do + let !(T.Text (TA.Array array) w16Off w16Len) = x + poke w16Len + pokeFromByteArray array (2 * w16Off) (2 * w16Len) + peek = do + w16Len <- peek + ByteArray array <- peekToByteArray "Data.Text.Text" (2 * w16Len) + return (T.Text (TA.Array array) 0 w16Len) + {-# INLINE size #-} + {-# INLINE peek #-} + {-# INLINE poke #-} + +{- +-- Gets a little tricky to compute size due to size of storing indices. + +instance (Store i, Store e) => Store (Array i e) where + size = combineSize' () () () $ + VarSize $ \t -> + case size :: Size e of + ConstSize n -> n * length x + VarSize f -> foldl' (\acc x -> acc + f x) 0 +-} + +------------------------------------------------------------------------ +-- Known size instances + +-- TODO: this doesn't scale nicely to 'Text'. Force it to be byte size? +-- 'StaticByteSize'? + +newtype StaticSize (n :: Nat) a = StaticSize { unStaticSize :: a } + deriving (Eq, Show, Ord, Data, Typeable, Generic) + +instance NFData a => NFData (StaticSize n a) + +class KnownNat n => IsStaticSize n a where + toStaticSize :: a -> Maybe (StaticSize n a) + +toStaticSizeEx :: IsStaticSize n a => a -> StaticSize n a +toStaticSizeEx x = + case toStaticSize x of + Just r -> r + Nothing -> error "Failed to assert a static size via toStaticSizeEx" + +instance KnownNat n => IsStaticSize n BS.ByteString where + toStaticSize bs + | BS.length bs == fromInteger (natVal (Proxy :: Proxy n)) = Just (StaticSize bs) + | otherwise = Nothing + +instance KnownNat n => Store (StaticSize n BS.ByteString) where + size = ConstSize (fromInteger (natVal (Proxy :: Proxy n))) + poke (StaticSize x) = do + -- TODO: worth it to put an assert here? + let (sourceFp, sourceOffset, sourceLength) = BS.toForeignPtr x + pokeFromForeignPtr sourceFp sourceOffset sourceLength + peek = do + let len = fromInteger (natVal (Proxy :: Proxy n)) + fp <- peekToPlainForeignPtr ("StaticSize " ++ show len ++ " Data.ByteString") len + return (StaticSize (BS.PS fp 0 len)) + {-# INLINE size #-} + {-# INLINE peek #-} + {-# INLINE poke #-} + +-- NOTE: this could be a 'Lift' instance, but we can't use type holes in +-- TH. Alternatively we'd need a (TypeRep -> Type) function and Typeable +-- constraint. +liftStaticSize :: forall n a. (KnownNat n, Lift a) => TypeQ -> StaticSize n a -> ExpQ +liftStaticSize tyq (StaticSize x) = do + let numTy = litT $ numTyLit $ natVal (Proxy :: Proxy n) + [| StaticSize $(lift x) :: StaticSize $(numTy) $(tyq) |] + +------------------------------------------------------------------------ +-- containers instances + +instance Store a => Store [a] where + size = sizeSequence + poke = pokeSequence + peek = peekSequence + {-# INLINE size #-} + {-# INLINE peek #-} + {-# INLINE poke #-} + +instance Store a => Store (NE.NonEmpty a) + +instance Store a => Store (Seq a) where + size = sizeSequence + poke = pokeSequence + peek = peekSequence + {-# INLINE size #-} + {-# INLINE peek #-} + {-# INLINE poke #-} + +instance (Store a, Ord a) => Store (Set a) where + size = sizeSet + poke = pokeSet + peek = peekSet + {-# INLINE size #-} + {-# INLINE peek #-} + {-# INLINE poke #-} + +instance Store IntSet where + size = sizeSet + poke = pokeSet + peek = peekSet + {-# INLINE size #-} + {-# INLINE peek #-} + {-# INLINE poke #-} + +instance Store a => Store (IntMap a) where + size = sizeMap + poke = pokeMap + peek = peekMap + {-# INLINE size #-} + {-# INLINE peek #-} + {-# INLINE poke #-} + +instance (Ord k, Store k, Store a) => Store (Map k a) where + size = sizeMap + poke = pokeMap + peek = peekMap + {-# INLINE size #-} + {-# INLINE peek #-} + {-# INLINE poke #-} + +instance (Eq k, Hashable k, Store k, Store a) => Store (HashMap k a) where + size = sizeMap + poke = pokeMap + peek = peekMap + {-# INLINE size #-} + {-# INLINE peek #-} + {-# INLINE poke #-} + +instance (Eq a, Hashable a, Store a) => Store (HashSet a) where + size = sizeSet + poke = pokeSet + peek = peekSet + {-# INLINE size #-} + {-# INLINE peek #-} + {-# INLINE poke #-} + +-- FIXME: implement +-- +-- instance (Ix i, Bounded i, Store a) => Store (Array ix a) where +-- +-- instance (Ix i, Bounded i, Store a) => Store (UA.UArray ix a) where + +instance Store Integer where +#if MIN_VERSION_integer_gmp(1,0,0) + size = VarSize $ \ x -> + sizeOf (undefined :: Word8) + case x of + I.S# _ -> sizeOf (undefined :: Int) + I.Jp# (I.BN# arr) -> sizeOf (undefined :: Int) + I# (sizeofByteArray# arr) + I.Jn# (I.BN# arr) -> sizeOf (undefined :: Int) + I# (sizeofByteArray# arr) + poke (I.S# x) = poke (0 :: Word8) >> poke (I# x) + poke (I.Jp# (I.BN# arr)) = do + let len = I# (sizeofByteArray# arr) + poke (1 :: Word8) + poke len + pokeFromByteArray arr 0 len + poke (I.Jn# (I.BN# arr)) = do + let len = I# (sizeofByteArray# arr) + poke (2 :: Word8) + poke len + pokeFromByteArray arr 0 len + peek = do + tag <- peek :: Peek Word8 + case tag of + 0 -> fromIntegral <$> (peek :: Peek Int) + 1 -> I.Jp# <$> peekBN + 2 -> I.Jn# <$> peekBN + _ -> peekException "Invalid Integer tag" + where + peekBN = do + len <- peek :: Peek Int + ByteArray arr <- peekToByteArray "GHC>Integer" len + return $ I.BN# arr +#else + -- May as well put in the extra effort to use the same encoding as + -- used for the newer integer-gmp. + size = VarSize $ \ x -> + sizeOf (undefined :: Word8) + case x of + I.S# _ -> sizeOf (undefined :: Int) + I.J# sz _ -> sizeOf (undefined :: Int) + (I# sz) * sizeOf (undefined :: Word) + poke (I.S# x) = poke (0 :: Word8) >> poke (I# x) + poke (I.J# sz arr) + | (I# sz) > 0 = do + let len = I# sz * sizeOf (undefined :: Word) + poke (1 :: Word8) + poke len + pokeFromByteArray arr 0 len + | (I# sz) < 0 = do + let len = negate (I# sz) * sizeOf (undefined :: Word) + poke (2 :: Word8) + poke len + pokeFromByteArray arr 0 len + | otherwise = do + poke (0 :: Word8) + poke (0 :: Int) + peek = do + tag <- peek :: Peek Word8 + case tag of + 0 -> fromIntegral <$> (peek :: Peek Int) + 1 -> peekJ False + 2 -> peekJ True + _ -> peekException "Invalid Integer tag" + where + peekJ neg = do + len <- peek :: Peek Int + ByteArray arr <- peekToByteArray "GHC>Integer" len + let (sz0, r) = len `divMod` (sizeOf (undefined :: Word)) + !(I# sz) = if neg then negate sz0 else sz0 + when (r /= 0) (peekException "Buffer size stored for encoded Integer not divisible by Word size (to get limb count).") + return (I.J# sz arr) +#endif + +-- instance Store GHC.Fingerprint.Types.Fingerprint where + +instance Store (Fixed a) where + size = contramapSize (\(MkFixed x) -> x) (size :: Size Integer) + poke (MkFixed x) = poke x + peek = MkFixed <$> peek + {-# INLINE size #-} + {-# INLINE peek #-} + {-# INLINE poke #-} + +-- instance Store a => Store (Tree a) where + +------------------------------------------------------------------------ +-- Other instances + +-- Manual implementation due to no Generic instance for Ratio. Also due +-- to the instance for Storable erroring when the denominator is 0. +-- Perhaps we should keep the behavior but instead a peekException? +-- +-- In that case it should also error on poke. +-- +-- I prefer being able to Store these, because they are constructable. + +instance Store a => Store (Ratio a) where + size = combineSize (\(x :% _) -> x) (\(_ :% y) -> y) + poke (x :% y) = poke (x, y) + peek = uncurry (:%) <$> peek + {-# INLINE size #-} + {-# INLINE peek #-} + {-# INLINE poke #-} + +instance Store Time.Day where + size = contramapSize Time.toModifiedJulianDay (size :: Size Integer) + poke = poke . Time.toModifiedJulianDay + peek = Time.ModifiedJulianDay <$> peek + {-# INLINE size #-} + {-# INLINE peek #-} + {-# INLINE poke #-} + +instance Store Time.DiffTime where + size = contramapSize (realToFrac :: Time.DiffTime -> Pico) (size :: Size Pico) + poke = (poke :: Pico -> Poke ()) . realToFrac + peek = Time.picosecondsToDiffTime <$> peek + {-# INLINE size #-} + {-# INLINE peek #-} + {-# INLINE poke #-} + +instance Store Time.UTCTime where + size = combineSize Time.utctDay Time.utctDayTime + poke (Time.UTCTime day time) = poke (day, time) + peek = uncurry Time.UTCTime <$> peek + {-# INLINE size #-} + {-# INLINE peek #-} + {-# INLINE poke #-} + +instance Store () +instance Store a => Store (Dual a) +instance Store a => Store (Sum a) +instance Store a => Store (Product a) +instance Store a => Store (First a) +instance Store a => Store (Last a) +instance Store a => Store (Maybe a) +instance (Store a, Store b) => Store (Either a b) + +-- Explicit definition needed because in base <= 4.7 (GHC 7.8), +-- Fingerprint is not Typeable. +instance Store Fingerprint where + size = sizeStorableTy "GHC.Fingerprint.Type.Fingerprint" + poke = pokeStorable + peek = peekStorableTy "GHC.Fingerprint.Type.Fingerprint" + +-- FIXME: have TH deriving handle unboxed fields? + +------------------------------------------------------------------------ +-- Instances generated by TH + +$($(derive [d| + -- TODO + -- instance Deriving (Store ()) + instance Deriving (Store All) + instance Deriving (Store Any) + instance Deriving (Store Void) + instance Deriving (Store Bool) + |])) + +-- TODO: higher arities? Limited now by Generics instances for tuples +$(return $ map deriveTupleStoreInstance [2..7]) + +$(deriveManyStoreUnboxVector) + +$(deriveManyStoreFromStorable (\_ -> True)) + +$(deriveManyStorePrimVector) + +$(reifyManyWithoutInstances ''Store [''ModName, ''NameSpace, ''PkgName] (const True) >>= +-- mapM (\name -> deriveStore [] (ConT name) .dtCons =<< reifyDataType name)) + mapM (\name -> return (deriveGenericInstance [] (ConT name)))) + +-- Explicit definition needed because in template-haskell <= 2.9 (GHC +-- 7.8), NameFlavour contains unboxed values, causing generic deriving +-- to fail. +#if !MIN_VERSION_template_haskell(2,10,0) +instance Store NameFlavour where + size = VarSize $ \x -> getSize (0 :: Word8) + case x of + NameS -> 0 + NameQ mn -> getSize mn + NameU i -> getSize (I# i) + NameL i -> getSize (I# i) + NameG ns pn mn -> getSize ns + getSize pn + getSize mn + poke NameS = poke (0 :: Word8) + poke (NameQ mn) = do + poke (1 :: Word8) + poke mn + poke (NameU i) = do + poke (2 :: Word8) + poke (I# i) + poke (NameL i) = do + poke (3 :: Word8) + poke (I# i) + poke (NameG ns pn mn) = do + poke (4 :: Word8) + poke ns + poke pn + poke mn + peek = do + tag <- peek + case tag :: Word8 of + 0 -> return NameS + 1 -> NameQ <$> peek + 2 -> do + !(I# i) <- peek + return (NameU i) + 3 -> do + !(I# i) <- peek + return (NameL i) + 4 -> NameG <$> peek <*> peek <*> peek + _ -> peekException "Invalid NameFlavour tag" +#endif + +$(reifyManyWithoutInstances ''Store [''Info] (const True) >>= +-- mapM (\name -> deriveStore [] (ConT name) .dtCons =<< reifyDataType name)) + mapM (\name -> return (deriveGenericInstance [] (ConT name)))) diff --git a/src/Data/Store/Streaming.hs b/src/Data/Store/Streaming.hs new file mode 100644 --- /dev/null +++ b/src/Data/Store/Streaming.hs @@ -0,0 +1,173 @@ +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE FlexibleContexts #-} +{-| +Module: Data.Store.Streaming +Description: A thin streaming layer that uses 'Store' for serialisation. + +For efficiency reasons, 'Store' does not provide facilities for +incrementally consuming input. In order to avoid partial input, this +module introduces 'Message's that wrap values of instances of 'Store'. + +In addition to the serialisation of a value, the serialised message +also contains the length of the serialisation. This way, instead of +consuming input incrementally, more input can be demanded before +serialisation is attempted in the first place. + +-} +module Data.Store.Streaming + ( -- * 'Message's to stream data using 'Store' for serialisation. + Message (..) + -- * Encoding 'Message's + , encodeMessage + -- * Decoding 'Message's + , PeekMessage (..) + , peekMessage + , decodeMessage + -- * Conduits for encoding and decoding + , conduitEncode + , conduitDecode + ) where + +import Control.Exception (assert) +import Control.Monad (liftM) +import Control.Monad.IO.Class +import Control.Monad.Trans.Resource (MonadResource) +import Data.ByteString (ByteString) +import qualified Data.ByteString.Internal as BS +import qualified Data.Conduit as C +import qualified Data.Conduit.List as C +import Data.Store +import Data.Store.Impl (Peek (..), Poke (..), tooManyBytes, getSize) +import Data.Word +import Foreign.Ptr +import qualified Foreign.Storable as Storable +import Prelude +import System.IO.ByteBuffer (ByteBuffer) +import qualified System.IO.ByteBuffer as BB + +-- | If @a@ is an instance of 'Store', @Message a@ can be serialised +-- and deserialised in a streaming fashion. +newtype Message a = Message { fromMessage :: a } deriving (Eq, Show) + +-- | Type used to store the length of a 'Message'. +type SizeTag = Int + +tagLength :: Int +tagLength = Storable.sizeOf (undefined :: SizeTag) +{-# INLINE tagLength #-} + +-- | Encode a 'Message' to a 'ByteString'. +encodeMessage :: Store a => Message a -> ByteString +encodeMessage (Message x) = + let l = getSize x + totalLength = tagLength + l + in BS.unsafeCreate + totalLength + (\p -> do (offset, ()) <- runPoke (poke l >> poke x) p 0 + assert (offset == totalLength) (return ())) +{-# INLINE encodeMessage #-} + +-- | The result of peeking at the next message can either be a +-- successfully deserialised 'Message', or a request for more input. +data PeekMessage m a = Done (Message a) + | NeedMoreInput (ByteString -> m (PeekMessage m a)) + +-- | Decode a 'Message' of known size from a 'ByteBuffer'. +peekSized :: (MonadIO m, Store a) => ByteBuffer -> Int -> m (PeekMessage m a) +peekSized bb n = + BB.unsafeConsume bb n >>= \case + Right ptr -> liftM (Done . Message) $ decodeFromPtr ptr n + Left _ -> return $ NeedMoreInput (\ bs -> BB.copyByteString bb bs + >> peekSized bb n) +{-# INLINE peekSized #-} + +-- | Decode a 'SizeTag' from a 'ByteBuffer'. +peekSizeTag :: MonadIO m => ByteBuffer -> m (PeekMessage m SizeTag) +peekSizeTag bb = peekSized bb tagLength +{-# INLINE peekSizeTag #-} + +-- | Decode some 'Message' from a 'ByteBuffer', by first reading its +-- size, and then the actual 'Message'. +peekMessage :: (MonadIO m, Store a) => ByteBuffer -> m (PeekMessage m a) +peekMessage bb = + peekSizeTag bb >>= \case + (Done (Message n)) -> peekSized bb n + NeedMoreInput _ -> + return $ NeedMoreInput (\ bs -> BB.copyByteString bb bs + >> peekMessage bb) +{-# INLINE peekMessage #-} + +-- | Decode a 'Message' from a 'ByteBuffer' and an action that can get +-- additional 'ByteString's to refill the buffer when necessary. +-- +-- The only conditions under which this function will give 'Nothing', +-- is when the 'ByteBuffer' contains zero bytes, and refilling yields +-- 'Nothing'. If there is some data available, but not enough to +-- decode the whole 'Message', a 'PeekException' will be thrown. +decodeMessage :: (MonadIO m, Store a) + => ByteBuffer -> m (Maybe ByteString) -> m (Maybe (Message a)) +decodeMessage bb getBs = + decodeSizeTag bb getBs >>= \case + Nothing -> return Nothing + Just n -> decodeSized bb getBs n +{-# INLINE decodeMessage #-} + +decodeSizeTag :: MonadIO m + => ByteBuffer + -> m (Maybe ByteString) + -> m (Maybe SizeTag) +decodeSizeTag bb getBs = + peekSizeTag bb >>= \case + (Done (Message n)) -> return (Just n) + (NeedMoreInput _) -> getBs >>= \case + Just bs -> BB.copyByteString bb bs >> decodeSizeTag bb getBs + Nothing -> BB.availableBytes bb >>= \case + 0 -> return Nothing + n -> liftIO $ tooManyBytes tagLength n "Data.Store.Message.SizeTag" +{-# INLINE decodeSizeTag #-} + +decodeSized :: (MonadIO m, Store a) + => ByteBuffer + -> m (Maybe ByteString) + -> Int + -> m (Maybe (Message a)) +decodeSized bb getBs n = + peekSized bb n >>= \case + Done message -> return (Just message) + NeedMoreInput _ -> getBs >>= \case + Just bs -> BB.copyByteString bb bs >> decodeSized bb getBs n + Nothing -> BB.availableBytes bb >>= \ available -> + liftIO $ tooManyBytes n available "Data.Store.Message.Message" +{-# INLINE decodeSized #-} + +-- | Decode a value, given a 'Ptr' and the number of bytes that make +-- up the encoded message. +decodeFromPtr :: (MonadIO m, Store a) => Ptr Word8 -> Int -> m a +decodeFromPtr ptr n = + liftIO (liftM snd $ runPeek peek (ptr `plusPtr` n) ptr) +{-# INLINE decodeFromPtr #-} + +-- | Conduit for encoding 'Message's to 'ByteString's. +conduitEncode :: (Monad m, Store a) => C.Conduit (Message a) m ByteString +conduitEncode = C.map encodeMessage +{-# INLINE conduitEncode #-} + +-- | Conduit for decoding 'Message's from 'ByteString's. +conduitDecode :: (MonadIO m, MonadResource m, Store a) + => Maybe Int + -- ^ Initial length of the 'ByteBuffer' used for + -- buffering the incoming 'ByteString's. If 'Nothing', + -- use the default value of 4MB. + -> C.Conduit ByteString m (Message a) +conduitDecode bufSize = + C.bracketP + (BB.new bufSize) + BB.free + go + where + go buffer = do + mmessage <- decodeMessage buffer C.await + case mmessage of + Nothing -> return () + Just message -> C.yield message >> go buffer +{-# INLINE conduitDecode #-} diff --git a/src/Data/Store/TH.hs b/src/Data/Store/TH.hs new file mode 100644 --- /dev/null +++ b/src/Data/Store/TH.hs @@ -0,0 +1,71 @@ +{-# LANGUAGE TemplateHaskell #-} + +-- | This module exports TH utilities intended to be useful to users. +-- +-- However, the visible exports do not show the main things that will be +-- useful, which is using TH to generate 'Store' instances, via +-- "TH.Derive". It's used like this: +-- +-- @ +-- data Foo = Foo Int | Bar Int +-- +-- $($(derive [d| +-- instance Deriving (Store Foo) +-- |])) +-- @ +-- +-- One advantage of using this Template Haskell definition of Store +-- instances is that in some cases they can be faster than the instances +-- defined via Generics. Specifically, sum types which can yield +-- 'ConstSize' from 'size' will be faster when used in array-like types. +-- The instances generated via generics always use 'VarSize' for sum +-- types. +module Data.Store.TH + ( + -- * Testing Store instances + smallcheckManyStore + , checkRoundtrip + , assertRoundtrip + ) where + +import Data.Complex () +import Data.Store.Impl +import Data.Typeable (Typeable, typeOf) +import Debug.Trace (trace) +import Language.Haskell.TH +import Prelude +import Test.Hspec +import Test.Hspec.SmallCheck (property) +import Test.SmallCheck + +------------------------------------------------------------------------ +-- Testing + +-- | Test a 'Store' instance using 'smallcheck' and 'hspec'. +smallcheckManyStore :: Bool -> Int -> [TypeQ] -> ExpQ +smallcheckManyStore verbose depth = smallcheckMany . map testRoundtrip + where + testRoundtrip tyq = do + ty <- tyq + expr <- [e| property $ changeDepth (\_ -> depth) $ \x -> checkRoundtrip verbose (x :: $(return ty)) |] + return ("Roundtrips (" ++ pprint ty ++ ")", expr) + +assertRoundtrip :: (Eq a, Show a, Store a, Monad m, Typeable a) => Bool -> a -> m () +assertRoundtrip verbose x + | checkRoundtrip verbose x = return () + | otherwise = fail $ "Failed to roundtrip " ++ show (typeOf x) + +-- | Check if a given value succeeds in decoding its encoded +-- representation. +checkRoundtrip :: (Eq a, Show a, Store a) => Bool -> a -> Bool +checkRoundtrip verbose x = decoded == Right x + where + encoded = verboseTrace verbose "encoded" (encode x) + decoded = verboseTrace verbose "decoded" (decode encoded) + +smallcheckMany :: [Q (String, Exp)] -> ExpQ +smallcheckMany = doE . map (\f -> f >>= \(name, expr) -> noBindS [e| it name $ $(return expr) |]) + +verboseTrace :: Show a => Bool -> String -> a -> a +verboseTrace True msg x = trace (show (msg, x)) x +verboseTrace False _ x = x diff --git a/src/Data/Store/TH/Internal.hs b/src/Data/Store/TH/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Data/Store/TH/Internal.hs @@ -0,0 +1,429 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE ParallelListComp #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE ViewPatterns #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} + +module Data.Store.TH.Internal + ( + -- * TH functions for generating Store instances + deriveManyStoreFromStorable + , deriveTupleStoreInstance + , deriveGenericInstance + , deriveManyStorePrimVector + , deriveManyStoreUnboxVector + , deriveStore + -- * Misc utilties used in Store test + , getAllInstanceTypes1 + , isMonoType + ) where + +import Control.Applicative +import Data.Complex () +import Data.Generics.Aliases (extT) +import Data.Generics.Schemes (listify, everywhere) +import Data.List (find, nub) +import qualified Data.Map as M +import Data.Maybe (fromMaybe) +import Data.Primitive.ByteArray +import Data.Primitive.Types +import Data.Store.Impl +import qualified Data.Text as T +import Data.Traversable (forM) +import Data.Typeable (Typeable) +import qualified Data.Vector.Primitive as PV +import qualified Data.Vector.Unboxed as UV +import Data.Word +import Foreign.Storable (Storable) +import GHC.Types (Int(..)) +import Language.Haskell.TH +import Language.Haskell.TH.ReifyMany.Internal (TypeclassInstance(..), getInstances, unAppsT) +import Language.Haskell.TH.Syntax (lift) +import Prelude +import Safe (headMay) +import TH.Derive (Deriver(..)) +import TH.ReifyDataType +import TH.Utilities (freeVarsT, expectTyCon1, dequalify) + +instance Deriver (Store a) where + runDeriver _ preds ty = do + argTy <- expectTyCon1 ''Store ty + dt <- reifyDataTypeSubstituted argTy + (:[]) <$> deriveStore preds argTy (dtCons dt) + +deriveStore :: Cxt -> Type -> [DataCon] -> Q Dec +deriveStore preds headTy cons0 = + makeStoreInstance preds headTy + <$> sizeExpr + <*> peekExpr + <*> pokeExpr + where + cons :: [(Name, [(Name, Type)])] + cons = + [ ( dcName dc + , [ (mkName ("c" ++ show ixc ++ "f" ++ show ixf), ty) + | ixf <- ints + | (_, ty) <- dcFields dc + ] + ) + | ixc <- ints + | dc <- cons0 + ] + -- NOTE: tag code duplicated in th-storable. + (tagType, _, tagSize) = + fromMaybe (error "Too many constructors") $ + find (\(_, maxN, _) -> maxN >= length cons) tagTypes + tagTypes :: [(Name, Int, Int)] + tagTypes = + [ ('(), 1, 0) + , (''Word8, fromIntegral (maxBound :: Word8), 1) + , (''Word16, fromIntegral (maxBound :: Word16), 2) + , (''Word32, fromIntegral (maxBound :: Word32), 4) + , (''Word64, fromIntegral (maxBound :: Word64), 8) + ] + fName ix = mkName ("f" ++ show ix) + ints = [0..] :: [Int] + fNames = map fName ints + sizeNames = zipWith (\_ -> mkName . ("sz" ++) . show) cons ints + tagName = mkName "tag" + valName = mkName "val" + sizeExpr = + caseE (tupE (concatMap (map sizeAtType . snd) cons)) + (if null sizeNames then [matchConstSize] else [matchConstSize, matchVarSize]) + where + sizeAtType :: (Name, Type) -> ExpQ + sizeAtType (_, ty) = [| size :: Size $(return ty) |] + matchConstSize :: MatchQ + matchConstSize = do + let sz0 = VarE (mkName "sz0") + sizeDecls = + if null sizeNames + then [valD (varP (mkName "sz0")) (normalB [| 0 |]) []] + else zipWith constSizeDec sizeNames cons + sameSizeExpr <- + case sizeNames of + (_ : tailSizeNames) -> + foldl (\l r -> [| $(l) && $(r) |]) [| True |] $ + map (\szn -> [| $(return sz0) == $(varE szn) |]) tailSizeNames + [] -> [| True |] + result <- [| ConstSize (tagSize + $(return sz0)) |] + match (tupP (map (\(n, _) -> conP 'ConstSize [varP n]) + (concatMap snd cons))) + (guardedB [return (NormalG sameSizeExpr, result)]) + sizeDecls + constSizeDec :: Name -> (Name, [(Name, Type)]) -> DecQ + constSizeDec szn (_, []) = + valD (varP szn) (normalB [| 0 |]) [] + constSizeDec szn (_, fields) = + valD (varP szn) body [] + where + body = normalB $ + foldl1 (\l r -> [| $(l) + $(r) |]) $ + map (\(sizeName, _) -> varE sizeName) fields + matchVarSize :: MatchQ + matchVarSize = do + match (tupP (map (\(n, _) -> varP n) (concatMap snd cons))) + (normalB [| VarSize $ \x -> tagSize + + $(caseE [| x |] (map matchVar cons)) |]) + [] + matchVar :: (Name, [(Name, Type)]) -> MatchQ + matchVar (cname, []) = + match (conP cname []) (normalB [| 0 |]) [] + matchVar (cname, fields) = + match (conP cname (zipWith (\_ fn -> varP fn) fields fNames)) + body + [] + where + body = normalB $ + foldl1 (\l r -> [| $(l) + $(r) |]) + (zipWith (\(sizeName, _) fn -> [| getSizeWith $(varE sizeName) $(varE fn) |]) + fields + fNames) + -- Choose a tag size large enough for this constructor count. + -- Expression used for the definition of peek. + peekExpr = case cons of + [] -> [| error ("Attempting to peek type with no constructors (" ++ $(lift (show headTy)) ++ ")") |] + [con] -> peekCon con + _ -> doE + [ bindS (varP tagName) [| peek |] + , noBindS (caseE (sigE (varE tagName) (conT tagType)) + (map peekMatch (zip [0..] cons) ++ [peekErr])) + ] + peekMatch (ix, con) = match (litP (IntegerL ix)) (normalB (peekCon con)) [] + peekErr = match wildP (normalB + [| peekException $ T.pack $ "Found invalid tag while peeking (" ++ $(lift (show headTy)) ++ ")" |]) [] + peekCon (cname, fields) = + case fields of + [] -> [| pure $(conE cname) |] + _ -> doE $ + map (\(fn, _) -> bindS (varP fn) [| peek |]) fields ++ + [noBindS $ appE (varE 'return) $ appsE $ conE cname : map (\(fn, _) -> varE fn) fields] + pokeExpr = lamE [varP valName] $ caseE (varE valName) $ zipWith pokeCon [0..] cons + pokeCon :: Int -> (Name, [(Name, Type)]) -> MatchQ + pokeCon ix (cname, fields) = + match (conP cname (map (\(fn, _) -> varP fn) fields)) body [] + where + body = normalB $ + case cons of + (_:_:_) -> doE (pokeTag ix : map pokeField fields) + _ -> doE (map pokeField fields) + pokeTag ix = noBindS [| poke (ix :: $(conT tagType)) |] + pokeField (fn, _) = noBindS [| poke $(varE fn) |] + +-- FIXME: make this work even when there are too many fields + +-- FIXME: make the ConstSize stuff explicit in the API. Have an option +-- that always errors at runtime if it isn't ConstSize? + +-- TODO: It would be really awesome, though a bit tricky, to know at +-- compile time if we have a static size. + +-- TODO: make sure that this tends to optimize even with tons of fields. +-- It should also optimize when some fields are known to be , but others +-- are unknown (determined by polymorphic var) + +{- What the generated code looks like + +data Foo = Foo Int Double Float + +instance Store Foo where + size = + case (size :: Size Int, size :: Size Double, size :: Size Float) of + (ConstSize c0f0, ConstSize c0f1, ConstSize c0f2) -> ConstSize (0 + sz0) + where + sz0 = c0f0 + c0f1 + c0f2 + (c0f0, c0f1, c0f2) + VarSize $ \(Foo f0 f1 f2) -> 0 + + getSizeWith c0f0 f0 + getSizeWith c0f1 f1 + getSizeWith c0f2 f2 + peek = do + f0 <- peek + f1 <- peek + f2 <- peek + return (Foo f0 f1 f2) + poke (Foo f0 f1 f2) = do + poke f0 + poke f1 + poke f2 + +data Bar = Bar Int | Baz Double + +instance Store Bar where + size = + case (size :: Size Int, size :: Size Double) of + (ConstSize c0f0, ConstSize c1f0) | sz0 == sz1 -> ConstSize (1 + sz0) + where + sz0 = c0f0 + sz1 = c1f0 + (c0f0, c1f0) -> VarSize $ \x -> 1 + + case x of + Bar f0 -> getSizeWith c0f0 f0 + Baz f0 -> getSizeWith c1f0 f0 + peek = do + tag <- peek + case (tag :: Word8) of + 0 -> do + f0 <- peek + return (Bar f0) + 1 -> do + f0 <- peek + return (Baz f0) + _ -> peekException "Found invalid tag while peeking (Bar)" + poke (Bar f0) = do + poke 0 + poke f0 + poke (Bar f0) = do + poke 1 + poke f0 +-} + +------------------------------------------------------------------------ +-- Generic + +deriveTupleStoreInstance :: Int -> Dec +deriveTupleStoreInstance n = + deriveGenericInstance (map storeCxt tvs) + (foldl1 AppT (TupleT n : tvs)) + where + tvs = take n (map (VarT . mkName . (:[])) ['a'..'z']) + storeCxt ty = +#if MIN_VERSION_template_haskell(2,10,0) + AppT (ConT ''Store) ty +#else + ClassP ''Store [ty] +#endif + +deriveGenericInstance :: Cxt -> Type -> Dec +deriveGenericInstance cs ty = InstanceD cs (AppT (ConT ''Store) ty) [] + +------------------------------------------------------------------------ +-- Storable + +-- TODO: Generate inline pragmas? Probably not necessary + +deriveManyStoreFromStorable :: (Type -> Bool) -> Q [Dec] +deriveManyStoreFromStorable p = do + storables <- postprocess . instancesMap <$> getInstances ''Storable + stores <- postprocess . instancesMap <$> getInstances ''Store + return $ M.elems $ flip M.mapMaybe (storables `M.difference` stores) $ + \(TypeclassInstance cs ty _) -> + let argTy = head (tail (unAppsT ty)) in + if p argTy + then Just $ makeStoreInstance (cs ++ everythingTypeable ty) + argTy + (VarE 'sizeStorable) + (VarE 'peekStorable) + (VarE 'pokeStorable) + else Nothing + +-- Necessitated by the Typeable constraint on peekStorable +everythingTypeable :: Type -> Cxt +everythingTypeable = +#if MIN_VERSION_template_haskell(2,10,0) + map (AppT (ConT ''Typeable) . VarT) . freeVarsT +#else + map (ClassP ''Typeable . (:[]) . VarT) . freeVarsT +#endif + +------------------------------------------------------------------------ +-- Vector + +deriveManyStorePrimVector :: Q [Dec] +deriveManyStorePrimVector = do + prims <- postprocess . instancesMap <$> getInstances ''PV.Prim + stores <- postprocess . instancesMap <$> getInstances ''Store + let primInsts = + M.mapKeys (map (AppT (ConT ''PV.Vector))) prims + `M.difference` + stores + forM (M.toList primInsts) $ \primInst -> case primInst of + ([instTy], TypeclassInstance cs ty _) -> do + let argTy = head (tail (unAppsT ty)) + sizeExpr <- [| + VarSize $ \x -> + I# $(primSizeOfExpr (ConT ''Int)) + + I# $(primSizeOfExpr argTy) * PV.length x + |] + peekExpr <- [| do + len <- peek + let sz = I# $(primSizeOfExpr argTy) + array <- peekToByteArray $(lift ("Primitive Vector (" ++ pprint argTy ++ ")")) + (len * sz) + return (PV.Vector 0 len array) + |] + pokeExpr <- [| \(PV.Vector offset len (ByteArray array)) -> do + let sz = I# $(primSizeOfExpr argTy) + poke len + pokeFromByteArray array (offset * sz) (len * sz) + |] + return $ makeStoreInstance cs instTy sizeExpr peekExpr pokeExpr + _ -> fail "Invariant violated in derivemanyStorePrimVector" + + +primSizeOfExpr :: Type -> ExpQ +primSizeOfExpr ty = [| $(varE 'sizeOf#) (error "sizeOf# evaluated its argument" :: $(return ty)) |] + +deriveManyStoreUnboxVector :: Q [Dec] +deriveManyStoreUnboxVector = do + unboxes <- getUnboxInfo + stores <- postprocess . instancesMap <$> getInstances ''Store + let unboxInsts = + M.fromList (map (\(preds, ty, cons) -> ([AppT (ConT ''UV.Vector) ty], (preds, cons))) unboxes) + `M.difference` + stores + -- TODO: ideally this would use a variant of 'deriveStore' which + -- assumes VarSize. + forM (M.toList unboxInsts) $ \case + ([ty], (preds, cons)) -> do + {- + -- While this approach is reasonable-ish, it ends up + -- requiring UndecidableInstances. + let extraPreds = + map (AppT (ConT ''Store)) $ + filter (not . isMonoType) $ + concatMap (map snd . dcFields) cons + -} + let extraPreds = + map (AppT (ConT ''Store) . AppT (ConT ''UV.Vector)) $ listify isVarT ty + deriveStore (nub (preds ++ extraPreds)) ty cons + _ -> fail "impossible case in deriveManyStoreUnboxVector" + +-- TODO: Add something for this purpose to TH.ReifyDataType + +getUnboxInfo :: Q [(Cxt, Type, [DataCon])] +getUnboxInfo = do + FamilyI _ insts <- reify ''UV.Vector + return (map (everywhere (id `extT` dequalVarT) . go) insts) + where + go (NewtypeInstD preds _ [ty] con _) = + (preds, ty, conToDataCons con) + go (DataInstD preds _ [ty] cons _) = + (preds, ty, concatMap conToDataCons cons) + go x = error ("Unexpected result from reifying Unboxed Vector instances: " ++ pprint x) + dequalVarT (VarT n) = VarT (dequalify n) + dequalVarT ty = ty + +------------------------------------------------------------------------ +-- Utilities + +-- Filters out overlapping instances and instances with more than one +-- type arg (should be impossible). +postprocess :: M.Map [Type] [a] -> M.Map [Type] a +postprocess = + M.mapMaybeWithKey $ \tys xs -> + case (tys, xs) of + ([_ty], [x]) -> Just x + _ -> Nothing + +makeStoreInstance :: Cxt -> Type -> Exp -> Exp -> Exp -> Dec +makeStoreInstance cs ty sizeExpr peekExpr pokeExpr = + InstanceD cs + (AppT (ConT ''Store) ty) + [ ValD (VarP 'size) (NormalB sizeExpr) [] + , PragmaD (InlineP 'size Inline FunLike AllPhases) + , ValD (VarP 'peek) (NormalB peekExpr) [] + , PragmaD (InlineP 'peek Inline FunLike AllPhases) + , ValD (VarP 'poke) (NormalB pokeExpr) [] + , PragmaD (InlineP 'poke Inline FunLike AllPhases) + ] + +-- TODO: either generate random types that satisfy instances with +-- variables in them, or have a check that there's at least a manual +-- check for polymorphic instances. + +getAllInstanceTypes :: Name -> Q [[Type]] +getAllInstanceTypes n = + map (\(TypeclassInstance _ ty _) -> drop 1 (unAppsT ty)) <$> + getInstances n + +getAllInstanceTypes1 :: Name -> Q [Type] +getAllInstanceTypes1 n = + fmap (fmap (fromMaybe (error "getAllMonoInstances1 expected only one type argument") . headMay)) + (getAllInstanceTypes n) + +isMonoType :: Type -> Bool +isMonoType = null . listify isVarT + +isVarT :: Type -> Bool +isVarT VarT{} = True +isVarT _ = False + +-- TOOD: move these to th-reify-many + +-- | Get a map from the 'getTyHead' type of instances to +-- 'TypeclassInstance'. +instancesMap :: [TypeclassInstance] -> M.Map [Type] [TypeclassInstance] +instancesMap = + M.fromListWith (++) . + map (\ti -> (map getTyHead (instanceArgTypes ti), [ti])) + +instanceArgTypes :: TypeclassInstance -> [Type] +instanceArgTypes (TypeclassInstance _ ty _) = drop 1 (unAppsT ty) + +getTyHead :: Type -> Type +getTyHead (SigT x _) = getTyHead x +getTyHead (ForallT _ _ x) = getTyHead x +getTyHead (AppT l _) = getTyHead l +getTyHead x = x diff --git a/src/Data/Store/TypeHash.hs b/src/Data/Store/TypeHash.hs new file mode 100644 --- /dev/null +++ b/src/Data/Store/TypeHash.hs @@ -0,0 +1,17 @@ +-- This module provides utilities for computing hashes based on the +-- structural definitions of datatypes. The purpose of this is to +-- provide a mechanism for tagging serialized data in such a way that +-- deserialization issues can be anticipated. +-- +-- This portion of the store package is still under development and will +-- likely change in interface and behavior. +module Data.Store.TypeHash + ( Tagged(..) + , TypeHash + , HasTypeHash(..) + -- * TH for generating HasTypeHash instances + , mkHasTypeHash + , mkManyHasTypeHash + ) where + +import Data.Store.TypeHash.Internal diff --git a/src/Data/Store/TypeHash/Internal.hs b/src/Data/Store/TypeHash/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Data/Store/TypeHash/Internal.hs @@ -0,0 +1,127 @@ +{-# LANGUAGE CPP #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE StandaloneDeriving #-} + +module Data.Store.TypeHash.Internal where + +import Control.Applicative +import Control.DeepSeq (NFData) +import Control.Monad (when, unless) +import qualified Crypto.Hash.SHA1 as SHA1 +import qualified Data.ByteString as BS +import Data.Char (isUpper, isLower) +import Data.Data (Data) +import Data.Generics (listify) +import Data.List (sortBy) +import Data.Monoid ((<>)) +import Data.Ord (comparing) +import Data.Proxy (Proxy(..)) +import Data.Store +import Data.Store.Internal +import Data.Typeable (Typeable) +import GHC.Generics (Generic) +import Instances.TH.Lift () +import Language.Haskell.TH +import Language.Haskell.TH.ReifyMany (reifyMany) +import Language.Haskell.TH.Syntax (Lift(lift)) +import Prelude + +newtype Tagged a = Tagged { unTagged :: a } + deriving (Eq, Ord, Show, Data, Typeable, Generic) + +instance NFData a => NFData (Tagged a) + +instance (Store a, HasTypeHash a) => Store (Tagged a) where + size = addSize 20 (contramapSize unTagged size) + peek = do + tag <- peek + let expected = typeHash (Proxy :: Proxy a) + when (tag /= expected) $ fail "Mismatched type hash" + Tagged <$> peek + poke (Tagged x) = do + poke (typeHash (Proxy :: Proxy a)) + poke x + +newtype TypeHash = TypeHash { unTypeHash :: StaticSize 20 BS.ByteString } + deriving (Eq, Ord, Show, Store, Generic) + +#if __GLASGOW_HASKELL__ >= 710 +deriving instance Typeable TypeHash +deriving instance Data TypeHash +#endif + +instance NFData TypeHash + +instance Lift TypeHash where + lift = liftStaticSize [t| BS.ByteString |] . unTypeHash + +-- TODO: move into th-reify-many +reifyManyTyDecls :: ((Name, Info) -> Q (Bool, [Name])) + -> [Name] + -> Q [(Name, Info)] +reifyManyTyDecls f = reifyMany go + where + go x@(_, TyConI{}) = f x + go x@(_, FamilyI{}) = f x + go x@(_, PrimTyConI{}) = f x + go x@(_, DataConI{}) = f x + go (_, ClassI{}) = return (False, []) + go (_, ClassOpI{}) = return (False, []) + go (_, VarI{}) = return (False, []) + go (_, TyVarI{}) = return (False, []) + +-- | At compiletime, this yields a hash of the specified datatypes. +-- Not only does this cover the datatypes themselves, but also all +-- transitive dependencies. +-- +-- The resulting expression is a literal of type 'Int'. +typeHashForNames :: [Name] -> Q Exp +typeHashForNames ns = do + infos <- getTypeInfosRecursively ns + [| TypeHash (BS.pack $(lift (BS.unpack (SHA1.hash (encode infos))))) |] + +-- | At compiletime, this yields a cryptographic hash of the specified 'Type', +-- including the definition of things it references (transitively). +-- +-- The resulting expression is a literal of type 'Int'. +hashOfType :: Type -> Q Exp +hashOfType ty = do + unless (null (getVarNames ty)) $ fail $ "hashOfType cannot handle polymorphic type " <> pprint ty + infos <- getTypeInfosRecursively (getConNames ty) + lift $ TypeHash $ toStaticSizeEx $ SHA1.hash $ encode (ty, infos) + +getTypeInfosRecursively :: [Name] -> Q [(Name, Info)] +getTypeInfosRecursively names = do + allInfos <- reifyManyTyDecls (\(_, info) -> return (True, getConNames info)) names + -- Sorting step probably unnecessary because this should be + -- deterministic, but hey why not. + return (sortBy (comparing fst) allInfos) + +getConNames :: Data a => a -> [Name] +getConNames = listify (isUpper . head . nameBase) + +getVarNames :: Data a => a -> [Name] +getVarNames = listify (isLower . head . nameBase) + +-- TODO: Generic instance for polymorphic types, or have TH generate +-- polymorphic instances. + +class HasTypeHash a where + typeHash :: Proxy a -> TypeHash + +mkHasTypeHash :: Type -> Q [Dec] +mkHasTypeHash ty = + [d| instance HasTypeHash $(return ty) where + typeHash _ = TypeHash $(hashOfType ty) + |] + +mkManyHasTypeHash :: [Q Type] -> Q [Dec] +mkManyHasTypeHash qtys = concat <$> mapM (mkHasTypeHash =<<) qtys + +combineTypeHashes :: [TypeHash] -> TypeHash +combineTypeHashes = TypeHash . toStaticSizeEx . SHA1.hash . BS.concat . map (unStaticSize . unTypeHash) diff --git a/src/System/IO/ByteBuffer.hs b/src/System/IO/ByteBuffer.hs new file mode 100644 --- /dev/null +++ b/src/System/IO/ByteBuffer.hs @@ -0,0 +1,249 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE FlexibleContexts #-} +{-| +Module: System.IO.ByteBuffer +Description: Provides an efficient buffering abstraction. + +A 'ByteBuffer' is a simple buffer for bytes. It supports two +operations: refilling with the contents of a 'ByteString', and +consuming a fixed number of bytes. + +It is implemented as a pointer, together with counters that keep track +of the offset and the number of bytes in the buffer. Note that the +counters are simple 'IORef's, so 'ByteBuffer's are not thread-safe! + +A 'ByteBuffer' is constructed by 'new' with a given starting length, +and will grow (by repeatedly multiplying its size by 1.5) whenever it +is being fed a 'ByteString' that is too large. +-} + +module System.IO.ByteBuffer + ( ByteBuffer + -- * Allocation and Deallocation + , new, free, with + -- * Query for number of available bytes + , totalSize, isEmpty, availableBytes + -- * Feeding new input + , copyByteString + -- * Consuming bytes from the buffer + , consume, unsafeConsume + ) where + +import Control.Applicative +import Control.Exception.Lifted (bracket) +import Control.Monad (when) +import Control.Monad.IO.Class +import Control.Monad.Trans.Control (MonadBaseControl) +import Data.ByteString (ByteString) +import qualified Data.ByteString.Internal as BS +import Data.IORef +import Data.Maybe (fromMaybe) +import Data.Word +import Foreign.ForeignPtr +import qualified Foreign.Marshal.Alloc as Alloc +import Foreign.Marshal.Utils hiding (new, with) +import GHC.Ptr +import Prelude + +-- | A buffer into which bytes can be written. +-- +-- Invariants: +-- +-- * @size >= containedBytes >= consumedBytes >= 0@ +-- +-- * The range from @ptr@ to @ptr `plusPtr` size@ will be allocated +-- +-- * The range from @ptr@ to @ptr `plusPtr` containedBytes@ will +-- contain bytes previously copied to the buffer +-- +-- * The buffer contains @containedBytes - consumedBytes@ bytes of +-- data that have been copied to it, but not yet read. They are in +-- the range from @ptr `plusPtr` consumedBytes@ to @ptr `plusPtr` +-- containedBytes@. + +data BBRef = BBRef { + size :: {-# UNPACK #-} !Int + -- ^ The amount of memory allocated. + , contained :: {-# UNPACK #-} !Int + -- ^ The number of bytes that the 'ByteBuffer' currently holds. + , consumed :: {-# UNPACK #-} !Int + -- ^ The number of bytes that have already been consumed. + , ptr :: {-# UNPACK #-} !(Ptr Word8) + -- ^ This points to the beginning of the memory allocated for + -- the 'ByteBuffer' + } + +type ByteBuffer = IORef BBRef + +totalSize :: MonadIO m => ByteBuffer -> m Int +totalSize bb = liftIO $ size <$> readIORef bb +{-# INLINE totalSize #-} + +isEmpty :: MonadIO m => ByteBuffer -> m Bool +isEmpty bb = liftIO $ (==0) <$> availableBytes bb +{-# INLINE isEmpty #-} + +-- | Number of available bytes in a 'ByteBuffer' (that is, bytes that +-- have been copied to, but not yet read from the 'ByteBuffer'. +availableBytes :: MonadIO m => ByteBuffer -> m Int +availableBytes bb = do + BBRef{..} <- liftIO $ readIORef bb + return $ contained - consumed +{-# INLINE availableBytes #-} + +-- | The number of bytes that can be appended to the buffer, without +-- resetting it. +freeCapacity :: MonadIO m => ByteBuffer -> m Int +freeCapacity bb = do + BBRef{..} <- liftIO $ readIORef bb + return $ size - contained +{-# INLINE freeCapacity #-} + +-- | Allocates a new ByteBuffer with a given buffer size filling from +-- the given FillBuffer. +-- +-- Note that 'ByteBuffer's created with 'new' have to be deallocated +-- explicitly using 'free'. For automatic deallocation, consider +-- using 'with' instead. +new :: MonadIO m + => Maybe Int + -- ^ Size of buffer to allocate. If 'Nothing', use the default + -- value of 4MB + -> m ByteBuffer + -- ^ The byte buffer. +new ml = liftIO $ do + let l = fromMaybe (4*1024*1024) ml + newPtr <- Alloc.mallocBytes l + newIORef BBRef { ptr = newPtr + , size = l + , contained = 0 + , consumed = 0 + } +{-# INLINE new #-} + +-- | Free a byte buffer. +free :: MonadIO m => ByteBuffer -> m () +free bb = liftIO $ readIORef bb >>= Alloc.free . ptr +{-# INLINE free #-} + +-- | Perform some action with a bytebuffer, with automatic allocation +-- and deallocation. +with :: (MonadIO m, MonadBaseControl IO m) + => Maybe Int + -- ^ Initial length of the 'ByteBuffer'. If 'Nothing', use the + -- default value of 4MB. + -> (ByteBuffer -> m a) + -> m a +with l action = + bracket + (new l) + free + action +{-# INLINE with #-} + +-- | Reset a 'ByteBuffer', i.e. copy all the bytes that have not yet +-- been consumed to the front of the buffer. +reset :: ByteBuffer -> IO () +reset bb = do + bbref@BBRef{..} <- readIORef bb + let available = contained - consumed + moveBytes ptr (ptr `plusPtr` consumed) available + writeIORef bb bbref { contained = available + , consumed = 0 + } +{-# INLINE reset #-} + +-- | Make sure the buffer is at least @minSize@ bytes long. +-- +-- In order to avoid havong to enlarge the buffer too often, we double +-- its size until it is at least @minSize@ bytes long. +enlargeByteBuffer :: ByteBuffer + -> Int + -- ^ minSize + -> IO () +enlargeByteBuffer bb minSize = do + bbref@BBRef{..} <- readIORef bb + when (size < minSize) $ do + let newSize = head . dropWhile ( ByteBuffer -> ByteString -> m () +copyByteString bb bs@(BS.PS _ _ bsSize) = liftIO $ do + BBRef{..} <- readIORef bb + -- if the byteBuffer is too small, resize it. + available <- availableBytes bb -- bytes not yet consumed + when (size < bsSize + available) (enlargeByteBuffer bb (bsSize + available)) + -- if it is currently too full, reset it + capacity <- freeCapacity bb + when (capacity < bsSize) (reset bb) + -- now we can safely copy. + unsafeCopyByteString bb bs +{-# INLINE copyByteString #-} + +-- | Copy the contents of a 'ByteString' to a 'ByteBuffer'. No bounds +-- checks are performed. +unsafeCopyByteString :: ByteBuffer -> ByteString -> IO () +unsafeCopyByteString bb (BS.PS bsFptr bsOffset bsSize) = do + bbref@BBRef{..} <- readIORef bb + withForeignPtr bsFptr $ \ bsPtr -> + copyBytes (ptr `plusPtr` contained) + (bsPtr `plusPtr` bsOffset) + bsSize + writeIORef bb bbref { contained = (contained + bsSize) } +{-# INLINE unsafeCopyByteString #-} + +-- | Try to get a pointer to @n@ bytes from the 'ByteBuffer'. +-- +-- Note that the pointer should be used before any other actions are +-- performed on the 'ByteBuffer'. It points to some address within the +-- buffer, so operations such as enlarging the buffer or feeding it +-- new data will change the data the pointer points to. This is why +-- this function is called unsafe. +unsafeConsume :: MonadIO m + => ByteBuffer + -> Int + -- ^ n + -> m (Either Int (Ptr Word8)) + -- ^ Will be @Left missing@ when there are only @n-missing@ + -- bytes left in the 'ByteBuffer'. +unsafeConsume bb n = liftIO $ do + available <- availableBytes bb + if available < n + then return $ Left (n - available) + else do + bbref@BBRef{..} <- readIORef bb + writeIORef bb bbref { consumed = (consumed + n) } + return $ Right (ptr `plusPtr` consumed) +{-# INLINE unsafeConsume #-} + +-- | As `unsafeConsume`, but instead of returning a `Ptr` into the +-- contents of the `ByteBuffer`, it returns a `ByteString` containing +-- the next @n@ bytes in the buffer. This involves allocating a new +-- 'ByteString' and copying the @n@ bytes to it. +consume :: MonadIO m + => ByteBuffer + -> Int + -> m (Either Int ByteString) +consume bb n = do + mPtr <- unsafeConsume bb n + case mPtr of + Right ptr -> do + bs <- liftIO . BS.create n $ \ bsPtr -> + copyBytes bsPtr ptr n + return (Right bs) + Left missing -> return (Left missing) +{-# INLINE consume #-} diff --git a/store.cabal b/store.cabal new file mode 100644 --- /dev/null +++ b/store.cabal @@ -0,0 +1,193 @@ +-- This file has been generated from package.yaml by hpack version 0.14.0. +-- +-- see: https://github.com/sol/hpack + +name: store +version: 0.1.0.0 +synopsis: Fast binary serialization +homepage: https://github.com/fpco/store#readme +bug-reports: https://github.com/fpco/store/issues +license: MIT +license-file: LICENSE +maintainer: Michael Sloan +copyright: 2016 FP Complete +category: Serialization, Data +build-type: Simple +cabal-version: >= 1.10 + +extra-source-files: + ChangeLog.md + README.md + +source-repository head + type: git + location: https://github.com/fpco/store + +flag comparison-bench + manual: True + default: False + +flag small-bench + default: False + manual: True + +library + hs-source-dirs: + src + exposed-modules: + Data.Store + Data.Store.Internal + Data.Store.Streaming + Data.Store.TH + Data.Store.TH.Internal + Data.Store.TypeHash + Data.Store.TypeHash.Internal + System.IO.ByteBuffer + other-modules: + Data.Store.Impl + build-depends: + base >= 4.7 && < 5 + , array + , base-orphans + , bytestring + , containers + , cryptohash + , deepseq + , fail + , ghc-prim + , hashable + , hspec + , hspec-smallcheck + , integer-gmp + , mono-traversable + , primitive + , safe + , smallcheck + , syb + , template-haskell + , text + , th-lift + , th-lift-instances >= 0.1.6 + , th-utilities + , th-reify-many + , time + , transformers + , unordered-containers + , vector + , conduit + , lifted-base + , monad-control + , resourcet + , semigroups + , void + , th-orphans + default-language: Haskell2010 + ghc-options: -Wall -fwarn-tabs -fwarn-incomplete-uni-patterns -fwarn-incomplete-record-updates -O2 + +test-suite store-test + type: exitcode-stdio-1.0 + hs-source-dirs: + test + main-is: Spec.hs + other-modules: + Data.Store.StreamingSpec + Data.StoreSpec + Data.StoreSpec.TH + System.IO.ByteBufferSpec + build-depends: + base >= 4.7 && < 5 + , array + , base-orphans + , bytestring + , containers + , cryptohash + , deepseq + , fail + , ghc-prim + , hashable + , hspec + , hspec-smallcheck + , integer-gmp + , mono-traversable + , primitive + , safe + , smallcheck + , syb + , template-haskell + , text + , th-lift + , th-lift-instances >= 0.1.6 + , th-utilities + , th-reify-many + , time + , transformers + , unordered-containers + , vector + , conduit + , lifted-base + , monad-control + , resourcet + , semigroups + , void + , th-orphans + , hspec + , smallcheck + , hspec-smallcheck + , store + ghc-options: -Wall -fwarn-tabs -fwarn-incomplete-uni-patterns -fwarn-incomplete-record-updates -O2 -threaded -rtsopts -with-rtsopts=-N + default-language: Haskell2010 + +benchmark store-bench + type: exitcode-stdio-1.0 + hs-source-dirs: + bench + ghc-options: -Wall -fwarn-tabs -fwarn-incomplete-uni-patterns -fwarn-incomplete-record-updates -O2 -threaded -rtsopts -with-rtsopts=-N1 -with-rtsopts=-s -with-rtsopts=-qg + main-is: Bench.hs + build-depends: + base >= 4.7 && < 5 + , array + , base-orphans + , bytestring + , containers + , cryptohash + , deepseq + , fail + , ghc-prim + , hashable + , hspec + , hspec-smallcheck + , integer-gmp + , mono-traversable + , primitive + , safe + , smallcheck + , syb + , template-haskell + , text + , th-lift + , th-lift-instances >= 0.1.6 + , th-utilities + , th-reify-many + , time + , transformers + , unordered-containers + , vector + , conduit + , lifted-base + , monad-control + , resourcet + , semigroups + , void + , th-orphans + , criterion + , store + if flag(comparison-bench) + cpp-options: -DCOMPARISON_BENCH + build-depends: + cereal + , binary + , vector-binary-instances + , cereal-vector + if flag(small-bench) + cpp-options: -DSMALL_BENCH + default-language: Haskell2010 diff --git a/test/Data/Store/StreamingSpec.hs b/test/Data/Store/StreamingSpec.hs new file mode 100644 --- /dev/null +++ b/test/Data/Store/StreamingSpec.hs @@ -0,0 +1,121 @@ +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE OverloadedStrings #-} +module Data.Store.StreamingSpec where + +import Control.Monad.Trans.Resource +import Control.Exception (try) +import qualified Data.ByteString as BS +import Data.Conduit ((=$=), ($$)) +import qualified Data.Conduit.List as C +import Data.List (unfoldr) +import Data.Monoid +import Data.Store (PeekException (..)) +import Data.Store.Streaming +import qualified System.IO.ByteBuffer as BB +import Test.Hspec +import Test.Hspec.SmallCheck +import Test.SmallCheck + +spec :: Spec +spec = do + describe "conduitEncode and conduitDecode" $ do + it "Roundtrips ([Int])." $ property roundtrip + it "Roundtrips ([Int]), with chunked transfer." $ property roundtripChunked + it "Throws an Exception on incomplete messages." conduitIncomplete + it "Throws an Exception on excess input." $ property conduitExcess + describe "peekMessage" $ do + it "demands more input when needed." $ property (askMore 8) + it "demands more input on incomplete SizeTag." $ property (askMore 1) + it "successfully decodes valid input." $ property peek + describe "decodeMessage" $ + it "Throws an Exception on incomplete messages." decodeIncomplete + +roundtrip :: [Int] -> Property IO +roundtrip xs = monadic $ do + xs' <- runResourceT $ C.sourceList xs + =$= C.map Message + =$= conduitEncode + =$= conduitDecode Nothing + =$= C.map fromMessage + $$ C.consume + return $ xs' == xs + +roundtripChunked :: [Int] -> Property IO +roundtripChunked input = monadic $ do + let (xs, chunkLengths) = splitAt (length input `div` 2) input + bs <- C.sourceList xs + =$= C.map Message + =$= conduitEncode + $$ C.fold (<>) mempty + let chunks = unfoldr takeChunk (bs, chunkLengths) + xs' <- runResourceT $ C.sourceList chunks + =$= conduitDecode (Just 10) + =$= C.map fromMessage + $$ C.consume + return $ xs' == xs + where + takeChunk (x, _) | BS.null x = Nothing + takeChunk (x, []) = Just (x, (BS.empty, [])) + takeChunk (x, l:ls) = + let (chunk, rest) = BS.splitAt l x + in Just (chunk, (rest, ls)) + +conduitIncomplete :: Expectation +conduitIncomplete = + (runResourceT (C.sourceList [incompleteInput] + =$= conduitDecode (Just 10) + $$ C.consume) + :: IO [Message Integer]) `shouldThrow` peekException + +conduitExcess :: [Int] -> Property IO +conduitExcess xs = monadic $ do + bs <- C.sourceList xs + =$= C.map Message + =$= conduitEncode + $$ C.fold (<>) mempty + res <- try (runResourceT (C.sourceList [bs `BS.append` "excess bytes"] + =$= conduitDecode (Just 10) + $$ C.consume) :: IO [Message Int]) + case res of + Right _ -> return False + Left (PeekException _ _) -> return True + +-- splits an encoded message after n bytes. Feeds the first part to +-- peekResult, expecting it to require more input. Then, feeds the +-- second part and checks if the decoded result is the original +-- message. +askMore :: Int -> Integer -> Property IO +askMore n x = monadic $ BB.with (Just 10) $ \ bb -> do + let bs = encodeMessage (Message x) + (start, end) = BS.splitAt n $ bs + BB.copyByteString bb start + peekResult <- peekMessage bb :: IO (PeekMessage IO Integer) + case peekResult of + NeedMoreInput cont -> + cont end >>= \case + Done (Message x') -> return $ x' == x + _ -> return False + _ -> return False + +peek :: Integer -> Property IO +peek x = monadic $ BB.with (Just 10) $ \ bb -> do + let bs = encodeMessage (Message x) + BB.copyByteString bb bs + peekResult <- peekMessage bb :: IO (PeekMessage IO Integer) + case peekResult of + NeedMoreInput _ -> return False + Done (Message x') -> return $ x' == x + +decodeIncomplete :: IO () +decodeIncomplete = BB.with (Just 0) $ \ bb -> do + BB.copyByteString bb (BS.take 1 incompleteInput) + (decodeMessage bb (return Nothing) :: IO (Maybe (Message Integer))) + `shouldThrow` peekException + +incompleteInput :: BS.ByteString +incompleteInput = + let bs = encodeMessage (Message (42 :: Integer)) + in BS.take (BS.length bs - 1) bs + +peekException :: Selector PeekException +peekException = const True diff --git a/test/Data/StoreSpec.hs b/test/Data/StoreSpec.hs new file mode 100644 --- /dev/null +++ b/test/Data/StoreSpec.hs @@ -0,0 +1,304 @@ +{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE OverloadedStrings #-} +module Data.StoreSpec where + +import Control.Applicative +import Control.Monad (unless) +import qualified Data.ByteString as BS +import qualified Data.ByteString.Lazy as LBS +import qualified Data.ByteString.Short as SBS +import Data.Complex (Complex(..)) +import Data.Containers (mapFromList, setFromList) +import Data.HashMap.Strict (HashMap) +import Data.HashSet (HashSet) +import Data.Hashable (Hashable) +import Data.Int +import Data.IntMap (IntMap) +import Data.IntSet (IntSet) +import qualified Data.List.NonEmpty as NE +import Data.Map (Map) +import Data.Monoid +import Data.Primitive.Types (Addr) +import Data.Sequence (Seq) +import Data.Sequences (fromList) +import Data.Set (Set) +import Data.Store +import Data.Store.Internal +import Data.Store.TH +import Data.Store.TH.Internal +import Data.Store.TypeHash +import Data.StoreSpec.TH +import Data.Text (Text) +import qualified Data.Text as T +import qualified Data.Time as Time +import qualified Data.Vector as V +import qualified Data.Vector.Primitive as PV +import qualified Data.Vector.Storable as SV +import qualified Data.Vector.Unboxed as UV +import Data.Void (Void) +import Data.Word +import Foreign.C.Types +import Foreign.Ptr +import Foreign.Storable (Storable) +import GHC.Fingerprint.Type (Fingerprint(..)) +import GHC.Generics +import GHC.Real (Ratio(..)) +import Language.Haskell.TH +import Language.Haskell.TH.ReifyMany +import Language.Haskell.TH.Syntax +import Prelude +import System.Posix.Types +import Test.Hspec hiding (runIO) +import Test.SmallCheck.Series + +------------------------------------------------------------------------ +-- Instances for base types + +-- TODO: should be possible to do something clever where it only defines +-- instances that don't already exist. For now, just doing it manually. + +addMinAndMaxBounds :: forall a. (Bounded a, Eq a, Num a) => [a] -> [a] +addMinAndMaxBounds xs = + (if (minBound :: a) `notElem` xs then [minBound] else []) ++ + (if (maxBound :: a) `notElem` xs && (maxBound :: a) /= minBound then maxBound : xs else xs) + +$(mkManyHasTypeHash [ [t| Int32 |] ]) + +-- Serial instances for (Num a, Bounded a) types. Only really +-- appropriate for the use here. + +$(do let ns = [ ''CWchar, ''CUid, ''CUShort, ''CULong, ''CULLong, ''CIntMax + , ''CUIntMax, ''CPtrdiff, ''CSChar, ''CShort, ''CUInt, ''CLLong + , ''CLong, ''CInt, ''CChar, ''CTcflag, ''CSsize, ''CRLim, ''CPid + , ''COff, ''CNlink, ''CMode, ''CIno, ''CGid, ''CDev + , ''Word8, ''Word16, ''Word32, ''Word64, ''Word + , ''Int8, ''Int16, ''Int32, ''Int64 + ] + f n = [d| instance Monad m => Serial m $(conT n) where + series = generate (\_ -> addMinAndMaxBounds [0, 1]) |] + concat <$> mapM f ns) + +-- Serial instances for (Num a) types. Only really appropriate for the +-- use here. + +$(do let ns = [ ''CUSeconds, ''CClock, ''CTime, ''CUChar, ''CSize, ''CSigAtomic + , ''CSUSeconds, ''CFloat, ''CDouble, ''CSpeed, ''CCc + ] + f n = [d| instance Monad m => Serial m $(conT n) where + series = generate (\_ -> [0, 1]) |] + concat <$> mapM f ns) + +-- Serial instances for Primitive vectors + +$(do tys <- getAllInstanceTypes1 ''PV.Prim + let f ty = [d| instance (Serial m $(return ty), Monad m) => Serial m (PV.Vector $(return ty)) where + series = fmap PV.fromList series |] + concat <$> mapM f tys) + +-- Instance needed for generic TH instances + +-- Needs to be done manually because in GHC 7.8's TH, NameFlavour uses +-- unboxed values and cannot use generic deriving. So we skip having an +-- instance for it. +instance Monad m => Serial m Name where series = fmap mkName series + +-- Serial instances for (Generic a) types. + +-- FIXME: generating for TH instances is probably just adding +-- unnecessary compiletime + runtime overhead. +$(do thNames <- reifyManyWithoutInstances + ''Serial + [''Info, ''Loc, ''ModName, ''PkgName, ''NameSpace, ''OccName] + (`notElem` [''NameFlavour]) + let ns = [ ''Any, ''All ] ++ thNames + f n = [d| instance Monad m => Serial m $(conT n) |] + concat <$> mapM f ns) + +$(do let ns = [ ''Dual, ''Sum, ''Product, ''First, ''Last ] + f n = [d| instance (Monad m, Serial m a) => Serial m ($(conT n) a) |] + concat <$> mapM f ns) + +instance Monad m => Serial m Fingerprint where + series = generate (\_ -> [Fingerprint 0 0, Fingerprint maxBound maxBound]) + +instance Monad m => Serial m BS.ByteString where + series = fmap BS.pack series + +instance Monad m => Serial m LBS.ByteString where + series = fmap LBS.pack series + +instance Monad m => Serial m SBS.ShortByteString where + series = fmap SBS.pack series + +instance (Monad m, Serial m a, Storable a) => Serial m (SV.Vector a) where + series = fmap SV.fromList series + +instance (Monad m, Serial m a) => Serial m (V.Vector a) where + series = fmap V.fromList series + +instance (Monad m, Serial m k, Serial m a, Ord k) => Serial m (Map k a) where + series = fmap mapFromList series + +instance (Monad m, Serial m a, Ord a) => Serial m (Set a) where + series = fmap setFromList series + +instance (Monad m, Serial m a) => Serial m (IntMap a) where + series = fmap mapFromList series + +instance Monad m => Serial m IntSet where + series = fmap setFromList series + +instance Monad m => Serial m Text where + series = fmap fromList series + +instance (Monad m, Serial m a) => Serial m (Seq a) where + series = fmap fromList series + +instance (Monad m, Serial m a) => Serial m (Complex a) where + series = uncurry (:+) <$> (series >< series) + +instance (Monad m, Serial m a, UV.Unbox a) => Serial m (UV.Vector a) where + series = fmap fromList series + +instance (Monad m, Serial m k, Serial m a, Hashable k, Eq k) => Serial m (HashMap k a) where + series = fmap mapFromList series + +instance (Monad m, Serial m a, Hashable a, Eq a) => Serial m (HashSet a) where + series = fmap setFromList series + +instance Monad m => Serial m Time.Day where + series = Time.ModifiedJulianDay <$> series + +instance Monad m => Serial m Time.DiffTime where + series = Time.picosecondsToDiffTime <$> series + +instance Monad m => Serial m Time.UTCTime where + series = uncurry Time.UTCTime <$> (series >< series) + +instance (Monad m, Serial m a) => Serial m (NE.NonEmpty a) + +instance (Monad m, Serial m a) => Serial m (Tagged a) + +-- Should probably get added to smallcheck :) +instance Monad m => Serial m Void where + series = generate (\_ -> []) + +deriving instance Show NameFlavour +deriving instance Show NameSpace + +------------------------------------------------------------------------ +-- Test datatypes for generics support + +data Test + = TestA Int64 Word32 + | TestB Bool + | TestC + | TestD BS.ByteString + deriving (Eq, Show, Generic) +-- $(return . (:[]) =<< deriveStore [] (ConT ''Test) . dtCons =<< reifyDataType ''Test) +instance Store Test +instance Monad m => Serial m Test + +data X = X + deriving (Eq, Show, Generic) +instance Monad m => Serial m X +instance Store X + +spec :: Spec +spec = do + describe "Store on all monomorphic instances" + $(do insts <- getAllInstanceTypes1 ''Store + omitTys <- sequence + [ [t| PV.Vector Addr |] + , [t| CUIntPtr |] + , [t| CIntPtr |] + , [t| IntPtr |] + , [t| WordPtr |] + , [t| TypeHash |] + , [t| Fd |] + , [t| NameFlavour |] + ] + let f ty = isMonoType ty && ty `notElem` omitTys + smallcheckManyStore verbose 2 . map return . filter f $ insts) + describe "Store on all custom generic instances" + $(smallcheckManyStore verbose 2 + [ [t| Test |] + , [t| X |] + ]) + describe "Manually listed polymorphic store instances" + $(smallcheckManyStore verbose 2 + [ [t| SV.Vector Int8 |] + , [t| V.Vector Int8 |] + , [t| Ratio Int8 |] + , [t| Complex Int8 |] + , [t| Dual Int8 |] + , [t| Sum Int8 |] + , [t| Product Int8 |] + , [t| First Int8 |] + , [t| Last Int8 |] + , [t| Maybe Int8 |] + , [t| Either Int8 Int8 |] + , [t| SV.Vector Int64 |] + , [t| V.Vector Int64 |] + , [t| Ratio Int64 |] + , [t| Complex Int64 |] + , [t| Dual Int64 |] + , [t| Sum Int64 |] + , [t| Product Int64 |] + , [t| First Int64 |] + , [t| Last Int64 |] + , [t| Maybe Int64 |] + , [t| Either Int64 Int64 |] + , [t| (Int8, Int16) |] + , [t| (Int8, Int16, Bool) |] + , [t| (Bool, (), (), ()) |] + , [t| (Bool, (), Int8, ()) |] + -- Container-ey types + , [t| [Int8] |] + , [t| [Int64] |] + , [t| Seq Int8 |] + , [t| Seq Int64 |] + , [t| Set Int8 |] + , [t| Set Int64 |] + , [t| IntMap Int8 |] + , [t| IntMap Int64 |] + , [t| Map Int8 Int8 |] + , [t| Map Int64 Int64 |] + , [t| HashMap Int8 Int8 |] + , [t| HashMap Int64 Int64 |] + , [t| HashSet Int8 |] + , [t| HashSet Int64 |] + , [t| NE.NonEmpty Int8 |] + , [t| NE.NonEmpty Int64 |] + , [t| Tagged Int32 |] + ]) + it "Slices roundtrip" $ do + assertRoundtrip False $ T.drop 3 $ T.take 3 "Hello world!" + assertRoundtrip False $ BS.drop 3 $ BS.take 3 "Hello world!" + assertRoundtrip False $ SV.drop 3 $ SV.take 3 (SV.fromList [1..10] :: SV.Vector Int32) + assertRoundtrip False $ UV.drop 3 $ UV.take 3 (UV.fromList [1..10] :: UV.Vector Word8) + (return () :: IO ()) + it "StaticSize roundtrips" $ do + let x :: StaticSize 17 BS.ByteString + x = toStaticSizeEx (BS.replicate 17 255) + unless (checkRoundtrip False x) $ + (fail "Failed to roundtrip StaticSize ByteString" :: IO ()) + it "Size of generic instance for single fieldless constructor is 0" $ do + case size :: Size X of + ConstSize 0 -> (return () :: IO ()) + _ -> fail "Empty datatype takes up space" + it "Printing out polymorphic store instances" $ do + putStrLn "" + putStrLn "Not really a test - printing out known polymorphic store instances (which should all be tested above)" + putStrLn "" + mapM_ putStrLn + $(do insts <- getAllInstanceTypes1 ''Store + lift $ map pprint $ filter (not . isMonoType) insts) diff --git a/test/Data/StoreSpec/TH.hs b/test/Data/StoreSpec/TH.hs new file mode 100644 --- /dev/null +++ b/test/Data/StoreSpec/TH.hs @@ -0,0 +1,13 @@ +{-# LANGUAGE CPP #-} + +-- Just exists due to TH stage restriction... The actual testing TH code +-- is in "Data.Store.TH". +module Data.StoreSpec.TH where + +verbose :: Bool +verbose = +#if VERBOSE_TEST + True +#else + False +#endif diff --git a/test/Spec.hs b/test/Spec.hs new file mode 100644 --- /dev/null +++ b/test/Spec.hs @@ -0,0 +1,1 @@ +{-# OPTIONS_GHC -F -pgmF hspec-discover #-} diff --git a/test/System/IO/ByteBufferSpec.hs b/test/System/IO/ByteBufferSpec.hs new file mode 100644 --- /dev/null +++ b/test/System/IO/ByteBufferSpec.hs @@ -0,0 +1,40 @@ +{-# LANGUAGE OverloadedStrings #-} +module System.IO.ByteBufferSpec where + +import qualified Data.ByteString as BS +import qualified System.IO.ByteBuffer as BB +import Test.Hspec + +spec :: Spec +spec = describe "ByteBuffer" $ do + it "can grow to store a value and return it." $ BB.with (Just 0) $ \ bb -> do + let bs = "some bytestring" + BB.copyByteString bb bs + bs' <- BB.consume bb (BS.length bs) + bs' `shouldBe` Right bs + bbIsEmpty bb + it "should request more input when needed." $ BB.with (Just 0) $ \ bb -> do + let bs = "some bytestring" + BB.copyByteString bb bs + bs' <- BB.consume bb (2 * BS.length bs) + bs' `shouldBe` Left (BS.length bs) + BB.copyByteString bb bs + bs'' <- BB.consume bb (2 * BS.length bs) + bs'' `shouldBe` Right (BS.concat [bs, bs]) + bbIsEmpty bb + it "should not grow if bytes can be freed." $ + let bs1 = "12345" + bs2 = "67810" -- what about nine? 7 8 9! + in BB.with (Just $ BS.length bs1) $ \ bb -> do + BB.copyByteString bb bs1 + bs1' <- BB.consume bb (BS.length bs1) + BB.copyByteString bb bs2 + bs2' <- BB.consume bb (BS.length bs2) + bs1' `shouldBe` Right bs1 + bs2' `shouldBe` Right bs2 + bbSize <- BB.totalSize bb + bbSize `shouldBe` BS.length bs1 + bbIsEmpty bb + +bbIsEmpty :: BB.ByteBuffer -> Expectation +bbIsEmpty bb = BB.isEmpty bb >>= (`shouldBe` True)