@@ -55,7 +55,9 @@ module Data.Binary.Bits.Get
5555 -- $bitget
5656
5757 BitGet
58+ , BitOrder(..)
5859 , runBitGet
60+ , selectBitOrder
5961
6062 -- ** Get bytes
6163 , getBool
@@ -158,33 +160,10 @@ import GHC.Word
158160-- @
159161
160162data S = S {-# UNPACK #-} !ByteString -- Input
161− {-# UNPACK #-} !Int -- Bit offset (0-7)
163+ {-# UNPACK #-} !Int -- Bit offset (0-7)
164+ !BitOrder -- Bit order
162165 deriving (Show)
163166
164−-- | A block that will be read with only one boundary check. Needs to know the
165−-- number of bits in advance.
166−data Block a = Block Int (S -> a)
167−
168−instance Functor Block where
169− fmap f (Block i p) = Block i (\s -> f (p s))
170−
171−instance Applicative Block where
172− pure a = Block 0 (\_ -> a)
173− (Block i p) <*> (Block j q) = Block (i+j) (\s -> p s $ q (incS i s))
174− (Block i _) *> (Block j q) = Block (i+j) (q . incS i)
175− (Block i p) <* (Block j _) = Block (i+j) p
176−
177−-- | Get a block. Will be read with one single boundry check, and
178−-- therefore requires a statically known number of bits.
179−-- Build blocks using 'bool', 'word8', 'word16be', 'word32be', 'word64be',
180−-- 'byteString' and 'Applicative'.
181−block :: Block a -> BitGet a
182−block (Block i p) = do
183− ensureBits i
184− s <- getState
185− putState $! (incS i s)
186− return $! p s
187−
188167-- | Compute bit offset (equivalent to x `mod` 8 but faster)
189168bit_offset :: Int -> Int
190169bit_offset n = make_mask 3 .&. n
@@ -197,12 +176,16 @@ byte_offset n = n `shiftR` 3
197176
198177-- | Increment the current bit offset
199178incS :: Int -> S -> S
200−incS o (S bs n) = S (unsafeDrop d bs) n'
179+incS o (S bs n bo) = S (unsafeDrop d bs) n' bo
201180 where
202181 !o' = (n+o)
203182 !d = byte_offset o'
204183 !n' = bit_offset o'
205184
185+-- | Select the bit ordering
186+selectBitOrderS :: BitOrder -> S -> S
187+selectBitOrderS bo (S bs o _) = S bs o bo
188+
206189-- | make_mask 3 = 00000111
207190make_mask :: (Bits a, Num a) => Int -> a
208191make_mask n = (1 `shiftL` fromIntegral n) - 1
@@ -218,33 +201,6 @@ mask :: (Bits a, Num a) => Int -> a -> a
218201mask n v = v .&. make_mask n
219202{-# INLINE mask #-}
220203
221−class FastBits a where
222− fastShiftR :: a -> Int -> a
223− fastShiftL :: a -> Int -> a
224−
225− {-# INLINE fastShift #-}
226− fastShift :: a -> Int -> a
227− fastShift x n
228− | n > 0 = fastShiftL x n
229− | n < 0 = fastShiftR x (negate n)
230− | otherwise = x
231−
232−instance FastBits Word8 where
233− fastShiftR = shiftr_w8
234− fastShiftL = shiftl_w8
235−
236−instance FastBits Word16 where
237− fastShiftR = shiftr_w16
238− fastShiftL = shiftl_w16
239−
240−instance FastBits Word32 where
241− fastShiftR = shiftr_w32
242− fastShiftL = shiftl_w32
243−
244−instance FastBits Word64 where
245− fastShiftR = shiftr_w64
246− fastShiftL = shiftl_w64
247−
248204-- | Bit order
249205--
250206-- E.g. two words of 5 bits: ABCDE, VWXYZ
@@ -252,13 +208,14 @@ instance FastBits Word64 where
252208-- - LB: XYZABCDE xxxxxxVW
253209-- - BL: EDCBAZYX WVxxxxxx -- not implemented
254210-- - LL: XWVEDCBA xxxxxxZY -- not implemented
255−data BitOrder = BB | LB
211+data BitOrder = BB | LB deriving (Show)
256212
257213
258214-- | Read a single bit
259−readBool :: BitOrder -> S -> Bool
260−readBool BB (S bs n) = testBit (unsafeHead bs) (7-n)
261−readBool LB (S bs n) = testBit (unsafeHead bs) n
215+readBool :: S -> Bool
216+readBool (S bs o bo) = case bo of
217+ BB -> testBit (unsafeHead bs) (7-o)
218+ LB -> testBit (unsafeHead bs) o
262219
263220-- | Extract a range of bits from (ws :: ByteString)
264221--
@@ -269,9 +226,9 @@ extract bo bs o n
269226 | B.length bs == 0 = error "Empty ByteString"
270227 | otherwise = mask n (foldlWithIndex' f 0 bs)
271228 where
272− -- foldl' with index
273− foldlWithIndex' f b = fst . B.foldl' g (b,0)
274− where g (b,i) w = (f b w i, (i+1))
229+ -- B.foldl' with index
230+ foldlWithIndex' op b = fst . B.foldl' g (b,0)
231+ where g (b',i) w = (op b' w i, (i+1))
275232
276233 -- 'or' correctly shifted words
277234 f b w i = b .|. (fromIntegral w `fastShift` off i)
@@ -286,8 +243,8 @@ extract bo bs o n
286243
287244
288245-- | Generic readWord
289−readWord :: (Num a, Bits a, FastBits a) => BitOrder -> Int -> S -> a
290−readWord bo n (S bs o)
246+readWord :: (Num a, Bits a, FastBits a) => Int -> S -> a
247+readWord n (S bs o bo)
291248 | n == 0 = 0
292249 | o+n <= 8 = extract bo (unsafeTake 1 bs) o n
293250 | o+n <= 16 = extract bo (unsafeTake 2 bs) o n
@@ -299,30 +256,20 @@ readWord bo n (S bs o)
299256 | o+n <= 64 = extract bo (unsafeTake 8 bs) o n
300257 | otherwise = extract bo (unsafeTake ((o+n+7) `div` 8) bs) o n
301258
302−readWordChecked :: (Num a, Bits a, FastBits a) => Int -> BitOrder -> Int -> S -> a
303−readWordChecked m bo n s
259+-- | Check that the number of bits to read is not greater than the first parameter
260+{-# INLINE readWordChecked #-}
261+readWordChecked :: (Num a, Bits a, FastBits a) => Int -> Int -> S -> a
262+readWordChecked m n s
304263 | n > m = error $ "Tried to read more than " ++ show m ++ " bits (" ++ show n ++")"
305− | otherwise = readWord bo n s
306−
307−{-# INLINE readWord8 #-}
308−readWord8 :: BitOrder -> Int -> S -> Word8
309−readWord8 = readWordChecked 8
310−
311−{-# INLINE readWord16 #-}
312−readWord16 :: BitOrder -> Int -> S -> Word16
313−readWord16 = readWordChecked 16
314−
315−{-# INLINE readWord32 #-}
316−readWord32 :: BitOrder -> Int -> S -> Word32
317−readWord32 = readWordChecked 32
318−
319−{-# INLINE readWord64 #-}
320−readWord64 :: BitOrder -> Int -> S -> Word64
321−readWord64 = readWordChecked 64
264+ | otherwise = readWord n s
322265
323266-- | Read the given number of bytes and return them in Big-Endian order
324−readByteString :: BitOrder -> Int -> S -> ByteString
325−readByteString bo n (S bs o)
267+--
268+-- Examples:
269+-- BB: xxxABCDE FGHIJKLM NOPQRxxx -> ABCDEFGH IJKLMNOP QRxxxxxx
270+-- LB: NOPQRxxx FGHIJKLM xxxABCDE -> ABCDEFGH IJKLMNOP QRxxxxxx
271+readByteString :: Int -> S -> ByteString
272+readByteString n (S bs o bo)
326273 | o == 0 = unsafeTake n bs
327274 | otherwise = unsafePerformIO $ do
328275 let len = n+1
@@ -357,7 +304,7 @@ newtype BitGet a = B { runState :: S -> Get (S,a) }
357304
358305instance Monad BitGet where
359306 return x = B $ \s -> return (s,x)
360− fail str = B $ \(S inp n) -> putBackState inp n >> fail str
307+ fail str = B $ \s -> putBackState s >> fail str
361308 (B f) >>= g = B $ \s -> do (s',a) <- f s
362309 runState (g a) s'
363310
@@ -373,20 +320,23 @@ instance Applicative BitGet where
373320runBitGet :: BitGet a -> Get a
374321runBitGet bg = do
375322 s <- mkInitState
376− ((S str' n),a) <- runState bg s
377− putBackState str' n
323+ (s',a) <- runState bg s
324+ putBackState s'
378325 return a
379326
380327mkInitState :: Get S
381328mkInitState = do
382− str <- get
329+ bs <- get
383330 put B.empty
384− return (S str 0)
331+ return (S bs 0 BB)
385332
386−putBackState :: B.ByteString -> Int -> Get ()
387−putBackState bs n = do
333+putBackState :: S -> Get ()
334+putBackState (S bs o _) = do
388335 remaining <- get
389− put (B.drop (if n==0 then 0 else 1) bs `B.append` remaining)
336+ let bs' = case o of
337+ 0 -> bs
338+ _ -> unsafeDrop 1 bs
339+ put (bs' `B.append` remaining)
390340
391341getState :: BitGet S
392342getState = B $ \s -> return (s,s)
@@ -394,32 +344,47 @@ getState = B $ \s -> return (s,s)
394344putState :: S -> BitGet ()
395345putState s = B $ \_ -> return (s,())
396346
347+withState :: (S -> S) -> BitGet ()
348+withState f = do
349+ s <- getState
350+ putState $! f s
351+
397352-- | Make sure there are at least @n@ bits.
398353ensureBits :: Int -> BitGet ()
399354ensureBits n = do
400− (S bs o) <- getState
355+ (S bs o bo) <- getState
401356 if n <= (B.length bs * 8 - o)
402357 then return ()
403358 else do let currentBits = B.length bs * 8 - o
404− let byteCount = (n - currentBits + 7) `div` 8
359+ let byteCount = byte_offset (n - currentBits + 7)
405360 B $ \_ -> do B.ensureN byteCount
406361 bs' <- B.get
407362 put B.empty
408− return (S (bs`append`bs') o, ())
363+ return (S (bs`append`bs') o bo, ())
409364
410365-- | Skip the given number of bits
411366skipBits :: Int -> BitGet ()
412367skipBits n = do
413368 ensureBits n
414− s <- getState
415− putState $! (incS n s)
369+ withState (incS n)
416370
417371-- | Skip bits if necessary to align to the next byte
418372alignByte :: BitGet ()
419373alignByte = do
420− (S _ n) <- getState
421− when (n /= 0) $
422− skipBits (8-n)
374+ (S _ o _) <- getState
375+ when (o /= 0) $
376+ skipBits (8-o)
377+
378+-- | Select the bit ordering
379+selectBitOrder :: BitOrder -> BitGet ()
380+selectBitOrder = withState . selectBitOrderS
381+
382+-- | Test whether all input has been consumed, i.e. there are no remaining
383+-- undecoded bytes.
384+isEmpty :: BitGet Bool
385+isEmpty = B $ \ (S bs o bo) -> if B.null bs
386+ then B.isEmpty >>= \e -> return (S bs o bo, e)
387+ else return (S bs o bo, False)
423388
424389-- | Get 1 bit as a 'Bool'.
425390getBool :: BitGet Bool
@@ -448,48 +413,99 @@ getByteString n = block (byteString n)
448413-- | Get @n@ bytes as a lazy ByteString.
449414getLazyByteString :: Int -> BitGet L.ByteString
450415getLazyByteString n = do
451− (S _ o) <- getState
416+ (S _ o bo) <- getState
452417 case o of
453− 0 -> B $ \ (S bs o') -> do
454− putBackState bs o'
418+ 0 -> B $ \s -> do
419+ putBackState s
455420 lbs <- B.getLazyByteString (fromIntegral n)
456− return (S B.empty 0, lbs)
421+ return (S B.empty 0 bo, lbs)
457422 _ -> L.fromChunks . (:[]) <$> Data.Binary.Bits.Get.getByteString n
458423
459−-- | Test whether all input has been consumed, i.e. there are no remaining
460−-- undecoded bytes.
461−isEmpty :: BitGet Bool
462−isEmpty = B $ \ (S bs o) -> if B.null bs
463− then B.isEmpty >>= \e -> return (S bs o, e)
464− else return (S bs o, False)
424+
425+
426+
427+
428+-- | A block that will be read with only one boundary check. Needs to know the
429+-- number of bits in advance.
430+data Block a = Block Int (S -> a)
431+
432+instance Functor Block where
433+ fmap f (Block i p) = Block i (\s -> f (p s))
434+
435+instance Applicative Block where
436+ pure a = Block 0 (\_ -> a)
437+ (Block i p) <*> (Block j q) = Block (i+j) (\s -> p s $ q (incS i s))
438+ (Block i _) *> (Block j q) = Block (i+j) (q . incS i)
439+ (Block i p) <* (Block j _) = Block (i+j) p
440+
441+-- | Get a block. Will be read with one single boundry check, and
442+-- therefore requires a statically known number of bits.
443+-- Build blocks using 'bool', 'word8', 'word16be', 'word32be', 'word64be',
444+-- 'byteString' and 'Applicative'.
445+block :: Block a -> BitGet a
446+block (Block i p) = do
447+ ensureBits i
448+ s <- getState
449+ putState $! (incS i s)
450+ return $! p s
465451
466452-- | Read a 1 bit 'Bool'.
467453bool :: Block Bool
468−bool = Block 1 (readBool BB)
454+bool = Block 1 readBool
469455
470456-- | Read @n@ bits as a 'Word8'. @n@ must be within @[0..8]@.
471457word8 :: Int -> Block Word8
472−word8 n = Block n (readWord8 BB n)
458+word8 n = Block n (readWordChecked 8 n)
473459
474460-- | Read @n@ bits as a 'Word16'. @n@ must be within @[0..16]@.
475461word16be :: Int -> Block Word16
476−word16be n = Block n (readWord16 BB n)
462+word16be n = Block n (readWordChecked 16 n)
477463
478464-- | Read @n@ bits as a 'Word32'. @n@ must be within @[0..32]@.
479465word32be :: Int -> Block Word32
480−word32be n = Block n (readWord32 BB n)
466+word32be n = Block n (readWordChecked 32 n)
481467
482468-- | Read @n@ bits as a 'Word64'. @n@ must be within @[0..64]@.
483469word64be :: Int -> Block Word64
484−word64be n = Block n (readWord64 BB n)
470+word64be n = Block n (readWordChecked 64 n)
485471
486472-- | Read @n@ bytes as a 'ByteString'.
487473byteString :: Int -> Block ByteString
488−byteString n | n > 0 = Block (n*8) (readByteString BB n)
474+byteString n | n > 0 = Block (n*8) (readByteString n)
489475 | otherwise = Block 0 (\_ -> B.empty)
490476
491−------------------------------------------------------------------------
492−-- Unchecked shifts, from the package binary
477+
478+---------------------------------------------------------------------
479+-- Unchecked shifts, from the "binary" package
480+
481+-- | Class for types supporting fast bit shifting
482+class FastBits a where
483+ fastShiftR :: a -> Int -> a
484+ fastShiftL :: a -> Int -> a
485+
486+ {-# INLINE fastShift #-}
487+ fastShift :: a -> Int -> a
488+ fastShift x n
489+ | n > 0 = fastShiftL x n
490+ | n < 0 = fastShiftR x (negate n)
491+ | otherwise = x
492+
493+instance FastBits Word8 where
494+ fastShiftR = shiftr_w8
495+ fastShiftL = shiftl_w8
496+
497+instance FastBits Word16 where
498+ fastShiftR = shiftr_w16
499+ fastShiftL = shiftl_w16
500+
501+instance FastBits Word32 where
502+ fastShiftR = shiftr_w32
503+ fastShiftL = shiftl_w32
504+
505+instance FastBits Word64 where
506+ fastShiftR = shiftr_w64
507+ fastShiftL = shiftl_w64
508+
493509
494510shiftl_w8 :: Word8 -> Int -> Word8
495511shiftl_w16 :: Word16 -> Int -> Word16