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                  → Word32     -- ^ 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 = fromIntegral $
 105                  if isMono
 106                  then nSamples0
 107                  else nSamples0 ⋅ 2
 108
 109       -- Hey, this is way tooooo long...
 110       loop ∷ v s Int32
 111            → Int
 112            → ContT Int (ST s) ()
 113            → ContT Int (ST s) ()
 114            → ContT Int (ST s) ()
 115       loop v n break continue
 116           = do let c | isMono        = fst $ wdEntropyData w
 117                      | n `testBit` 0 = fst $ wdEntropyData w
 118                      | otherwise     = snd $ wdEntropyData w
 119                med00   ← lift $ readSTRef (edMedian0 $ fst $ wdEntropyData w)
 120                hldZero ← lift $ readSTRef (wdHoldingZero w)
 121                hldOne  ← lift $ readSTRef (wdHoldingOne  w)
 122                med10   ← lift $ readSTRef (edMedian0 $ snd $ wdEntropyData w)
 123                when (med00 < 2 ∧ hldZero ≡ False ∧ hldOne ≡ 0 ∧ med10 < 2) $
 124                     do zAcc ← lift $ readSTRef (wdZeroesAcc w)
 125                        if zAcc > 0 then
 126                            do lift $ modifySTRef (wdZeroesAcc w) ((-) 1)
 127                               when (zAcc > 1) $
 128                                    do lift $ MV.unsafeWrite v n 0
 129                                       continue
 130                        else
 131                            do cBits ← lift $ takeWhileLessThan id 33 bs
 132
 133                               when (cBits ≡ 33) $
 134                                   break
 135
 136                               if cBits < 2 then
 137                                   lift $ writeSTRef (wdZeroesAcc w) cBits
 138                               else
 139                                   do lift $ writeSTRef (wdZeroesAcc w) 0
 140                                      (mask, _)
 141                                          ← for (1, cBits)
 142                                                ((> 1) ∘ snd)
 143                                                (\(m, cb) → (m `shiftL` 1, cb - 1)) $ \(mask, _) _ _ →
 144                                                    do b ← lift $ takeHead bs
 145                                                       when b $
 146                                                           lift $ modifySTRef (wdZeroesAcc w) (.|. mask)
 147                                      lift $ modifySTRef (wdZeroesAcc w) (.|. mask)
 148
 149                               zAcc' ← lift$ readSTRef (wdZeroesAcc w)
 150                               when (zAcc' > 0) $
 151                                   do lift $ clearMedians $ fst $ wdEntropyData w
 152                                      lift $ clearMedians $ snd $ wdEntropyData w
 153                                      lift $ MV.unsafeWrite v n 0
 154                                      continue
 155
 156                onesCount ← lift $ newSTRef (⊥)
 157                if hldZero then
 158                    do lift $ writeSTRef onesCount 0
 159                       lift $ writeSTRef (wdHoldingZero w) False
 160                else
 161                    do next8 ← lift $ readBits (8 ∷ Word8) bs
 162                       if next8 ≡ 0xFF then
 163                           do lift $ dropBits (8 ∷ Word8) bs
 164                              oc ← for 8 (< limitOnes + 1) (+ 1) $ \_ break' _ →
 165                                       do h ← lift $ takeHead bs
 166                                          unless h $
 167                                              break'
 168                              lift $ writeSTRef onesCount oc
 169
 170                              when (oc ≡ limitOnes + 1) $
 171                                  break
 172
 173                              when (oc ≡ limitOnes) $
 174                                  do cBits ← for 0 (< 33) (+ 1) $ \_ break' _ →
 175                                                 do h ← lift $ takeHead bs
 176                                                    unless h $
 177                                                        break'
 178
 179                                     when (cBits ≡ 33) $
 180                                         break
 181
 182                                     if cBits < 2 then
 183                                         lift $ writeSTRef onesCount cBits
 184                                     else
 185                                         do lift $ writeSTRef onesCount 0
 186                                            (mask, _)
 187                                                ← for (1, cBits)
 188                                                      ((> 1) ∘ snd)
 189                                                      (\(m, cb) → (m `shiftL` 1, cb - 1)) $ \(mask, _) _ _ →
 190                                                          do b ← lift $ takeHead bs
 191                                                             when b $
 192                                                                 lift $ modifySTRef onesCount (.|. mask)
 193                                            lift $ modifySTRef onesCount (.|. mask)
 194
 195                                     lift $ modifySTRef onesCount (+ limitOnes)
 196                       else
 197                           do let oc = getOnesCount next8
 198                              lift $ writeSTRef onesCount oc
 199                              lift $ dropBits (oc + 1) bs
 200
 201                       oc ← lift $ readSTRef onesCount
 202                       let hldOne' = oc .&. 1
 203                       lift $ writeSTRef (wdHoldingOne w) hldOne'
 204                       if hldOne > 0 then
 205                           lift $ writeSTRef onesCount ((oc `shiftR` 1) + 1)
 206                       else
 207                           lift $ writeSTRef onesCount  (oc `shiftR` 1)
 208
 209                       lift $ writeSTRef (wdHoldingZero w)
 210                            $ ((complement hldOne') .&. 1) ≢ 0
 211
 212                oc ← lift $ readSTRef onesCount
 213                (low, high)
 214                   ← if oc ≡ 0 then
 215                         do high ← fmap ((-) 1) $ lift $ getMedian0 c
 216                            lift $ decMedian0 c
 217                            return (0, high)
 218                     else
 219                         do low ← lift $ getMedian0 c
 220                            lift $ incMedian0 c
 221
 222                            if oc ≡ 1 then
 223                                do high ← fmap (((-) 1) ∘ (+ low)) $ lift $ getMedian1 c
 224                                   lift $ decMedian1 c
 225                                   return (low, high)
 226                            else
 227                                do low' ← fmap (+ low) $ lift $ getMedian1 c
 228                                   lift $ incMedian1 c
 229
 230                                   if oc ≡ 2 then
 231                                       do high ← fmap (((-) 1) ∘ (+ low')) $ lift $ getMedian2 c
 232                                          lift $ decMedian2 c
 233                                          return (low', high)
 234                                   else
 235                                       do med2 ← lift $ getMedian2 c
 236                                          let low'' = low' + (oc - 2) ⋅ med2
 237                                              high  = low'' + med2 - 1
 238                                          lift $ incMedian2 c
 239                                          return (low'', high)
 240
 241                code ← lift $ readCode bs (high - low)
 242                b    ← lift $ takeHead bs
 243                let word = if b then
 244                               complement (low + code)
 245                           else
 246                               low + code
 247                lift $ MV.unsafeWrite v n (fromIntegral word)
 248
 249 -- | Read a single unsigned value from the specified bitstream with a
 250 -- value from 0 to maxCode. If there are exactly a power of two number
 251 -- of possible codes then this will read a fixed number of bits;
 252 -- otherwise it reads the minimum number of bits and then determines
 253 -- whether another bit is needed to define the code.
 254 readCode ∷ Bitstream bs ⇒ STRef s bs → Word32 → ST s Word32
 255 {-# INLINEABLE readCode #-}
 256 readCode _  0       = return 0
 257 readCode bs 1       = fmap b2n $ takeHead bs
 258 readCode bs maxCode
 259     = do let bitCount = countBits maxCode
 260              extras   = bit bitCount - maxCode - 1
 261          code ← takeBits (bitCount - 1) bs
 262          if code ≥ extras then
 263              do nextBit ← takeHead bs
 264                 return $ (code `shiftL` 1) - extras + b2n nextBit
 265            else
 266              return code
 267
 268 takeHead ∷ Bitstream bs ⇒ STRef s bs → ST s Bool
 269 {-# INLINEABLE takeHead #-}
 270 takeHead bsr
 271     = do bs ← readSTRef bsr
 272          writeSTRef bsr (B.tail bs)
 273          return (B.head bs)
 274
 275 takeWhileLessThan ∷ (Integral n, Bitstream bs)
 276                   ⇒ (Bool → Bool)
 277                   → n
 278                   → STRef s bs
 279                   → ST s n
 280 {-# INLINEABLE takeWhileLessThan #-}
 281 takeWhileLessThan f n bsr = go 0
 282     where
 283       {-# INLINE go #-}
 284       go i | i < n
 285                = do b ← takeHead bsr
 286                     if f b then
 287                         go (i + 1)
 288                       else
 289                         return i
 290            | otherwise
 291                = return i
 292
 293 readBits ∷ (Integral n, Bitstream bs, Bits a) ⇒ n → STRef s bs → ST s a
 294 {-# INLINEABLE readBits #-}
 295 readBits n bsr
 296     = do bs ← readSTRef bsr
 297          return (B.toBits (B.take n bs))
 298
 299 takeBits ∷ (Integral n, Bitstream bs, Bits a) ⇒ n → STRef s bs → ST s a
 300 {-# INLINEABLE takeBits #-}
 301 takeBits n bsr
 302     = do bs ← readSTRef bsr
 303          writeSTRef bsr (B.drop n bs)
 304          return (B.toBits (B.take n bs))
 305
 306 dropBits ∷ (Integral n, Bitstream bs) ⇒ n → STRef s bs → ST s ()
 307 {-# INLINEABLE dropBits #-}
 308 dropBits n bsr
 309     = do bs ← readSTRef bsr
 310          writeSTRef bsr (B.drop n bs)
 311
 312 -- | C style /for/ loop with /break/ and /continue/.
 313 for ∷ ∀m α. MonadCont m
 314     ⇒ α          -- ^ Initial state
 315     → (α → Bool) -- ^ Continue-the-loop predicate
 316     → (α → α)    -- ^ State modifier
 317     → (α → m () → m () → m ()) -- ^ Loop body taking breaker and
 318                                -- continuer
 319     → m α        -- ^ Final state
 320 for α0 contLoop next body
 321     = callCC $ \break → loop break α0
 322     where
 323       loop ∷ (α → m ()) → α → m α
 324       loop break α
 325           | contLoop α
 326               = do callCC $ \continue → body α (break α) (continue ())
 327                    loop break (next α)
 328           | otherwise
 329               = return α