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