@@ -93,11 +93,15 @@ import Data.Binary.Get.Internal as B ( get, put, ensureN )
93import Data.ByteString as B
93import Data.ByteString as B
94import qualified Data.ByteString.Lazy as L
94import qualified Data.ByteString.Lazy as L
95import Data.ByteString.Unsafe
95import Data.ByteString.Unsafe
96+import Foreign.Marshal.Alloc (mallocBytes)
98+import Foreign.Storable (poke)
99+import System.IO.Unsafe (unsafePerformIO)
97import Data.Bits
101import Data.Bits
98import Data.Word
102import Data.Word
99import Control.Applicative
103import Control.Applicative
100−import Control.Monad (when)
104+import Control.Monad (when,foldM_)
102import Prelude as P
106import Prelude as P
@@ -157,7 +161,7 @@ data S = S {-# UNPACK #-} !ByteString -- Input
157 {-# UNPACK #-} !Int -- Bit offset (0-7)
161 {-# UNPACK #-} !Int -- Bit offset (0-7)
158 deriving (Show)
162 deriving (Show)
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
161-- number of bits in advance.
165-- number of bits in advance.
162data Block a = Block Int (S -> a)
166data Block a = Block Int (S -> a)
@@ -181,12 +185,23 @@ block (Block i p) = do
181 putState $! (incS i s)
185 putState $! (incS i s)
182 return $! p s
186 return $! p s
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 #-}
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 #-}
198+-- | Increment the current bit offset
184incS :: Int -> S -> S
199incS :: Int -> S -> S
185−incS o (S bs n) =
200+incS o (S bs n) = S (unsafeDrop d bs) n'
186− let !o' = (n+o)
202+ !o' = (n+o)
187− !d = o' `shiftR` 3
203+ !d = byte_offset o'
188− !n' = o' .&. make_mask 3
204+ !n' = bit_offset o'
189− in S (unsafeDrop d bs) n'
191-- | make_mask 3 = 00000111
206-- | make_mask 3 = 00000111
192make_mask :: (Bits a, Num a) => Int -> a
207make_mask :: (Bits a, Num a) => Int -> a
@@ -198,151 +213,142 @@ make_mask n = (1 `shiftL` fromIntegral n) - 1
198{-# SPECIALIZE make_mask :: Int -> Word32 #-}
213{-# SPECIALIZE make_mask :: Int -> Word32 #-}
199{-# SPECIALIZE make_mask :: Int -> Word64 #-}
214{-# SPECIALIZE make_mask :: Int -> Word64 #-}
201−bit_offset :: Int -> Int
216+-- | Keep only the n least-significant bits of the given value
202−bit_offset n = make_mask 3 .&. n
217+mask :: (Bits a, Num a) => Int -> a -> a
204−byte_offset :: Int -> Int
205−byte_offset n = n `shiftR` 3
207−readBool :: S -> Bool
208−readBool (S bs n) = testBit (unsafeHead bs) (7-n)
210−{-# INLINE readWord8 #-}
211−readWord8 :: Int -> S -> Word8
213− -- no bits at all, return 0
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
218+mask n v = v .&. make_mask n
220− w' = (w `shiftr_w8` (8 - o - n)) .&. m
219+{-# INLINE mask #-}
223− -- the bits are in two different bytes
224− -- make a word16 using both bytes, and then shift and mask
227− w = (fromIntegral (unsafeHead bs) `shiftl_w16` 8) .|.
228− (fromIntegral (unsafeIndex bs 1))
230− w' = (w `shiftr_w16` (16 - o - n)) .&. m
233−{-# INLINE readWord16be #-}
234−readWord16be :: Int -> S -> Word16
235−readWord16be n s@(S bs o)
237− -- 8 or fewer bits, use readWord8
238− | n <= 8 = fromIntegral (readWord8 n s)
240− -- handle 9 or more bits, stored in two bytes
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
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))
254− -- with offset, and n=9-16
221+class FastBits a where
255− | n <= 16 = readWithOffset s shiftl_w16 shiftr_w16 n
222+ fastShiftR :: a -> Int -> a
223+ fastShiftL :: a -> Int -> a
257− | otherwise = error "readWord16be: tried to read more than 16 bits"
225+ {-# INLINE fastShift #-}
226+ fastShift :: a -> Int -> a
228+ | n > 0 = fastShiftL x n
229+ | n < 0 = fastShiftR x (negate n)
259−{-# INLINE readWord32be #-}
232+instance FastBits Word8 where
260−readWord32be :: Int -> S -> Word32
233+ fastShiftR = shiftr_w8
261−readWord32be n s@(S _ o)
234+ fastShiftL = shiftl_w8
262− -- 8 or fewer bits, use readWord8
263− | n <= 8 = fromIntegral (readWord8 n s)
265− -- 16 or fewer bits, use readWord16be
236+instance FastBits Word16 where
266− | n <= 16 = fromIntegral (readWord16be n s)
237+ fastShiftR = shiftr_w16
238+ fastShiftL = shiftl_w16
268− | o == 0 = readWithoutOffset s shiftl_w32 shiftr_w32 n
240+instance FastBits Word32 where
241+ fastShiftR = shiftr_w32
242+ fastShiftL = shiftl_w32
270− | n <= 32 = readWithOffset s shiftl_w32 shiftr_w32 n
244+instance FastBits Word64 where
245+ fastShiftR = shiftr_w64
246+ fastShiftL = shiftl_w64
272− | otherwise = error "readWord32be: tried to read more than 32 bits"
248+-- | Bit order
274−
250+-- E.g. two words of 5 bits: ABCDE, VWXYZ
275−{-# INLINE readWord64be #-}
251+-- - BB: ABCDEVWX YZxxxxxx
276−readWord64be :: Int -> S -> Word64
252+-- - LB: XYZABCDE xxxxxxVW
277−readWord64be n s@(S _ o)
253+-- - BL: EDCBAZYX WVxxxxxx -- not implemented
278− -- 8 or fewer bits, use readWord8
254+-- - LL: XWVEDCBA xxxxxxZY -- not implemented
279− | n <= 8 = fromIntegral (readWord8 n s)
255+data BitOrder = BB | LB
281− -- 16 or fewer bits, use readWord16be
282− | n <= 16 = fromIntegral (readWord16be n s)
284− | o == 0 = readWithoutOffset s shiftl_w64 shiftr_w64 n
286− | n <= 64 = readWithOffset s shiftl_w64 shiftr_w64 n
288− | otherwise = error "readWord64be: tried to read more than 64 bits"
291−readByteString :: Int -> S -> ByteString
292−readByteString n s@(S bs o)
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)))
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"
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)
311− | otherwise = let segs = byte_offset n
312− o' = bit_offset (n - 8 + o)
314− bn 0 = fromIntegral (unsafeHead bs)
315− bn m = (bn (m-1) `shifterL` 8) .|. fromIntegral (unsafeIndex bs m)
317− msegs = bn (segs-1) `shifterL` o'
319− lst = (fromIntegral (unsafeIndex bs segs)) `shifterR` (8 - o')
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')
331− segs = byte_offset n'
334− bn m = (bn (m-1) `shifterL` 8) .|. fromIntegral (unsafeIndex bs m)
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
340− lst | o' > 0 = (fromIntegral (unsafeIndex bs (segs + 1))) `shifterR` (8 - o')
263+-- | Extract a range of bits from (ws :: ByteString)
265+-- Constraint: 8 * (length ws -1 ) < o+n <= 8 * length ws
266+extract :: (Num a, Bits a, FastBits a) => BitOrder -> ByteString -> Int -> Int -> a
269+ | B.length bs == 0 = error "Empty ByteString"
270+ | otherwise = mask n (foldlWithIndex' f 0 bs)
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))
276+ -- 'or' correctly shifted words
277+ f b w i = b .|. (fromIntegral w `fastShift` off i)
280+ r = B.length bs * 8 - (o + n)
282+ -- shift offset depending on the byte position (0..B.length-1)
285+ BB -> (B.length bs -1 - i) * 8 - r
288+-- | Generic readWord
289+readWord :: (Num a, Bits a, FastBits a) => BitOrder -> Int -> S -> a
290+readWord bo n (S bs o)
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
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 ++")"
341− | otherwise = 0
305+ | otherwise = readWord bo n s
343− w = top .|. mseg .|. lst
307+{-# INLINE readWord8 #-}
344− in w
308+readWord8 :: BitOrder -> Int -> S -> Word8
345− | otherwise = error $ "Invalid offset (" ++ show n ++ " > 64)"
309+readWord8 = readWordChecked 8
311+{-# INLINE readWord16 #-}
312+readWord16 :: BitOrder -> Int -> S -> Word16
313+readWord16 = readWordChecked 16
315+{-# INLINE readWord32 #-}
316+readWord32 :: BitOrder -> Int -> S -> Word32
317+readWord32 = readWordChecked 32
319+{-# INLINE readWord64 #-}
320+readWord64 :: BitOrder -> Int -> S -> Word64
321+readWord64 = readWordChecked 64
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
329+ ptr <- mallocBytes len
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'
340+ unsafeInit <$> unsafePackMallocCStringLen (ptr,len)
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'
351+ unsafeTail <$> unsafePackMallocCStringLen (ptr,len)
347------------------------------------------------------------------------
353------------------------------------------------------------------------
348-- | 'BitGet' is a monad, applicative and a functor. See 'runBitGet'
354-- | 'BitGet' is a monad, applicative and a functor. See 'runBitGet'
@@ -408,7 +414,7 @@ skipBits n = do
408 s <- getState
414 s <- getState
409 putState $! (incS n s)
415 putState $! (incS n s)
411−-- | Skip bits if necessary to align on a byte
417+-- | Skip bits if necessary to align to the next byte
412alignByte :: BitGet ()
418alignByte :: BitGet ()
413alignByte = do
419alignByte = do
414 (S _ n) <- getState
420 (S _ n) <- getState
@@ -459,32 +465,33 @@ isEmpty = B $ \ (S bs o) -> if B.null bs
460-- | Read a 1 bit 'Bool'.
466-- | Read a 1 bit 'Bool'.
461bool :: Block Bool
467bool :: Block Bool
462−bool = Block 1 readBool
468+bool = Block 1 (readBool BB)
464-- | Read @n@ bits as a 'Word8'. @n@ must be within @[0..8]@.
470-- | Read @n@ bits as a 'Word8'. @n@ must be within @[0..8]@.
465word8 :: Int -> Block Word8
471word8 :: Int -> Block Word8
466−word8 n = Block n (readWord8 n)
472+word8 n = Block n (readWord8 BB n)
468-- | Read @n@ bits as a 'Word16'. @n@ must be within @[0..16]@.
474-- | Read @n@ bits as a 'Word16'. @n@ must be within @[0..16]@.
469word16be :: Int -> Block Word16
475word16be :: Int -> Block Word16
470−word16be n = Block n (readWord16be n)
476+word16be n = Block n (readWord16 BB n)
472-- | Read @n@ bits as a 'Word32'. @n@ must be within @[0..32]@.
478-- | Read @n@ bits as a 'Word32'. @n@ must be within @[0..32]@.
473word32be :: Int -> Block Word32
479word32be :: Int -> Block Word32
474−word32be n = Block n (readWord32be n)
480+word32be n = Block n (readWord32 BB n)
476-- | Read @n@ bits as a 'Word64'. @n@ must be within @[0..64]@.
482-- | Read @n@ bits as a 'Word64'. @n@ must be within @[0..64]@.
477word64be :: Int -> Block Word64
483word64be :: Int -> Block Word64
478−word64be n = Block n (readWord64be n)
484+word64be n = Block n (readWord64 BB n)
480-- | Read @n@ bytes as a 'ByteString'.
486-- | Read @n@ bytes as a 'ByteString'.
481byteString :: Int -> Block ByteString
487byteString :: Int -> Block ByteString
482−byteString n | n > 0 = Block (n*8) (readByteString n)
488+byteString n | n > 0 = Block (n*8) (readByteString BB n)
483 | otherwise = Block 0 (\_ -> B.empty)
489 | otherwise = Block 0 (\_ -> B.empty)
485------------------------------------------------------------------------
491------------------------------------------------------------------------
486-- Unchecked shifts, from the package binary
492-- Unchecked shifts, from the package binary
494+shiftl_w8 :: Word8 -> Int -> Word8
488shiftl_w16 :: Word16 -> Int -> Word16
495shiftl_w16 :: Word16 -> Int -> Word16
489shiftl_w32 :: Word32 -> Int -> Word32
496shiftl_w32 :: Word32 -> Int -> Word32
490shiftl_w64 :: Word64 -> Int -> Word64
497shiftl_w64 :: Word64 -> Int -> Word64
@@ -495,6 +502,7 @@ shiftr_w32 :: Word32 -> Int -> Word32
495shiftr_w64 :: Word64 -> Int -> Word64
502shiftr_w64 :: Word64 -> Int -> Word64
497#if defined(__GLASGOW_HASKELL__) && !defined(__HADDOCK__)
504#if defined(__GLASGOW_HASKELL__) && !defined(__HADDOCK__)
505+shiftl_w8 (W8# w) (I# i) = W8# (w `uncheckedShiftL#` i)
498shiftl_w16 (W16# w) (I# i) = W16# (w `uncheckedShiftL#` i)
506shiftl_w16 (W16# w) (I# i) = W16# (w `uncheckedShiftL#` i)
499shiftl_w32 (W32# w) (I# i) = W32# (w `uncheckedShiftL#` i)
507shiftl_w32 (W32# w) (I# i) = W32# (w `uncheckedShiftL#` i)
@@ -521,6 +529,7 @@ shiftr_w64 (W64# w) (I# i) = W64# (w `uncheckedShiftRL#` i)
524shiftl_w16 = shiftL
533shiftl_w16 = shiftL
525shiftl_w32 = shiftL
534shiftl_w32 = shiftL
526shiftl_w64 = shiftL
535shiftl_w64 = shiftL