git.iko.soy/binary-bits
on master help sign in
Make ensureBits pull just enough bytes from binary
771f4b2Lennart Kolmodin
BitsGet.hs+7−10
compare with master
unifiedsplit
@@ -19,7 +19,7 @@ module BitsGet
19
19
20 ) where
20 ) where
21
21
22−import Data.Binary.Get ( runGet, Get, needMore, getByteString, getS, putS )
22+import Data.Binary.Get ( runGet, Get, getByteString )
23
23
24import Data.ByteString as B
24import Data.ByteString as B
25import Data.ByteString.Internal
25import Data.ByteString.Internal
@@ -254,15 +254,11 @@ runBitGetSimple bg = do
254
254
255runBitGet :: BitGet a -> Get (Either String a)
255runBitGet :: BitGet a -> Get (Either String a)
256runBitGet bg = do
256runBitGet bg = do
257− bs <- getS -- get the current chunk
257+ let bs = B.empty -- initial state
258− putS B.empty
259 v <- runCont bg (S bs 0) (\str -> return (Left str)) (\s a -> return (Right (s, a)))
258 v <- runCont bg (S bs 0) (\str -> return (Left str)) (\s a -> return (Right (s, a)))
260 case v of
259 case v of
261 Left err -> return (Left err)
260 Left err -> return (Left err)
262− Right ((S bs' n), x) -> do
261+ Right ((S _bs' _n), x) -> do
263− let bs'' | n > 0 = B.tail bs'
264− | otherwise = bs'
265− putS bs''
266 return (Right x)
262 return (Right x)
267
263
268-- | Make sure there are at lest @n@ bits.
264-- | Make sure there are at lest @n@ bits.
@@ -270,9 +266,10 @@ ensureBits :: Int -> BitGet ()
270ensureBits n = C $ \s kf ks ->
266ensureBits n = C $ \s kf ks ->
271 let loop (S bs o) | n <= (B.length bs * 8 - o) = ks (S bs o) ()
267 let loop (S bs o) | n <= (B.length bs * 8 - o) = ks (S bs o) ()
272 | otherwise = do
268 | otherwise = do
273− needMore
269+ let currentBits = B.length bs * 8 - o
274− bs' <- getS
270+ let byteCount = (n - currentBits + 7) `div` 8
275− putS B.empty
271+ bs' <- getByteString byteCount
272+ -- should always be enough going once throguh the loop
276 loop (S (bs`append`bs') o)
273 loop (S (bs`append`bs') o)
277 in loop s
274 in loop s
278
275