- = let !bitCount = countBits maxCode
- !extras = bit bitCount - maxCode - 1
- !code = B.toBits (B.take (bitCount - 1) bs)
- (# code', bitCount' #)
- = if code ≥ extras then
- (# (code `shiftL` 1)
- - extras
- + b2n (bs B.!! bitCount)
- , bitCount #)
- else
- (# code, bitCount - 1 #)
- !bs' = B.drop bitCount' bs
- in
- (# code', bs' #)
+ = do let bitCount = countBits maxCode
+ extras = bit bitCount - maxCode - 1
+ code ← takeBits bs (bitCount - 1)
+ if code ≥ extras then
+ do nextBit ← takeHead bs
+ return $ (code `shiftL` 1) - extras + b2n nextBit
+ else
+ return code
+
+takeHead ∷ Bitstream bs ⇒ STRef s bs → ST s Bool
+{-# INLINEABLE takeHead #-}
+takeHead bsr
+ = do bs ← readSTRef bsr
+ writeSTRef bsr (B.tail bs)
+ return (B.head bs)
+
+takeBits ∷ (Integral n, Bitstream bs, Bits a) ⇒ STRef s bs → n → ST s a
+{-# INLINEABLE takeBits #-}
+takeBits bsr n
+ = do bs ← readSTRef bsr
+ writeSTRef bsr (B.drop n bs)
+ return (B.toBits (B.take n bs))
+
+-- | C style /for/ loop with /break/ and /continue/.
+for ∷ ∀m α. MonadCont m
+ ⇒ α -- ^ Initial state
+ → (α → Bool) -- ^ Continue-the-loop predicate
+ → (α → α) -- ^ State modifier
+ → (α → m () → m () → m ()) -- ^ Loop body taking breaker and
+ -- continuer
+ → m α -- ^ Final state
+for α0 contLoop next body
+ = callCC $ \break → loop break α0
+ where
+ loop ∷ (α → m ()) → α → m α
+ loop break α
+ | contLoop α
+ = do callCC $ \continue → body α (break α) (continue ())
+ loop break (next α)
+ | otherwise
+ = return α