still working on decorrStereoPass

[wavpack.git] / Codec / Audio / WavPack / Decorrelation.hs

diff --git a/Codec/Audio/WavPack/Decorrelation.hs b/Codec/Audio/WavPack/Decorrelation.hs
index 0230244b90ff045da22c8579e475aab0c59519c1..3e2271e0ab1ae874ffaa087e114c1e8c295c2bbb 100644 (file)
--- a/Codec/Audio/WavPack/Decorrelation.hs
+++ b/Codec/Audio/WavPack/Decorrelation.hs
@@ -14,7 +14,6 @@ import Control.Monad.ST
 import Data.Bits
 import Data.Int
 import Data.STRef
-import qualified Data.Strict as S
 import qualified Data.Vector.Generic.Mutable as MV
 import Prelude.Unicode
 
@@ -36,11 +35,11 @@ data DecorrPass v s
       , dpWeightA  ∷ !(STRef s Int16)
       -- | The decorrelation weight for channel B: @-1024 <= weight <=
       --   1024@
-      , dpWeightB  ∷ !(S.Maybe (STRef s Int16))
+      , dpWeightB  ∷ !(STRef s Int16)
       -- | The decorrelation samples for channel A.
       , dpSamplesA ∷ !(v s Int32)
       -- | The decorrelation samples for channel B.
-      , dpSamplesB ∷ !(S.Maybe (v s Int32))
+      , dpSamplesB ∷ !(v s Int32)
       }
 
 -- |General function to perform stereo decorrelation pass on specified
@@ -55,14 +54,31 @@ decorrStereoPass ∷ ∀v s. MV.MVector v Int32
                  → ST s ()
 {-# INLINEABLE decorrStereoPass #-}
 decorrStereoPass dp buffer
+    | MV.length buffer `rem` 2 ≢ 0
+        = fail "decorrStereoPass: assertion failed: MV.length buffer `rem` 2 ≡ 0"
     | dpTerm dp ≡ 17
         = flip mapM_ [0 .. (MV.length buffer `div` 2) - 1] $ \n →
               do a0 ← MV.unsafeRead (dpSamplesA dp) 0
                  a1 ← MV.unsafeRead (dpSamplesA dp) 1
-                 let samA = 2 ⋅ a0 - a1
                  MV.unsafeWrite (dpSamplesA dp) 1 a0
                  tmpA ← MV.unsafeRead buffer (n ⋅ 2)
-                 fail "FIXME"
+                 weiA ← readSTRef (dpWeightA dp)
+                 let samA  = 2 ⋅ a0 - a1
+                     samA' = applyWeight weiA samA + tmpA
+                 MV.unsafeWrite (dpSamplesA dp) 0 samA'
+                 MV.unsafeWrite buffer (n ⋅ 2) samA'
+                 writeSTRef (dpWeightA dp) $ updateWeight weiA (dpDelta dp) tmpA samA
+
+                 b0 ← MV.unsafeRead (dpSamplesB dp) 0
+                 b1 ← MV.unsafeRead (dpSamplesB dp) 1
+                 MV.unsafeWrite (dpSamplesB dp) 1 b0
+                 tmpB ← MV.unsafeRead buffer (n ⋅ 2 + 1)
+                 weiB ← readSTRef (dpWeightB dp)
+                 let samB  = 2 ⋅ b0 - b1
+                     samB' = applyWeight weiB samB + tmpB
+                 MV.unsafeWrite (dpSamplesB dp) 0 samB'
+                 MV.unsafeWrite buffer (n ⋅ 2 + 1) samB'
+                 writeSTRef (dpWeightB dp) $ updateWeight weiB (dpDelta dp) tmpB samB
     | otherwise
         = fail "FIXME"
 
@@ -76,3 +92,13 @@ applyWeight weight sample
             (((sample .&. complement 0xFFFF) `shiftR` 9) ⋅ fromIntegral weight) +
             1
           ) `shiftR` 1
+
+updateWeight ∷ Int16 → Int8 → Int32 → Int32 → Int16
+{-# INLINE updateWeight #-}
+updateWeight weight δ source result
+    | source ≢ 0 ∧ result ≢ 0
+        = let s = fromIntegral $ (source `xor` result) `shiftR` 31
+          in
+            (fromIntegral δ `xor` s) + (weight - s)
+    | otherwise
+        = weight