Codec/Audio/WavPack/Words.hs

   1 {-# LANGUAGE
   2     BangPatterns
   3   , DoAndIfThenElse
   4   , FlexibleContexts
   5   , ScopedTypeVariables
   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.Cont
  32 import Control.Monad.ST
  33 import Data.Bits
  34 import Data.Bitstream.Generic (Bitstream)
  35 import qualified Data.Bitstream.Generic as B
  36 import Data.Int
  37 import Data.STRef
  38 import qualified Data.Vector.Generic.Mutable as MV
  39 import qualified Data.Vector.Unboxed as UV
  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 -- | Maximum consecutive 1s sent for /div/ data.
  58 limitOnes ∷ Num n ⇒ n
  59 {-# INLINE limitOnes #-}
  60 limitOnes = 16
  61
  62 getOnesCount ∷ Num a ⇒ Word8 → a
  63 {-# INLINE getOnesCount #-}
  64 getOnesCount = fromIntegral ∘ UV.unsafeIndex oct ∘ fromIntegral
  65     where
  66       oct ∷ UV.Vector Word8
  67       {-# NOINLINE oct #-}
  68       oct = UV.fromList
  69             [ 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 4 --   0 -  15
  70             , 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 5 --  16 -  31
  71             , 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 4 --  32 -  47
  72             , 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 6 --  48 -  63
  73             , 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 4 --  64 -  79
  74             , 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 5 --  80 -  95
  75             , 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 4 --  96 - 111
  76             , 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 7 -- 112 - 127
  77             , 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 4 -- 128 - 143
  78             , 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 5 -- 144 - 159
  79             , 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 4 -- 160 - 175
  80             , 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 6 -- 176 - 191
  81             , 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 4 -- 192 - 207
  82             , 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 5 -- 208 - 223
  83             , 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 4 -- 124 - 239
  84             , 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 8 -- 240 - 255
  85             ]
  86
  87 -- | This is an optimized version of 'getWord' that is used for
  88 -- lossless only ('edErrorLimit' ≡ 0). Also, rather than obtaining a
  89 -- single sample, it can be used to obtain an entire buffer of either
  90 -- mono or stereo samples.
  91 getWordsLossless ∷ ∀bs v s. (Bitstream bs, MV.MVector v Int32)
  92                  ⇒ Bool       -- ^ Is the stream monaural?
  93                  → WordsData s
  94                  → STRef s bs -- ^ WV bitstream
  95                  → Int        -- ^ Number of samples to get
  96                  → ST s (v s Int32)
  97 {-# INLINEABLE getWordsLossless #-}
  98 getWordsLossless isMono w bs nSamples0
  99     = do v ← MV.new nSamples
 100          n ← runContT (for 0 (< nSamples) (+ 1) (loop v)) return
 101          return $ MV.take n v
 102     where
 103       nSamples ∷ Int
 104       nSamples = if isMono
 105                  then nSamples0
 106                  else nSamples0 ⋅ 2
 107
 108       -- Hey, this is way tooooo long...
 109       loop ∷ v s Int32
 110            → Int
 111            → ContT Int (ST s) ()
 112            → ContT Int (ST s) ()
 113            → ContT Int (ST s) ()
 114       loop v n break continue
 115           = do let c | isMono        = fst $ wdEntropyData w
 116                      | n `testBit` 0 = fst $ wdEntropyData w
 117                      | otherwise     = snd $ wdEntropyData w
 118                med00   ← lift $ readSTRef (edMedian0 $ fst $ wdEntropyData w)
 119                hldZero ← lift $ readSTRef (wdHoldingZero w)
 120                hldOne  ← lift $ readSTRef (wdHoldingOne  w)
 121                med10   ← lift $ readSTRef (edMedian0 $ snd $ wdEntropyData w)
 122                when (med00 < 2 ∧ hldZero ≡ False ∧ hldOne ≡ 0 ∧ med10 < 2) $
 123                     do zAcc ← lift $ readSTRef (wdZeroesAcc w)
 124                        if zAcc > 0 then
 125                            do lift $ modifySTRef (wdZeroesAcc w) ((-) 1)
 126                               when (zAcc > 1) $
 127                                    do lift $ MV.unsafeWrite v n 0
 128                                       continue
 129                        else
 130                            do cBits ← lift $ takeWhileLessThan id 33 bs
 131
 132                               when (cBits ≡ 33) $
 133                                   break
 134
 135                               if cBits < 2 then
 136                                   lift $ writeSTRef (wdZeroesAcc w) cBits
 137                               else
 138                                   do lift $ writeSTRef (wdZeroesAcc w) 0
 139                                      (mask, _)
 140                                          ← for (1, cBits)
 141                                                ((> 1) ∘ snd)
 142                                                (\(m, cb) → (m `shiftL` 1, cb - 1)) $ \(mask, _) _ _ →
 143                                                    do b ← lift $ takeHead bs
 144                                                       when b $
 145                                                           lift $ modifySTRef (wdZeroesAcc w) (.|. mask)
 146                                      lift $ modifySTRef (wdZeroesAcc w) (.|. mask)
 147
 148                               zAcc' ← lift$ readSTRef (wdZeroesAcc w)
 149                               when (zAcc' > 0) $
 150                                   do lift $ clearMedians $ fst $ wdEntropyData w
 151                                      lift $ clearMedians $ snd $ wdEntropyData w
 152                                      lift $ MV.unsafeWrite v n 0
 153                                      continue
 154
 155                onesCount ← lift $ newSTRef (⊥)
 156                if hldZero then
 157                    do lift $ writeSTRef onesCount 0
 158                       lift $ writeSTRef (wdHoldingZero w) False
 159                else
 160                    do next8 ← lift $ readBits (8 ∷ Word8) bs
 161                       if next8 ≡ 0xFF then
 162                           do lift $ dropBits (8 ∷ Word8) bs
 163                              oc ← for 8 (< limitOnes + 1) (+ 1) $ \_ break' _ →
 164                                       do h ← lift $ takeHead bs
 165                                          unless h $
 166                                              break'
 167                              lift $ writeSTRef onesCount oc
 168
 169                              when (oc ≡ limitOnes + 1) $
 170                                  break
 171
 172                              when (oc ≡ limitOnes) $
 173                                  do cBits ← for 0 (< 33) (+ 1) $ \_ break' _ →
 174                                                 do h ← lift $ takeHead bs
 175                                                    unless h $
 176                                                        break'
 177
 178                                     when (cBits ≡ 33) $
 179                                         break
 180
 181                                     if cBits < 2 then
 182                                         lift $ writeSTRef onesCount cBits
 183                                     else
 184                                         do lift $ writeSTRef onesCount 0
 185                                            (mask, _)
 186                                                ← for (1, cBits)
 187                                                      ((> 1) ∘ snd)
 188                                                      (\(m, cb) → (m `shiftL` 1, cb - 1)) $ \(mask, _) _ _ →
 189                                                          do b ← lift $ takeHead bs
 190                                                             when b $
 191                                                                 lift $ modifySTRef onesCount (.|. mask)
 192                                            lift $ modifySTRef onesCount (.|. mask)
 193
 194                                     lift $ modifySTRef onesCount (+ limitOnes)
 195                       else
 196                           do let oc = getOnesCount next8
 197                              lift $ writeSTRef onesCount oc
 198                              lift $ dropBits (oc + 1) bs
 199
 200                       oc ← lift $ readSTRef onesCount
 201                       let hldOne' = oc .&. 1
 202                       lift $ writeSTRef (wdHoldingOne w) hldOne'
 203                       if hldOne > 0 then
 204                           lift $ writeSTRef onesCount ((oc `shiftR` 1) + 1)
 205                       else
 206                           lift $ writeSTRef onesCount  (oc `shiftR` 1)
 207
 208                       lift $ writeSTRef (wdHoldingZero w)
 209                            $ ((complement hldOne') .&. 1) ≢ 0
 210
 211                oc ← lift $ readSTRef onesCount
 212                (low, high)
 213                   ← if oc ≡ 0 then
 214                         do high ← fmap ((-) 1) $ lift $ getMedian0 c
 215                            lift $ decMedian0 c
 216                            return (0, high)
 217                     else
 218                         do low ← lift $ getMedian0 c
 219                            lift $ incMedian0 c
 220
 221                            if oc ≡ 1 then
 222                                do high ← fmap (((-) 1) ∘ (+ low)) $ lift $ getMedian1 c
 223                                   lift $ decMedian1 c
 224                                   return (low, high)
 225                            else
 226                                do low' ← fmap (+ low) $ lift $ getMedian1 c
 227                                   lift $ incMedian1 c
 228
 229                                   if oc ≡ 2 then
 230                                       do high ← fmap (((-) 1) ∘ (+ low')) $ lift $ getMedian2 c
 231                                          lift $ decMedian2 c
 232                                          return (low', high)
 233                                   else
 234                                       do med2 ← lift $ getMedian2 c
 235                                          let low'' = low' + (oc - 2) ⋅ med2
 236                                              high  = low'' + med2 - 1
 237                                          lift $ incMedian2 c
 238                                          return (low'', high)
 239
 240                code ← lift $ readCode bs (high - low)
 241                b    ← lift $ takeHead bs
 242                let word = if b then
 243                               complement (low + code)
 244                           else
 245                               low + code
 246                lift $ MV.unsafeWrite v n (fromIntegral word)
 247
 248 -- | Read a single unsigned value from the specified bitstream with a
 249 -- value from 0 to maxCode. If there are exactly a power of two number
 250 -- of possible codes then this will read a fixed number of bits;
 251 -- otherwise it reads the minimum number of bits and then determines
 252 -- whether another bit is needed to define the code.
 253 readCode ∷ Bitstream bs ⇒ STRef s bs → Word32 → ST s Word32
 254 {-# INLINEABLE readCode #-}
 255 readCode _  0       = return 0
 256 readCode bs 1       = fmap b2n $ takeHead bs
 257 readCode bs maxCode
 258     = do let bitCount = countBits maxCode
 259              extras   = bit bitCount - maxCode - 1
 260          code ← takeBits (bitCount - 1) bs
 261          if code ≥ extras then
 262              do nextBit ← takeHead bs
 263                 return $ (code `shiftL` 1) - extras + b2n nextBit
 264            else
 265              return code
 266
 267 takeHead ∷ Bitstream bs ⇒ STRef s bs → ST s Bool
 268 {-# INLINEABLE takeHead #-}
 269 takeHead bsr
 270     = do bs ← readSTRef bsr
 271          writeSTRef bsr (B.tail bs)
 272          return (B.head bs)
 273
 274 takeWhileLessThan ∷ (Integral n, Bitstream bs)
 275                   ⇒ (Bool → Bool)
 276                   → n
 277                   → STRef s bs
 278                   → ST s n
 279 {-# INLINEABLE takeWhileLessThan #-}
 280 takeWhileLessThan f n bsr = go 0
 281     where
 282       {-# INLINE go #-}
 283       go i | i < n
 284                = do b ← takeHead bsr
 285                     if f b then
 286                         go (i + 1)
 287                       else
 288                         return i
 289            | otherwise
 290                = return i
 291
 292 readBits ∷ (Integral n, Bitstream bs, Bits a) ⇒ n → STRef s bs → ST s a
 293 {-# INLINEABLE readBits #-}
 294 readBits n bsr
 295     = do bs ← readSTRef bsr
 296          return (B.toBits (B.take n bs))
 297
 298 takeBits ∷ (Integral n, Bitstream bs, Bits a) ⇒ n → STRef s bs → ST s a
 299 {-# INLINEABLE takeBits #-}
 300 takeBits n bsr
 301     = do bs ← readSTRef bsr
 302          writeSTRef bsr (B.drop n bs)
 303          return (B.toBits (B.take n bs))
 304
 305 dropBits ∷ (Integral n, Bitstream bs) ⇒ n → STRef s bs → ST s ()
 306 {-# INLINEABLE dropBits #-}
 307 dropBits n bsr
 308     = do bs ← readSTRef bsr
 309          writeSTRef bsr (B.drop n bs)
 310
 311 -- | C style /for/ loop with /break/ and /continue/.
 312 for ∷ ∀m α. MonadCont m
 313     ⇒ α          -- ^ Initial state
 314     → (α → Bool) -- ^ Continue-the-loop predicate
 315     → (α → α)    -- ^ State modifier
 316     → (α → m () → m () → m ()) -- ^ Loop body taking breaker and
 317                                -- continuer
 318     → m α        -- ^ Final state
 319 for α0 contLoop next body
 320     = callCC $ \break → loop break α0
 321     where
 322       loop ∷ (α → m ()) → α → m α
 323       loop break α
 324           | contLoop α
 325               = do callCC $ \continue → body α (break α) (continue ())
 326                    loop break (next α)
 327           | otherwise
 328               = return α