@@ -1,9 +1,8 @@
1−{-# LANGUAGE RankNTypes, MagicHash, CPP #-}
1+{-# LANGUAGE RankNTypes, MagicHash, BangPatterns, CPP #-}
22
33module Data.Binary.Bits.Get
44 ( T(..)
55 , runBitGet
6− , runBitGetSimple
76 , BitGet
87
98 , getBool
@@ -12,6 +11,14 @@ module Data.Binary.Bits.Get
1211 , getWord32be
1312 , getWord64be
1413
14+ , bool
15+ , word8
16+ , word16be
17+ , word32be
18+ , word64be
19+
20+ , block
21+
1522 , Data.Binary.Bits.Get.getByteString
1623
1724 , readBool
@@ -34,6 +41,8 @@ import Data.List as List ( reverse )
3441
3542import Control.Applicative
3643
44+import Prelude as P
45+
3746#if defined(__GLASGOW_HASKELL__) && !defined(__HADDOCK__)
3847import GHC.Base
3948import GHC.Word
@@ -46,45 +55,61 @@ data S = S !ByteString -- ^ Input
4655 !Int -- ^ Bit offset (0-7)
4756 deriving (Show)
4857
58+-- | A block that will be read with only one boundry check.
59+data Block a = Block Int (S -> a)
60+
61+instance Functor Block where
62+ fmap f (Block i p) = Block i (\s -> f (p s))
63+
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
69+
70+block :: Block a -> BitGet a
71+block (Block i p) = do
72+ ensureBits i
73+ modifyState (\s -> (incS i s, p s))
74+
4975-- | make_mask 3 = 00000111
5076make_mask :: Bits a => Int -> a
5177make_mask n = (1 `shiftL` fromIntegral n) - 1
78+{-# SPECIALIZE make_mask :: Int -> Int #-}
5279{-# SPECIALIZE make_mask :: Int -> Word #-}
5380{-# SPECIALIZE make_mask :: Int -> Word8 #-}
5481{-# SPECIALIZE make_mask :: Int -> Word16 #-}
5582{-# SPECIALIZE make_mask :: Int -> Word32 #-}
5683{-# SPECIALIZE make_mask :: Int -> Word64 #-}
5784
58−readBool :: S -> T Bool S
59−readBool s@(S bs n) = testBit (unsafeHead bs) (7-n) :*: incS s 1
60−
61−incS :: S -> Int -> S
62−incS (S bs n) 1 =
63− let
64− n' = n + 1
65− s' | n' == 8 = S (unsafeTail bs) 0
66− | otherwise = S bs n'
67− in s'
68−incS (S bs n) o =
69− let (d,n') = divMod (n+o) 8
85+readBool :: S -> Bool
86+readBool (S bs n) = testBit (unsafeHead bs) (7-n)
87+
88+incS :: Int -> S -> S
89+incS o (S bs n) =
90+ let !o' = (n+o)
91+ !d = o' `shiftR` 3
92+ !n' = o' .&. make_mask 3
7093 in S (unsafeDrop d bs) n'
7194
72−readByteString :: S -> Int -> T ByteString S
73−readByteString s@(S bs o) n
74− | o == 0 = unsafeTake n bs :*: (S (unsafeDrop n bs) 0)
75− | otherwise = unsafeTake n (unsafeTail bs) :*: (S (unsafeDrop (n+1) bs) 0)
95+readByteString :: Int -> S -> ByteString
96+readByteString n s@(S bs o)
97+ -- no offset, easy.
98+ | o == 0 = unsafeTake n bs
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)))
76101
77−readWord8 :: S -> Int -> T Word8 S
78−readWord8 s@(S bs o) n
102+readWord8 :: Int -> S -> Word8
103+readWord8 n s@(S bs o)
79104 -- no bits at all, return 0
80− | n == 0 = 0 :*: s
105+ | n == 0 = 0
81106
82107 -- all bits are in the same byte
83108 -- we just need to shift and mask them right
84109 | n <= 8 - o = let w = unsafeHead bs
85110 m = make_mask n
86111 w' = (w `shiftr_w8` (8 - o - n)) .&. m
87− in w' :*: incS s n
112+ in w'
88113
89114 -- the bits are in two different bytes
90115 -- make a word16 using both bytes, and then shift and mask
@@ -92,7 +117,7 @@ readWord8 s@(S bs o) n
92117 (fromIntegral (unsafeIndex bs 1))
93118 m = make_mask n
94119 w' = (w `shiftr_w16` (16 - o - n)) .&. m
95− in fromIntegral w' :*: incS s n
120+ in fromIntegral w'
96121
97122bit_offset :: Int -> Int
98123bit_offset n = make_mask 3 .&. n
@@ -100,12 +125,11 @@ bit_offset n = make_mask 3 .&. n
100125byte_offset :: Int -> Int
101126byte_offset n = n `shiftR` 3
102127
103−readWord16be :: S -> Int -> T Word16 S
104−readWord16be s@(S bs o) n
128+readWord16be :: Int -> S -> Word16
129+readWord16be n s@(S bs o)
105130
106131 -- 8 or fewer bits, use readWord8
107− | n <= 8 = let w :*: s' = readWord8 s n
108− in fromIntegral w :*: s'
132+ | n <= 8 = fromIntegral (readWord8 n s)
109133
110134 -- handle 9 or more bits, stored in two bytes
111135
@@ -113,13 +137,13 @@ readWord16be s@(S bs o) n
113137 | o == 0 && n == 16 = let msb = fromIntegral (unsafeHead bs)
114138 lsb = fromIntegral (unsafeIndex bs 1)
115139 w = (msb `shiftl_w16` 8) .|. lsb
116− in w :*: incS s n
140+ in w
117141
118142 -- no offset, but not full 16 bytes
119143 | o == 0 = let msb = fromIntegral (unsafeHead bs)
120144 lsb = fromIntegral (unsafeIndex bs 1)
121145 w = (msb `shiftl_w16` (n-8)) .|. (lsb `shiftr_w16` (16-n))
122− in w :*: incS s n
146+ in w
123147
124148 -- with offset, and n=9-16
125149 | n <= 16 = readWithOffset s shiftl_w16 shiftr_w16 n
@@ -127,7 +151,7 @@ readWord16be s@(S bs o) n
127151 | otherwise = error "readWord16be: tried to read more than 16 bits"
128152
129153readWithoutOffset :: (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
131155readWithoutOffset s@(S bs o) shifterL shifterR n
132156 | o /= 0 = error "readWithoutOffset: there is an offset"
133157
@@ -136,7 +160,7 @@ readWithoutOffset s@(S bs o) shifterL shifterR n
136160 bn 0 = fromIntegral (unsafeHead bs)
137161 bn n = (bn (n-1) `shifterL` 8) .|. fromIntegral (unsafeIndex bs n)
138162
139− in (bn (segs-1) :*: incS s n)
163+ in bn (segs-1)
140164
141165 | n <= 64 = let segs = byte_offset n
142166 o' = bit_offset (n - 8 + o)
@@ -149,10 +173,10 @@ readWithoutOffset s@(S bs o) shifterL shifterR n
149173 last = (fromIntegral (unsafeIndex bs segs)) `shifterR` (8 - o')
150174
151175 w = msegs .|. last
152− in w :*: incS s n
176+ in w
153177
154178readWithOffset :: (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
156180readWithOffset s@(S bs o) shifterL shifterR n
157181 | n <= 64 = let bits_in_msb = 8 - o
158182 (n',top) = (n - bits_in_msb
@@ -161,7 +185,6 @@ readWithOffset s@(S bs o) shifterL shifterR n
161185 segs = byte_offset n'
162186
163187 bn 0 = 0
164− bn 1 = fromIntegral (unsafeIndex bs 1)
165188 bn n = (bn (n-1) `shifterL` 8) .|. fromIntegral (unsafeIndex bs n)
166189
167190 o' = bit_offset n'
@@ -172,18 +195,16 @@ readWithOffset s@(S bs o) shifterL shifterR n
172195 | otherwise = 0
173196
174197 w = top .|. mseg .|. last
175− in w :*: incS s n
198+ in w
176199
177200
178−readWord32be :: S -> Int -> T Word32 S
179−readWord32be s@(S bs o) n
201+readWord32be :: Int -> S -> Word32
202+readWord32be n s@(S bs o)
180203 -- 8 or fewer bits, use readWord8
181− | n <= 8 = let w :*: s' = readWord8 s n
182− in fromIntegral w :*: s'
204+ | n <= 8 = fromIntegral (readWord8 n s)
183205
184206 -- 16 or fewer bits, use readWord16be
185− | n <= 16 = let w :*: s' = readWord16be s n
186− in fromIntegral w :*: s'
207+ | n <= 16 = fromIntegral (readWord16be n s)
187208
188209 | o == 0 = readWithoutOffset s shiftl_w32 shiftr_w32 n
189210
@@ -192,15 +213,13 @@ readWord32be s@(S bs o) n
192213 | otherwise = error "readWord32be: tried to read more than 32 bits"
193214
194215
195−readWord64be :: S -> Int -> T Word64 S
196−readWord64be s@(S bs o) n
216+readWord64be :: Int -> S -> Word64
217+readWord64be n s@(S bs o)
197218 -- 8 or fewer bits, use readWord8
198− | n <= 8 = let w :*: s' = readWord8 s n
199− in fromIntegral w :*: s'
219+ | n <= 8 = fromIntegral (readWord8 n s)
200220
201221 -- 16 or fewer bits, use readWord16be
202− | n <= 16 = let w :*: s' = readWord16be s n
203− in fromIntegral w :*: s'
222+ | n <= 16 = fromIntegral (readWord16be n s)
204223
205224 | o == 0 = readWithoutOffset s shiftl_w64 shiftr_w64 n
206225
@@ -210,17 +229,13 @@ readWord64be s@(S bs o) n
210229
211230------------------------------------------------------------------------
212231-- unrolled codensity/state monad
213−newtype BitGet a = C { runCont :: forall r.
214− S ->
215− Success a r ->
216− Get r }
217−
218−type Success a r = S -> a -> Get r
232+newtype BitGet a = B { runState :: S -> Get (S,a) }
219233
220234instance Monad BitGet where
221− return x = C $ \s ks -> ks s x
222− fail str = C $ \s ks -> fail str
223− (C c) >>= f = C $ \s ks -> c s (\s' a -> runCont (f a) s' ks)
235+ return x = B $ \s -> return (s,x)
236+ fail str = B $ \_s -> fail str
237+ (B f) >>= g = B $ \s -> do (s',a) <- f s
238+ runState (g a) s'
224239
225240instance Functor BitGet where
226241 fmap f m = m >>= \a -> return (f a)
@@ -229,59 +244,66 @@ instance Applicative BitGet where
229244 pure x = return x
230245 fm <*> m = fm >>= \f -> m >>= \v -> return (f v)
231246
232−runBitGetSimple :: BitGet a -> Get a
233−runBitGetSimple bg = runBitGet bg
234−
235247runBitGet :: 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)
250+ return a
237251
238−-- | Make sure there are at lest @n@ bits.
252+getState :: BitGet S
253+getState = B $ \s -> return (s,s)
254+
255+putState :: S -> BitGet ()
256+putState s = B $ \_ -> return (s,())
257+
258+modifyState :: (S -> (S,a)) -> BitGet a
259+modifyState f = B $ \s -> return (f s)
260+
261+-- | Make sure there are at least @n@ bits.
239262ensureBits :: Int -> BitGet ()
240−ensureBits n = C $ \s ks ->
241− let loop (S bs o) | n <= (B.length bs * 8 - o) = ks (S bs o) ()
242− | otherwise = do
243− let currentBits = B.length bs * 8 - o
244− let byteCount = (n - currentBits + 7) `div` 8
245− bs' <- B.getByteString byteCount
246− -- should always be enough going once throguh the loop
247− loop (S (bs`append`bs') o)
248− in loop s
263+ensureBits n = do
264+ s@(S bs o) <- getState
265+ if n <= (B.length bs * 8 - o)
266+ then return ()
267+ else do let currentBits = B.length bs * 8 - o
268+ let byteCount = (n - currentBits + 7) `div` 8
269+ B $ \s -> do bs' <- B.getByteString byteCount
270+ return (S (bs`append`bs') o, ())
249271
250272getBool :: BitGet Bool
251−getBool = ensureBits 1 >> modifyState readBool
273+getBool = block bool
252274
253275getWord8 :: Int -> BitGet Word8
254−getWord8 n = ensureBits n >> modifyState (flip readWord8 n)
276+getWord8 n = block (word8 n)
255277
256278getWord16be :: Int -> BitGet Word16
257−getWord16be n = ensureBits n >> modifyState (flip readWord16be n)
279+getWord16be n = block (word16be n)
258280
259281getWord32be :: Int -> BitGet Word32
260−getWord32be n = ensureBits n >> modifyState (flip readWord32be n)
282+getWord32be n = block (word32be n)
261283
262284getWord64be :: Int -> BitGet Word64
263−getWord64be n = ensureBits n >> modifyState (flip readWord64be n)
285+getWord64be n = block (word64be n)
264286
265−-- | Get a @ByteString@. Drops the current byte if we've started to read
266−-- from it.
267287getByteString :: Int -> BitGet ByteString
268−getByteString n = do
269− ensureBits (n*8)
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
288+getByteString n = block (byteString n)
275289
276−getState :: BitGet S
277−getState = C $ \s ks -> ks s s
290+bool :: Block Bool
291+bool = Block 1 readBool
278292
279−putState :: S -> BitGet ()
280−putState s = C $ \_ ks -> ks s ()
293+word8 :: Int -> Block Word8
294+word8 n = Block n (readWord8 n)
295+
296+word16be :: Int -> Block Word16
297+word16be n = Block n (readWord16be n)
298+
299+word32be :: Int -> Block Word32
300+word32be n = Block n (readWord32be n)
301+
302+word64be :: Int -> Block Word64
303+word64be n = Block n (readWord64be n)
281304
282−modifyState :: (S -> (T a S)) -> BitGet a
283−modifyState f = C $ \s ks -> case f s of
284− w :*: s' -> ks s' w
305+byteString :: Int -> Block ByteString
306+byteString n = Block (n*8) (readByteString n)
285307
286308------------------------------------------------------------------------
287309-- Unchecked shifts, from the package binary