diff --git a/binary-strict.cabal b/binary-strict.cabal --- a/binary-strict.cabal +++ b/binary-strict.cabal @@ -1,5 +1,5 @@ name: binary-strict -version: 0.4.1 +version: 0.4.2 license: BSD3 license-file: LICENSE author: Lennart Kolmodin @@ -17,13 +17,15 @@ build-depends: base, containers, array, bytestring>=0.9 stability: provisional tested-with: GHC == 6.8.2 -exposed-modules: Data.Binary.Strict.Get, - Data.Binary.Strict.IncrementalGet, - Data.Binary.Strict.BitGet, - Data.Binary.Strict.Util +exposed-modules: Data.Binary.BitPut Data.Binary.Strict.ByteSet - Data.Binary.Strict.Class + Data.Binary.Strict.BitGet, Data.Binary.Strict.BitUtil + Data.Binary.Strict.Get, + Data.Binary.Strict.Class + Data.Binary.Strict.IncrementalGet, + Data.Binary.Strict.Util + Data.Binary.BitBuilder extensions: CPP, FlexibleContexts hs-source-dirs: src extra-source-files: tests/BitGetTest.hs, src/Data/Binary/Strict/Common.h diff --git a/src/Data/Binary/BitBuilder.hs b/src/Data/Binary/BitBuilder.hs new file mode 100644 --- /dev/null +++ b/src/Data/Binary/BitBuilder.hs @@ -0,0 +1,291 @@ +{-# LANGUAGE CPP #-} +{-# OPTIONS_GHC -fno-warn-name-shadowing #-} + +----------------------------------------------------------------------------- +-- | +-- Module : Data.Binary.BitBuilder +-- Copyright : Lennart Kolmodin, Ross Paterson, Adam Langley +-- License : BSD3-style (see LICENSE) +-- +-- Maintainer : Adam Langley +-- Stability : experimental +-- Portability : portable to Hugs and GHC +-- +-- Efficient construction of lazy bytestrings, bit by bit. +----------------------------------------------------------------------------- + +#if defined(__GLASGOW_HASKELL__) && !defined(__HADDOCK__) +#include "MachDeps.h" +#endif + +module Data.Binary.BitBuilder ( + -- * The Builder type + BitBuilder + , toLazyByteString + + -- * Constructing Builders + , empty + , singleton + , append + , fromByteString -- :: S.ByteString -> BitBuilder + , fromLazyByteString -- :: L.ByteString -> BitBuilder + , fromBits + + -- * Flushing the buffer state + , flush + ) where + +import Foreign +import Data.Monoid +import Data.Word +import qualified Data.ByteString as S +import qualified Data.ByteString.Lazy as L + +#ifdef BYTESTRING_IN_BASE +import Data.ByteString.Base (inlinePerformIO) +import qualified Data.ByteString.Base as S +#else +import Data.ByteString.Internal (inlinePerformIO) +import qualified Data.ByteString.Internal as S +import qualified Data.ByteString.Lazy.Internal as L +#endif + +import Data.Binary.Strict.BitUtil + +#if defined(__GLASGOW_HASKELL__) && !defined(__HADDOCK__) +import GHC.Base +#endif + +------------------------------------------------------------------------ + +-- | A 'BitBuilder' is an efficient way to build lazy 'L.ByteString's. +-- There are several functions for constructing 'BitBuilder's, but only one +-- to inspect them: to extract any data, you have to turn them into lazy +-- 'L.ByteString's using 'toLazyByteString'. +-- +-- Internally, a 'BitBuilder' constructs a lazy 'L.Bytestring' by filling byte +-- arrays piece by piece. As each buffer is filled, it is \'popped\' +-- off, to become a new chunk of the resulting lazy 'L.ByteString'. +-- All this is hidden from the user of the 'BitBuilder'. +-- +-- This is closely based on the Builder monad, but this one deals with +-- single bits at a time. + +newtype BitBuilder = BitBuilder { + -- Invariant (from Data.ByteString.Lazy): + -- The lists include no null ByteStrings. + runBitBuilder :: (Buffer -> [S.ByteString]) -> Buffer -> [S.ByteString] + } + +instance Show BitBuilder where + show = const "" + +instance Monoid BitBuilder where + mempty = empty + mappend = append + +------------------------------------------------------------------------ + +-- | /O(1)./ The empty BitBuilder, satisfying +-- +-- * @'toLazyByteString' 'empty' = 'L.empty'@ +-- +empty :: BitBuilder +empty = BitBuilder id + +-- | /O(1)./ A BitBuilder taking a single bit, satisfying +-- +-- * @'toLazyByteString' ('singleton' b) = 'L.singleton' b@ +-- +singleton :: Bool -> BitBuilder +singleton bit = writeN 1 $ \p phase -> do + byte <- peek p + let mask = complement (0x80 `shiftR` phase) + value = if not bit then 0 else 0x80 `shiftR` phase + poke p $ (byte .&. mask) .|. value +{-# INLINE singleton #-} + +fromByteString :: (S.ByteString, Int) -> BitBuilder +fromByteString (bs, bsPhase) = withPhase f where + f phase + -- the trival case, a no-op + | S.length bs == 0 = empty + | phase == 0 && bsPhase == 0 = + -- if we are the start of a byte, and the bytestring is an exact number + -- of bytes long, we can just include it in our output + flush `append` (BitBuilder $ \k buf -> bs : k buf) + | phase == 0 = + -- if we are at the beginning of a byte, in general, we can still just + -- include most of the string in our output, we just need to handle the + -- partial byte at the end + flush `append` (BitBuilder $ \k buf -> S.init bs : + (runBitBuilder (writeN bsPhase (\p _ -> poke p (S.last bs))) k buf)) + | otherwise = + -- the fully general case. We take the first n bits from the bytestring, + -- phase shift the rest and recurse. + writeN (8 - phase) (mergeByte $ S.head bs) `mappend` fromByteString shiftedBS where + mergeByte nextByte p phase = do + byte <- peek p + let takingBits = 8 - phase + mask = topNBits phase + a = topNBits takingBits .&. nextByte + b = a `shiftR` phase + c = (byte .&. mask) .|. b + poke p c + shiftedBS = (S.take newLength shifted, bsPhase') + shifted = leftShift (8 - phase) bs + oldBitLength = + if bsPhase == 0 + then 8 * S.length bs + else (S.length bs - 1) * 8 + bsPhase + newLength = ((oldBitLength - (8 - phase)) + 7) `div` 8 + bsPhase' = (bsPhase - (8 - phase)) `mod` 8 + +-- | Construct a BitBuilder by taking the bottom n bits of a Bits instance. If +-- the instance has less than n bits, this acts as if there was an infinite +-- zero filled prefix +fromBits :: (Integral a, Bits a) => Int -> a -> BitBuilder +fromBits n v + | n == 0 = empty + | otherwise = writeN n $ f n where + f n p phase = do + let space = 8 - phase + if n <= space + then g p phase v n + else g p phase (v `shiftR` (n - space)) space >> f (n - space) (p `plusPtr` 1) 0 + g p phase v n = do + byte <- peek p + let mask = topNBits phase + bits = ((fromIntegral v) .&. bottomNBits n) `shiftL` ((8 - phase) - n) + poke p $ (byte .&. mask) .|. bits +{-# INLINE fromBits #-} + +------------------------------------------------------------------------ + +-- | /O(1)./ The concatenation of two BitBuilders, an associative operation +-- with identity 'empty', satisfying +-- +-- * @'toLazyByteString' ('append' x y) = 'L.append' ('toLazyByteString' x) ('toLazyByteString' y)@ +-- +append :: BitBuilder -> BitBuilder -> BitBuilder +append (BitBuilder f) (BitBuilder g) = BitBuilder (f . g) + +-- | /O(1)./ A BitBuilder taking a lazy 'L.ByteString', satisfying +-- +-- * @'toLazyByteString' ('fromLazyByteString' bs) = bs@ +-- +fromLazyByteString :: L.ByteString -> BitBuilder +fromLazyByteString = foldr (append . fromByteString . flip (,) 0) empty . L.toChunks + +------------------------------------------------------------------------ + +-- Our internal buffer type +-- The pointer points to the start of the buffer. This never changes +-- for a given buffer. We may 'flush' a partial buffer, in which case +-- the new Buffer has the same pointer, but starts at a different byte +-- offset. +-- +-- The bit offset gives the number of valid bits (from the MSB +-- downwards) in the current byte. This ranges from 0 to 8. The number +-- of used bytes does not include the current byte and the number of +-- bytes left is as if the current byte was empty. +data Buffer = Buffer {-# UNPACK #-} !(ForeignPtr Word8) + {-# UNPACK #-} !Int -- byte offset + {-# UNPACK #-} !Int -- bit offset (0..7) + {-# UNPACK #-} !Int -- used bytes + {-# UNPACK #-} !Int -- bytes left + +------------------------------------------------------------------------ + +-- | /O(n)./ Extract a lazy 'L.ByteString' from a 'BitBuilder'. +-- The construction work takes place if and when the relevant part of +-- the lazy 'L.ByteString' is demanded. +-- +toLazyByteString :: BitBuilder -> L.ByteString +toLazyByteString m = L.fromChunks $ unsafePerformIO $ do + fp <- S.mallocByteString (defaultSize `div` 8) + let buf = Buffer fp 0 0 0 (defaultSize `div` 8) + return (runBitBuilder (m `append` zeroExtendFinalByte `append` flush) (const []) buf) + +-- | /O(1)./ Pop the 'S.ByteString' we have constructed so far, if any, +-- yielding a new chunk in the result lazy 'L.ByteString'. +flush :: BitBuilder +flush = BitBuilder $ \ k buf@(Buffer p bo phase u l) -> + if u == 0 + then k buf + else S.PS p bo u : k (Buffer p (bo+u) phase 0 l) + +------------------------------------------------------------------------ + +-- | The default size of a new chunk, in bits +defaultSize :: Int +defaultSize = 8 * (512 - overhead) where + overhead = 2 * sizeOf (undefined :: Int) + +------------------------------------------------------------------------ + +-- | Sequence an IO operation on the buffer +unsafeLiftIO :: (Buffer -> IO Buffer) -> BitBuilder +unsafeLiftIO f = BitBuilder $ \ k buf -> inlinePerformIO $ do + buf' <- f buf + return (k buf') +{-# INLINE unsafeLiftIO #-} + +-- | Get the size of the buffer, in bits +withSize :: (Int -> BitBuilder) -> BitBuilder +withSize f = BitBuilder $ \ k buf@(Buffer _ _ phase _ l) -> + runBitBuilder (f $ l*8 - phase) k buf + +withPhase :: (Int -> BitBuilder) -> BitBuilder +withPhase f = BitBuilder $ \ k buf@(Buffer _ _ phase _ _) -> + runBitBuilder (f phase) k buf + +------------------------------------------------------------------------ + +zeroExtendFinalByte :: BitBuilder +zeroExtendFinalByte = withPhase $ \phase -> + if phase == 0 + then empty + else writeN (8 - phase) (\p phase -> do + byte <- peek p + poke p $ byte .&. topNBits phase) + +-- | Ensure that there are at least @n@ many bits available. +ensureFree :: Int -> BitBuilder +ensureFree n = n `seq` withSize $ \ l -> + if n <= l then empty else + flush `append` unsafeLiftIO (newBuffer (max n defaultSize)) +{-# INLINE ensureFree #-} + +-- | Ensure that @n@ many bits are available, and then use @f@ to write some +-- bytes into the memory. +writeN :: Int -> (Ptr Word8 -> Int -> IO ()) -> BitBuilder +writeN n f = ensureFree n `append` unsafeLiftIO (writeNBuffer n f) +{-# INLINE [1] writeN #-} + +writeNBuffer :: Int -> (Ptr Word8 -> Int -> IO ()) -> Buffer -> IO Buffer +writeNBuffer n f (Buffer fp bo phase u l) = do + withForeignPtr fp (\p -> f (p `plusPtr` (bo+u)) phase) + let (bytesUsed, phase') = divMod (phase + n) 8 + return (Buffer fp bo phase' (u+bytesUsed) (l-bytesUsed)) +{-# INLINE writeNBuffer #-} + +-- | Create a new buffer of, at least, the given bit size. The current +-- buffer is passed in. If it has a partial byte in progress, +-- that byte is merged in. The given buffer must be in the form as +-- produced by 'flush' +newBuffer :: Int -> Buffer -> IO Buffer +newBuffer size (Buffer p bo phase u _) = + if phase == 0 + then do + let byteSize = (size + 7) `div` 8 + fp <- S.mallocByteString byteSize + return $! Buffer fp 0 0 0 byteSize + else do + let byteSize = (size + 15) `div` 8 + fp <- S.mallocByteString byteSize + withForeignPtr fp (\fp -> + withForeignPtr p (\p -> do + byte <- peek (p `plusPtr` (bo+u)) + poke fp byte)) + return $! Buffer fp 0 phase 0 byteSize diff --git a/src/Data/Binary/BitPut.hs b/src/Data/Binary/BitPut.hs new file mode 100644 --- /dev/null +++ b/src/Data/Binary/BitPut.hs @@ -0,0 +1,82 @@ +----------------------------------------------------------------------------- +-- | +-- Module : Data.Binary.BitPut +-- Copyright : Dominic Steinitz +-- License : BSD3-style (see LICENSE) +-- +-- Maintainer : Dominic Steinitz +-- Stability : experimental +-- +-- This is the writer dual to BitGet. It allows one to append bits in a monad +-- and get a strict ByteString as a result. Bits are appended from the MSB of +-- the first byte towards the LSB of the last byte. +-- +-- This is best suited to small bit-fields because it accumulates bytes using +-- snoc, so large results will cause a lot of copying. It would be possible +-- to switch to using something similar to the Builder monad if need arises. +-- However, since most protocols only have small bit fields, this should +-- suffice for many cases. +----------------------------------------------------------------------------- +module Data.Binary.BitPut + ( BitPut + , runBitPut + , putBit + , putNBits + , putBits + , putByteString + , putLeftByteString + ) where + +import Data.Bits (bitSize, Bits) +import Control.Monad +import qualified Data.ByteString as B +import qualified Data.ByteString.Lazy as BL +import qualified Data.Binary.BitBuilder as BB + +newtype BitPut' a = BitPut' { unPut :: (a, BB.BitBuilder) } + +type BitPut = BitPut' () + +instance Functor BitPut' where + fmap f m = BitPut' (let (a, w) = unPut m in (f a, w)) + +instance Monad BitPut' where + return a = BitPut' (a,BB.empty) + m >>= k = BitPut' (let (a, w) = unPut m + (b, w') = unPut (k a) + in (b, w `BB.append` w')) + + m >> k = BitPut' (let (_, w) = unPut m + (b, w') = unPut k + in (b, w `BB.append` w')) + {-# INLINE (>>) #-} + +-- | Append a single bit +putBit :: Bool -> BitPut +putBit bit = BitPut' ((), BB.singleton bit) + +-- | Append the bottom n bits of the given bits value. In the case that more +-- bits are requested than the value provides, this acts as if the value +-- has as unlimited number of leading 0 bits. +putNBits :: (Integral a, Bits a) => Int -> a -> BitPut +putNBits n v = BitPut' ((), BB.fromBits n v) + +-- | Append a value. Note that this function is undefined for instances of Bits +-- which have no fixed bitsize (like Integer) +putBits :: (Integral a, Bits a) => a -> BitPut +putBits v = putNBits (bitSize v) v + +-- | Append a ByteString +putByteString :: B.ByteString -> BitPut +putByteString bs = BitPut' ((), BB.fromByteString (bs, 0)) + +-- | Append a left aligned ByteString where ByteString has a partial byte +-- with the given number of valid bits, from the MSB downwards. The number +-- of such bits must be 0..7. (A normal ByteString, which all bytes full +-- would use 0) +putLeftByteString :: (B.ByteString, Int) -> BitPut +putLeftByteString bs = BitPut' ((), BB.fromByteString bs) + +runBitPut :: BitPut -> BL.ByteString +runBitPut m = let (_, w) = unPut m + in BB.toLazyByteString w