{-# LANGUAGE
    UnboxedTuples
  , UnicodeSyntax
  #-}
-- | Lazy bitstrings based on 'L.ByteString' which treats a byte
-- stream as a bit sequence in the same way as WavPack's manner.
module Codec.Audio.WavPack.BitString
    ( -- * The BitString Type
      BitString

      -- * Construction
    , (∅)
    , singleton

      -- * Basic Interface
    , cons
    , snoc
    , head
    , uncons
    , tail
    , null
    , length

      -- * Conversion from/to 'L.ByteString'
    , fromByteString
    , toByteString
    )
    where
import qualified Data.ByteString.Lazy as L
import Data.Bits
import Data.Int
import qualified Data.Strict as S
import Data.Word
import Prelude hiding (head, length, null, tail)
import Prelude.Unicode

-- | The BitString Type.
data BitString
    = BitString {
        leftRem    ∷ !Remnant
      , leftBytes  ∷ !L.ByteString
      -- | reversed
      , rightBytes ∷ !L.ByteString
      , rightRem   ∷ !Remnant
      }
    deriving (Eq, Show)

data Remnant
    = Remnant {
      -- | bit length from 0 to 8
        remLen  ∷ !Int
      -- | current byte
      , remByte ∷ !Word8
      }
    deriving (Eq, Show)

remEmpty ∷ Remnant
remEmpty = Remnant 0 0

remNull ∷ Remnant → Bool
remNull = (0 ≡) ∘ remLen

remSingleton ∷ Bool -> Remnant
remSingleton b = Remnant {
                   remLen  = 1
                 , remByte = if b then 1 else 0
                 }

byteToRem ∷ Word8 → Remnant
byteToRem w = Remnant {
                remLen  = 8
              , remByte = w
              }

remToStr ∷ Remnant → L.ByteString
remToStr r
    | remNull r = L.empty
    | otherwise = L.singleton $ remByte r

consRem ∷ Bool → Remnant → (# Remnant, S.Maybe Word8 #)
consRem b r
    | remLen r ≡ 8 = (# remSingleton b, S.Just $ remByte r #)
    | remLen r ≡ 7 = (# remEmpty, S.Just w' #)
    | otherwise     = let !r' = r {
                                  remLen  = remLen r + 1
                                , remByte = w'
                                }
                      in
                        (# r', S.Nothing #)
    where
      w' = (remByte r `shiftR` 1)
           .|.
           if b then 1 else 0

snocRem ∷ Remnant → Bool → (# S.Maybe Word8, Remnant #)
snocRem r b
    | remLen r ≡ 8 = (# S.Just $ remByte r, remSingleton b #)
    | remLen r ≡ 7 = (# S.Just w', remEmpty #)
    | otherwise     = let !r' = r {
                                  remLen  = remLen r + 1
                                , remByte = w'
                                }
                      in
                        (# S.Nothing, r' #)
    where
      w' | b         = bit (remLen r) .|. remByte r
         | otherwise = remByte r

unconsRem ∷ Remnant → (# S.Maybe Bool, Remnant #)
unconsRem r
    | remNull r = (# S.Nothing, remEmpty #)
    | otherwise = let !b  = remByte r `testBit` 1
                      !r' = Remnant {
                              remLen  = remLen r - 1
                            , remByte = remByte r `shiftR` 1
                            }
                  in
                    (# S.Just b, r' #)

unconsBytes ∷ L.ByteString → (# S.Maybe Remnant, L.ByteString #)
unconsBytes bs
    = case L.uncons bs of
        Just (w, bs')
            → (# S.Just (byteToRem w), bs' #)
        Nothing
            → (# S.Nothing, L.empty #)

-- | /O(1)/ The empty 'BitString'.
(∅) ∷ BitString
(∅) = BitString {
        leftRem    = remEmpty
      , leftBytes  = L.empty
      , rightBytes = L.empty
      , rightRem   = remEmpty
      }

-- | /O(1)/ Convert a 'Bool' into a 'BitString'.
singleton ∷ Bool → BitString
singleton b = BitString {
                leftRem    = remSingleton b
              , leftBytes  = L.empty
              , rightBytes = L.empty
              , rightRem   = remEmpty
              }

-- | /O(1)/ Prepend a bit to the beginning of a 'BitString'.
cons ∷ Bool → BitString → BitString
cons b bs
    = case consRem b $ leftRem bs of
        (# lr', S.Just w #)
            → bs {
                 leftRem   = lr'
               , leftBytes = L.cons' w $ leftBytes bs
               }

        (# lr', S.Nothing #)
            → bs { leftRem = lr' }

-- | /O(1)/ Append a bit to the end of a 'BitString'.
snoc ∷ BitString → Bool → BitString
snoc bs b
    = case snocRem (rightRem bs) b of
        (# S.Just w, rr' #)
            → bs {
                 rightBytes = L.cons' w $ rightBytes bs
               , rightRem   = rr'
               }

        (# S.Nothing, rr' #)
            → bs { rightRem = rr' }

-- | /amortized O(1)/ Extract the first bit of a 'BitString', which
-- must be non-empty.
head ∷ BitString → Bool
head bs = case uncons bs of
            Just (b, _ )
                → b
            Nothing
                → error "head: empty BitString"

-- | /amortized O(1)/ Extract the bits after the head of a
-- 'BitString', which must be non-empty.
tail ∷ BitString → BitString
tail bs = case uncons bs of
            Just (_, bs')
                → bs'
            Nothing
                → error "tail: empty BitString"

-- | /amortized O(1)/ Extract the the head and tail of a 'BitString',
-- returning 'Nothing' if it's empty.
uncons ∷ BitString → Maybe (Bool, BitString)
uncons bs
    = case unconsRem $ leftRem bs of
        (# S.Just b, lr' #)
            → Just (b, bs { leftRem = lr' })

        (# S.Nothing, _ #)
            → case unconsBytes $ leftBytes bs of
                 (# S.Just lr, lb' #)
                     → let !bs' = bs { leftRem   = lr
                                      , leftBytes = lb'
                                      }
                        in
                          uncons bs'

                 (# S.Nothing, _ #)
                     → if L.null $ rightBytes bs then
                            case unconsRem $ rightRem bs of
                              (# S.Just b, rr' #)
                                  → Just (b, bs { rightRem = rr' })

                              (# S.Nothing, _ #)
                                  → Nothing
                        else
                            let !bs' = bs { leftBytes = L.reverse
                                                        $ rightBytes bs
                                          }
                            in
                              uncons bs'

-- | /O(1)/ Test whether a 'BitString' is empty.
null ∷ BitString → Bool
null bs
    = remLen (leftRem bs) ≡ 0
      ∧
      L.null (leftBytes bs)
      ∧
      L.null (rightBytes bs)
      ∧
      remLen (rightRem bs) ≡ 0

-- | /O(n)/ Return the number of bits in a 'BitString'.
length ∷ Integral n ⇒ BitString → n
length bs
    = fromIntegral $ ( fromIntegral (remLen $ leftRem bs)
                       +
                       L.length (leftBytes bs) ⋅ 8
                       +
                       L.length (rightBytes bs) ⋅ 8
                       +
                       fromIntegral (remLen $ rightRem bs)
                     )
{-# SPECIALISE length ∷ BitString → Int64 #-}

-- | /O(1)/ Convert a 'L.ByteString' into a 'BitString'.
fromByteString ∷ L.ByteString → BitString
fromByteString bs = BitString {
                      leftRem    = remEmpty
                    , leftBytes  = bs
                    , rightBytes = L.empty
                    , rightRem   = remEmpty
                    }

-- | /O(n)/ Convert a 'BitString' into 'L.ByteString', padding
-- incomplete bytes to single bytes with necessary 0's at their
-- MSB. Thus the following equation does not hold when the length of
-- 'BitString' isn't multiple of 8.
--
-- > fromByteString . toByteString = id
--
-- But the following always holds true.
--
-- > toByteString . fromByteString = id
toByteString ∷ BitString → L.ByteString
toByteString bs
    = L.concat [ remToStr $ leftRem bs
               , leftBytes bs
               , L.reverse $ rightBytes bs
               , remToStr $ rightRem bs
               ]