git.iko.soy/binary-bits
on d259c3a help sign in
binary-bits/BitsPut.hs
115 lines 3.7 kB history raw
1module BitsPut
2 ( BitPut
3 , runBitPut
4 , joinPut
6 , putBool
7 , putWord8
8 , putWord16be
9 , putWord32be
10 , putWord64be
11 )
12 where
14import qualified Data.Binary.Builder as B
15import Data.Binary.Builder ( Builder )
16import qualified Data.Binary.Put as Put
17import Data.Binary.Put ( Put )
19import Data.ByteString
21import Data.Bits
22import Data.Monoid
23import Data.Word
25data BitPut a = BitPut { run :: (S -> PairS a) }
27data PairS a = PairS a {-# UNPACK #-} !S
29data S = S !Builder !Word8 !Int
31instance Show S where
32 show (S _ w o) = "S( " ++ show w ++ " : " ++ show o ++ ")"
34putBool :: Bool -> BitPut ()
35putBool b = putWord8 1 (if b then 1 else 0)
37putWord8 :: Int -> Word8 -> BitPut ()
38putWord8 n w = BitPut $ \s -> PairS () $
39 case s of
40 -- a whole word8, no offset
41 (S b t o) | n == 8 && o == 0 -> flush $ S b w n
42 -- less than a word8, will fit in the current word8
43 | n <= 8 - o -> flush $ S b (t .|. (w `shiftL` (8 - n - o))) (o+n)
44 -- will finish this word8, and spill into the next one
45 | otherwise -> flush $
46 let o' = o + n - 8
47 w' = t .|. (w `shiftR` o')
48 t' = w `shiftL` (8 - o')
49 in S (b `mappend` B.singleton w') t' o'
51putWord16be :: Int -> Word16 -> BitPut ()
52putWord16be n w
53 | n <= 8 = putWord8 n (fromIntegral w)
54 | otherwise =
55 BitPut $ \s -> PairS () $
56 case s of
57 -- as n>=9, it's too big to fit into one single byte
58 -- it'll either use 2 or 3 bytes
59 -- it'll fit in 2 bytes
60 (S b t o) | o + n <= 16 -> flush $
61 let o' = o + n - 8
62 w' = t .|. fromIntegral (w `shiftR` o')
63 t' = fromIntegral (w `shiftL` (8-o'))
64 in (S (b `mappend` B.singleton w') t' o')
65 -- 3 bytes required
66 | otherwise -> flush $
67 let o' = o + n - 16
68 w' = t .|. fromIntegral (w `shiftR` (o' + 8))
69 w'' = fromIntegral ((w `shiftR` o') .&. 0xff)
70 t' = fromIntegral (w `shiftL` (8-o'))
71 in (S (b `mappend` B.singleton w' `mappend` B.singleton w'') t' o')
73putWord32be :: Int -> Word32 -> BitPut ()
74putWord32be n w
75 | n <= 16 = putWord16be n (fromIntegral w)
76 | otherwise = do
77 putWord32be (n-16) (w`shiftR`16)
78 putWord32be 16 (w .&. 0x0000ffff)
80putWord64be :: Int -> Word64 -> BitPut ()
81putWord64be n w
82 | n <= 32 = putWord32be n (fromIntegral w)
83 | otherwise = do
84 putWord64be (n-32) (w`shiftR`32)
85 putWord64be 32 (w .&. 0xffffffff)
87joinPut :: Put -> BitPut ()
88joinPut m = BitPut $ \s0 -> PairS () $
89 let (S b0 _ _) = flushIncomplete s0
90 b = Put.execPut m
91 in (S (b0`mappend`b) 0 0)
93flush :: S -> S
94flush s@(S b w o)
95 | o > 8 = error "flush: offset > 8"
96 | o == 8 = S (b `mappend` B.singleton w) 0 0
97 | otherwise = s
99flushIncomplete :: S -> S
100flushIncomplete s@(S b w o)
101 | o == 0 = s
102 | otherwise = (S (b `mappend` B.singleton w) 0 0)
104runBitPut :: BitPut () -> Put.Put
105runBitPut m = Put.putBuilder b'
106 where
107 PairS _ s@(S b t o) = run m (S mempty 0 0)
108 (S b' _ _) = flushIncomplete s
110instance Monad BitPut where
111 m >>= k = BitPut $ \s ->
112 let PairS a s' = run m s
113 PairS b s'' = run (k a) s'
114 in PairS b s''
115 return x = BitPut $ \s -> PairS x s