Codec/Audio/WavPack/Decorrelation.hs

   1 {-# LANGUAGE
   2     DoAndIfThenElse
   3   , FlexibleContexts
   4   , ScopedTypeVariables
   5   , UnicodeSyntax
   6   #-}
   7 -- | FIXME
   8 module Codec.Audio.WavPack.Decorrelation
   9     ( DecorrPass(..)
  10     , decorrStereoPass
  11     )
  12     where
  13 import Control.Monad.ST
  14 import Data.Bits
  15 import Data.Int
  16 import Data.STRef
  17 import qualified Data.Vector.Generic.Mutable as MV
  18 import Prelude.Unicode
  19
  20 {-
  21 maxTerm ∷ Num a ⇒ a
  22 maxTerm = 8
  23 -}
  24
  25 -- | FIXME
  26 data DecorrPass v s
  27     = DecorrPass {
  28       -- | The decorrelation term: @(term /= 0) && ((-3 <= term <= 8)
  29       --   || (term == 17) || (term <= 18))@
  30         dpTerm     ∷ !Int8
  31       -- | The decorrelation delta: @0 <= delta <= 8@
  32       , dpDelta    ∷ !Int8
  33       -- | The decorrelation weight for channel A: @-1024 <= weight <=
  34       --   1024@
  35       , dpWeightA  ∷ !(STRef s Int16)
  36       -- | The decorrelation weight for channel B: @-1024 <= weight <=
  37       --   1024@
  38       , dpWeightB  ∷ !(STRef s Int16)
  39       -- | The decorrelation samples for channel A.
  40       , dpSamplesA ∷ !(v s Int32)
  41       -- | The decorrelation samples for channel B.
  42       , dpSamplesB ∷ !(v s Int32)
  43       }
  44
  45 -- |General function to perform stereo decorrelation pass on specified
  46 -- buffer (although since this is the reverse function it might
  47 -- technically be called /correlation/ instead). This version handles
  48 -- all sample resolutions and weight deltas. The dpSamplesX data is
  49 -- *not* returned normalized for term values 1-8, so it should be
  50 -- normalized if it is going to be used to call this function again.
  51 decorrStereoPass ∷ ∀v s. MV.MVector v Int32
  52                  ⇒ DecorrPass v s
  53                  → v s Int32
  54                  → ST s ()
  55 {-# INLINEABLE decorrStereoPass #-}
  56 decorrStereoPass dp buffer
  57     | MV.length buffer `rem` 2 ≢ 0
  58         = fail "decorrStereoPass: assertion failed: MV.length buffer `rem` 2 ≡ 0"
  59     | dpTerm dp ≡ 17
  60         = flip mapM_ [0 .. (MV.length buffer `div` 2) - 1] $ \n →
  61               do a0 ← MV.unsafeRead (dpSamplesA dp) 0
  62                  a1 ← MV.unsafeRead (dpSamplesA dp) 1
  63                  MV.unsafeWrite (dpSamplesA dp) 1 a0
  64                  tmpA ← MV.unsafeRead buffer (n ⋅ 2)
  65                  weiA ← readSTRef (dpWeightA dp)
  66                  let samA  = 2 ⋅ a0 - a1
  67                      samA' = applyWeight weiA samA + tmpA
  68                  MV.unsafeWrite (dpSamplesA dp) 0 samA'
  69                  MV.unsafeWrite buffer (n ⋅ 2) samA'
  70                  writeSTRef (dpWeightA dp) $ updateWeight weiA (dpDelta dp) tmpA samA
  71
  72                  b0 ← MV.unsafeRead (dpSamplesB dp) 0
  73                  b1 ← MV.unsafeRead (dpSamplesB dp) 1
  74                  MV.unsafeWrite (dpSamplesB dp) 1 b0
  75                  tmpB ← MV.unsafeRead buffer (n ⋅ 2 + 1)
  76                  weiB ← readSTRef (dpWeightB dp)
  77                  let samB  = 2 ⋅ b0 - b1
  78                      samB' = applyWeight weiB samB + tmpB
  79                  MV.unsafeWrite (dpSamplesB dp) 0 samB'
  80                  MV.unsafeWrite buffer (n ⋅ 2 + 1) samB'
  81                  writeSTRef (dpWeightB dp) $ updateWeight weiB (dpDelta dp) tmpB samB
  82     | otherwise
  83         = fail "FIXME"
  84
  85 applyWeight ∷ Int16 → Int32 → Int32
  86 {-# INLINE applyWeight #-}
  87 applyWeight weight sample
  88     | sample `shiftR` 0xFFFF ≡ 0
  89         = (fromIntegral weight ⋅ sample + 512) `shiftR` 10
  90     | otherwise
  91         = ( (((sample .&. 0xFFFF) ⋅ fromIntegral weight) `shiftR` 9) +
  92             (((sample .&. complement 0xFFFF) `shiftR` 9) ⋅ fromIntegral weight) +
  93             1
  94           ) `shiftR` 1
  95
  96 updateWeight ∷ Int16 → Int8 → Int32 → Int32 → Int16
  97 {-# INLINE updateWeight #-}
  98 updateWeight weight δ source result
  99     | source ≢ 0 ∧ result ≢ 0
 100         = let s = fromIntegral $ (source `xor` result) `shiftR` 31
 101           in
 102             (fromIntegral δ `xor` s) + (weight - s)
 103     | otherwise
 104         = weight