-----------------------------------------------------------------------------
-- |
-- Module      : Data.Binary.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.BitPut
  ( BitPut
  , runBitPut
  , putBit
  , putNBits
  , putBits
  , putByteString
  , putLeftByteString
  ) where

import Data.Bits (bitSize, Bits)
import Control.Monad
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL
import qualified Data.Binary.BitBuilder as BB

newtype BitPut' a = BitPut' { unPut :: (a, BB.BitBuilder) }

type BitPut = BitPut' ()

instance Functor BitPut' where
   fmap f m = BitPut' (let (a, w) = unPut m in (f a, w))

instance Monad BitPut' where
   return a = BitPut' (a,BB.empty)
   m >>= k = BitPut' (let (a, w) = unPut m
                          (b, w') = unPut (k a)
                       in (b, w `BB.append` w'))

   m >> k = BitPut' (let (_, w) = unPut m
                         (b, w') = unPut k
                      in (b, w `BB.append` w'))
   {-# INLINE (>>) #-}

-- | Append a single bit
putBit :: Bool -> BitPut
putBit bit = BitPut' ((), BB.singleton bit)

-- | 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 = BitPut' ((), BB.fromBits n 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 a ByteString
putByteString :: B.ByteString -> BitPut
putByteString bs = BitPut' ((), BB.fromByteString (bs, 0))

-- | Append a left aligned ByteString where ByteString has a partial byte
--   with the given number of valid bits, from the MSB downwards. The number
--   of such bits must be 0..7. (A normal ByteString, which all bytes full
--   would use 0)
putLeftByteString :: (B.ByteString, Int) -> BitPut
putLeftByteString bs = BitPut' ((), BB.fromByteString bs)

runBitPut :: BitPut -> BL.ByteString
runBitPut m = let (_, w) = unPut m
               in BB.toLazyByteString w