BitPut.hs

-----------------------------------------------------------------------------
-- |
-- Module      : Data.Binary.Strict.BitPut
-- Copyright   : Dominic Steinitz
-- License     : BSD3-style (see LICENSE)
--
-- Maintainer  : Dominic Steinitz <dominic.steinitz@blueyonder.co.uk>
-- Stability   : experimental
--
-- This is the writer dual to BitGet. It allows one to append bits in a monad
-- and get a strict ByteString as a result. Bits are appended from the MSB of
-- the first byte towards the LSB of the last byte.
--
-- This is best suited to small bit-fields because it accumulates bytes using
-- snoc, so large results will cause a lot of copying. It would be possible
-- to switch to using something similar to the Builder monad if need arises.
-- However, since most protocols only have small bit fields, this should
-- suffice for many cases.
-----------------------------------------------------------------------------
module Data.Binary.Strict.BitPut
  ( BitPut
  , runBitPut
  , putBit
  , putNBits
  , putBits
  , putLeftByteString
  ) where

import Control.Monad
import Data.Word (Word8)
import Data.Bits (shiftL, shiftR, (.&.), (.|.), bitSize, Bits)
import qualified Data.ByteString as B

import Data.Binary.Strict.BitUtil

-- | The state of the BitPut.
--   The current offset is in [0..7]. The current byte is packed from MSB, downwards
data S = S {-# UNPACK #-} !B.ByteString  -- ^ output
           {-# UNPACK #-} !Word8  -- ^ bit offset in current byte
           {-# UNPACK #-} !Word8  -- ^ current byte
           deriving (Show)

newtype BitPut' a = BitPut' { unPut :: S -> (a,S) }
type BitPut = BitPut' ()

instance Functor BitPut' where
   fmap f m = BitPut' (\s -> let (a,s') = unPut m s in (f a,s'))

instance Monad BitPut' where
   return a = BitPut' (\s -> (a,s))
   m >>= k = BitPut' (\s -> let (a,s')  = unPut m s in unPut (k a) s')

get :: BitPut' S
get = BitPut' (\s -> (s,s))

put :: S -> BitPut
put s = BitPut' (const ((), s))

-- | Append a single bit
putBit :: Bool -> BitPut
putBit bit = do
  S bytes boff curr <- get
  let v = if bit then 1 else 0
      newCurr = curr .|. (shiftL (fromIntegral v) (fromIntegral (7 - boff)))
      newBoff = boff + 1
  if newBoff == 8
     then put (S (bytes `B.snoc` (curr .|. v)) 0 0)
     else put (S bytes newBoff newCurr)

-- | Append the bottom n bits of the given bits value. In the case that more
--   bits are requested than the value provides, this acts as if the value
--   has as unlimited number of leading 0 bits.
putNBits :: (Integral a, Bits a) => Int -> a -> BitPut
putNBits n v
  | n == 0 = return ()
  | otherwise = do
      S bytes boff curr <- get
      let space = 8 - boff
      if n < fromIntegral space
         then do
           let boff'   = boff + fromIntegral n
               shifted = (fromIntegral v .&. bottomNBits n) `shiftL` (8 - (fromIntegral boff + n))
               curr'  = curr .|. shifted
           put $ S bytes boff' curr'
         else do
           let bits = v `shiftR` remainingBits
               remainingBits = n - fromIntegral space
               mask = bottomNBits $ fromIntegral space
               bytes' = bytes `B.snoc` (curr .|. (mask .&. fromIntegral bits))
           put $ S bytes' 0 0
           putNBits remainingBits v

-- | Append a value. Note that this function is undefined for instances of Bits
--   which have no fixed bitsize (like Integer)
putBits :: (Integral a, Bits a) => a -> BitPut
putBits v = putNBits (bitSize v) v

-- | Append the first n bits of a left aligned ByteString.
putLeftByteString :: Int -> B.ByteString -> BitPut
putLeftByteString bits bs
  | bits < 8 = putNBits bits $ B.head bs
  | otherwise = putBits (B.head bs) >> putLeftByteString (bits - 8) (B.tail bs)

runBitPut :: BitPut -> B.ByteString
runBitPut m = r where
   (_, (S bytes boff curr)) = unPut m (S B.empty 0 0)
   r = if boff > 0
          then bytes `B.snoc` (topNBits (fromIntegral boff) .&. curr)
          else bytes