Codec/Audio/WavPack/Words.hs

   1 {-# LANGUAGE
   2     BangPatterns
   3   , FlexibleContexts
   4   , ScopedTypeVariables
   5   , UnboxedTuples
   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 Data.Bits
  32 import Data.Bitstream.Generic (Bitstream)
  33 import qualified Data.Bitstream.Generic as B
  34 import Data.Int
  35 import qualified Data.Vector.Generic as GV
  36 import qualified Data.Vector.Generic.Mutable as MV
  37 import Data.Vector.Generic.New (New)
  38 import qualified Data.Vector.Generic.New as New
  39 import Data.Word
  40 import Prelude.Unicode
  41
  42 -- | FIXME
  43 data WordsData
  44     = WordsData {
  45         wdBitrateDelta ∷ !(Word32, Word32)
  46       , wdBitrateAcc   ∷ !(Word32, Word32)
  47       , wdPendingData  ∷ !Word32
  48       , wdHoldingOne   ∷ !Word32
  49       , wdZeroesAcc    ∷ !Word32
  50       , wdHoldingZero  ∷ !Bool
  51       , wdPendingCount ∷ !Int
  52       , wdEntropyData  ∷ !(EntropyData, EntropyData)
  53       }
  54     deriving (Eq, Show)
  55
  56 -- | This is an optimized version of 'getWord' that is used for
  57 -- lossless only ('edErrorLimit' ≡ 0). Also, rather than obtaining a
  58 -- single sample, it can be used to obtain an entire buffer of either
  59 -- mono or stereo samples.
  60 getWordsLossless ∷ ∀bs v. (Bitstream bs, GV.Vector v Int32)
  61                  ⇒ Bool -- ^ Is the stream monaural?
  62                  → WordsData
  63                  → bs   -- ^ WV bitstream.
  64                  → Int   -- ^ Number of samples to get.
  65                  → (# WordsData, bs, v Int32 #)
  66 {-# INLINEABLE getWordsLossless #-}
  67 getWordsLossless isMono w0 bs0 nSamples0
  68     = let v0  = New.create $ MV.new nSamples
  69           (# w1, bs1, n1, v1 #)
  70               = go0 w0 bs0 0 v0
  71           v2  = GV.new $ New.take n1 v1
  72       in
  73         (# w1, bs1, v2 #)
  74     where
  75       nSamples ∷ Int
  76       nSamples = if isMono
  77                  then nSamples0
  78                  else nSamples0 ⋅ 2
  79
  80       go0 ∷ WordsData → bs → Int → New v Int32
  81           → (# WordsData, bs, Int, New v Int32 #)
  82       go0 w bs n v
  83           | n ≥ nSamples
  84               = (# w, bs, n, v #)
  85           | edMedian0 (fst $ wdEntropyData w) < 2 ∧
  86             wdHoldingZero w ≡ False               ∧
  87             wdHoldingOne  w ≡ 0                   ∧
  88             edMedian1 (fst $ wdEntropyData w) < 2
  89               = if wdZeroesAcc w > 0 then
  90                     let w' = w { wdZeroesAcc = wdZeroesAcc w - 1 }
  91                     in
  92                       if wdZeroesAcc w' > 0 then
  93                           let v' = New.modify (\mv → MV.unsafeWrite mv n 0) v
  94                               n' = n + 1
  95                           in
  96                             go0 w' bs n' v'
  97                       else
  98                           go1 w' bs n v
  99                 else
 100                     let cBits = min 33 $ B.length (B.takeWhile id bs)
 101                         bs'   = B.drop cBits bs
 102                     in
 103                       if cBits ≡ 33 then
 104                           (# w, bs', n, v #)
 105                       else
 106                           let (# w', bs'' #) = go0' cBits w bs'
 107                           in
 108                             if wdZeroesAcc w' > 0 then
 109                                 let w'' = w' {
 110                                             wdEntropyData =
 111                                                 ( clearMedian $ fst $ wdEntropyData w'
 112                                                 , clearMedian $ snd $ wdEntropyData w' )
 113                                           }
 114                                     v'  = New.modify (\mv → MV.unsafeWrite mv n 0) v
 115                                     n'  = n + 1
 116                                 in
 117                                   go0 w'' bs'' n' v'
 118                             else
 119                                 go1 w' bs'' n v
 120           | otherwise
 121               = go1 w bs n v
 122
 123       go0' ∷ Word32 → WordsData → bs → (# WordsData, bs #)
 124       go0' cBits w bs
 125           | cBits < 2
 126               = let w' = w { wdZeroesAcc = cBits }
 127                 in
 128                   (# w', bs #)
 129           | otherwise
 130               = let w' = w { wdZeroesAcc = 0 }
 131                 in
 132                   go0'' 1 cBits w' bs
 133
 134       go0'' ∷ Word32 → Word32 → WordsData → bs → (# WordsData, bs #)
 135       go0'' mask cBits w bs
 136           | cBits ≡ 1
 137               = let w' = w { wdZeroesAcc = wdZeroesAcc w .|. mask }
 138                 in
 139                   (# w', bs #)
 140           | otherwise
 141               = let cBits' = cBits - 1
 142                     w'     = if B.head bs then
 143                                  w { wdZeroesAcc = wdZeroesAcc w .|. mask }
 144                              else
 145                                  w
 146                     mask'  = mask `shiftL` 1
 147                     bs'    = B.tail bs
 148                 in
 149                   go0'' mask' cBits' w' bs'
 150
 151       go1 ∷ WordsData → bs → Int → New v Int32
 152           → (# WordsData, bs, Int, New v Int32 #)
 153       go1 w bs n v
 154           | wdHoldingZero w
 155               = let w' = w { wdHoldingZero = False }
 156                 in
 157                   go2 0 w' bs n v
 158           | otherwise
 159               = error "FIXME"
 160
 161       go2 ∷ Word32 → WordsData → bs → Int → New v Int32
 162           → (# WordsData, bs, Int, New v Int32 #)
 163       go2 0 w bs n v
 164           = let ent  = getEntropy n w
 165                 low  = 0
 166                 high = getMedian0 ent
 167                 ent' = decMedian0 ent
 168                 w'   = setEntropy ent' n w
 169             in
 170               go3 low high w' bs n v
 171       go2 1 w bs n v
 172           = let ent  = getEntropy n w
 173                 low  = getMedian0 ent
 174                 high = low + getMedian1 ent - 1
 175                 ent' = (incMedian0 ∘ decMedian1) ent
 176                 w'   = setEntropy ent' n w
 177             in
 178               go3 low high w' bs n v
 179       go2 2 w bs n v
 180           = let ent   = getEntropy n w
 181                 low   = getMedian0 ent + getMedian1 ent
 182                 high  = low + getMedian2 ent - 1
 183                 ent'  = (incMedian0 ∘ incMedian1 ∘ decMedian2) ent
 184                 w'    = setEntropy ent' n w
 185             in
 186               go3 low high w' bs n v
 187       go2 onesCount w bs n v
 188           = let ent   = getEntropy n w
 189                 low   = getMedian0 ent + getMedian1 ent + (onesCount-2) ⋅ getMedian2 ent
 190                 high  = low + getMedian2 ent - 1
 191                 ent'  = (incMedian0 ∘ incMedian1 ∘ incMedian2) ent
 192                 w'    = setEntropy ent' n w
 193             in
 194               go3 low high w' bs n v
 195
 196       go3 ∷ Word32 → Word32 → WordsData → bs → Int → New v Int32
 197           → (# WordsData, bs, Int, New v Int32 #)
 198       go3 low high w bs n v
 199           = let (# code, bs' #)
 200                      = readCode bs (high - low)
 201                 low' = low + code
 202                 a    = if B.head bs' then
 203                            fromIntegral $ complement low'
 204                        else
 205                            fromIntegral low'
 206                 bs'' = B.tail bs'
 207                 v'   = New.modify (\mv → MV.unsafeWrite mv n a) v
 208                 n'   = n + 1
 209             in
 210               go0 w bs'' n' v'
 211
 212       getEntropy ∷ Int → WordsData → EntropyData
 213       getEntropy n w
 214           | isMono        = fst $ wdEntropyData w
 215           | n `testBit` 0 = fst $ wdEntropyData w
 216           | otherwise     = snd $ wdEntropyData w
 217
 218       setEntropy ∷ EntropyData → Int → WordsData → WordsData
 219       setEntropy e n w
 220           | isMono        = w { wdEntropyData = (e, snd $ wdEntropyData w) }
 221           | n `testBit` 0 = w { wdEntropyData = (e, snd $ wdEntropyData w) }
 222           | otherwise     = w { wdEntropyData = (fst $ wdEntropyData w, e) }
 223
 224 -- | Read a single unsigned value from the specified bitstream with a
 225 -- value from 0 to maxCode. If there are exactly a power of two number
 226 -- of possible codes then this will read a fixed number of bits;
 227 -- otherwise it reads the minimum number of bits and then determines
 228 -- whether another bit is needed to define the code.
 229 readCode ∷ Bitstream bs ⇒ bs → Word32 → (# Word32, bs #)
 230 {-# INLINEABLE readCode #-}
 231 readCode bs 0       = (# 0, bs #)
 232 readCode bs 1       = (# b2n (B.head bs), B.tail bs #)
 233 readCode bs maxCode
 234     = let !bitCount = countBits maxCode
 235           !extras   = bit bitCount - maxCode - 1
 236           !code     = B.toBits (B.take (bitCount - 1) bs)
 237           (# code', bitCount' #)
 238                     = if code ≥ extras then
 239                           (# (code `shiftL` 1)
 240                              - extras
 241                              + b2n (bs B.!! bitCount)
 242                            , bitCount #)
 243                       else
 244                           (# code, bitCount - 1 #)
 245           !bs'      = B.drop bitCount' bs
 246       in
 247         (# code', bs' #)