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                            error "FIXME"
 100                error "FIXME"
 101
 102 {-
 103 getWordsLossless ∷ ∀bs v. (Bitstream bs, GV.Vector v Int32)
 104                  ⇒ Bool -- ^ Is the stream monaural?
 105                  → WordsData
 106                  → bs   -- ^ WV bitstream.
 107                  → Int   -- ^ Number of samples to get.
 108                  → (# WordsData, bs, v Int32 #)
 109 {-# INLINEABLE getWordsLossless #-}
 110 getWordsLossless isMono w0 bs0 nSamples0
 111     = let v0  = New.create $ MV.new nSamples
 112           (# w1, bs1, n1, v1 #)
 113               = go0 w0 bs0 0 v0
 114           v2  = GV.new $ New.take n1 v1
 115       in
 116         (# w1, bs1, v2 #)
 117     where
 118       go0 ∷ WordsData → bs → Int → New v Int32
 119           → (# WordsData, bs, Int, New v Int32 #)
 120       go0 w bs n v
 121           | n ≥ nSamples
 122               = (# w, bs, n, v #)
 123           | edMedian0 (fst $ wdEntropyData w) < 2 ∧
 124             wdHoldingZero w ≡ False               ∧
 125             wdHoldingOne  w ≡ 0                   ∧
 126             edMedian1 (fst $ wdEntropyData w) < 2
 127               = if wdZeroesAcc w > 0 then
 128                     let w' = w { wdZeroesAcc = wdZeroesAcc w - 1 }
 129                     in
 130                       if wdZeroesAcc w' > 0 then
 131                           let (# n', v' #) = appendWord 0 n v
 132                           in
 133                             go0 w' bs n' v'
 134                       else
 135                           go1 w' bs n v
 136                 else
 137                     let cBits = min 33 $ B.length (B.takeWhile id bs)
 138                         bs'   = B.drop cBits bs
 139                     in
 140                       if cBits ≡ 33 then
 141                           (# w, bs', n, v #)
 142                       else
 143                           let (# w', bs'' #) = go0' cBits w bs'
 144                           in
 145                             if wdZeroesAcc w' > 0 then
 146                                 let w'' = w' {
 147                                             wdEntropyData =
 148                                                 ( clearMedian $ fst $ wdEntropyData w'
 149                                                 , clearMedian $ snd $ wdEntropyData w' )
 150                                           }
 151                                     (# n', v' #)
 152                                         = appendWord 0 n v
 153                                 in
 154                                   go0 w'' bs'' n' v'
 155                             else
 156                                 go1 w' bs'' n v
 157           | otherwise
 158               = go1 w bs n v
 159
 160       go0' ∷ Word32 → WordsData → bs → (# WordsData, bs #)
 161       go0' cBits w bs
 162           | cBits < 2
 163               = let w' = w { wdZeroesAcc = cBits }
 164                 in
 165                   (# w', bs #)
 166           | otherwise
 167               = let w' = w { wdZeroesAcc = 0 }
 168                 in
 169                   go0'' 1 cBits w' bs
 170
 171       go0'' ∷ Word32 → Word32 → WordsData → bs → (# WordsData, bs #)
 172       go0'' mask cBits w bs
 173           | cBits ≡ 1
 174               = let w' = w { wdZeroesAcc = wdZeroesAcc w .|. mask }
 175                 in
 176                   (# w', bs #)
 177           | otherwise
 178               = let cBits' = cBits - 1
 179                     w'     = if B.head bs then
 180                                  w { wdZeroesAcc = wdZeroesAcc w .|. mask }
 181                              else
 182                                  w
 183                     mask'  = mask `shiftL` 1
 184                     bs'    = B.tail bs
 185                 in
 186                   go0'' mask' cBits' w' bs'
 187
 188       go1 ∷ WordsData → bs → Int → New v Int32
 189           → (# WordsData, bs, Int, New v Int32 #)
 190       go1 w bs n v
 191           | wdHoldingZero w
 192               = let w' = w { wdHoldingZero = False }
 193                 in
 194                   go2 0 w' bs n v
 195           | otherwise
 196               = let next8 ∷ Word8
 197                     next8 = B.toBits (B.take (8 ∷ Int) bs)
 198                 in
 199                   if next8 ≡ 0xFF then
 200                       error "FIXME"
 201                   else
 202                       error "FIXME"
 203
 204       go2 ∷ Word32 → WordsData → bs → Int → New v Int32
 205           → (# WordsData, bs, Int, New v Int32 #)
 206       go2 0 w bs n v
 207           = let ent  = getEntropy n w
 208                 low  = 0
 209                 high = getMedian0 ent - 1
 210                 ent' = decMedian0 ent
 211                 w'   = setEntropy ent' n w
 212             in
 213               go3 low high w' bs n v
 214       go2 1 w bs n v
 215           = let ent  = getEntropy n w
 216                 low  = getMedian0 ent
 217                 high = low + getMedian1 ent - 1
 218                 ent' = (incMedian0 ∘ decMedian1) ent
 219                 w'   = setEntropy ent' n w
 220             in
 221               go3 low high w' bs n v
 222       go2 2 w bs n v
 223           = let ent   = getEntropy n w
 224                 low   = getMedian0 ent + getMedian1 ent
 225                 high  = low + getMedian2 ent - 1
 226                 ent'  = (incMedian0 ∘ incMedian1 ∘ decMedian2) ent
 227                 w'    = setEntropy ent' n w
 228             in
 229               go3 low high w' bs n v
 230       go2 onesCount w bs n v
 231           = let ent   = getEntropy n w
 232                 low   = getMedian0 ent + getMedian1 ent + (onesCount-2) ⋅ getMedian2 ent
 233                 high  = low + getMedian2 ent - 1
 234                 ent'  = (incMedian0 ∘ incMedian1 ∘ incMedian2) ent
 235                 w'    = setEntropy ent' n w
 236             in
 237               go3 low high w' bs n v
 238
 239       go3 ∷ Word32 → Word32 → WordsData → bs → Int → New v Int32
 240           → (# WordsData, bs, Int, New v Int32 #)
 241       go3 low high w bs n v
 242           = let (# code, bs' #)
 243                      = readCode bs (high - low)
 244                 low' = low + code
 245                 word = if B.head bs' then
 246                            fromIntegral $ complement low'
 247                        else
 248                            fromIntegral low'
 249                 bs'' = B.tail bs'
 250                 (# n', v' #)
 251                      = appendWord word n v
 252             in
 253               go0 w bs'' n' v'
 254
 255       appendWord ∷ Int32 → Int → New v Int32 → (# Int, New v Int32 #)
 256       appendWord word n v
 257           = let v' = New.modify (\mv → MV.unsafeWrite mv n word) v
 258                 n' = n + 1
 259             in
 260               (# n', v' #)
 261
 262       getEntropy ∷ Int → WordsData → EntropyData
 263       getEntropy n w
 264           | isMono        = fst $ wdEntropyData w
 265           | n `testBit` 0 = fst $ wdEntropyData w
 266           | otherwise     = snd $ wdEntropyData w
 267
 268       setEntropy ∷ EntropyData → Int → WordsData → WordsData
 269       setEntropy e n w
 270           | isMono        = w { wdEntropyData = (e, snd $ wdEntropyData w) }
 271           | n `testBit` 0 = w { wdEntropyData = (e, snd $ wdEntropyData w) }
 272           | otherwise     = w { wdEntropyData = (fst $ wdEntropyData w, e) }
 273 -}
 274
 275 -- | Read a single unsigned value from the specified bitstream with a
 276 -- value from 0 to maxCode. If there are exactly a power of two number
 277 -- of possible codes then this will read a fixed number of bits;
 278 -- otherwise it reads the minimum number of bits and then determines
 279 -- whether another bit is needed to define the code.
 280 readCode ∷ Bitstream bs ⇒ STRef s bs → Word32 → ST s Word32
 281 {-# INLINEABLE readCode #-}
 282 readCode bs 0       = return 0
 283 readCode bs 1       = fmap b2n $ takeHead bs
 284 readCode bs maxCode
 285     = do let bitCount = countBits maxCode
 286              extras   = bit bitCount - maxCode - 1
 287          code ← takeBits bs (bitCount - 1)
 288          if code ≥ extras then
 289              do nextBit ← takeHead bs
 290                 return $ (code `shiftL` 1) - extras + b2n nextBit
 291            else
 292              return code
 293
 294 takeHead ∷ Bitstream bs ⇒ STRef s bs → ST s Bool
 295 {-# INLINEABLE takeHead #-}
 296 takeHead bsr
 297     = do bs ← readSTRef bsr
 298          writeSTRef bsr (B.tail bs)
 299          return (B.head bs)
 300
 301 takeBits ∷ (Integral n, Bitstream bs, Bits a) ⇒ STRef s bs → n → ST s a
 302 {-# INLINEABLE takeBits #-}
 303 takeBits bsr n
 304     = do bs ← readSTRef bsr
 305          writeSTRef bsr (B.drop n bs)
 306          return (B.toBits (B.take n bs))
 307
 308 -- | C style /for/ loop with /break/ and /continue/.
 309 for ∷ ∀m α. MonadCont m
 310     ⇒ α          -- ^ Initial state
 311     → (α → Bool) -- ^ Continue-the-loop predicate
 312     → (α → α)    -- ^ State modifier
 313     → (α → m () → m () → m ()) -- ^ Loop body taking breaker and
 314                                -- continuer
 315     → m α        -- ^ Final state
 316 for α0 contLoop next body
 317     = callCC $ \break → loop break α0
 318     where
 319       loop ∷ (α → m ()) → α → m α
 320       loop break α
 321           | contLoop α
 322               = do callCC $ \continue → body α (break α) (continue ())
 323                    loop break (next α)
 324           | otherwise
 325               = return α