Codec/Audio/WavPack/Words.hs

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