- = 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))