@@ -93,11 +93,15 @@ import Data.Binary.Get.Internal as B ( get, put, ensureN )
9393import Data.ByteString as B
9494import qualified Data.ByteString.Lazy as L
9595import Data.ByteString.Unsafe
96+import Foreign.Marshal.Alloc (mallocBytes)
97+import Foreign.Ptr
98+import Foreign.Storable (poke)
99+import System.IO.Unsafe (unsafePerformIO)
96100
97101import Data.Bits
98102import Data.Word
99103import Control.Applicative
100−import Control.Monad (when)
104+import Control.Monad (when,foldM_)
101105
102106import Prelude as P
103107
@@ -157,7 +161,7 @@ data S = S {-# UNPACK #-} !ByteString -- Input
157161 {-# UNPACK #-} !Int -- Bit offset (0-7)
158162 deriving (Show)
159163
160−-- | A block that will be read with only one boundry check. Needs to know the
164+-- | A block that will be read with only one boundary check. Needs to know the
161165-- number of bits in advance.
162166data Block a = Block Int (S -> a)
163167
@@ -181,12 +185,23 @@ block (Block i p) = do
181185 putState $! (incS i s)
182186 return $! p s
183187
188+-- | Compute bit offset (equivalent to x `mod` 8 but faster)
189+bit_offset :: Int -> Int
190+bit_offset n = make_mask 3 .&. n
191+{-# INLINE bit_offset #-}
192+
193+-- | Compute byte offset (equivalent to x `div` 8 but faster)
194+byte_offset :: Int -> Int
195+byte_offset n = n `shiftR` 3
196+{-# INLINE byte_offset #-}
197+
198+-- | Increment the current bit offset
184199incS :: Int -> S -> S
185−incS o (S bs n) =
186− let !o' = (n+o)
187− !d = o' `shiftR` 3
188− !n' = o' .&. make_mask 3
189− in S (unsafeDrop d bs) n'
200+incS o (S bs n) = S (unsafeDrop d bs) n'
201+ where
202+ !o' = (n+o)
203+ !d = byte_offset o'
204+ !n' = bit_offset o'
190205
191206-- | make_mask 3 = 00000111
192207make_mask :: (Bits a, Num a) => Int -> a
@@ -198,151 +213,142 @@ make_mask n = (1 `shiftL` fromIntegral n) - 1
198213{-# SPECIALIZE make_mask :: Int -> Word32 #-}
199214{-# SPECIALIZE make_mask :: Int -> Word64 #-}
200215
201−bit_offset :: Int -> Int
202−bit_offset n = make_mask 3 .&. n
203−
204−byte_offset :: Int -> Int
205−byte_offset n = n `shiftR` 3
206−
207−readBool :: S -> Bool
208−readBool (S bs n) = testBit (unsafeHead bs) (7-n)
209−
210−{-# INLINE readWord8 #-}
211−readWord8 :: Int -> S -> Word8
212−readWord8 n (S bs o)
213− -- no bits at all, return 0
214− | n == 0 = 0
215−
216− -- all bits are in the same byte
217− -- we just need to shift and mask them right
218− | n <= 8 - o = let w = unsafeHead bs
219− m = make_mask n
220− w' = (w `shiftr_w8` (8 - o - n)) .&. m
221− in w'
222−
223− -- the bits are in two different bytes
224− -- make a word16 using both bytes, and then shift and mask
225− -- n <= 8
226− | otherwise = let
227− w = (fromIntegral (unsafeHead bs) `shiftl_w16` 8) .|.
228− (fromIntegral (unsafeIndex bs 1))
229− m = make_mask n
230− w' = (w `shiftr_w16` (16 - o - n)) .&. m
231− in fromIntegral w'
232−
233−{-# INLINE readWord16be #-}
234−readWord16be :: Int -> S -> Word16
235−readWord16be n s@(S bs o)
236−
237− -- 8 or fewer bits, use readWord8
238− | n <= 8 = fromIntegral (readWord8 n s)
239−
240− -- handle 9 or more bits, stored in two bytes
241−
242− -- no offset, plain and simple 16 bytes
243− | o == 0 && n == 16 = let msb = fromIntegral (unsafeHead bs)
244− lsb = fromIntegral (unsafeIndex bs 1)
245− w = (msb `shiftl_w16` 8) .|. lsb
246− in w
247−
248− -- no offset, but not full 16 bytes
249− | o == 0 = let msb = fromIntegral (unsafeHead bs)
250− lsb = fromIntegral (unsafeIndex bs 1)
251− w = (msb `shiftl_w16` (n-8)) .|. (lsb `shiftr_w16` (16-n))
252− in w
216+-- | Keep only the n least-significant bits of the given value
217+mask :: (Bits a, Num a) => Int -> a -> a
218+mask n v = v .&. make_mask n
219+{-# INLINE mask #-}
253220
254− -- with offset, and n=9-16
255− | n <= 16 = readWithOffset s shiftl_w16 shiftr_w16 n
221+class FastBits a where
222+ fastShiftR :: a -> Int -> a
223+ fastShiftL :: a -> Int -> a
256224
257− | otherwise = error "readWord16be: tried to read more than 16 bits"
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
258231
259−{-# INLINE readWord32be #-}
260−readWord32be :: Int -> S -> Word32
261−readWord32be n s@(S _ o)
262− -- 8 or fewer bits, use readWord8
263− | n <= 8 = fromIntegral (readWord8 n s)
232+instance FastBits Word8 where
233+ fastShiftR = shiftr_w8
234+ fastShiftL = shiftl_w8
264235
265− -- 16 or fewer bits, use readWord16be
266− | n <= 16 = fromIntegral (readWord16be n s)
236+instance FastBits Word16 where
237+ fastShiftR = shiftr_w16
238+ fastShiftL = shiftl_w16
267239
268− | o == 0 = readWithoutOffset s shiftl_w32 shiftr_w32 n
240+instance FastBits Word32 where
241+ fastShiftR = shiftr_w32
242+ fastShiftL = shiftl_w32
269243
270− | n <= 32 = readWithOffset s shiftl_w32 shiftr_w32 n
244+instance FastBits Word64 where
245+ fastShiftR = shiftr_w64
246+ fastShiftL = shiftl_w64
271247
272− | otherwise = error "readWord32be: tried to read more than 32 bits"
273−
274−
275−{-# INLINE readWord64be #-}
276−readWord64be :: Int -> S -> Word64
277−readWord64be n s@(S _ o)
278− -- 8 or fewer bits, use readWord8
279− | n <= 8 = fromIntegral (readWord8 n s)
280−
281− -- 16 or fewer bits, use readWord16be
282− | n <= 16 = fromIntegral (readWord16be n s)
283−
284− | o == 0 = readWithoutOffset s shiftl_w64 shiftr_w64 n
285−
286− | n <= 64 = readWithOffset s shiftl_w64 shiftr_w64 n
287−
288− | otherwise = error "readWord64be: tried to read more than 64 bits"
289−
290−
291−readByteString :: Int -> S -> ByteString
292−readByteString n s@(S bs o)
293− -- no offset, easy.
294− | o == 0 = unsafeTake n bs
295− -- offset. ugg. this is really naive and slow. but also pretty easy :)
296− | otherwise = B.pack (P.map (readWord8 8) (P.take n (iterate (incS 8) s)))
297−
298−readWithoutOffset :: (Bits a, Num a)
299− => S -> (a -> Int -> a) -> (a -> Int -> a) -> Int -> a
300−readWithoutOffset (S bs o) shifterL shifterR n
301− | o /= 0 = error "readWithoutOffset: there is an offset"
302−
303− | bit_offset n == 0 && byte_offset n <= 4 =
304− let segs = byte_offset n
305− bn 0 = fromIntegral (unsafeHead bs)
306− bn m = (bn (m-1) `shifterL` 8) .|. fromIntegral (unsafeIndex bs m)
307−
308− in bn (segs-1)
309−
310− -- n <= 64
311− | otherwise = let segs = byte_offset n
312− o' = bit_offset (n - 8 + o)
313−
314− bn 0 = fromIntegral (unsafeHead bs)
315− bn m = (bn (m-1) `shifterL` 8) .|. fromIntegral (unsafeIndex bs m)
316−
317− msegs = bn (segs-1) `shifterL` o'
318−
319− lst = (fromIntegral (unsafeIndex bs segs)) `shifterR` (8 - o')
320−
321− w = msegs .|. lst
322− in w
323−
324−readWithOffset :: (Bits a, Num a)
325− => S -> (a -> Int -> a) -> (a -> Int -> a) -> Int -> a
326−readWithOffset (S bs o) shifterL shifterR n
327− | n <= 64 = let bits_in_msb = 8 - o
328− (n',top) = (n - bits_in_msb
329− , (fromIntegral (unsafeHead bs) .&. make_mask bits_in_msb) `shifterL` n')
330−
331− segs = byte_offset n'
332−
333− bn 0 = 0
334− bn m = (bn (m-1) `shifterL` 8) .|. fromIntegral (unsafeIndex bs m)
248+-- | Bit order
249+--
250+-- E.g. two words of 5 bits: ABCDE, VWXYZ
251+-- - BB: ABCDEVWX YZxxxxxx
252+-- - LB: XYZABCDE xxxxxxVW
253+-- - BL: EDCBAZYX WVxxxxxx -- not implemented
254+-- - LL: XWVEDCBA xxxxxxZY -- not implemented
255+data BitOrder = BB | LB
335256
336− o' = bit_offset n'
337257
338− mseg = bn segs `shifterL` o'
258+-- | 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
339262
340− lst | o' > 0 = (fromIntegral (unsafeIndex bs (segs + 1))) `shifterR` (8 - o')
341− | otherwise = 0
263+-- | Extract a range of bits from (ws :: ByteString)
264+--
265+-- Constraint: 8 * (length ws -1 ) < o+n <= 8 * length ws
266+extract :: (Num a, Bits a, FastBits a) => BitOrder -> ByteString -> Int -> Int -> a
267+extract bo bs o n
268+ | n == 0 = 0
269+ | B.length bs == 0 = error "Empty ByteString"
270+ | otherwise = mask n (foldlWithIndex' f 0 bs)
271+ 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))
275+
276+ -- 'or' correctly shifted words
277+ f b w i = b .|. (fromIntegral w `fastShift` off i)
278+
279+ -- co-offset
280+ r = B.length bs * 8 - (o + n)
281+
282+ -- shift offset depending on the byte position (0..B.length-1)
283+ off i = case bo of
284+ LB -> 8*i - o
285+ BB -> (B.length bs -1 - i) * 8 - r
286+
287+
288+-- | Generic readWord
289+readWord :: (Num a, Bits a, FastBits a) => BitOrder -> Int -> S -> a
290+readWord bo n (S bs o)
291+ | n == 0 = 0
292+ | o+n <= 8 = extract bo (unsafeTake 1 bs) o n
293+ | o+n <= 16 = extract bo (unsafeTake 2 bs) o n
294+ | o+n <= 24 = extract bo (unsafeTake 3 bs) o n
295+ | o+n <= 32 = extract bo (unsafeTake 4 bs) o n
296+ | o+n <= 40 = extract bo (unsafeTake 5 bs) o n
297+ | o+n <= 48 = extract bo (unsafeTake 6 bs) o n
298+ | o+n <= 56 = extract bo (unsafeTake 7 bs) o n
299+ | o+n <= 64 = extract bo (unsafeTake 8 bs) o n
300+ | otherwise = extract bo (unsafeTake ((o+n+7) `div` 8) bs) o n
301+
302+readWordChecked :: (Num a, Bits a, FastBits a) => Int -> BitOrder -> Int -> S -> a
303+readWordChecked m bo n s
304+ | n > m = error $ "Tried to read more than " ++ show m ++ " bits (" ++ show n ++")"
305+ | otherwise = readWord bo n s
342306
343− w = top .|. mseg .|. lst
344− in w
345− | otherwise = error $ "Invalid offset (" ++ show n ++ " > 64)"
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
322+
323+-- | 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)
326+ | o == 0 = unsafeTake n bs
327+ | otherwise = unsafePerformIO $ do
328+ let len = n+1
329+ ptr <- mallocBytes len
330+ case bo of
331+ BB -> do
332+ let f r i = do
333+ let
334+ w = unsafeIndex bs (len-i)
335+ w' = (w `fastShiftL` o) .|. r
336+ r' = w `fastShiftR` (8-o)
337+ poke (ptr `plusPtr` (len-i)) w'
338+ return r'
339+ foldM_ f 0 [1..len]
340+ unsafeInit <$> unsafePackMallocCStringLen (ptr,len)
341+
342+ LB -> do
343+ let f r i = do
344+ let
345+ w = unsafeIndex bs (i-1)
346+ w' = (w `fastShiftL` (8-o)) .|. r
347+ r' = w `fastShiftR` o
348+ poke (ptr `plusPtr` (len-i)) w'
349+ return r'
350+ foldM_ f 0 [1..len]
351+ unsafeTail <$> unsafePackMallocCStringLen (ptr,len)
346352
347353------------------------------------------------------------------------
348354-- | 'BitGet' is a monad, applicative and a functor. See 'runBitGet'
@@ -408,7 +414,7 @@ skipBits n = do
408414 s <- getState
409415 putState $! (incS n s)
410416
411−-- | Skip bits if necessary to align on a byte
417+-- | Skip bits if necessary to align to the next byte
412418alignByte :: BitGet ()
413419alignByte = do
414420 (S _ n) <- getState
@@ -459,32 +465,33 @@ isEmpty = B $ \ (S bs o) -> if B.null bs
459465
460466-- | Read a 1 bit 'Bool'.
461467bool :: Block Bool
462−bool = Block 1 readBool
468+bool = Block 1 (readBool BB)
463469
464470-- | Read @n@ bits as a 'Word8'. @n@ must be within @[0..8]@.
465471word8 :: Int -> Block Word8
466−word8 n = Block n (readWord8 n)
472+word8 n = Block n (readWord8 BB n)
467473
468474-- | Read @n@ bits as a 'Word16'. @n@ must be within @[0..16]@.
469475word16be :: Int -> Block Word16
470−word16be n = Block n (readWord16be n)
476+word16be n = Block n (readWord16 BB n)
471477
472478-- | Read @n@ bits as a 'Word32'. @n@ must be within @[0..32]@.
473479word32be :: Int -> Block Word32
474−word32be n = Block n (readWord32be n)
480+word32be n = Block n (readWord32 BB n)
475481
476482-- | Read @n@ bits as a 'Word64'. @n@ must be within @[0..64]@.
477483word64be :: Int -> Block Word64
478−word64be n = Block n (readWord64be n)
484+word64be n = Block n (readWord64 BB n)
479485
480486-- | Read @n@ bytes as a 'ByteString'.
481487byteString :: Int -> Block ByteString
482−byteString n | n > 0 = Block (n*8) (readByteString n)
488+byteString n | n > 0 = Block (n*8) (readByteString BB n)
483489 | otherwise = Block 0 (\_ -> B.empty)
484490
485491------------------------------------------------------------------------
486492-- Unchecked shifts, from the package binary
487493
494+shiftl_w8 :: Word8 -> Int -> Word8
488495shiftl_w16 :: Word16 -> Int -> Word16
489496shiftl_w32 :: Word32 -> Int -> Word32
490497shiftl_w64 :: Word64 -> Int -> Word64
@@ -495,6 +502,7 @@ shiftr_w32 :: Word32 -> Int -> Word32
495502shiftr_w64 :: Word64 -> Int -> Word64
496503
497504#if defined(__GLASGOW_HASKELL__) && !defined(__HADDOCK__)
505+shiftl_w8 (W8# w) (I# i) = W8# (w `uncheckedShiftL#` i)
498506shiftl_w16 (W16# w) (I# i) = W16# (w `uncheckedShiftL#` i)
499507shiftl_w32 (W32# w) (I# i) = W32# (w `uncheckedShiftL#` i)
500508
@@ -521,6 +529,7 @@ shiftr_w64 (W64# w) (I# i) = W64# (w `uncheckedShiftRL#` i)
521529#endif
522530
523531#else
532+shiftl_w8 = shiftL
524533shiftl_w16 = shiftL
525534shiftl_w32 = shiftL
526535shiftl_w64 = shiftL