@@ -0,0 +1,66 @@
1+module BitsPut
2+ ( putBool
3+ , putWord8
4+ , runBitPut
5+ )
6+ where
7+
8+import qualified Data.Binary.Builder as B
9+import Data.Binary.Builder ( Builder )
10+import qualified Data.Binary.Put as Put
11+import Data.ByteString
12+
13+import Data.Bits
14+import Data.Monoid
15+import Data.Word
16+
17+data BitPut a = BitPut { run :: (S -> PairS a) }
18+
19+data PairS a = PairS a {-# UNPACK #-} !S
20+
21+data S = S !Builder !Word8 !Int
22+
23+putBool :: Bool -> BitPut ()
24+putBool b = putWord8 1 (if b then 1 else 0)
25+
26+putWord8 :: Int -> Word8 -> BitPut ()
27+putWord8 n w = BitPut $ \s -> PairS () $
28+ case s of
29+ -- a whole word8, no offset
30+ (S b t o) | n == 8 && o == 0 -> flush $ S b w n
31+ -- less than a word8, will fit in the current word8
32+ | n <= 8 - o -> flush $ S b (t .|. (w `shiftL` (8 - n - o))) (o+n)
33+ -- will finish this word8, and spill into the next one
34+ | otherwise -> flush $
35+ let o' = o + n - 8
36+ w' = t .|. (w `shiftR` o')
37+ t' = w `shiftL` (8 - o')
38+ in S (b `mappend` B.singleton w') t' o'
39+
40+putWord16be :: Int -> Word16 -> BitPut ()
41+putWord16be n w
42+ | n <= 8 = putWord8 n (fromIntegral w)
43+ | otherwise =
44+ BitPut $ \s -> PairS () $
45+ case s of
46+ (S b t o) | otherwise -> undefined
47+
48+flush :: S -> S
49+flush s@(S b w o)
50+ | o > 8 = error "flush: offset > 8"
51+ | o == 8 = S (b `mappend` B.singleton w) 0 0
52+ | otherwise = s
53+
54+runBitPut :: BitPut () -> Put.Put
55+runBitPut m = Put.putBuilder b'
56+ where
57+ PairS _ (S b t o) = run m (S mempty 0 0)
58+ b' | o == 0 = b
59+ | otherwise = b `mappend` B.singleton t
60+
61+instance Monad BitPut where
62+ m >>= k = BitPut $ \s ->
63+ let PairS a s' = run m s
64+ PairS b s'' = run (k a) s'
65+ in PairS b s''
66+ return x = BitPut $ \s -> PairS x s