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