@@ -1,9 +1,8 @@
1−{-# LANGUAGE RankNTypes, MagicHash, CPP #-}
1+{-# LANGUAGE RankNTypes, MagicHash, BangPatterns, CPP #-}
3module Data.Binary.Bits.Get
3module Data.Binary.Bits.Get
5 , runBitGet
5 , runBitGet
@@ -12,6 +11,14 @@ module Data.Binary.Bits.Get
12 , getWord32be
11 , getWord32be
13 , getWord64be
12 , getWord64be
15 , Data.Binary.Bits.Get.getByteString
22 , Data.Binary.Bits.Get.getByteString
17 , readBool
24 , readBool
@@ -34,6 +41,8 @@ import Data.List as List ( reverse )
35import Control.Applicative
42import Control.Applicative
37#if defined(__GLASGOW_HASKELL__) && !defined(__HADDOCK__)
46#if defined(__GLASGOW_HASKELL__) && !defined(__HADDOCK__)
38import GHC.Base
47import GHC.Base
39import GHC.Word
48import GHC.Word
@@ -46,45 +55,61 @@ data S = S !ByteString -- ^ Input
46 !Int -- ^ Bit offset (0-7)
55 !Int -- ^ Bit offset (0-7)
47 deriving (Show)
56 deriving (Show)
58+-- | A block that will be read with only one boundry check.
59+data Block a = Block Int (S -> a)
61+instance Functor Block where
62+ fmap f (Block i p) = Block i (\s -> f (p s))
64+instance Applicative Block where
65+ pure a = Block 0 (\_ -> a)
66+ (Block i p) <*> (Block j q) = Block (i+j) (\s -> p s $ q (incS i s))
67+ (Block i _) *> (Block j q) = Block (i+j) (q . incS i)
68+ (Block i p) <* (Block j _) = Block (i+j) p
70+block :: Block a -> BitGet a
71+block (Block i p) = do
73+ modifyState (\s -> (incS i s, p s))
49-- | make_mask 3 = 00000111
75-- | make_mask 3 = 00000111
50make_mask :: Bits a => Int -> a
76make_mask :: Bits a => Int -> a
51make_mask n = (1 `shiftL` fromIntegral n) - 1
77make_mask n = (1 `shiftL` fromIntegral n) - 1
78+{-# SPECIALIZE make_mask :: Int -> Int #-}
52{-# SPECIALIZE make_mask :: Int -> Word #-}
79{-# SPECIALIZE make_mask :: Int -> Word #-}
53{-# SPECIALIZE make_mask :: Int -> Word8 #-}
80{-# SPECIALIZE make_mask :: Int -> Word8 #-}
54{-# SPECIALIZE make_mask :: Int -> Word16 #-}
81{-# SPECIALIZE make_mask :: Int -> Word16 #-}
55{-# SPECIALIZE make_mask :: Int -> Word32 #-}
82{-# SPECIALIZE make_mask :: Int -> Word32 #-}
56{-# SPECIALIZE make_mask :: Int -> Word64 #-}
83{-# SPECIALIZE make_mask :: Int -> Word64 #-}
58−readBool :: S -> T Bool S
85+readBool :: S -> Bool
59−readBool s@(S bs n) = testBit (unsafeHead bs) (7-n) :*: incS s 1
86+readBool (S bs n) = testBit (unsafeHead bs) (7-n)
61−incS :: S -> Int -> S
88+incS :: Int -> S -> S
62−incS (S bs n) 1 =
89+incS o (S bs n) =
63− let
90+ let !o' = (n+o)
64− n' = n + 1
91+ !d = o' `shiftR` 3
65− s' | n' == 8 = S (unsafeTail bs) 0
92+ !n' = o' .&. make_mask 3
66− | otherwise = S bs n'
69− let (d,n') = divMod (n+o) 8
70 in S (unsafeDrop d bs) n'
93 in S (unsafeDrop d bs) n'
72−readByteString :: S -> Int -> T ByteString S
95+readByteString :: Int -> S -> ByteString
73−readByteString s@(S bs o) n
96+readByteString n s@(S bs o)
74− | o == 0 = unsafeTake n bs :*: (S (unsafeDrop n bs) 0)
98+ | o == 0 = unsafeTake n bs
75− | otherwise = unsafeTake n (unsafeTail bs) :*: (S (unsafeDrop (n+1) bs) 0)
99+ -- offset. ugg. this is really naive and slow. but also pretty easy :)
100+ | otherwise = B.pack (P.map (readWord8 8) (P.take n (iterate (incS 8) s)))
77−readWord8 :: S -> Int -> T Word8 S
102+readWord8 :: Int -> S -> Word8
78−readWord8 s@(S bs o) n
103+readWord8 n s@(S bs o)
79 -- no bits at all, return 0
104 -- no bits at all, return 0
80− | n == 0 = 0 :*: s
105+ | n == 0 = 0
82 -- all bits are in the same byte
107 -- all bits are in the same byte
83 -- we just need to shift and mask them right
108 -- we just need to shift and mask them right
84 | n <= 8 - o = let w = unsafeHead bs
109 | n <= 8 - o = let w = unsafeHead bs
85 m = make_mask n
110 m = make_mask n
86 w' = (w `shiftr_w8` (8 - o - n)) .&. m
111 w' = (w `shiftr_w8` (8 - o - n)) .&. m
87− in w' :*: incS s n
112+ in w'
89 -- the bits are in two different bytes
114 -- the bits are in two different bytes
90 -- make a word16 using both bytes, and then shift and mask
115 -- make a word16 using both bytes, and then shift and mask
@@ -92,7 +117,7 @@ readWord8 s@(S bs o) n
92 (fromIntegral (unsafeIndex bs 1))
117 (fromIntegral (unsafeIndex bs 1))
93 m = make_mask n
118 m = make_mask n
94 w' = (w `shiftr_w16` (16 - o - n)) .&. m
119 w' = (w `shiftr_w16` (16 - o - n)) .&. m
95− in fromIntegral w' :*: incS s n
120+ in fromIntegral w'
97bit_offset :: Int -> Int
122bit_offset :: Int -> Int
98bit_offset n = make_mask 3 .&. n
123bit_offset n = make_mask 3 .&. n
@@ -100,12 +125,11 @@ bit_offset n = make_mask 3 .&. n
100byte_offset :: Int -> Int
125byte_offset :: Int -> Int
101byte_offset n = n `shiftR` 3
126byte_offset n = n `shiftR` 3
103−readWord16be :: S -> Int -> T Word16 S
128+readWord16be :: Int -> S -> Word16
104−readWord16be s@(S bs o) n
129+readWord16be n s@(S bs o)
106 -- 8 or fewer bits, use readWord8
131 -- 8 or fewer bits, use readWord8
107− | n <= 8 = let w :*: s' = readWord8 s n
132+ | n <= 8 = fromIntegral (readWord8 n s)
108− in fromIntegral w :*: s'
110 -- handle 9 or more bits, stored in two bytes
134 -- handle 9 or more bits, stored in two bytes
@@ -113,13 +137,13 @@ readWord16be s@(S bs o) n
113 | o == 0 && n == 16 = let msb = fromIntegral (unsafeHead bs)
137 | o == 0 && n == 16 = let msb = fromIntegral (unsafeHead bs)
114 lsb = fromIntegral (unsafeIndex bs 1)
138 lsb = fromIntegral (unsafeIndex bs 1)
115 w = (msb `shiftl_w16` 8) .|. lsb
139 w = (msb `shiftl_w16` 8) .|. lsb
116− in w :*: incS s n
140+ in w
118 -- no offset, but not full 16 bytes
142 -- no offset, but not full 16 bytes
119 | o == 0 = let msb = fromIntegral (unsafeHead bs)
143 | o == 0 = let msb = fromIntegral (unsafeHead bs)
120 lsb = fromIntegral (unsafeIndex bs 1)
144 lsb = fromIntegral (unsafeIndex bs 1)
121 w = (msb `shiftl_w16` (n-8)) .|. (lsb `shiftr_w16` (16-n))
145 w = (msb `shiftl_w16` (n-8)) .|. (lsb `shiftr_w16` (16-n))
122− in w :*: incS s n
146+ in w
124 -- with offset, and n=9-16
148 -- with offset, and n=9-16
125 | n <= 16 = readWithOffset s shiftl_w16 shiftr_w16 n
149 | n <= 16 = readWithOffset s shiftl_w16 shiftr_w16 n
@@ -127,7 +151,7 @@ readWord16be s@(S bs o) n
127 | otherwise = error "readWord16be: tried to read more than 16 bits"
151 | otherwise = error "readWord16be: tried to read more than 16 bits"
129readWithoutOffset :: (Bits a, Num a)
153readWithoutOffset :: (Bits a, Num a)
130− => S -> (a -> Int -> a) -> (a -> Int -> a) -> Int -> T a S
154+ => S -> (a -> Int -> a) -> (a -> Int -> a) -> Int -> a
131readWithoutOffset s@(S bs o) shifterL shifterR n
155readWithoutOffset s@(S bs o) shifterL shifterR n
132 | o /= 0 = error "readWithoutOffset: there is an offset"
156 | o /= 0 = error "readWithoutOffset: there is an offset"
@@ -136,7 +160,7 @@ readWithoutOffset s@(S bs o) shifterL shifterR n
136 bn 0 = fromIntegral (unsafeHead bs)
160 bn 0 = fromIntegral (unsafeHead bs)
137 bn n = (bn (n-1) `shifterL` 8) .|. fromIntegral (unsafeIndex bs n)
161 bn n = (bn (n-1) `shifterL` 8) .|. fromIntegral (unsafeIndex bs n)
139− in (bn (segs-1) :*: incS s n)
163+ in bn (segs-1)
141 | n <= 64 = let segs = byte_offset n
165 | n <= 64 = let segs = byte_offset n
142 o' = bit_offset (n - 8 + o)
166 o' = bit_offset (n - 8 + o)
@@ -149,10 +173,10 @@ readWithoutOffset s@(S bs o) shifterL shifterR n
149 last = (fromIntegral (unsafeIndex bs segs)) `shifterR` (8 - o')
173 last = (fromIntegral (unsafeIndex bs segs)) `shifterR` (8 - o')
151 w = msegs .|. last
175 w = msegs .|. last
152− in w :*: incS s n
176+ in w
154readWithOffset :: (Bits a, Num a)
178readWithOffset :: (Bits a, Num a)
155− => S -> (a -> Int -> a) -> (a -> Int -> a) -> Int -> T a S
179+ => S -> (a -> Int -> a) -> (a -> Int -> a) -> Int -> a
156readWithOffset s@(S bs o) shifterL shifterR n
180readWithOffset s@(S bs o) shifterL shifterR n
157 | n <= 64 = let bits_in_msb = 8 - o
181 | n <= 64 = let bits_in_msb = 8 - o
158 (n',top) = (n - bits_in_msb
182 (n',top) = (n - bits_in_msb
@@ -161,7 +185,6 @@ readWithOffset s@(S bs o) shifterL shifterR n
161 segs = byte_offset n'
185 segs = byte_offset n'
164− bn 1 = fromIntegral (unsafeIndex bs 1)
165 bn n = (bn (n-1) `shifterL` 8) .|. fromIntegral (unsafeIndex bs n)
188 bn n = (bn (n-1) `shifterL` 8) .|. fromIntegral (unsafeIndex bs n)
167 o' = bit_offset n'
190 o' = bit_offset n'
@@ -172,18 +195,16 @@ readWithOffset s@(S bs o) shifterL shifterR n
172 | otherwise = 0
195 | otherwise = 0
174 w = top .|. mseg .|. last
197 w = top .|. mseg .|. last
175− in w :*: incS s n
198+ in w
178−readWord32be :: S -> Int -> T Word32 S
201+readWord32be :: Int -> S -> Word32
179−readWord32be s@(S bs o) n
202+readWord32be n s@(S bs o)
180 -- 8 or fewer bits, use readWord8
203 -- 8 or fewer bits, use readWord8
181− | n <= 8 = let w :*: s' = readWord8 s n
204+ | n <= 8 = fromIntegral (readWord8 n s)
182− in fromIntegral w :*: s'
184 -- 16 or fewer bits, use readWord16be
206 -- 16 or fewer bits, use readWord16be
185− | n <= 16 = let w :*: s' = readWord16be s n
207+ | n <= 16 = fromIntegral (readWord16be n s)
186− in fromIntegral w :*: s'
188 | o == 0 = readWithoutOffset s shiftl_w32 shiftr_w32 n
209 | o == 0 = readWithoutOffset s shiftl_w32 shiftr_w32 n
@@ -192,15 +213,13 @@ readWord32be s@(S bs o) n
192 | otherwise = error "readWord32be: tried to read more than 32 bits"
213 | otherwise = error "readWord32be: tried to read more than 32 bits"
195−readWord64be :: S -> Int -> T Word64 S
216+readWord64be :: Int -> S -> Word64
196−readWord64be s@(S bs o) n
217+readWord64be n s@(S bs o)
197 -- 8 or fewer bits, use readWord8
218 -- 8 or fewer bits, use readWord8
198− | n <= 8 = let w :*: s' = readWord8 s n
219+ | n <= 8 = fromIntegral (readWord8 n s)
199− in fromIntegral w :*: s'
201 -- 16 or fewer bits, use readWord16be
221 -- 16 or fewer bits, use readWord16be
202− | n <= 16 = let w :*: s' = readWord16be s n
222+ | n <= 16 = fromIntegral (readWord16be n s)
203− in fromIntegral w :*: s'
205 | o == 0 = readWithoutOffset s shiftl_w64 shiftr_w64 n
224 | o == 0 = readWithoutOffset s shiftl_w64 shiftr_w64 n
@@ -210,17 +229,13 @@ readWord64be s@(S bs o) n
211------------------------------------------------------------------------
230------------------------------------------------------------------------
212-- unrolled codensity/state monad
231-- unrolled codensity/state monad
213−newtype BitGet a = C { runCont :: forall r.
232+newtype BitGet a = B { runState :: S -> Get (S,a) }
218−type Success a r = S -> a -> Get r
220instance Monad BitGet where
234instance Monad BitGet where
221− return x = C $ \s ks -> ks s x
235+ return x = B $ \s -> return (s,x)
222− fail str = C $ \s ks -> fail str
236+ fail str = B $ \_s -> fail str
223− (C c) >>= f = C $ \s ks -> c s (\s' a -> runCont (f a) s' ks)
237+ (B f) >>= g = B $ \s -> do (s',a) <- f s
225instance Functor BitGet where
240instance Functor BitGet where
226 fmap f m = m >>= \a -> return (f a)
241 fmap f m = m >>= \a -> return (f a)
@@ -229,59 +244,66 @@ instance Applicative BitGet where
229 pure x = return x
244 pure x = return x
230 fm <*> m = fm >>= \f -> m >>= \v -> return (f v)
245 fm <*> m = fm >>= \f -> m >>= \v -> return (f v)
232−runBitGetSimple :: BitGet a -> Get a
233−runBitGetSimple bg = runBitGet bg
235runBitGet :: BitGet a -> Get a
247runBitGet :: BitGet a -> Get a
236−runBitGet bg = runCont bg (S B.empty 0) (\_s a -> return a)
248+runBitGet bg = do
249+ (_,a) <- runState bg (S B.empty 0)
253+getState = B $ \s -> return (s,s)
255+putState :: S -> BitGet ()
256+putState s = B $ \_ -> return (s,())
258+modifyState :: (S -> (S,a)) -> BitGet a
259+modifyState f = B $ \s -> return (f s)
238−-- | Make sure there are at lest @n@ bits.
261+-- | Make sure there are at least @n@ bits.
239ensureBits :: Int -> BitGet ()
262ensureBits :: Int -> BitGet ()
240−ensureBits n = C $ \s ks ->
263+ensureBits n = do
264+ s@(S bs o) <- getState
241− let loop (S bs o) | n <= (B.length bs * 8 - o) = ks (S bs o) ()
265+ if n <= (B.length bs * 8 - o)
242− | otherwise = do
266+ then return ()
243− let currentBits = B.length bs * 8 - o
267+ else do let currentBits = B.length bs * 8 - o
244− let byteCount = (n - currentBits + 7) `div` 8
268+ let byteCount = (n - currentBits + 7) `div` 8
245− bs' <- B.getByteString byteCount
269+ B $ \s -> do bs' <- B.getByteString byteCount
246− -- should always be enough going once throguh the loop
247− loop (S (bs`append`bs') o)
270+ return (S (bs`append`bs') o, ())
250getBool :: BitGet Bool
272getBool :: BitGet Bool
251−getBool = ensureBits 1 >> modifyState readBool
273+getBool = block bool
253getWord8 :: Int -> BitGet Word8
275getWord8 :: Int -> BitGet Word8
254−getWord8 n = ensureBits n >> modifyState (flip readWord8 n)
276+getWord8 n = block (word8 n)
256getWord16be :: Int -> BitGet Word16
278getWord16be :: Int -> BitGet Word16
257−getWord16be n = ensureBits n >> modifyState (flip readWord16be n)
279+getWord16be n = block (word16be n)
259getWord32be :: Int -> BitGet Word32
281getWord32be :: Int -> BitGet Word32
260−getWord32be n = ensureBits n >> modifyState (flip readWord32be n)
282+getWord32be n = block (word32be n)
262getWord64be :: Int -> BitGet Word64
284getWord64be :: Int -> BitGet Word64
263−getWord64be n = ensureBits n >> modifyState (flip readWord64be n)
285+getWord64be n = block (word64be n)
265−-- | Get a @ByteString@. Drops the current byte if we've started to read
267getByteString :: Int -> BitGet ByteString
287getByteString :: Int -> BitGet ByteString
268−getByteString n = do
288+getByteString n = block (byteString n)
270− -- drop the first byte if we've started to use it
271− modifyState $ \s@(S bs o) ->
272− let bs' | o == 0 = bs
273− | otherwise = B.tail bs
274− in B.take n bs' :*: S (B.drop n bs') 0
276−getState :: BitGet S
290+bool :: Block Bool
277−getState = C $ \s ks -> ks s s
291+bool = Block 1 readBool
279−putState :: S -> BitGet ()
293+word8 :: Int -> Block Word8
280−putState s = C $ \_ ks -> ks s ()
294+word8 n = Block n (readWord8 n)
296+word16be :: Int -> Block Word16
297+word16be n = Block n (readWord16be n)
299+word32be :: Int -> Block Word32
300+word32be n = Block n (readWord32be n)
302+word64be :: Int -> Block Word64
303+word64be n = Block n (readWord64be n)
282−modifyState :: (S -> (T a S)) -> BitGet a
305+byteString :: Int -> Block ByteString
283−modifyState f = C $ \s ks -> case f s of
306+byteString n = Block (n*8) (readByteString n)
286------------------------------------------------------------------------
308------------------------------------------------------------------------
287-- Unchecked shifts, from the package binary
309-- Unchecked shifts, from the package binary