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