Codec/Audio/WavPack/BitString.hs

   1 {-# LANGUAGE
   2     BangPatterns
   3   , UnboxedTuples
   4   , UnicodeSyntax
   5   #-}
   6 -- | Lazy bitstrings based on 'L.ByteString' which treats a byte
   7 -- stream as a bit sequence in the same way as WavPack's manner.
   8 module Codec.Audio.WavPack.BitString
   9     ( -- * The BitString Type
  10       BitString
  11
  12       -- * Construction
  13     , (∅)
  14     , singleton
  15
  16       -- * Basic Interface
  17     , cons
  18     , snoc
  19     , head
  20     , uncons
  21     , tail
  22     , null
  23     , length
  24
  25       -- * Conversion from/to 'L.ByteString'
  26     , fromByteString
  27     , toByteString
  28
  29       -- * Substrings
  30     , take
  31     , drop
  32     , splitAt
  33     )
  34     where
  35 import qualified Data.ByteString.Lazy as L
  36 import Data.Bits
  37 import Data.Int
  38 import qualified Data.Strict as S
  39 import Data.Word
  40 import Prelude hiding (drop, head, length, null, splitAt, tail, take)
  41 import Prelude.Unicode
  42
  43 -- | The BitString Type.
  44 data BitString
  45     = BitString {
  46         leftRem    ∷ !Remnant
  47       , leftBytes  ∷ !L.ByteString
  48       -- | reversed
  49       , rightBytes ∷ !L.ByteString
  50       , rightRem   ∷ !Remnant
  51       }
  52     deriving (Eq, Show)
  53
  54 data Remnant
  55     = Remnant {
  56       -- | bit length from 0 to 8
  57         remLen  ∷ !Int
  58       -- | current byte @(0 <= remByte <= (1 `'shiftL'` 'remLen') - 1)@
  59       , remByte ∷ !Word8
  60       }
  61     deriving (Eq, Show)
  62
  63 remEmpty ∷ Remnant
  64 remEmpty = Remnant 0 0
  65
  66 remNull ∷ Remnant → Bool
  67 remNull = (0 ≡) ∘ remLen
  68
  69 remSingleton ∷ Bool -> Remnant
  70 remSingleton b = Remnant {
  71                    remLen  = 1
  72                  , remByte = if b then 1 else 0
  73                  }
  74
  75 byteToRem ∷ Word8 → Remnant
  76 byteToRem w = Remnant {
  77                 remLen  = 8
  78               , remByte = w
  79               }
  80
  81 remToStr ∷ Remnant → L.ByteString
  82 remToStr r
  83     | remNull r = L.empty
  84     | otherwise = L.singleton $ remByte r
  85
  86 consRem ∷ Bool → Remnant → (# Remnant, S.Maybe Word8 #)
  87 consRem b r
  88     | remLen r ≡ 8 = (# remSingleton b, S.Just $ remByte r #)
  89     | remLen r ≡ 7 = (# remEmpty, S.Just w' #)
  90     | otherwise     = let !r' = r {
  91                                   remLen  = remLen r + 1
  92                                 , remByte = w'
  93                                 }
  94                       in
  95                         (# r', S.Nothing #)
  96     where
  97       w' = (remByte r `shiftR` 1)
  98            .|.
  99            if b then 1 else 0
 100
 101 snocRem ∷ Remnant → Bool → (# S.Maybe Word8, Remnant #)
 102 snocRem r b
 103     | remLen r ≡ 8 = (# S.Just $ remByte r, remSingleton b #)
 104     | remLen r ≡ 7 = (# S.Just w', remEmpty #)
 105     | otherwise     = let !r' = r {
 106                                   remLen  = remLen r + 1
 107                                 , remByte = w'
 108                                 }
 109                       in
 110                         (# S.Nothing, r' #)
 111     where
 112       w' | b         = bit (remLen r) .|. remByte r
 113          | otherwise = remByte r
 114
 115 unconsRem ∷ Remnant → (# S.Maybe Bool, Remnant #)
 116 unconsRem r
 117     | remNull r = (# S.Nothing, remEmpty #)
 118     | otherwise = let !b  = remByte r `testBit` 1
 119                       !r' = Remnant {
 120                               remLen  = remLen r - 1
 121                             , remByte = remByte r `shiftR` 1
 122                             }
 123                   in
 124                     (# S.Just b, r' #)
 125
 126 unconsBytes ∷ L.ByteString → (# S.Maybe Remnant, L.ByteString #)
 127 unconsBytes bs
 128     = case L.uncons bs of
 129         Just (!w, bs')
 130             → (# S.Just (byteToRem w), bs' #)
 131         Nothing
 132             → (# S.Nothing, L.empty #)
 133
 134 appendRem ∷ Remnant → Remnant → Remnant
 135 appendRem a b = Remnant {
 136                   remLen  = remLen a + remLen b
 137                 , remByte = (remByte a `shiftL` remLen b) .|. remByte b
 138                 }
 139
 140 splitRemAt ∷ Integral n ⇒ n → Remnant → (# S.Maybe Remnant, Remnant #)
 141 splitRemAt n r
 142     | remLen r ≥ fromIntegral n
 143         = let !h  = Remnant {
 144                       remLen  = fromIntegral n
 145                     , remByte = remByte r .&. ((1 `shiftL` fromIntegral n) - 1)
 146                     }
 147               !r' = Remnant {
 148                       remLen  = remLen r - fromIntegral n
 149                     , remByte = remByte r `shiftR` fromIntegral n
 150                     }
 151           in
 152             (# S.Just h, r' #)
 153     | otherwise
 154         = (# S.Nothing, remEmpty #)
 155
 156 -- | /O(1)/ The empty 'BitString'.
 157 (∅) ∷ BitString
 158 (∅) = BitString {
 159         leftRem    = remEmpty
 160       , leftBytes  = L.empty
 161       , rightBytes = L.empty
 162       , rightRem   = remEmpty
 163       }
 164
 165 -- | /O(1)/ Convert a 'Bool' into a 'BitString'.
 166 singleton ∷ Bool → BitString
 167 singleton b = BitString {
 168                 leftRem    = remSingleton b
 169               , leftBytes  = L.empty
 170               , rightBytes = L.empty
 171               , rightRem   = remEmpty
 172               }
 173
 174 -- | /O(1)/ Prepend a bit to the beginning of a 'BitString'.
 175 cons ∷ Bool → BitString → BitString
 176 cons b bs
 177     = case consRem b $ leftRem bs of
 178         (# lr', S.Just w #)
 179             → bs {
 180                  leftRem   = lr'
 181                , leftBytes = L.cons' w $ leftBytes bs
 182                }
 183
 184         (# lr', S.Nothing #)
 185             → bs { leftRem = lr' }
 186
 187 -- | /O(1)/ Append a bit to the end of a 'BitString'.
 188 snoc ∷ BitString → Bool → BitString
 189 snoc bs b
 190     = case snocRem (rightRem bs) b of
 191         (# S.Just w, rr' #)
 192             → bs {
 193                  rightBytes = L.cons' w $ rightBytes bs
 194                , rightRem   = rr'
 195                }
 196
 197         (# S.Nothing, rr' #)
 198             → bs { rightRem = rr' }
 199
 200 -- | /amortized O(1)/ Extract the first bit of a 'BitString', which
 201 -- must be non-empty.
 202 head ∷ BitString → Bool
 203 head bs = case uncons bs of
 204             Just (b, _ )
 205                 → b
 206             Nothing
 207                 → error "head: empty BitString"
 208
 209 -- | /amortized O(1)/ Extract the bits after the head of a
 210 -- 'BitString', which must be non-empty.
 211 tail ∷ BitString → BitString
 212 tail bs = case uncons bs of
 213             Just (_, bs')
 214                 → bs'
 215             Nothing
 216                 → error "tail: empty BitString"
 217
 218 -- | /amortized O(1)/ Extract the the head and tail of a 'BitString',
 219 -- returning 'Nothing' if it's empty.
 220 uncons ∷ BitString → Maybe (Bool, BitString)
 221 uncons bs
 222     = case unconsRem $ leftRem bs of
 223         (# S.Just b, lr' #)
 224             → Just (b, bs { leftRem = lr' })
 225
 226         (# S.Nothing, _ #)
 227             → case unconsBytes $ leftBytes bs of
 228                  (# S.Just lr, lb' #)
 229                      → let !bs' = bs { leftRem   = lr
 230                                       , leftBytes = lb'
 231                                       }
 232                         in
 233                           uncons bs'
 234
 235                  (# S.Nothing, _ #)
 236                      → if L.null $ rightBytes bs then
 237                             case unconsRem $ rightRem bs of
 238                               (# S.Just b, rr' #)
 239                                   → Just (b, bs { rightRem = rr' })
 240
 241                               (# S.Nothing, _ #)
 242                                   → Nothing
 243                         else
 244                             let !bs' = bs { leftBytes = L.reverse
 245                                                         $ rightBytes bs
 246                                           }
 247                             in
 248                               uncons bs'
 249
 250 -- | /O(1)/ Test whether a 'BitString' is empty.
 251 null ∷ BitString → Bool
 252 null bs
 253     = remLen (leftRem bs) ≡ 0
 254       ∧
 255       L.null (leftBytes bs)
 256       ∧
 257       L.null (rightBytes bs)
 258       ∧
 259       remLen (rightRem bs) ≡ 0
 260
 261 -- | /O(n)/ Return the number of bits in a 'BitString'.
 262 length ∷ Integral n ⇒ BitString → n
 263 length bs
 264     = fromIntegral $ ( fromIntegral (remLen $ leftRem bs)
 265                        +
 266                        L.length (leftBytes bs) ⋅ 8
 267                        +
 268                        L.length (rightBytes bs) ⋅ 8
 269                        +
 270                        fromIntegral (remLen $ rightRem bs)
 271                      )
 272 {-# SPECIALISE length ∷ BitString → Int64 #-}
 273
 274 -- | /O(1)/ Convert a 'L.ByteString' into a 'BitString'.
 275 fromByteString ∷ L.ByteString → BitString
 276 fromByteString bs = BitString {
 277                       leftRem    = remEmpty
 278                     , leftBytes  = bs
 279                     , rightBytes = L.empty
 280                     , rightRem   = remEmpty
 281                     }
 282
 283 -- | /O(n)/ Convert a 'BitString' into 'L.ByteString', padding
 284 -- incomplete bytes to single bytes with necessary 0's at their
 285 -- MSB. Thus the following equation does not hold when the length of
 286 -- 'BitString' isn't multiple of 8.
 287 --
 288 -- > fromByteString . toByteString = id
 289 --
 290 -- But the following always holds true.
 291 --
 292 -- > toByteString . fromByteString = id
 293 toByteString ∷ BitString → L.ByteString
 294 toByteString bs
 295     = L.concat [ remToStr $ leftRem bs
 296                , leftBytes bs
 297                , L.reverse $ rightBytes bs
 298                , remToStr $ rightRem bs
 299                ]
 300
 301 -- | /O(n)/ @'take' n bs@ returns the prefix of @bs@ of length
 302 -- @n@. @bs@ must have at least @n@ bits.
 303 take ∷ Integral n ⇒ n → BitString → BitString
 304 take n xs = case splitAt n xs of
 305               (taken, _) → taken
 306
 307 -- | /O(n)/ @'drop' n bs@ returns the suffix of @bs@ after the first
 308 -- @n@ bits. @bs@ must have at least @n@ bits.
 309 drop ∷ Integral n ⇒ n → BitString → BitString
 310 drop n xs = case splitAt n xs of
 311               (_, dropped) → dropped
 312
 313 -- | /O(n)/ @'splitAt' n bs@ is equivalent to @('take' n bs, 'drop' n
 314 -- bs)@.
 315 splitAt ∷ Integral n ⇒ n → BitString → (BitString, BitString)
 316 splitAt n bs
 317     = case splitRemAt n $ leftRem bs of
 318         (# S.Just h, lr' #)
 319             -- The leftRem has enough bits.
 320             → let !h' = BitString {
 321                            leftRem    = h
 322                          , leftBytes  = L.empty
 323                          , rightBytes = L.empty
 324                          , rightRem   = remEmpty
 325                          }
 326                in
 327                  (h', bs { leftRem = lr' })
 328
 329         (# S.Nothing, _ #)
 330             → let !bytesToTake  = (n - fromIntegral (remLen $ leftRem bs)) `div` 8
 331                   !amortised    = leftBytes bs `L.append` (L.reverse $ rightBytes bs)
 332                   -- 0 ≤ bitsToTake < 8 - remLen (leftRem bs)
 333                   -- ∴ remLen (leftRem bs) + bitsToTake < 8
 334                   !bitsToTake   = n - bytesToTake ⋅ 8 - fromIntegral (remLen $ leftRem bs)
 335               in
 336                 case (# bytesToTake, L.uncons amortised #) of
 337                   (#_, Just (!w, amor') #)
 338                       -- There is at least one byte in the byte pool
 339                       -- (amortised).
 340                       → let !h1  = leftRem bs
 341                             (# h2, lr' #)
 342                                  = case splitRemAt bitsToTake $ byteToRem w of
 343                                     (# S.Just h2', lr'' #) → (# h2', lr'' #)
 344                                     (# S.Nothing , _    #) → error "internal error"
 345                             !h   = h1 `appendRem` h2
 346                             (bytes, lb')
 347                                  = L.splitAt (fromIntegral bytesToTake) amor'
 348                         in
 349                           if L.length bytes ≡ fromIntegral bytesToTake then
 350                               -- The byte pool actuall has at least
 351                               -- bytesToTake bytes.
 352                               let !h'  = BitString {
 353                                            leftRem    = h
 354                                          , leftBytes  = bytes
 355                                          , rightBytes = L.empty
 356                                          , rightRem   = remEmpty
 357                                          }
 358                                   !bs' = BitString {
 359                                            leftRem    = lr'
 360                                          , leftBytes  = lb'
 361                                          , rightBytes = L.empty
 362                                          , rightRem   = rightRem bs
 363                                          }
 364                               in
 365                                 (h', bs')
 366                           else
 367                               error "splitAt: not enough bits"
 368
 369                   (# 0, Nothing #)
 370                       -- No bytes are in the byte pool but the
 371                       -- rightRem may have enough bits.
 372                       → case splitRemAt bitsToTake $ rightRem bs of
 373                            (# S.Just h, rr' #)
 374                                → let !h'  = BitString {
 375                                               leftRem    = leftRem bs
 376                                             , leftBytes  = L.empty
 377                                             , rightBytes = L.empty
 378                                             , rightRem   = h
 379                                             }
 380                                      !bs' = BitString {
 381                                               leftRem    = remEmpty
 382                                             , leftBytes  = L.empty
 383                                             , rightBytes = L.empty
 384                                             , rightRem   = rr'
 385                                             }
 386                                  in
 387                                    (h', bs')
 388
 389                            (# S.Nothing, _ #)
 390                                → error "splitAt: not enough bits"
 391
 392                   (# _, Nothing #)
 393                       -- No bytes are in the byte pool but more than 8
 394                       -- bits are requested. The rightRem can't have
 395                       -- that many bits.
 396                       → error "splitAt: not enough bits"