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