diff --git a/CHANGELOG.markdown b/CHANGELOG.markdown new file mode 100644 --- /dev/null +++ b/CHANGELOG.markdown @@ -0,0 +1,21 @@ +# Change log + +## 0.5 + +- Released on 2015-01-09. + +## 0.4 + +- Released on 2015-01-09. + +## 0.3 + +- Released on 2013-03-14. + +## 0.2 + +- Released on 2012-10-28. + +## 0.1 + +- Initially released on 2011-10-22. diff --git a/LICENSE.txt b/LICENSE.txt new file mode 100644 --- /dev/null +++ b/LICENSE.txt @@ -0,0 +1,30 @@ +Copyright (c) Lennart Kolmodin + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions +are met: + +1. Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in the + documentation and/or other materials provided with the distribution. + +3. Neither the name of the author nor the names of his contributors + may be used to endorse or promote products derived from this software + without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE CONTRIBUTORS ``AS IS'' AND ANY EXPRESS +OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE AUTHORS OR CONTRIBUTORS BE LIABLE FOR +ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS +OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) +HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, +STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN +ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +POSSIBILITY OF SUCH DAMAGE. diff --git a/README.markdown b/README.markdown new file mode 100644 --- /dev/null +++ b/README.markdown @@ -0,0 +1,3 @@ +# caerbannog + +Bit parsing/writing on top of `binary`. diff --git a/caerbannog.cabal b/caerbannog.cabal new file mode 100644 --- /dev/null +++ b/caerbannog.cabal @@ -0,0 +1,50 @@ +cabal-version: >= 1.10 + +name: caerbannog +version: 0.6.0.2 +synopsis: Bit parsing/writing on top of binary. +description: + Bit parsing/writing on top of binary. Provides functions to read and write + bits to and from 8\/16\/32\/64 words. + +author: Lennart Kolmodin +build-type: Simple +category: Data, Parsing +extra-source-files: + CHANGELOG.markdown + README.markdown +license-file: LICENSE.txt +license: BSD3 +maintainer: Taylor Fausak + +source-repository head + location: https://github.com/tfausak/caerbannog + type: git + +library + build-depends: + base >= 4.12.0 && < 4.15 + , binary >= 0.8.6 && < 0.9 + , bytestring >= 0.10.8 && < 0.11 + default-language: Haskell98 + exposed-modules: + Data.Binary.Bits + Data.Binary.Bits.Get + Data.Binary.Bits.Put + ghc-options: + -Wall + hs-source-dirs: src/lib + +test-suite test + build-depends: + base + , binary + , caerbannog + , bytestring + , hspec >= 2.7.1 && < 2.8 + , QuickCheck >= 2.13.2 && < 2.15 + , random >= 1.1 && < 1.3 + default-language: Haskell98 + hs-source-dirs: src/test + main-is: Main.hs + type: exitcode-stdio-1.0 diff --git a/src/lib/Data/Binary/Bits.hs b/src/lib/Data/Binary/Bits.hs new file mode 100644 --- /dev/null +++ b/src/lib/Data/Binary/Bits.hs @@ -0,0 +1,33 @@ +-- | Parse and write bits easily. Parsing can be done either in a monadic +-- style, or more efficiently, using the 'Applicative' style. Writing is +-- monadic style only. See "Data.Binary.Bits.Get" and "Data.Binary.Bits.Put", +-- respectively. +module Data.Binary.Bits ( BinaryBit(putBits, getBits) ) where + +import qualified Data.Binary.Bits.Get as Get +import qualified Data.Binary.Bits.Put as Put +import qualified Data.Word as Word + +class BinaryBit a where + putBits :: Int -> a -> Put.BitPut () + getBits :: Int -> Get.BitGet a + +instance BinaryBit Bool where + putBits = const Put.putBool + getBits = const Get.getBool + +instance BinaryBit Word.Word8 where + putBits = Put.putWord8 + getBits = Get.getWord8 + +instance BinaryBit Word.Word16 where + putBits = Put.putWord16be + getBits = Get.getWord16be + +instance BinaryBit Word.Word32 where + putBits = Put.putWord32be + getBits = Get.getWord32be + +instance BinaryBit Word.Word64 where + putBits = Put.putWord64be + getBits = Get.getWord64be diff --git a/src/lib/Data/Binary/Bits/Get.hs b/src/lib/Data/Binary/Bits/Get.hs new file mode 100644 --- /dev/null +++ b/src/lib/Data/Binary/Bits/Get.hs @@ -0,0 +1,473 @@ +{-# LANGUAGE BangPatterns, CPP #-} + +-- | Parse bits easily. Parsing can be done either in a monadic style, or more +-- efficiently, using the 'Applicative' style. +-- +-- For the monadic style, write your parser as a 'BitGet' monad using the +-- +-- * 'getBool' +-- +-- * 'getWord8' +-- +-- * 'getWord16be' +-- +-- * 'getWord32be' +-- +-- * 'getWord64be' +-- +-- * 'getByteString' +-- +-- functions and run it with 'runBitGet'. +-- +-- For the applicative style, compose the fuctions +-- +-- * 'bool' +-- +-- * 'word8' +-- +-- * 'word16be' +-- +-- * 'word32be' +-- +-- * 'word64be' +-- +-- * 'byteString' +-- +-- to make a 'Block'. +-- Use 'block' to turn it into the 'BitGet' monad to be able to run it with +-- 'runBitGet'. + +module Data.Binary.Bits.Get + ( + -- * BitGet monad + + -- $bitget + + BitGet + , runBitGet + + -- ** Get bytes + , getBool + , getWord8 + , getWord16be + , getWord32be + , getWord64be + + -- * Blocks + + -- $blocks + , Block + , block + + -- ** Read in Blocks + , bool + , word8 + , word16be + , word32be + , word64be + , byteString + , Data.Binary.Bits.Get.getByteString + , Data.Binary.Bits.Get.getLazyByteString + , Data.Binary.Bits.Get.isEmpty + + ) where + +import qualified Control.Monad.Fail as Fail + +import Data.Binary.Get as B ( Get, getLazyByteString, isEmpty ) +import Data.Binary.Get.Internal as B ( get, put, ensureN ) + +import Data.ByteString as B +import qualified Data.ByteString.Lazy as L +import Data.ByteString.Unsafe + +import Data.Bits +import Data.Word +import Control.Applicative + +import Prelude as P + + +-- $bitget +-- Parse bits using a monad. +-- +-- @ +--myBitParser :: 'Get' ('Word8', 'Word8') +--myBitParser = 'runGetBit' parse4by4 +-- +--parse4by4 :: 'BitGet' ('Word8', 'Word8') +--parse4by4 = do +-- bits <- 'getWord8' 4 +-- more <- 'getWord8' 4 +-- return (bits,more) +-- @ + +-- $blocks +-- Parse more efficiently in blocks. Each block is read with only one boundry +-- check (checking that there is enough input) as the size of the block can be +-- calculated statically. This is somewhat limiting as you cannot make the +-- parsing depend on the input being parsed. +-- +-- @ +--data IPV6Header = IPV6Header { +-- ipv6Version :: 'Word8' +-- , ipv6TrafficClass :: 'Word8' +-- , ipv6FlowLabel :: 'Word32 +-- , ipv6PayloadLength :: 'Word16' +-- , ipv6NextHeader :: 'Word8' +-- , ipv6HopLimit :: 'Word8' +-- , ipv6SourceAddress :: 'ByteString' +-- , ipv6DestinationAddress :: 'ByteString' +-- } +-- +-- ipv6headerblock = +-- IPV6Header '<$>' 'word8' 4 +-- '<*>' 'word8' 8 +-- '<*>' 'word32be' 24 +-- '<*>' 'word16be' 16 +-- '<*>' 'word8' 8 +-- '<*>' 'word8' 8 +-- '<*>' 'byteString' 16 +-- '<*>' 'byteString' 16 +-- +--ipv6Header :: 'Get' IPV6Header +--ipv6Header = 'runBitGet' ('block' ipv6headerblock) +-- @ + +data S = S {-# UNPACK #-} !ByteString -- Input + {-# UNPACK #-} !Int -- Bit offset (0-7) + deriving (Show) + +-- | A block that will be read with only one boundry check. Needs to know the +-- number of bits in advance. +data Block a = Block Int (S -> a) + +instance Functor Block where + fmap f (Block i p) = Block i (\s -> f (p s)) + +instance Applicative Block where + pure a = Block 0 (\_ -> a) + (Block i p) <*> (Block j q) = Block (i+j) (\s -> p s $ q (incS i s)) + (Block i _) *> (Block j q) = Block (i+j) (q . incS i) + (Block i p) <* (Block j _) = Block (i+j) p + +-- | Get a block. Will be read with one single boundry check, and +-- therefore requires a statically known number of bits. +-- Build blocks using 'bool', 'word8', 'word16be', 'word32be', 'word64be', +-- 'byteString' and 'Applicative'. +block :: Block a -> BitGet a +block (Block i p) = do + ensureBits i + s <- getState + putState $! (incS i s) + return $! p s + +incS :: Int -> S -> S +incS o (S bs n) = + let !o' = (n+o) + !d = o' `shiftR` 3 + !n' = o' .&. make_mask 3 + in S (unsafeDrop d bs) n' + +-- | make_mask 3 = 00000111 +make_mask :: (Bits a, Num a) => Int -> a +make_mask n = (1 `shiftL` fromIntegral n) - 1 +{-# SPECIALIZE make_mask :: Int -> Int #-} +{-# SPECIALIZE make_mask :: Int -> Word #-} +{-# SPECIALIZE make_mask :: Int -> Word8 #-} +{-# SPECIALIZE make_mask :: Int -> Word16 #-} +{-# SPECIALIZE make_mask :: Int -> Word32 #-} +{-# SPECIALIZE make_mask :: Int -> Word64 #-} + +bit_offset :: Int -> Int +bit_offset n = make_mask 3 .&. n + +byte_offset :: Int -> Int +byte_offset n = n `shiftR` 3 + +readBool :: S -> Bool +readBool (S bs n) = testBit (unsafeHead bs) (7-n) + +{-# INLINE readWord8 #-} +readWord8 :: Int -> S -> Word8 +readWord8 n (S bs o) + -- no bits at all, return 0 + | n == 0 = 0 + + -- all bits are in the same byte + -- we just need to shift and mask them right + | n <= 8 - o = let w = unsafeHead bs + m = make_mask n + w' = (w `shiftr_w8` (8 - o - n)) .&. m + in w' + + -- the bits are in two different bytes + -- make a word16 using both bytes, and then shift and mask + | n <= 8 = let w = (fromIntegral (unsafeHead bs) `shiftl_w16` 8) .|. + (fromIntegral (unsafeIndex bs 1)) + m = make_mask n + w' = (w `shiftr_w16` (16 - o - n)) .&. m + in fromIntegral w' + | otherwise = error "readWord8: tried to read more than 8 bits" + +{-# INLINE readWord16be #-} +readWord16be :: Int -> S -> Word16 +readWord16be n s@(S bs o) + + -- 8 or fewer bits, use readWord8 + | n <= 8 = fromIntegral (readWord8 n s) + + -- handle 9 or more bits, stored in two bytes + + -- no offset, plain and simple 16 bytes + | o == 0 && n == 16 = let msb = fromIntegral (unsafeHead bs) + lsb = fromIntegral (unsafeIndex bs 1) + w = (msb `shiftl_w16` 8) .|. lsb + in w + + -- no offset, but not full 16 bytes + | o == 0 = let msb = fromIntegral (unsafeHead bs) + lsb = fromIntegral (unsafeIndex bs 1) + w = (msb `shiftl_w16` (n-8)) .|. (lsb `shiftr_w16` (16-n)) + in w + + -- with offset, and n=9-16 + | n <= 16 = readWithOffset s shiftl_w16 shiftr_w16 n + + | otherwise = error "readWord16be: tried to read more than 16 bits" + +{-# INLINE readWord32be #-} +readWord32be :: Int -> S -> Word32 +readWord32be n s@(S _ o) + -- 8 or fewer bits, use readWord8 + | n <= 8 = fromIntegral (readWord8 n s) + + -- 16 or fewer bits, use readWord16be + | n <= 16 = fromIntegral (readWord16be n s) + + | o == 0 = readWithoutOffset s shiftl_w32 shiftr_w32 n + + | n <= 32 = readWithOffset s shiftl_w32 shiftr_w32 n + + | otherwise = error "readWord32be: tried to read more than 32 bits" + + +{-# INLINE readWord64be #-} +readWord64be :: Int -> S -> Word64 +readWord64be n s@(S _ o) + -- 8 or fewer bits, use readWord8 + | n <= 8 = fromIntegral (readWord8 n s) + + -- 16 or fewer bits, use readWord16be + | n <= 16 = fromIntegral (readWord16be n s) + + | o == 0 = readWithoutOffset s shiftl_w64 shiftr_w64 n + + | n <= 64 = readWithOffset s shiftl_w64 shiftr_w64 n + + | otherwise = error "readWord64be: tried to read more than 64 bits" + + +readByteString :: Int -> S -> ByteString +readByteString n s@(S bs o) + -- no offset, easy. + | o == 0 = unsafeTake n bs + -- offset. ugg. this is really naive and slow. but also pretty easy :) + | otherwise = B.pack (P.map (readWord8 8) (P.take n (iterate (incS 8) s))) + +readWithoutOffset :: (Bits a, Num a) + => S -> (a -> Int -> a) -> (a -> Int -> a) -> Int -> a +readWithoutOffset (S bs o) shifterL shifterR n + | o /= 0 = error "readWithoutOffset: there is an offset" + + | bit_offset n == 0 && byte_offset n <= 4 = + let segs = byte_offset n + bn 0 = fromIntegral (unsafeHead bs) + bn x = (bn (x-1) `shifterL` 8) .|. fromIntegral (unsafeIndex bs x) + + in bn (segs-1) + + | n <= 64 = let segs = byte_offset n + o' = bit_offset (n - 8 + o) + + bn 0 = fromIntegral (unsafeHead bs) + bn x = (bn (x-1) `shifterL` 8) .|. fromIntegral (unsafeIndex bs x) + + msegs = bn (segs-1) `shifterL` o' + + lst = (fromIntegral (unsafeIndex bs segs)) `shifterR` (8 - o') + + w = msegs .|. lst + in w + | otherwise = error "readWithoutOffset: tried to read more than 64 bits" + +readWithOffset :: (Bits a, Num a) + => S -> (a -> Int -> a) -> (a -> Int -> a) -> Int -> a +readWithOffset (S bs o) shifterL shifterR n + | n <= 64 = let bits_in_msb = 8 - o + (n',top) = (n - bits_in_msb + , (fromIntegral (unsafeHead bs) .&. make_mask bits_in_msb) `shifterL` n') + + segs = byte_offset n' + + bn 0 = 0 + bn x = (bn (x-1) `shifterL` 8) .|. fromIntegral (unsafeIndex bs x) + + o' = bit_offset n' + + mseg = bn segs `shifterL` o' + + lst | o' > 0 = (fromIntegral (unsafeIndex bs (segs + 1))) `shifterR` (8 - o') + | otherwise = 0 + + w = top .|. mseg .|. lst + in w + | otherwise = error "readWithOffset: tried to read more than 64 bits" + +-- | 'BitGet' is a monad, applicative and a functor. See 'runBitGet' +-- for how to run it. +newtype BitGet a = B { runState :: S -> Get (S,a) } + +instance Monad BitGet where + return = pure + (B f) >>= g = B $ \s -> do (s',a) <- f s + runState (g a) s' + +#if !MIN_VERSION_GLASGOW_HASKELL(8, 8, 1, 0) + fail = Fail.fail +#endif + +instance Fail.MonadFail BitGet where + fail str = B $ \(S inp n) -> putBackState inp n >> fail str + +instance Functor BitGet where + fmap f m = m >>= \a -> return (f a) + +instance Applicative BitGet where + pure x = B $ \s -> return (s,x) + fm <*> m = fm >>= \f -> m >>= \v -> return (f v) + +-- | Run a 'BitGet' within the Binary packages 'Get' monad. If a byte has +-- been partially consumed it will be discarded once 'runBitGet' is finished. +runBitGet :: BitGet a -> Get a +runBitGet bg = do + s <- mkInitState + ((S str' n),a) <- runState bg s + putBackState str' n + return a + +mkInitState :: Get S +mkInitState = do + str <- get + put B.empty + return (S str 0) + +putBackState :: B.ByteString -> Int -> Get () +putBackState bs n = do + remaining <- get + put (B.drop (if n==0 then 0 else 1) bs `B.append` remaining) + +getState :: BitGet S +getState = B $ \s -> return (s,s) + +putState :: S -> BitGet () +putState s = B $ \_ -> return (s,()) + +-- | Make sure there are at least @n@ bits. +ensureBits :: Int -> BitGet () +ensureBits n = do + (S bs o) <- getState + if n <= (B.length bs * 8 - o) + then return () + else do let currentBits = B.length bs * 8 - o + let byteCount = (n - currentBits + 7) `div` 8 + B $ \_ -> do B.ensureN byteCount + bs' <- B.get + put B.empty + return (S (bs`append`bs') o, ()) + +-- | Get 1 bit as a 'Bool'. +getBool :: BitGet Bool +getBool = block bool + +-- | Get @n@ bits as a 'Word8'. @n@ must be within @[0..8]@. +getWord8 :: Int -> BitGet Word8 +getWord8 n = block (word8 n) + +-- | Get @n@ bits as a 'Word16'. @n@ must be within @[0..16]@. +getWord16be :: Int -> BitGet Word16 +getWord16be n = block (word16be n) + +-- | Get @n@ bits as a 'Word32'. @n@ must be within @[0..32]@. +getWord32be :: Int -> BitGet Word32 +getWord32be n = block (word32be n) + +-- | Get @n@ bits as a 'Word64'. @n@ must be within @[0..64]@. +getWord64be :: Int -> BitGet Word64 +getWord64be n = block (word64be n) + +-- | Get @n@ bytes as a 'ByteString'. +getByteString :: Int -> BitGet ByteString +getByteString n = block (byteString n) + +-- | Get @n@ bytes as a lazy ByteString. +getLazyByteString :: Int -> BitGet L.ByteString +getLazyByteString n = do + (S _ o) <- getState + case o of + 0 -> B $ \ (S bs o') -> do + putBackState bs o' + lbs <- B.getLazyByteString (fromIntegral n) + return (S B.empty 0, lbs) + _ -> L.fromChunks . (:[]) <$> Data.Binary.Bits.Get.getByteString n + +-- | Test whether all input has been consumed, i.e. there are no remaining +-- undecoded bytes. +isEmpty :: BitGet Bool +isEmpty = B $ \ (S bs o) -> if B.null bs + then B.isEmpty >>= \e -> return (S bs o, e) + else return (S bs o, False) + +-- | Read a 1 bit 'Bool'. +bool :: Block Bool +bool = Block 1 readBool + +-- | Read @n@ bits as a 'Word8'. @n@ must be within @[0..8]@. +word8 :: Int -> Block Word8 +word8 n = Block n (readWord8 n) + +-- | Read @n@ bits as a 'Word16'. @n@ must be within @[0..16]@. +word16be :: Int -> Block Word16 +word16be n = Block n (readWord16be n) + +-- | Read @n@ bits as a 'Word32'. @n@ must be within @[0..32]@. +word32be :: Int -> Block Word32 +word32be n = Block n (readWord32be n) + +-- | Read @n@ bits as a 'Word64'. @n@ must be within @[0..64]@. +word64be :: Int -> Block Word64 +word64be n = Block n (readWord64be n) + +-- | Read @n@ bytes as a 'ByteString'. +byteString :: Int -> Block ByteString +byteString n | n > 0 = Block (n*8) (readByteString n) + | otherwise = Block 0 (\_ -> B.empty) + +-- Unchecked shifts, from the package binary + +shiftl_w16 :: Word16 -> Int -> Word16 +shiftl_w32 :: Word32 -> Int -> Word32 +shiftl_w64 :: Word64 -> Int -> Word64 +shiftr_w8 :: Word8 -> Int -> Word8 +shiftr_w16 :: Word16 -> Int -> Word16 +shiftr_w32 :: Word32 -> Int -> Word32 +shiftr_w64 :: Word64 -> Int -> Word64 + +shiftl_w16 = unsafeShiftL +shiftl_w32 = unsafeShiftL +shiftl_w64 = unsafeShiftL + +shiftr_w8 = unsafeShiftR +shiftr_w16 = unsafeShiftR +shiftr_w32 = unsafeShiftR +shiftr_w64 = unsafeShiftR diff --git a/src/lib/Data/Binary/Bits/Put.hs b/src/lib/Data/Binary/Bits/Put.hs new file mode 100644 --- /dev/null +++ b/src/lib/Data/Binary/Bits/Put.hs @@ -0,0 +1,160 @@ +-- | Put bits easily. + +module Data.Binary.Bits.Put + ( BitPut + , runBitPut + , joinPut + + -- * Data types + -- ** Bool + , putBool + + -- ** Words + , putWord8 + , putWord16be + , putWord32be + , putWord64be + + -- ** ByteString + , putByteString + ) + where + +import Data.Bits ((.&.), (.|.)) + +import qualified Data.Binary.Builder as B +import qualified Data.Binary.Put as Put +import qualified Data.Bits as Bits +import qualified Data.ByteString as ByteString +import qualified Data.Word as Word + +data BitPut a = BitPut { run :: (S -> PairS a) } + +data PairS a = PairS a {-# UNPACK #-} !S + +data S = S !B.Builder !Word.Word8 !Int + +-- | Put a 1 bit 'Bool'. +putBool :: Bool -> BitPut () +putBool b = putWord8 1 (if b then 0xff else 0x00) + +-- | make_mask 3 = 00000111 +make_mask :: (Bits.Bits a, Num a) => Int -> a +make_mask n = (1 `Bits.shiftL` fromIntegral n) - 1 +{-# SPECIALIZE make_mask :: Int -> Int #-} +{-# SPECIALIZE make_mask :: Int -> Word #-} +{-# SPECIALIZE make_mask :: Int -> Word.Word8 #-} +{-# SPECIALIZE make_mask :: Int -> Word.Word16 #-} +{-# SPECIALIZE make_mask :: Int -> Word.Word32 #-} +{-# SPECIALIZE make_mask :: Int -> Word.Word64 #-} + +-- | Put the @n@ lower bits of a 'Word8'. +putWord8 :: Int -> Word.Word8 -> BitPut () +putWord8 n w = BitPut $ \s -> PairS () $ + let w' = make_mask n .&. w in + case s of + -- a whole word8, no offset + (S b t o) | n == 8 && o == 0 -> flush $ S b w n + -- less than a word8, will fit in the current word8 + | n <= 8 - o -> flush $ S b (t .|. (w' `Bits.shiftL` (8 - n - o))) (o+n) + -- will finish this word8, and spill into the next one + | otherwise -> flush $ + let o' = o + n - 8 + b' = t .|. (w' `Bits.shiftR` o') + t' = w `Bits.shiftL` (8 - o') + in S (b `mappend` B.singleton b') t' o' + +-- | Put the @n@ lower bits of a 'Word16'. +putWord16be :: Int -> Word.Word16 -> BitPut () +putWord16be n w + | n <= 8 = putWord8 n (fromIntegral w) + | otherwise = + BitPut $ \s -> PairS () $ + let w' = make_mask n .&. w in + case s of + -- as n>=9, it's too big to fit into one single byte + -- it'll either use 2 or 3 bytes + -- it'll fit in 2 bytes + (S b t o) | o + n <= 16 -> flush $ + let o' = o + n - 8 + b' = t .|. fromIntegral (w' `Bits.shiftR` o') + t' = fromIntegral (w `Bits.shiftL` (8-o')) + in (S (b `mappend` B.singleton b') t' o') + -- 3 bytes required + | otherwise -> flush $ + let o' = o + n - 16 + b' = t .|. fromIntegral (w' `Bits.shiftR` (o' + 8)) + b'' = fromIntegral ((w `Bits.shiftR` o') .&. 0xff) + t' = fromIntegral (w `Bits.shiftL` (8-o')) + in (S (b `mappend` B.singleton b' `mappend` B.singleton b'') t' o') + +-- | Put the @n@ lower bits of a 'Word32'. +putWord32be :: Int -> Word.Word32 -> BitPut () +putWord32be n w + | n <= 16 = putWord16be n (fromIntegral w) + | otherwise = do + putWord32be (n-16) (w`Bits.shiftR`16) + putWord32be 16 (w .&. 0x0000ffff) + +-- | Put the @n@ lower bits of a 'Word64'. +putWord64be :: Int -> Word.Word64 -> BitPut () +putWord64be n w + | n <= 32 = putWord32be n (fromIntegral w) + | otherwise = do + putWord64be (n-32) (w`Bits.shiftR`32) + putWord64be 32 (w .&. 0xffffffff) + +-- | Put a 'ByteString'. +putByteString :: ByteString.ByteString -> BitPut () +putByteString bs = do + offset <- hasOffset + if offset + then mapM_ (putWord8 8) (ByteString.unpack bs) -- naive + else joinPut (Put.putByteString bs) + where + hasOffset = BitPut $ \ s@(S _ _ o) -> PairS (o /= 0) s + +-- | Run a 'Put' inside 'BitPut'. Any partially written bytes will be flushed +-- before 'Put' executes to ensure byte alignment. +joinPut :: Put.Put -> BitPut () +joinPut m = BitPut $ \s0 -> PairS () $ + let (S b0 _ _) = flushIncomplete s0 + b = Put.execPut m + in (S (b0`mappend`b) 0 0) + +flush :: S -> S +flush s@(S b w o) + | o > 8 = error "flush: offset > 8" + | o == 8 = S (b `mappend` B.singleton w) 0 0 + | otherwise = s + +flushIncomplete :: S -> S +flushIncomplete s@(S b w o) + | o == 0 = s + | otherwise = (S (b `mappend` B.singleton w) 0 0) + +-- | Run the 'BitPut' monad inside 'Put'. +runBitPut :: BitPut () -> Put.Put +runBitPut m = Put.putBuilder b + where + PairS _ s = run m (S mempty 0 0) + (S b _ _) = flushIncomplete s + +instance Functor BitPut where + fmap f (BitPut k) = BitPut $ \s -> + let PairS x s' = k s + in PairS (f x) s' + +instance Applicative BitPut where + pure a = BitPut (\s -> PairS a s) + (BitPut f) <*> (BitPut g) = BitPut $ \s -> + let PairS a s' = f s + PairS b s'' = g s' + in PairS (a b) s'' + +instance Monad BitPut where + m >>= k = BitPut $ \s -> + let PairS a s' = run m s + PairS b s'' = run (k a) s' + in PairS b s'' + return x = BitPut $ \s -> PairS x s diff --git a/src/test/Main.hs b/src/test/Main.hs new file mode 100644 --- /dev/null +++ b/src/test/Main.hs @@ -0,0 +1,469 @@ +{-# OPTIONS_GHC -Wno-orphans #-} +{-# LANGUAGE FlexibleInstances, FlexibleContexts #-} + +module Main ( main ) where + +import Data.Bits ((.|.)) +import Test.QuickCheck ((==>)) + +import qualified Data.Binary as Binary +import qualified Data.Binary.Bits as BB +import qualified Data.Binary.Bits.Get as BB +import qualified Data.Binary.Bits.Put as BB +import qualified Data.Binary.Get as Binary +import qualified Data.Binary.Put as Binary +import qualified Data.Bits as Bits +import qualified Data.ByteString as B +import qualified Data.ByteString.Lazy as L +import qualified Data.Word as Word +import qualified Foreign +import qualified Test.Hspec as Hspec +import qualified Test.QuickCheck as QC + +main :: IO () +main = Hspec.hspec $ do + Hspec.describe "Internal test functions" $ do + Hspec.it "prop_bitreq" $ QC.property prop_bitreq + + Hspec.describe "Custom test cases" $ do + Hspec.it "prop_composite_case" $ QC.property prop_composite_case + + Hspec.describe "getByteString" $ do + Hspec.it "prop_getByteString_negative" $ QC.property prop_getByteString_negative + + Hspec.describe "getLazyByteString" $ do + Hspec.it "getLazyByteString == getByteString" $ QC.property + prop_getLazyByteString_equal_to_ByteString + Hspec.it "getLazyByteString == getByteString (with shift)" $ QC.property + prop_getLazyByteString_equal_to_ByteString2 + + Hspec.describe "isEmpty" $ do + Hspec.it "prop_isEmptyOfEmptyEmpty" $ QC.property prop_isEmptyOfEmptyEmpty + Hspec.it "prop_isEmptyOfNonEmptyEmpty" $ QC.property prop_isEmptyOfNonEmptyEmpty + Hspec.it "prop_isEmptyOfConsumedEmpty" $ QC.property prop_isEmptyOfConsumedEmpty + Hspec.it "prop_isEmptyOfNotConsumedNotEmpty" $ QC.property prop_isEmptyOfNotConsumedNotEmpty + + Hspec.describe "Fail" $ do + Hspec.it "monadic fail" $ QC.property prop_fail + + Hspec.describe "prop_bitput_with_get_from_binary" $ do + Hspec.it "Word8" $ QC.property (prop_bitput_with_get_from_binary :: W [Word.Word8] -> QC.Property) + Hspec.it "Word16" $ QC.property (prop_bitput_with_get_from_binary :: W [Word.Word16] -> QC.Property) + Hspec.it "Word32" $ QC.property (prop_bitput_with_get_from_binary :: W [Word.Word32] -> QC.Property) + Hspec.it "Word64" $ QC.property (prop_bitput_with_get_from_binary :: W [Word.Word64] -> QC.Property) + + Hspec.describe "prop_bitget_with_put_from_binary" $ do + Hspec.it "Word8" $ QC.property (prop_bitget_with_put_from_binary :: W [Word.Word8] -> QC.Property) + Hspec.it "Word16" $ QC.property (prop_bitget_with_put_from_binary :: W [Word.Word16] -> QC.Property) + Hspec.it "Word32" $ QC.property (prop_bitget_with_put_from_binary :: W [Word.Word32] -> QC.Property) + Hspec.it "Word64" $ QC.property (prop_bitget_with_put_from_binary :: W [Word.Word64] -> QC.Property) + + Hspec.describe "prop_compare_put_with_naive" $ do + Hspec.it "Word8" $ QC.property (prop_compare_put_with_naive :: W [Word.Word8] -> QC.Property) + Hspec.it "Word16" $ QC.property (prop_compare_put_with_naive :: W [Word.Word16] -> QC.Property) + Hspec.it "Word32" $ QC.property (prop_compare_put_with_naive :: W [Word.Word32] -> QC.Property) + Hspec.it "Word64" $ QC.property (prop_compare_put_with_naive :: W [Word.Word64] -> QC.Property) + + Hspec.describe "prop_compare_get_with_naive" $ do + Hspec.it "Word8" $ QC.property (prop_compare_get_with_naive:: W [Word.Word8] -> QC.Property) + Hspec.it "Word16" $ QC.property (prop_compare_get_with_naive:: W [Word.Word16] -> QC.Property) + Hspec.it "Word32" $ QC.property (prop_compare_get_with_naive:: W [Word.Word32] -> QC.Property) + Hspec.it "Word64" $ QC.property (prop_compare_get_with_naive:: W [Word.Word64] -> QC.Property) + + Hspec.describe "prop_put_with_bitreq" $ do + Hspec.it "Word8" $ QC.property (prop_putget_with_bitreq :: W Word.Word8 -> QC.Property) + Hspec.it "Word16" $ QC.property (prop_putget_with_bitreq :: W Word.Word16 -> QC.Property) + Hspec.it "Word32" $ QC.property (prop_putget_with_bitreq :: W Word.Word32 -> QC.Property) + Hspec.it "Word64" $ QC.property (prop_putget_with_bitreq :: W Word.Word64 -> QC.Property) + + Hspec.describe "prop_putget_list_simple" $ do + Hspec.it "Bool" $ QC.property (prop_putget_list_simple :: W [Bool] -> QC.Property) + Hspec.it "Word8" $ QC.property (prop_putget_list_simple :: W [Word.Word8] -> QC.Property) + Hspec.it "Word16" $ QC.property (prop_putget_list_simple :: W [Word.Word16] -> QC.Property) + Hspec.it "Word32" $ QC.property (prop_putget_list_simple :: W [Word.Word32] -> QC.Property) + Hspec.it "Word64" $ QC.property (prop_putget_list_simple :: W [Word.Word64] -> QC.Property) + + Hspec.describe "prop_putget_applicative_with_bitreq" $ do + Hspec.it "Word8" $ QC.property (prop_putget_applicative_with_bitreq :: W [(Word.Word8,Word.Word8,Word.Word8)] -> QC.Property) + Hspec.it "Word16" $ QC.property (prop_putget_applicative_with_bitreq :: W [(Word.Word16,Word.Word16,Word.Word16)] -> QC.Property) + Hspec.it "Word32" $ QC.property (prop_putget_applicative_with_bitreq :: W [(Word.Word32,Word.Word32,Word.Word32)] -> QC.Property) + Hspec.it "Word64" $ QC.property (prop_putget_applicative_with_bitreq :: W [(Word.Word64,Word.Word64,Word.Word64)] -> QC.Property) + + Hspec.describe "prop_putget_list_with_bitreq" $ do + Hspec.it "Word8" $ QC.property (prop_putget_list_with_bitreq :: W [Word.Word8] -> QC.Property) + Hspec.it "Word16" $ QC.property (prop_putget_list_with_bitreq :: W [Word.Word16] -> QC.Property) + Hspec.it "Word32" $ QC.property (prop_putget_list_with_bitreq :: W [Word.Word32] -> QC.Property) + Hspec.it "Word64" $ QC.property (prop_putget_list_with_bitreq :: W [Word.Word64] -> QC.Property) + Hspec.describe "prop_bitget_bytestring_interspersed" $ do + Hspec.it "Word8" $ QC.property (prop_bitget_bytestring_interspersed :: W Word.Word8 -> [B.ByteString] -> QC.Property) + Hspec.it "Word16" $ QC.property (prop_bitget_bytestring_interspersed :: W Word.Word16 -> [B.ByteString] -> QC.Property) + Hspec.it "Word32" $ QC.property (prop_bitget_bytestring_interspersed :: W Word.Word32 -> [B.ByteString] -> QC.Property) + Hspec.it "Word64" $ QC.property (prop_bitget_bytestring_interspersed :: W Word.Word64 -> [B.ByteString] -> QC.Property) + Hspec.describe "Simulate programs" $ do + Hspec.it "primitive" $ QC.property prop_primitive + Hspec.it "many primitives in sequence" $ QC.property prop_program + +prop_isEmptyOfEmptyEmpty :: Bool +prop_isEmptyOfEmptyEmpty = Binary.runGet (BB.runBitGet BB.isEmpty) L.empty + +prop_isEmptyOfNonEmptyEmpty :: L.ByteString -> QC.Property +prop_isEmptyOfNonEmptyEmpty bs = + not (L.null bs) ==> not (Binary.runGet (BB.runBitGet BB.isEmpty) bs) + +prop_isEmptyOfConsumedEmpty :: L.ByteString -> QC.Property +prop_isEmptyOfConsumedEmpty bs = + not (L.null bs) ==> + Binary.runGet (BB.runBitGet (BB.getByteString n >> BB.isEmpty)) bs + where n = fromIntegral $ L.length bs + +prop_isEmptyOfNotConsumedNotEmpty :: L.ByteString -> Int -> QC.Property +prop_isEmptyOfNotConsumedNotEmpty bs n = + (fromIntegral n) < L.length bs && not (L.null bs) ==> + not (Binary.runGet (BB.runBitGet (BB.getByteString n >> BB.isEmpty)) bs) + +prop_getLazyByteString_equal_to_ByteString :: L.ByteString -> Int -> QC.Property +prop_getLazyByteString_equal_to_ByteString bs n = + (fromIntegral n) <= L.length bs ==> + Binary.runGet (BB.runBitGet (BB.getLazyByteString (fromIntegral n))) bs == + (L.fromChunks . (:[]) $ Binary.runGet (BB.runBitGet (BB.getByteString n)) bs) + +prop_getLazyByteString_equal_to_ByteString2 :: L.ByteString -> Int -> QC.Property +prop_getLazyByteString_equal_to_ByteString2 bs n = + (L.length bs > 1) && (fromIntegral n) < L.length bs ==> + Binary.runGet (BB.runBitGet (BB.getWord8 2 >> BB.getLazyByteString (fromIntegral n))) bs == + (L.fromChunks . (:[]) $ Binary.runGet (BB.runBitGet (BB.getWord8 2 >> BB.getByteString n)) bs) + +prop_getByteString_negative :: Int -> QC.Property +prop_getByteString_negative n = + n < 1 ==> + Binary.runGet (BB.runBitGet (BB.getByteString n)) L.empty == B.empty + +prop_putget_with_bitreq :: (BB.BinaryBit a, Num a, Bits.Bits a, Ord a) => W a -> QC.Property +prop_putget_with_bitreq (W w) = QC.property $ + -- write all words with as many bits as it's required + let p = BB.putBits (bitreq w) w + g = BB.getBits (bitreq w) + lbs = Binary.runPut (BB.runBitPut p) + w' = Binary.runGet (BB.runBitGet g) lbs + in w == w' + +-- | Write a list of items. Each item is written with the maximum amount of +-- bits, i.e. 8 for Word8, 16 for Word16, etc. +prop_putget_list_simple :: (BB.BinaryBit a, Eq a, Foreign.Storable a) => W [a] -> QC.Property +prop_putget_list_simple (W ws) = QC.property $ + let s = Foreign.sizeOf (head ws) * 8 + p = mapM_ (BB.putBits s) ws + g = mapM (const (BB.getBits s)) ws + lbs = Binary.runPut (BB.runBitPut p) + ws' = Binary.runGet (BB.runBitGet g) lbs + in ws == ws' + +-- | Write a list of items. Each item is written with exactly as many bits +-- as required. Then read it back. +prop_putget_list_with_bitreq :: (BB.BinaryBit a, Num a, Bits.Bits a, Ord a) => W [a] -> QC.Property +prop_putget_list_with_bitreq (W ws) = QC.property $ + -- write all words with as many bits as it's required + let p = mapM_ (\v -> BB.putBits (bitreq v) v) ws + g = mapM BB.getBits bitlist + lbs = Binary.runPut (BB.runBitPut p) + ws' = Binary.runGet (BB.runBitGet g) lbs + in ws == ws' + where + bitlist = map bitreq ws + +prop_putget_applicative_with_bitreq :: (BB.BinaryBit a, Num a, Bits.Bits a, Ord a) => W [(a,a,a)] -> QC.Property +prop_putget_applicative_with_bitreq (W ts) = QC.property $ + let p = mapM_ (\(a,b,c) -> do BB.putBits (bitreq a) a + BB.putBits (bitreq b) b + BB.putBits (bitreq c) c) ts + g = mapM (\(a,b,c) -> (,,) <$> BB.getBits a <*> BB.getBits b <*> BB.getBits c) bitlist + lbs = Binary.runPut (BB.runBitPut p) + ts' = Binary.runGet (BB.runBitGet g) lbs + in ts == ts' + where + bitlist = map (\(a,b,c) -> (bitreq a, bitreq b, bitreq c)) ts + +-- | Write bits using this library, and read them back using the binary +-- library. +prop_bitput_with_get_from_binary :: (BB.BinaryBit a, Binary.Binary a, Foreign.Storable a, Eq a) => W [a] -> QC.Property +prop_bitput_with_get_from_binary (W ws) = QC.property $ + let s = Foreign.sizeOf (head ws) * 8 + p = mapM_ (BB.putBits s) ws + g = mapM (const Binary.get) ws + lbs = Binary.runPut (BB.runBitPut p) + ws' = Binary.runGet g lbs + in ws == ws' + +-- | Write bits using the binary library, and read them back using this +-- library. +prop_bitget_with_put_from_binary :: (BB.BinaryBit a, Binary.Binary a, Foreign.Storable a, Eq a) => W [a] -> QC.Property +prop_bitget_with_put_from_binary (W ws) = QC.property $ + let s = Foreign.sizeOf (head ws) * 8 + p = mapM_ Binary.put ws + g = mapM (const (BB.getBits s)) ws + lbs = Binary.runPut p + ws' = Binary.runGet (BB.runBitGet g) lbs + in ws == ws' + +-- | Write each 'ByteString' with a variable sized value as a separator. +prop_bitget_bytestring_interspersed :: (BB.BinaryBit a, Binary.Binary a, Num a, Ord a, Bits.Bits a) => W a -> [B.ByteString] -> QC.Property +prop_bitget_bytestring_interspersed (W ws) bss = QC.property $ + let p = mapM_ (\bs -> BB.putBits (bitreq ws) ws >> BB.putByteString bs) bss + g = mapM (\bs -> (,) <$> BB.getBits (bitreq ws) <*> BB.getByteString (B.length bs)) bss + lbs = Binary.runPut (BB.runBitPut p) + r = Binary.runGet (BB.runBitGet g) lbs + in map ((,) ws) bss == r + +-- | Test failing. +prop_fail :: L.ByteString -> String -> QC.Property +prop_fail lbs errMsg0 = QC.forAll (QC.choose (0, 8 * L.length lbs)) $ \len -> + let (bytes,bits) = len `divMod` 8 + expectedBytesConsumed + | bits == 0 = bytes + | otherwise = bytes + 1 + p = do _ <- BB.getByteString (fromIntegral bytes) + _ <- BB.getBits (fromIntegral bits) :: BB.BitGet Word.Word8 + fail errMsg0 + r = Binary.runGetIncremental (BB.runBitGet p) `Binary.pushChunks` lbs + in case r of + Binary.Fail remainingBS pos errMsg -> + and [ L.fromChunks [remainingBS] == L.drop expectedBytesConsumed lbs + , pos == expectedBytesConsumed + , errMsg == errMsg0 + ] + _ -> False + +-- | number of bits required to write @v@ +bitreq :: (Num b, Num a, Bits.Bits a, Ord a) => a -> b +bitreq v = fromIntegral . head $ [ req | (req, top) <- bittable, v <= top ] + +bittable :: (Bits.Bits a, Num a) => [(Integer, a)] +bittable = [ (fromIntegral x, (1 `Bits.shiftL` x) - 1) | x <- [1..64] ] + +prop_bitreq :: W Word.Word64 -> QC.Property +prop_bitreq (W w) = QC.property $ + ( w == 0 && bitreq w == (1 :: Integer) ) + || bitreq w == bitreq (w `Bits.shiftR` 1) + (1 :: Integer) + +prop_composite_case :: Bool -> W Word.Word16 -> QC.Property +prop_composite_case b (W w) = w < 0x8000 ==> + let p = do BB.putBool b + BB.putWord16be 15 w + g = do v <- BB.getBool + case v of + True -> BB.getWord16be 15 + False -> do + msb <- BB.getWord8 7 + lsb <- BB.getWord8 8 + return ((fromIntegral msb `Bits.shiftL` 8) .|. fromIntegral lsb) + lbs = Binary.runPut (BB.runBitPut p) + w' = Binary.runGet (BB.runBitGet g) lbs + in w == w' + +prop_compare_put_with_naive :: (Bits.Bits a, BB.BinaryBit a, Ord a, Num a) => W [a] -> QC.Property +prop_compare_put_with_naive (W ws) = QC.property $ + let pn = mapM_ (\v -> naive_put (bitreq v) v) ws + p = mapM_ (\v -> BB.putBits (bitreq v) v) ws + lbs_n = Binary.runPut (BB.runBitPut pn) + lbs = Binary.runPut (BB.runBitPut p) + in lbs_n == lbs + +prop_compare_get_with_naive :: (Bits.Bits a, BB.BinaryBit a, Ord a, Num a) => W [a] -> QC.Property +prop_compare_get_with_naive (W ws) = QC.property $ + let gn = mapM (\v -> naive_get (bitreq v)) ws + g = mapM (\v -> BB.getBits (bitreq v)) ws + p = mapM_ (\v -> naive_put (bitreq v) v) ws + lbs = Binary.runPut (BB.runBitPut p) + rn = Binary.runGet (BB.runBitGet gn) lbs + r = Binary.runGet (BB.runBitGet g ) lbs + -- we must help our compiler to resolve the types of 'gn' and 'g' + _types = rn == ws && r == ws + in rn == r + +-- | Write one bit at a time until the full word has been written +naive_put :: (Bits.Bits a) => Int -> a -> BB.BitPut () +naive_put n w = mapM_ (\b -> BB.putBool (Bits.testBit w b)) [n-1,n-2..0] + +-- | Read one bit at a time until we've reconstructed the whole word +naive_get :: (Bits.Bits a, Num a) => Int -> BB.BitGet a +naive_get n0 = + let loop 0 acc = return acc + loop n acc = do + b <- BB.getBool + case b of + False -> loop (n-1) (acc `Bits.shiftL` 1) + True -> loop (n-1) ((acc `Bits.shiftL` 1) + 1) + in loop n0 0 + +shrinker :: (Num a, Ord a, Bits.Bits a) => a -> [a] +shrinker 0 = [] +shrinker w = [ w `Bits.shiftR` 1 -- try to make everything roughly half size + ] ++ [ w' -- flip bits to zero, left->right + | m <- [n, n-1..1] + , let w' = w `Bits.clearBit` m + , w /= w' + ] ++ [w-1] -- just make it a little smaller + where + n = bitreq w + +data W a = W { unW :: a } deriving (Show, Eq, Ord) + +arbitraryW :: (QC.Arbitrary (W a)) => QC.Gen a +arbitraryW = unW <$> QC.arbitrary + +shrinkW :: (QC.Arbitrary (W a)) => a -> [a] +shrinkW x = unW <$> QC.shrink (W x) + +instance QC.Arbitrary (W Bool) where + arbitrary = W <$> QC.arbitrary + shrink = map W <$> QC.shrink . unW + +instance QC.Arbitrary (W Word.Word8) where + arbitrary = W <$> QC.choose (minBound, maxBound) + shrink = map W . shrinker . unW + +instance QC.Arbitrary (W Word.Word16) where + arbitrary = W <$> QC.choose (minBound, maxBound) + shrink = map W . shrinker . unW + +instance QC.Arbitrary (W Word.Word32) where + arbitrary = W <$> QC.choose (minBound, maxBound) + shrink = map W . shrinker . unW + +instance QC.Arbitrary (W Word.Word64) where + arbitrary = W <$> QC.choose (minBound, maxBound) + shrink = map W . shrinker . unW + +instance QC.Arbitrary B.ByteString where + arbitrary = B.pack <$> QC.arbitrary + shrink bs = B.pack <$> QC.shrink (B.unpack bs) + +instance QC.Arbitrary L.ByteString where + arbitrary = L.fromChunks <$> QC.arbitrary + shrink bs = L.fromChunks <$> QC.shrink (L.toChunks bs) + +instance (QC.Arbitrary (W a)) => QC.Arbitrary (W [a]) where + arbitrary = W . map unW <$> QC.arbitrary + shrink = map (W . map unW) <$> mapM QC.shrink . map W . unW + +instance (QC.Arbitrary (W a), QC.Arbitrary (W b)) => QC.Arbitrary (W (a,b)) where + arbitrary = (W .) . (,) <$> arbitraryW <*> arbitraryW + shrink (W (a,b)) = (W .) . (,) <$> shrinkW a <*> shrinkW b + +instance (QC.Arbitrary (W a), QC.Arbitrary (W b), QC.Arbitrary (W c)) => QC.Arbitrary (W (a,b,c)) where + arbitrary = ((W .) .) . (,,) <$> arbitraryW <*> arbitraryW <*> arbitraryW + shrink (W (a,b,c)) = ((W .) .) . (,,) <$> shrinkW a <*> shrinkW b <*> shrinkW c + +data Primitive + = Bool Bool + | W8 Int Word.Word8 + | W16 Int Word.Word16 + | W32 Int Word.Word32 + | W64 Int Word.Word64 + | BS Int B.ByteString + | LBS Int L.ByteString + | IsEmpty + deriving (Eq, Show) + +type Program = [Primitive] + +instance QC.Arbitrary Primitive where + arbitrary = do + let gen c = do + let (maxBits, _) = (\w -> (Bits.finiteBitSize w, c undefined w)) undefined + bits <- QC.choose (0, maxBits) + n <- QC.choose (0, fromIntegral (2^bits-1 :: Integer)) + return (c bits n) + QC.oneof + [ Bool <$> QC.arbitrary + , gen W8 + , gen W16 + , gen W32 + , gen W64 + , do n <- QC.choose (0,10) + cs <- QC.vector n + return (BS n (B.pack cs)) + , do n <- QC.choose (0,10) + cs <- QC.vector n + return (LBS n (L.pack cs)) + , return IsEmpty + ] + shrink p = + let snk c x = map (\x' -> c (bitreq x') x') (shrinker x) in + case p of + Bool b -> if b then [Bool False] else [] + W8 _ x -> snk W8 x + W16 _ x -> snk W16 x + W32 _ x -> snk W32 x + W64 _ x -> snk W64 x + BS _ bs -> let ws = B.unpack bs in map (\ws' -> BS (length ws') (B.pack ws')) (QC.shrink ws) + LBS _ lbs -> let ws = L.unpack lbs in map (\ws' -> LBS (length ws') (L.pack ws')) (QC.shrink ws) + IsEmpty -> [] + +prop_primitive :: Primitive -> QC.Property +prop_primitive prim = QC.property $ + let p = putPrimitive prim + g = getPrimitive prim + lbs = Binary.runPut (BB.runBitPut p) + r = Binary.runGet (BB.runBitGet g) lbs + in r == prim + +prop_program :: Program -> QC.Property +prop_program program = QC.property $ + let p = mapM_ putPrimitive program + g = verifyProgram (8 * fromIntegral (L.length lbs)) program + lbs = Binary.runPut (BB.runBitPut p) + r = Binary.runGet (BB.runBitGet g) lbs + in r + +putPrimitive :: Primitive -> BB.BitPut () +putPrimitive p = + case p of + Bool b -> BB.putBool b + W8 n x -> BB.putWord8 n x + W16 n x -> BB.putWord16be n x + W32 n x -> BB.putWord32be n x + W64 n x -> BB.putWord64be n x + BS _ bs -> BB.putByteString bs + LBS _ lbs -> mapM_ BB.putByteString (L.toChunks lbs) + IsEmpty -> return () + +getPrimitive :: Primitive -> BB.BitGet Primitive +getPrimitive p = + case p of + Bool _ -> Bool <$> BB.getBool + W8 n _ -> W8 n <$> BB.getWord8 n + W16 n _ -> W16 n <$> BB.getWord16be n + W32 n _ -> W32 n <$> BB.getWord32be n + W64 n _ -> W64 n <$> BB.getWord64be n + BS n _ -> BS n <$> BB.getByteString n + LBS n _ -> LBS n <$> BB.getLazyByteString n + IsEmpty -> BB.isEmpty >> return IsEmpty + + +verifyProgram :: Int -> Program -> BB.BitGet Bool +verifyProgram totalLength ps0 = go 0 ps0 + where + go _ [] = return True + go pos (p:ps) = + case p of + Bool x -> check x BB.getBool >> go (pos+1) ps + W8 n x -> check x (BB.getWord8 n) >> go (pos+n) ps + W16 n x -> check x (BB.getWord16be n) >> go (pos+n) ps + W32 n x -> check x (BB.getWord32be n) >> go (pos+n) ps + W64 n x -> check x (BB.getWord64be n) >> go (pos+n) ps + BS n x -> check x (BB.getByteString n) >> go (pos+(8*n)) ps + LBS n x -> check x (BB.getLazyByteString n) >> go (pos+(8*n)) ps + IsEmpty -> do + let expected = pos == totalLength + actual <- BB.isEmpty + if expected == actual + then go pos ps + else error $ "isEmpty returned wrong value, expected " + ++ show expected ++ " but got " ++ show actual + check x g = do + y <- g + if x == y + then return () + else error $ "Roundtrip error: Expected " + ++ show x ++ " but got " ++ show y