Codec/Audio/WavPack/Words.hs

   1 {-# LANGUAGE
   2     BangPatterns
   3   , DoAndIfThenElse
   4   , FlexibleContexts
   5   , ScopedTypeVariables
   6   , UnicodeSyntax
   7   #-}
   8 {-| This module provides entropy word encoding and decoding functions
   9 using a variation on the Rice method.  This was introduced in wavpack
  10 3.93 because it allows splitting the data into a \"lossy\" stream and
  11 a \"correction\" stream in a very efficient manner and is therefore
  12 ideal for the "hybrid" mode.  For 4.0, the efficiency of this method
  13 was significantly improved by moving away from the normal Rice
  14 restriction of using powers of two for the modulus divisions and now
  15 the method can be used for both hybrid and pure lossless encoding.
  16
  17 Samples are divided by median probabilities at 5\/7 (71.43%), 10\/49
  18 (20.41%), and 20\/343 (5.83%). Each zone has 3.5 times fewer samples
  19 than the previous. Using standard Rice coding on this data would
  20 result in 1.4 bits per sample average (not counting sign
  21 bit). However, there is a very simple encoding that is over 99%
  22 efficient with this data and results in about 1.22 bits per sample. -}
  23 module Codec.Audio.WavPack.Words
  24     ( WordsData(..)
  25
  26     , getWordsLossless
  27     )
  28     where
  29 import Codec.Audio.WavPack.Entropy
  30 import Codec.Audio.WavPack.Internal
  31 import Control.Monad.Cont
  32 import Control.Monad.ST
  33 import Control.Monad.Trans
  34 import Control.Monad.Unicode
  35 import Data.Bits
  36 import Data.Bitstream.Generic (Bitstream)
  37 import qualified Data.Bitstream.Generic as B
  38 import Data.Int
  39 import Data.STRef
  40 import qualified Data.Vector.Generic.Mutable as MV
  41 import qualified Data.Vector.Unboxed as UV
  42 import Data.Word
  43 import Prelude hiding (break)
  44 import Prelude.Unicode
  45
  46 -- | FIXME
  47 data WordsData s
  48     = WordsData {
  49         wdBitrateDelta ∷ !(STRef s (Word32, Word32))
  50       , wdBitrateAcc   ∷ !(STRef s (Word32, Word32))
  51       , wdPendingData  ∷ !(STRef s Word32)
  52       , wdHoldingOne   ∷ !(STRef s Word32)
  53       , wdZeroesAcc    ∷ !(STRef s Word32)
  54       , wdHoldingZero  ∷ !(STRef s Bool)
  55       , wdPendingCount ∷ !(STRef s Int)
  56       , wdEntropyData  ∷ !(EntropyData s, EntropyData s)
  57       }
  58
  59 -- | Maximum consecutive 1s sent for /div/ data.
  60 limitOnes ∷ Num n ⇒ n
  61 {-# INLINE limitOnes #-}
  62 limitOnes = 16
  63
  64 getOnesCount ∷ Num a ⇒ Word8 → a
  65 {-# INLINE getOnesCount #-}
  66 getOnesCount = fromIntegral ∘ UV.unsafeIndex oct ∘ fromIntegral
  67     where
  68       oct ∷ UV.Vector Word8
  69       {-# NOINLINE oct #-}
  70       oct = UV.fromList
  71             [ 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 4 --   0 -  15
  72             , 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 5 --  16 -  31
  73             , 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 4 --  32 -  47
  74             , 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 6 --  48 -  63
  75             , 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 4 --  64 -  79
  76             , 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 5 --  80 -  95
  77             , 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 4 --  96 - 111
  78             , 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 7 -- 112 - 127
  79             , 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 4 -- 128 - 143
  80             , 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 5 -- 144 - 159
  81             , 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 4 -- 160 - 175
  82             , 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 6 -- 176 - 191
  83             , 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 4 -- 192 - 207
  84             , 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 5 -- 208 - 223
  85             , 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 4 -- 124 - 239
  86             , 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 8 -- 240 - 255
  87             ]
  88
  89 -- | This is an optimized version of 'getWord' that is used for
  90 -- lossless only ('edErrorLimit' ≡ 0). Also, rather than obtaining a
  91 -- single sample, it can be used to obtain an entire buffer of either
  92 -- mono or stereo samples.
  93 getWordsLossless ∷ ∀bs v s. (Bitstream bs, MV.MVector v Int32)
  94                  ⇒ Bool       -- ^ Is the stream monaural?
  95                  → WordsData s
  96                  → STRef s bs -- ^ WV bitstream
  97                  → Int        -- ^ Number of samples to get
  98                  → ST s (v s Int32)
  99 {-# INLINEABLE getWordsLossless #-}
 100 getWordsLossless isMono w bs nSamples0
 101     = do v ← MV.new nSamples
 102          n ← runContT (for 0 (< nSamples) (+ 1) (loop v)) return
 103          return $ MV.take n v
 104     where
 105       nSamples ∷ Int
 106       nSamples = if isMono
 107                  then nSamples0
 108                  else nSamples0 ⋅ 2
 109
 110       loop ∷ v s Int32
 111            → Int
 112            → ContT Int (ST s) ()
 113            → ContT Int (ST s) ()
 114            → ContT Int (ST s) ()
 115       loop v n break continue
 116           = do let c | isMono        = fst $ wdEntropyData w
 117                      | n `testBit` 0 = fst $ wdEntropyData w
 118                      | otherwise     = snd $ wdEntropyData w
 119                med00   ← lift $ readSTRef (edMedian0 $ fst $ wdEntropyData w)
 120                hldZero ← lift $ readSTRef (wdHoldingZero w)
 121                hldOne  ← lift $ readSTRef (wdHoldingOne  w)
 122                med10   ← lift $ readSTRef (edMedian0 $ snd $ wdEntropyData w)
 123                when (med00 < 2 ∧ hldZero ≡ False ∧ hldOne ≡ 0 ∧ med10 < 2) $
 124                     do zAcc ← lift $ readSTRef (wdZeroesAcc w)
 125                        if zAcc > 0 then
 126                            do lift $ modifySTRef (wdZeroesAcc w) ((-) 1)
 127                               when (zAcc > 1) $
 128                                    do lift $ MV.unsafeWrite v n 0
 129                                       continue
 130                        else
 131                            do cBits ← lift $ takeWhileLessThan id 33 bs
 132
 133                               when (cBits ≡ 33) $
 134                                   break
 135
 136                               if cBits < 2 then
 137                                   lift $ writeSTRef (wdZeroesAcc w) cBits
 138                               else
 139                                   do lift $ writeSTRef (wdZeroesAcc w) 0
 140                                      (mask, _)
 141                                          ← for (1, cBits)
 142                                                ((> 1) ∘ snd)
 143                                                (\(m, cb) → (m `shiftL` 1, cb - 1)) $ \(mask, _) _ _ →
 144                                                    do b ← lift $ takeHead bs
 145                                                       when b $
 146                                                           lift $ modifySTRef (wdZeroesAcc w) (.|. mask)
 147                                      lift $ modifySTRef (wdZeroesAcc w) (.|. mask)
 148
 149                               zAcc' ← lift$ readSTRef (wdZeroesAcc w)
 150                               when (zAcc' > 0) $
 151                                   do lift $ clearMedians $ fst $ wdEntropyData w
 152                                      lift $ clearMedians $ snd $ wdEntropyData w
 153                                      lift $ MV.unsafeWrite v n 0
 154                                      continue
 155
 156                onesCount ← lift $ newSTRef (⊥)
 157                if hldZero then
 158                    do lift $ writeSTRef onesCount 0
 159                       lift $ writeSTRef (wdHoldingZero w) False
 160                else
 161                    do next8 ← lift $ readBits (8 ∷ Word8) bs
 162                       if next8 ≡ 0xFF then
 163                           do lift $ dropBits (8 ∷ Word8) bs
 164                              oc ← for 8 (< limitOnes + 1) (+ 1) $ \oc break' _ →
 165                                       do h ← lift $ takeHead bs
 166                                          unless h $
 167                                              break'
 168                              lift $ writeSTRef onesCount oc
 169
 170                              when (oc ≡ limitOnes + 1) $
 171                                  break
 172
 173                              when (oc ≡ limitOnes) $
 174                                  do cBits ← for 0 (< 33) (+ 1) $ \cBits break' _ →
 175                                                 do h ← lift $ takeHead bs
 176                                                    unless h $
 177                                                        break'
 178
 179                                     when (cBits ≡ 33) $
 180                                         break
 181
 182                                     if cBits < 2 then
 183                                         lift $ writeSTRef onesCount cBits
 184                                     else
 185                                         do lift $ writeSTRef onesCount 0
 186                                            (mask, _)
 187                                                ← for (1, cBits)
 188                                                      ((> 1) ∘ snd)
 189                                                      (\(m, cb) → (m `shiftL` 1, cb - 1)) $ \(mask, _) _ _ →
 190                                                          do b ← lift $ takeHead bs
 191                                                             when b $
 192                                                                 lift $ modifySTRef onesCount (.|. mask)
 193                                            lift $ modifySTRef onesCount (.|. mask)
 194
 195                                     lift $ modifySTRef onesCount (+ limitOnes)
 196                       else
 197                           do let oc ∷ Word32
 198                                  oc = getOnesCount next8
 199                              lift $ writeSTRef onesCount oc
 200                              lift $ dropBits (oc + 1) bs
 201
 202                       if hldOne > 0 then
 203                           error "FIXME"
 204                       else
 205                           error "FIXME"
 206
 207 {-
 208 getWordsLossless ∷ ∀bs v. (Bitstream bs, GV.Vector v Int32)
 209                  ⇒ Bool -- ^ Is the stream monaural?
 210                  → WordsData
 211                  → bs   -- ^ WV bitstream.
 212                  → Int   -- ^ Number of samples to get.
 213                  → (# WordsData, bs, v Int32 #)
 214 {-# INLINEABLE getWordsLossless #-}
 215 getWordsLossless isMono w0 bs0 nSamples0
 216     = let v0  = New.create $ MV.new nSamples
 217           (# w1, bs1, n1, v1 #)
 218               = go0 w0 bs0 0 v0
 219           v2  = GV.new $ New.take n1 v1
 220       in
 221         (# w1, bs1, v2 #)
 222     where
 223       go0 ∷ WordsData → bs → Int → New v Int32
 224           → (# WordsData, bs, Int, New v Int32 #)
 225       go0 w bs n v
 226           | n ≥ nSamples
 227               = (# w, bs, n, v #)
 228           | edMedian0 (fst $ wdEntropyData w) < 2 ∧
 229             wdHoldingZero w ≡ False               ∧
 230             wdHoldingOne  w ≡ 0                   ∧
 231             edMedian1 (fst $ wdEntropyData w) < 2
 232               = if wdZeroesAcc w > 0 then
 233                     let w' = w { wdZeroesAcc = wdZeroesAcc w - 1 }
 234                     in
 235                       if wdZeroesAcc w' > 0 then
 236                           let (# n', v' #) = appendWord 0 n v
 237                           in
 238                             go0 w' bs n' v'
 239                       else
 240                           go1 w' bs n v
 241                 else
 242                     let cBits = min 33 $ B.length (B.takeWhile id bs)
 243                         bs'   = B.drop cBits bs
 244                     in
 245                       if cBits ≡ 33 then
 246                           (# w, bs', n, v #)
 247                       else
 248                           let (# w', bs'' #) = go0' cBits w bs'
 249                           in
 250                             if wdZeroesAcc w' > 0 then
 251                                 let w'' = w' {
 252                                             wdEntropyData =
 253                                                 ( clearMedian $ fst $ wdEntropyData w'
 254                                                 , clearMedian $ snd $ wdEntropyData w' )
 255                                           }
 256                                     (# n', v' #)
 257                                         = appendWord 0 n v
 258                                 in
 259                                   go0 w'' bs'' n' v'
 260                             else
 261                                 go1 w' bs'' n v
 262           | otherwise
 263               = go1 w bs n v
 264
 265       go0' ∷ Word32 → WordsData → bs → (# WordsData, bs #)
 266       go0' cBits w bs
 267           | cBits < 2
 268               = let w' = w { wdZeroesAcc = cBits }
 269                 in
 270                   (# w', bs #)
 271           | otherwise
 272               = let w' = w { wdZeroesAcc = 0 }
 273                 in
 274                   go0'' 1 cBits w' bs
 275
 276       go0'' ∷ Word32 → Word32 → WordsData → bs → (# WordsData, bs #)
 277       go0'' mask cBits w bs
 278           | cBits ≡ 1
 279               = let w' = w { wdZeroesAcc = wdZeroesAcc w .|. mask }
 280                 in
 281                   (# w', bs #)
 282           | otherwise
 283               = let cBits' = cBits - 1
 284                     w'     = if B.head bs then
 285                                  w { wdZeroesAcc = wdZeroesAcc w .|. mask }
 286                              else
 287                                  w
 288                     mask'  = mask `shiftL` 1
 289                     bs'    = B.tail bs
 290                 in
 291                   go0'' mask' cBits' w' bs'
 292
 293       go1 ∷ WordsData → bs → Int → New v Int32
 294           → (# WordsData, bs, Int, New v Int32 #)
 295       go1 w bs n v
 296           | wdHoldingZero w
 297               = let w' = w { wdHoldingZero = False }
 298                 in
 299                   go2 0 w' bs n v
 300           | otherwise
 301               = let next8 ∷ Word8
 302                     next8 = B.toBits (B.take (8 ∷ Int) bs)
 303                 in
 304                   if next8 ≡ 0xFF then
 305                       error "FIXME"
 306                   else
 307                       error "FIXME"
 308
 309       go2 ∷ Word32 → WordsData → bs → Int → New v Int32
 310           → (# WordsData, bs, Int, New v Int32 #)
 311       go2 0 w bs n v
 312           = let ent  = getEntropy n w
 313                 low  = 0
 314                 high = getMedian0 ent - 1
 315                 ent' = decMedian0 ent
 316                 w'   = setEntropy ent' n w
 317             in
 318               go3 low high w' bs n v
 319       go2 1 w bs n v
 320           = let ent  = getEntropy n w
 321                 low  = getMedian0 ent
 322                 high = low + getMedian1 ent - 1
 323                 ent' = (incMedian0 ∘ decMedian1) ent
 324                 w'   = setEntropy ent' n w
 325             in
 326               go3 low high w' bs n v
 327       go2 2 w bs n v
 328           = let ent   = getEntropy n w
 329                 low   = getMedian0 ent + getMedian1 ent
 330                 high  = low + getMedian2 ent - 1
 331                 ent'  = (incMedian0 ∘ incMedian1 ∘ decMedian2) ent
 332                 w'    = setEntropy ent' n w
 333             in
 334               go3 low high w' bs n v
 335       go2 onesCount w bs n v
 336           = let ent   = getEntropy n w
 337                 low   = getMedian0 ent + getMedian1 ent + (onesCount-2) ⋅ getMedian2 ent
 338                 high  = low + getMedian2 ent - 1
 339                 ent'  = (incMedian0 ∘ incMedian1 ∘ incMedian2) ent
 340                 w'    = setEntropy ent' n w
 341             in
 342               go3 low high w' bs n v
 343
 344       go3 ∷ Word32 → Word32 → WordsData → bs → Int → New v Int32
 345           → (# WordsData, bs, Int, New v Int32 #)
 346       go3 low high w bs n v
 347           = let (# code, bs' #)
 348                      = readCode bs (high - low)
 349                 low' = low + code
 350                 word = if B.head bs' then
 351                            fromIntegral $ complement low'
 352                        else
 353                            fromIntegral low'
 354                 bs'' = B.tail bs'
 355                 (# n', v' #)
 356                      = appendWord word n v
 357             in
 358               go0 w bs'' n' v'
 359
 360       appendWord ∷ Int32 → Int → New v Int32 → (# Int, New v Int32 #)
 361       appendWord word n v
 362           = let v' = New.modify (\mv → MV.unsafeWrite mv n word) v
 363                 n' = n + 1
 364             in
 365               (# n', v' #)
 366
 367       getEntropy ∷ Int → WordsData → EntropyData
 368       getEntropy n w
 369           | isMono        = fst $ wdEntropyData w
 370           | n `testBit` 0 = fst $ wdEntropyData w
 371           | otherwise     = snd $ wdEntropyData w
 372
 373       setEntropy ∷ EntropyData → Int → WordsData → WordsData
 374       setEntropy e n w
 375           | isMono        = w { wdEntropyData = (e, snd $ wdEntropyData w) }
 376           | n `testBit` 0 = w { wdEntropyData = (e, snd $ wdEntropyData w) }
 377           | otherwise     = w { wdEntropyData = (fst $ wdEntropyData w, e) }
 378 -}
 379
 380 -- | Read a single unsigned value from the specified bitstream with a
 381 -- value from 0 to maxCode. If there are exactly a power of two number
 382 -- of possible codes then this will read a fixed number of bits;
 383 -- otherwise it reads the minimum number of bits and then determines
 384 -- whether another bit is needed to define the code.
 385 readCode ∷ Bitstream bs ⇒ STRef s bs → Word32 → ST s Word32
 386 {-# INLINEABLE readCode #-}
 387 readCode bs 0       = return 0
 388 readCode bs 1       = fmap b2n $ takeHead bs
 389 readCode bs maxCode
 390     = do let bitCount = countBits maxCode
 391              extras   = bit bitCount - maxCode - 1
 392          code ← takeBits (bitCount - 1) bs
 393          if code ≥ extras then
 394              do nextBit ← takeHead bs
 395                 return $ (code `shiftL` 1) - extras + b2n nextBit
 396            else
 397              return code
 398
 399 takeHead ∷ Bitstream bs ⇒ STRef s bs → ST s Bool
 400 {-# INLINEABLE takeHead #-}
 401 takeHead bsr
 402     = do bs ← readSTRef bsr
 403          writeSTRef bsr (B.tail bs)
 404          return (B.head bs)
 405
 406 takeWhileLessThan ∷ (Integral n, Bitstream bs)
 407                   ⇒ (Bool → Bool)
 408                   → n
 409                   → STRef s bs
 410                   → ST s n
 411 {-# INLINEABLE takeWhileLessThan #-}
 412 takeWhileLessThan f n bsr = go 0
 413     where
 414       {-# INLINE go #-}
 415       go i | i < n
 416                = do b ← takeHead bsr
 417                     if b then
 418                         go (i + 1)
 419                       else
 420                         return i
 421            | otherwise
 422                = return i
 423
 424 readBits ∷ (Integral n, Bitstream bs, Bits a) ⇒ n → STRef s bs → ST s a
 425 {-# INLINEABLE readBits #-}
 426 readBits n bsr
 427     = do bs ← readSTRef bsr
 428          return (B.toBits (B.take n bs))
 429
 430 takeBits ∷ (Integral n, Bitstream bs, Bits a) ⇒ n → STRef s bs → ST s a
 431 {-# INLINEABLE takeBits #-}
 432 takeBits n bsr
 433     = do bs ← readSTRef bsr
 434          writeSTRef bsr (B.drop n bs)
 435          return (B.toBits (B.take n bs))
 436
 437 dropBits ∷ (Integral n, Bitstream bs) ⇒ n → STRef s bs → ST s ()
 438 {-# INLINEABLE dropBits #-}
 439 dropBits n bsr
 440     = do bs ← readSTRef bsr
 441          writeSTRef bsr (B.drop n bs)
 442
 443 -- | C style /for/ loop with /break/ and /continue/.
 444 for ∷ ∀m α. MonadCont m
 445     ⇒ α          -- ^ Initial state
 446     → (α → Bool) -- ^ Continue-the-loop predicate
 447     → (α → α)    -- ^ State modifier
 448     → (α → m () → m () → m ()) -- ^ Loop body taking breaker and
 449                                -- continuer
 450     → m α        -- ^ Final state
 451 for α0 contLoop next body
 452     = callCC $ \break → loop break α0
 453     where
 454       loop ∷ (α → m ()) → α → m α
 455       loop break α
 456           | contLoop α
 457               = do callCC $ \continue → body α (break α) (continue ())
 458                    loop break (next α)
 459           | otherwise
 460               = return α