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 (# n', v' #) = appendWord 0 n v
  94                           in
  95                             go0 w' bs n' v'
  96                       else
  97                           go1 w' bs n v
  98                 else
  99                     let cBits = min 33 $ B.length (B.takeWhile id bs)
 100                         bs'   = B.drop cBits bs
 101                     in
 102                       if cBits ≡ 33 then
 103                           (# w, bs', n, v #)
 104                       else
 105                           let (# w', bs'' #) = go0' cBits w bs'
 106                           in
 107                             if wdZeroesAcc w' > 0 then
 108                                 let w'' = w' {
 109                                             wdEntropyData =
 110                                                 ( clearMedian $ fst $ wdEntropyData w'
 111                                                 , clearMedian $ snd $ wdEntropyData w' )
 112                                           }
 113                                     (# n', v' #)
 114                                         = appendWord 0 n v
 115                                 in
 116                                   go0 w'' bs'' n' v'
 117                             else
 118                                 go1 w' bs'' n v
 119           | otherwise
 120               = go1 w bs n v
 121
 122       go0' ∷ Word32 → WordsData → bs → (# WordsData, bs #)
 123       go0' cBits w bs
 124           | cBits < 2
 125               = let w' = w { wdZeroesAcc = cBits }
 126                 in
 127                   (# w', bs #)
 128           | otherwise
 129               = let w' = w { wdZeroesAcc = 0 }
 130                 in
 131                   go0'' 1 cBits w' bs
 132
 133       go0'' ∷ Word32 → Word32 → WordsData → bs → (# WordsData, bs #)
 134       go0'' mask cBits w bs
 135           | cBits ≡ 1
 136               = let w' = w { wdZeroesAcc = wdZeroesAcc w .|. mask }
 137                 in
 138                   (# w', bs #)
 139           | otherwise
 140               = let cBits' = cBits - 1
 141                     w'     = if B.head bs then
 142                                  w { wdZeroesAcc = wdZeroesAcc w .|. mask }
 143                              else
 144                                  w
 145                     mask'  = mask `shiftL` 1
 146                     bs'    = B.tail bs
 147                 in
 148                   go0'' mask' cBits' w' bs'
 149
 150       go1 ∷ WordsData → bs → Int → New v Int32
 151           → (# WordsData, bs, Int, New v Int32 #)
 152       go1 w bs n v
 153           | wdHoldingZero w
 154               = let w' = w { wdHoldingZero = False }
 155                 in
 156                   go2 0 w' bs n v
 157           | otherwise
 158               = let next8 ∷ Word8
 159                     next8 = B.toBits (B.take (8 ∷ Int) bs)
 160                 in
 161                   if next8 ≡ 0xFF then
 162                       error "FIXME"
 163                   else
 164                       error "FIXME"
 165
 166       go2 ∷ Word32 → WordsData → bs → Int → New v Int32
 167           → (# WordsData, bs, Int, New v Int32 #)
 168       go2 0 w bs n v
 169           = let ent  = getEntropy n w
 170                 low  = 0
 171                 high = getMedian0 ent - 1
 172                 ent' = decMedian0 ent
 173                 w'   = setEntropy ent' n w
 174             in
 175               go3 low high w' bs n v
 176       go2 1 w bs n v
 177           = let ent  = getEntropy n w
 178                 low  = getMedian0 ent
 179                 high = low + getMedian1 ent - 1
 180                 ent' = (incMedian0 ∘ decMedian1) ent
 181                 w'   = setEntropy ent' n w
 182             in
 183               go3 low high w' bs n v
 184       go2 2 w bs n v
 185           = let ent   = getEntropy n w
 186                 low   = getMedian0 ent + getMedian1 ent
 187                 high  = low + getMedian2 ent - 1
 188                 ent'  = (incMedian0 ∘ incMedian1 ∘ decMedian2) ent
 189                 w'    = setEntropy ent' n w
 190             in
 191               go3 low high w' bs n v
 192       go2 onesCount w bs n v
 193           = let ent   = getEntropy n w
 194                 low   = getMedian0 ent + getMedian1 ent + (onesCount-2) ⋅ getMedian2 ent
 195                 high  = low + getMedian2 ent - 1
 196                 ent'  = (incMedian0 ∘ incMedian1 ∘ incMedian2) ent
 197                 w'    = setEntropy ent' n w
 198             in
 199               go3 low high w' bs n v
 200
 201       go3 ∷ Word32 → Word32 → WordsData → bs → Int → New v Int32
 202           → (# WordsData, bs, Int, New v Int32 #)
 203       go3 low high w bs n v
 204           = let (# code, bs' #)
 205                      = readCode bs (high - low)
 206                 low' = low + code
 207                 word = if B.head bs' then
 208                            fromIntegral $ complement low'
 209                        else
 210                            fromIntegral low'
 211                 bs'' = B.tail bs'
 212                 (# n', v' #)
 213                      = appendWord word n v
 214             in
 215               go0 w bs'' n' v'
 216
 217       appendWord ∷ Int32 → Int → New v Int32 → (# Int, New v Int32 #)
 218       appendWord word n v
 219           = let v' = New.modify (\mv → MV.unsafeWrite mv n word) v
 220                 n' = n + 1
 221             in
 222               (# n', v' #)
 223
 224       getEntropy ∷ Int → WordsData → EntropyData
 225       getEntropy n w
 226           | isMono        = fst $ wdEntropyData w
 227           | n `testBit` 0 = fst $ wdEntropyData w
 228           | otherwise     = snd $ wdEntropyData w
 229
 230       setEntropy ∷ EntropyData → Int → WordsData → WordsData
 231       setEntropy e n w
 232           | isMono        = w { wdEntropyData = (e, snd $ wdEntropyData w) }
 233           | n `testBit` 0 = w { wdEntropyData = (e, snd $ wdEntropyData w) }
 234           | otherwise     = w { wdEntropyData = (fst $ wdEntropyData w, e) }
 235
 236 -- | Read a single unsigned value from the specified bitstream with a
 237 -- value from 0 to maxCode. If there are exactly a power of two number
 238 -- of possible codes then this will read a fixed number of bits;
 239 -- otherwise it reads the minimum number of bits and then determines
 240 -- whether another bit is needed to define the code.
 241 readCode ∷ Bitstream bs ⇒ bs → Word32 → (# Word32, bs #)
 242 {-# INLINEABLE readCode #-}
 243 readCode bs 0       = (# 0, bs #)
 244 readCode bs 1       = (# b2n (B.head bs), B.tail bs #)
 245 readCode bs maxCode
 246     = let !bitCount = countBits maxCode
 247           !extras   = bit bitCount - maxCode - 1
 248           !code     = B.toBits (B.take (bitCount - 1) bs)
 249           (# code', bitCount' #)
 250                     = if code ≥ extras then
 251                           (# (code `shiftL` 1)
 252                              - extras
 253                              + b2n (bs B.!! bitCount)
 254                            , bitCount #)
 255                       else
 256                           (# code, bitCount - 1 #)
 257           !bs'      = B.drop bitCount' bs
 258       in
 259         (# code', bs' #)