@@ -52,20 +52,31 @@ data S = S !Builder !Word8 !Int
52putBool :: Bool -> BitPut ()
52putBool :: Bool -> BitPut ()
53putBool b = putWord8 1 (if b then 0xff else 0x00)
53putBool b = putWord8 1 (if b then 0xff else 0x00)
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 #-}
55-- | Put the @n@ lower bits of a 'Word8'.
65-- | Put the @n@ lower bits of a 'Word8'.
56putWord8 :: Int -> Word8 -> BitPut ()
66putWord8 :: Int -> Word8 -> BitPut ()
57putWord8 n w = BitPut $ \s -> PairS () $
67putWord8 n w = BitPut $ \s -> PairS () $
68+ let w' = make_mask n .&. w in
59 -- a whole word8, no offset
70 -- a whole word8, no offset
60 (S b t o) | n == 8 && o == 0 -> flush $ S b w n
71 (S b t o) | n == 8 && o == 0 -> flush $ S b w n
61 -- less than a word8, will fit in the current word8
72 -- 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)
63 -- will finish this word8, and spill into the next one
74 -- will finish this word8, and spill into the next one
64 | otherwise -> flush $
75 | otherwise -> flush $
65 let o' = o + n - 8
76 let o' = o + n - 8
66− w' = t .|. (w `shiftR` o')
77+ b' = t .|. (w' `shiftR` o')
67 t' = w `shiftL` (8 - o')
78 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'
70-- | Put the @n@ lower bits of a 'Word16'.
81-- | Put the @n@ lower bits of a 'Word16'.
71putWord16be :: Int -> Word16 -> BitPut ()
82putWord16be :: Int -> Word16 -> BitPut ()
@@ -73,22 +84,23 @@ putWord16be n w
73 | n <= 8 = putWord8 n (fromIntegral w)
84 | n <= 8 = putWord8 n (fromIntegral w)
74 | otherwise =
85 | otherwise =
75 BitPut $ \s -> PairS () $
86 BitPut $ \s -> PairS () $
87+ let w' = make_mask n .&. w in
77 -- as n>=9, it's too big to fit into one single byte
89 -- as n>=9, it's too big to fit into one single byte
78 -- it'll either use 2 or 3 bytes
90 -- it'll either use 2 or 3 bytes
79 -- it'll fit in 2 bytes
91 -- it'll fit in 2 bytes
80 (S b t o) | o + n <= 16 -> flush $
92 (S b t o) | o + n <= 16 -> flush $
81 let o' = o + n - 8
93 let o' = o + n - 8
82− w' = t .|. fromIntegral (w `shiftR` o')
94+ b' = t .|. fromIntegral (w' `shiftR` o')
83 t' = fromIntegral (w `shiftL` (8-o'))
95 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')
85 -- 3 bytes required
97 -- 3 bytes required
86 | otherwise -> flush $
98 | otherwise -> flush $
87 let o' = o + n - 16
99 let o' = o + n - 16
88− w' = t .|. fromIntegral (w `shiftR` (o' + 8))
100+ b' = t .|. fromIntegral (w' `shiftR` (o' + 8))
89− w'' = fromIntegral ((w `shiftR` o') .&. 0xff)
101+ b'' = fromIntegral ((w `shiftR` o') .&. 0xff)
90 t' = fromIntegral (w `shiftL` (8-o'))
102 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')
93-- | Put the @n@ lower bits of a 'Word32'.
105-- | Put the @n@ lower bits of a 'Word32'.
94putWord32be :: Int -> Word32 -> BitPut ()
106putWord32be :: Int -> Word32 -> BitPut ()