@@ -82,6 +82,7 @@ module Data.Binary.Bits.Get
8282 ) where
8383
8484import Data.Binary.Get as B ( runGet, Get, getByteString )
85+import Data.Binary.Get.Internal as B ( get, put, ensureN )
8586
8687import Data.ByteString as B
8788import Data.ByteString.Unsafe
@@ -335,7 +336,7 @@ newtype BitGet a = B { runState :: S -> Get (S,a) }
335336
336337instance Monad BitGet where
337338 return x = B $ \s -> return (s,x)
338− fail str = B $ \_s -> fail str
339+ fail str = B $ \(S inp n) -> putBackState inp n >> fail str
339340 (B f) >>= g = B $ \s -> do (s',a) <- f s
340341 runState (g a) s'
341342
@@ -350,9 +351,22 @@ instance Applicative BitGet where
350351-- been partially consumed it will be discarded once 'runBitGet' is finished.
351352runBitGet :: BitGet a -> Get a
352353runBitGet bg = do
353− (_,a) <- runState bg (S B.empty 0)
354+ s <- mkInitState
355+ ((S str' n),a) <- runState bg s
356+ putBackState str' n
354357 return a
355358
359+mkInitState :: Get S
360+mkInitState = do
361+ str <- get
362+ put B.empty
363+ return (S str 0)
364+
365+putBackState :: B.ByteString -> Int -> Get ()
366+putBackState bs n = do
367+ remaining <- get
368+ put (B.drop (if n==0 then 0 else 1) bs `B.append` remaining)
369+
356370getState :: BitGet S
357371getState = B $ \s -> return (s,s)
358372
@@ -370,7 +384,9 @@ ensureBits n = do
370384 then return ()
371385 else do let currentBits = B.length bs * 8 - o
372386 let byteCount = (n - currentBits + 7) `div` 8
373− B $ \_ -> do bs' <- B.getByteString byteCount
387+ B $ \_ -> do B.ensureN byteCount
388+ bs' <- B.get
389+ put B.empty
374390 return (S (bs`append`bs') o, ())
375391
376392-- | Get 1 bit as a 'Bool'.