@@ -1,3 +1,5 @@
1+{-# LANGUAGE MagicHash #-}
1-----------------------------------------------------------------------------
3-----------------------------------------------------------------------------
3-- Module : Data.Binary.Bits.Put
5-- Module : Data.Binary.Bits.Put
@@ -20,8 +22,10 @@ module Data.Binary.Bits.Put
20 -- * Data types
22 -- * Data types
21 -- ** Bool
23 -- ** Bool
24 -- ** Words
27 -- ** Words
25 , putWord8
29 , putWord8
26 , putWord16
30 , putWord16
27 , putWord32
31 , putWord32
@@ -41,14 +45,16 @@ import qualified Data.Binary.Put as Put
41import Data.Binary.Put ( Put )
45import Data.Binary.Put ( Put )
42import Data.Binary.Bits.Internal
46import Data.Binary.Bits.Internal
43import Data.Binary.Bits.BitOrder
47import Data.Binary.Bits.BitOrder
48+import Data.Binary.Bits.Alignment
45import Data.ByteString as BS
50import Data.ByteString as BS
46import Data.ByteString.Unsafe as BS
51import Data.ByteString.Unsafe as BS
48−import Control.Applicative
53+import Control.Monad (when)
49import Data.Bits
54import Data.Bits
51import Data.Word
55import Data.Word
56+import GHC.Base (Word#, shiftRL#, Int (..), Int#, (-#), (<#), isTrue#, (+#), shiftL#, wordToWord8#, orWord8#)
57+import GHC.Word (Word8(..))
53data BitPut a = BitPut { run :: (S -> PairS a) }
59data BitPut a = BitPut { run :: (S -> PairS a) }
@@ -60,11 +66,18 @@ data S = S !Builder !Word8 !Int !BitOrder
60putBool :: Bool -> BitPut ()
66putBool :: Bool -> BitPut ()
61putBool b = putWord8 1 (if b then 0xff else 0x00)
67putBool b = putWord8 1 (if b then 0xff else 0x00)
69+-- | Put a 1 bit 'Bool'.
70+putBool# :: Word# -> BitPut ()
71+putBool# b = putWord# 1# b
64-- | Generic putWord
74-- | Generic putWord
65putWord :: (Num a, FastBits a, Integral a) => Int -> a -> BitPut ()
75putWord :: (Num a, FastBits a, Integral a) => Int -> a -> BitPut ()
66putWord n w = BitPut $ \s -> PairS () (putWordS n w s)
76putWord n w = BitPut $ \s -> PairS () (putWordS n w s)
78+putWord# :: Int# -> Word# -> BitPut ()
79+putWord# n w = BitPut $ \s -> PairS () (putWordS# n w s)
68putWordS :: (Num a, FastBits a, Integral a) => Int -> a -> S -> S
81putWordS :: (Num a, FastBits a, Integral a) => Int -> a -> S -> S
69putWordS n w s@(S builder b o bo) = s'
82putWordS n w s@(S builder b o bo) = s'
@@ -83,8 +96,8 @@ putWordS n w s@(S builder b o bo) = s'
83 w' = case bo of
96 w' = case bo of
85 BL -> w `fastShiftR` cn
98 BL -> w `fastShiftR` cn
86− LB -> w `fastShiftR` cn
99+ LL -> w `fastShiftR` cn
89 -- Select bits to store in the current byte.
102 -- Select bits to store in the current byte.
90 -- Put them in the correct order and return them in the least-significant
103 -- Put them in the correct order and return them in the least-significant
@@ -92,8 +105,8 @@ putWordS n w s@(S builder b o bo) = s'
92 selectBits :: (Num a, FastBits a, Integral a) => a -> Word8
105 selectBits :: (Num a, FastBits a, Integral a) => a -> Word8
93 selectBits x = fromIntegral $ case bo of
106 selectBits x = fromIntegral $ case bo of
94 BB -> mask cn $ x `fastShiftR` (n-cn)
107 BB -> mask cn $ x `fastShiftR` (n-cn)
95− LL -> reverseBits cn $ mask cn $ x `fastShiftR` (n-cn)
108+ LB -> reverseBits cn $ mask cn $ x `fastShiftR` (n-cn)
96− LB -> mask cn x
109+ LL -> mask cn x
97 BL -> reverseBits cn $ mask cn x
110 BL -> reverseBits cn $ mask cn x
99 -- shift left at the correct position
112 -- shift left at the correct position
@@ -101,8 +114,55 @@ putWordS n w s@(S builder b o bo) = s'
101 shl x = case bo of
114 shl x = case bo of
102 BB -> x `fastShiftL` (8-o-cn)
115 BB -> x `fastShiftL` (8-o-cn)
103 BL -> x `fastShiftL` (8-o-cn)
116 BL -> x `fastShiftL` (8-o-cn)
104− LB -> x `fastShiftL` o
105 LL -> x `fastShiftL` o
117 LL -> x `fastShiftL` o
118+ LB -> x `fastShiftL` o
120+ flush s2@(S b2 w2 o2 bo2)
121+ | o2 == 8 = S (b2 `mappend` B.singleton w2) 0 0 bo2
124+putWordS# :: Int# -> Word# -> S -> S
125+putWordS# n w s@(S builder (W8# b) (I# o) bo) = s'
127+ min# :: Int# -> Int# -> Int#
128+ min# x y | isTrue# (x <# y) = x
131+ -- number of bits that will be stored in the current byte
132+ cn = min# (8# -# o) n
137+ _ -> putWordS# (n -# cn) w' (flush (S builder (W8# b') (I# (o +# cn)) bo))
140+ b' = wordToWord8# (shl (selectBits w)) `orWord8#` b
142+ -- Word containing the remaining (n-cn) bits to store in its LSB
145+ BL -> w `shiftRL#` cn
146+ LL -> w `shiftRL#` cn
149+ -- Select bits to store in the current byte.
150+ -- Put them in the correct order and return them in the least-significant
151+ -- bits of the returned value
152+ selectBits :: Word# -> Word#
153+ selectBits x = case bo of
154+ BB -> mask# cn (x `shiftRL#` (n -# cn))
155+ LB -> reverseBits# cn (mask# cn (x `shiftRL#` (n -# cn)))
157+ BL -> reverseBits# cn (mask# cn x)
159+ -- shift left at the correct position
160+ shl :: Word# -> Word#
162+ BB -> x `shiftL#` (8# -# o -# cn)
163+ BL -> x `shiftL#` (8# -# o -# cn)
107 flush s2@(S b2 w2 o2 bo2)
167 flush s2@(S b2 w2 o2 bo2)
108 | o2 == 8 = S (b2 `mappend` B.singleton w2) 0 0 bo2
168 | o2 == 8 = S (b2 `mappend` B.singleton w2) 0 0 bo2
@@ -141,9 +201,9 @@ putWord64be = putWord
142-- Examples: 3 bits are already written in the current byte
202-- Examples: 3 bits are already written in the current byte
143-- BB: ABCDEFGH IJKLMNOP -> xxxABCDE FGHIJKLM NOPxxxxx
203-- BB: ABCDEFGH IJKLMNOP -> xxxABCDE FGHIJKLM NOPxxxxx
144−-- LB: ABCDEFGH IJKLMNOP -> LMNOPxxx DEFGHIJK xxxxxABC
204+-- LL: ABCDEFGH IJKLMNOP -> LMNOPxxx DEFGHIJK xxxxxABC
145-- BL: ABCDEFGH IJKLMNOP -> xxxPONML KJIHGFED CBAxxxxx
205-- BL: ABCDEFGH IJKLMNOP -> xxxPONML KJIHGFED CBAxxxxx
146−-- LL: ABCDEFGH IJKLMNOP -> EDCBAxxx MLKJIHGF xxxxxPON
206+-- LB: ABCDEFGH IJKLMNOP -> EDCBAxxx MLKJIHGF xxxxxPON
147putByteString :: ByteString -> BitPut ()
207putByteString :: ByteString -> BitPut ()
148putByteString bs = BitPut $ \s -> PairS () (putByteStringS bs s)
208putByteString bs = BitPut $ \s -> PairS () (putByteStringS bs s)
@@ -152,13 +212,13 @@ putByteStringS bs s
152 | BS.null bs = s
212 | BS.null bs = s
153 | otherwise = case s of
213 | otherwise = case s of
154 (S builder b 0 BB) -> S (builder `mappend` B.fromByteString bs) b 0 BB
214 (S builder b 0 BB) -> S (builder `mappend` B.fromByteString bs) b 0 BB
155− (S builder b 0 LB) -> S (builder `mappend` B.fromByteString (BS.reverse bs)) b 0 LB
215+ (S builder b 0 LL) -> S (builder `mappend` B.fromByteString (BS.reverse bs)) b 0 LL
156− (S builder b 0 LL) -> S (builder `mappend` B.fromByteString (rev bs)) b 0 LL
216+ (S builder b 0 LB) -> S (builder `mappend` B.fromByteString (rev bs)) b 0 LB
157 (S builder b 0 BL) -> S (builder `mappend` B.fromByteString (rev (BS.reverse bs))) b 0 BL
217 (S builder b 0 BL) -> S (builder `mappend` B.fromByteString (rev (BS.reverse bs))) b 0 BL
158 (S _ _ _ BB) -> putByteStringS (BS.unsafeTail bs) (putWordS 8 (BS.unsafeHead bs) s)
218 (S _ _ _ BB) -> putByteStringS (BS.unsafeTail bs) (putWordS 8 (BS.unsafeHead bs) s)
159− (S _ _ _ LB) -> putByteStringS (BS.unsafeInit bs) (putWordS 8 (BS.unsafeLast bs) s)
219+ (S _ _ _ LL) -> putByteStringS (BS.unsafeInit bs) (putWordS 8 (BS.unsafeLast bs) s)
160 (S _ _ _ BL) -> putByteStringS (BS.unsafeInit bs) (putWordS 8 (BS.unsafeLast bs) s)
220 (S _ _ _ BL) -> putByteStringS (BS.unsafeInit bs) (putWordS 8 (BS.unsafeLast bs) s)
161− (S _ _ _ LL) -> putByteStringS (BS.unsafeTail bs) (putWordS 8 (BS.unsafeHead bs) s)
221+ (S _ _ _ LB) -> putByteStringS (BS.unsafeTail bs) (putWordS 8 (BS.unsafeHead bs) s)
163 rev = BS.map (reverseBits 8)
223 rev = BS.map (reverseBits 8)
@@ -178,6 +238,9 @@ flushIncomplete s@(S b w o bo)
178 | o == 0 = s
238 | o == 0 = s
179 | otherwise = (S (b `mappend` B.singleton w) 0 0 bo)
239 | otherwise = (S (b `mappend` B.singleton w) 0 0 bo)
241+getOffset :: BitPut Int
242+getOffset = BitPut $ \s@(S _ _ o _) -> PairS o s
181-- | Run the 'BitPut' monad inside 'Put'.
244-- | Run the 'BitPut' monad inside 'Put'.
182runBitPut :: BitPut () -> Put.Put
245runBitPut :: BitPut () -> Put.Put
183runBitPut m = Put.putBuilder b
246runBitPut m = Put.putBuilder b
@@ -208,3 +271,15 @@ instance BitOrderable BitPut where
208 setBitOrder bo = BitPut $ \(S bu b o _) -> PairS () (S bu b o bo)
271 setBitOrder bo = BitPut $ \(S bu b o _) -> PairS () (S bu b o bo)
210 getBitOrder = BitPut $ \s@(S _ _ _ bo) -> PairS bo s
273 getBitOrder = BitPut $ \s@(S _ _ _ bo) -> PairS bo s
275+instance Alignable BitPut where
276+ -- | Skip the given number of bits
278+ | n <= 64 = putWord64 n 0
279+ | otherwise = putWord64 64 0 >> skipBits (n-64)
281+ -- | Skip bits if necessary to align to the next byte