@@ -52,20 +52,31 @@ data S = S !Builder !Word8 !Int
5252putBool :: Bool -> BitPut ()
5353putBool b = putWord8 1 (if b then 0xff else 0x00)
5454
55+-- | make_mask 3 = 00000111
56+make_mask :: (Bits a, Num a) => Int -> a
57+make_mask n = (1 `shiftL` fromIntegral n) - 1
58+{-# SPECIALIZE make_mask :: Int -> Int #-}
59+{-# SPECIALIZE make_mask :: Int -> Word #-}
60+{-# SPECIALIZE make_mask :: Int -> Word8 #-}
61+{-# SPECIALIZE make_mask :: Int -> Word16 #-}
62+{-# SPECIALIZE make_mask :: Int -> Word32 #-}
63+{-# SPECIALIZE make_mask :: Int -> Word64 #-}
64+
5565-- | Put the @n@ lower bits of a 'Word8'.
5666putWord8 :: Int -> Word8 -> BitPut ()
5767putWord8 n w = BitPut $ \s -> PairS () $
68+ let w' = make_mask n .&. w in
5869 case s of
5970 -- a whole word8, no offset
6071 (S b t o) | n == 8 && o == 0 -> flush $ S b w n
6172 -- less than a word8, will fit in the current word8
62− | n <= 8 - o -> flush $ S b (t .|. (w `shiftL` (8 - n - o))) (o+n)
73+ | n <= 8 - o -> flush $ S b (t .|. (w' `shiftL` (8 - n - o))) (o+n)
6374 -- will finish this word8, and spill into the next one
6475 | otherwise -> flush $
6576 let o' = o + n - 8
66− w' = t .|. (w `shiftR` o')
77+ b' = t .|. (w' `shiftR` o')
6778 t' = w `shiftL` (8 - o')
68− in S (b `mappend` B.singleton w') t' o'
79+ in S (b `mappend` B.singleton b') t' o'
6980
7081-- | Put the @n@ lower bits of a 'Word16'.
7182putWord16be :: Int -> Word16 -> BitPut ()
@@ -73,22 +84,23 @@ putWord16be n w
7384 | n <= 8 = putWord8 n (fromIntegral w)
7485 | otherwise =
7586 BitPut $ \s -> PairS () $
87+ let w' = make_mask n .&. w in
7688 case s of
7789 -- as n>=9, it's too big to fit into one single byte
7890 -- it'll either use 2 or 3 bytes
7991 -- it'll fit in 2 bytes
8092 (S b t o) | o + n <= 16 -> flush $
8193 let o' = o + n - 8
82− w' = t .|. fromIntegral (w `shiftR` o')
94+ b' = t .|. fromIntegral (w' `shiftR` o')
8395 t' = fromIntegral (w `shiftL` (8-o'))
84− in (S (b `mappend` B.singleton w') t' o')
96+ in (S (b `mappend` B.singleton b') t' o')
8597 -- 3 bytes required
8698 | otherwise -> flush $
8799 let o' = o + n - 16
88− w' = t .|. fromIntegral (w `shiftR` (o' + 8))
89− w'' = fromIntegral ((w `shiftR` o') .&. 0xff)
100+ b' = t .|. fromIntegral (w' `shiftR` (o' + 8))
101+ b'' = fromIntegral ((w `shiftR` o') .&. 0xff)
90102 t' = fromIntegral (w `shiftL` (8-o'))
91− in (S (b `mappend` B.singleton w' `mappend` B.singleton w'') t' o')
103+ in (S (b `mappend` B.singleton b' `mappend` B.singleton b'') t' o')
92104
93105-- | Put the @n@ lower bits of a 'Word32'.
94106putWord32be :: Int -> Word32 -> BitPut ()