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 b72909fd3b0162e46d59f1e073a2ed4942d207b3..3e2271e0ab1ae874ffaa087e114c1e8c295c2bbb 100644 (file)
--- a/Codec/Audio/WavPack/Decorrelation.hs
+++ b/Codec/Audio/WavPack/Decorrelation.hs
@@ -1,13 +1,21 @@
 {-# LANGUAGE
-    UnicodeSyntax
+    DoAndIfThenElse
+  , FlexibleContexts
+  , ScopedTypeVariables
+  , UnicodeSyntax
   #-}
 -- | FIXME
 module Codec.Audio.WavPack.Decorrelation
     ( DecorrPass(..)
+    , decorrStereoPass
     )
     where
+import Control.Monad.ST
+import Data.Bits
 import Data.Int
-import qualified Data.Vector.Unboxed as UV
+import Data.STRef
+import qualified Data.Vector.Generic.Mutable as MV
+import Prelude.Unicode
 
 {-
 maxTerm ∷ Num a ⇒ a
@@ -15,22 +23,82 @@ maxTerm = 8
 -}
 
 -- | FIXME
-data DecorrPass
+data DecorrPass v s
     = DecorrPass {
-      -- | The decorrelation term: (term /= 0) && ((-3 <= term <= 8)
-      --   || (term == 17) || (term <= 18))
+      -- | The decorrelation term: @(term /= 0) && ((-3 <= term <= 8)
+      --   || (term == 17) || (term <= 18))@
         dpTerm     ∷ !Int8
-      -- | The decorrelation delta: 0 <= delta <= 8
+      -- | The decorrelation delta: @0 <= delta <= 8@
       , dpDelta    ∷ !Int8
-      -- | The decorrelation weight for channel A: -1024 <= weight <=
-      --   1024
-      , dpWeightA  ∷ !Int16
-      -- | The decorrelation weight for channel B: -1024 <= weight <=
-      --   1024
-      , dpWeightB  ∷ !Int16
+      -- | The decorrelation weight for channel A: @-1024 <= weight <=
+      --   1024@
+      , dpWeightA  ∷ !(STRef s Int16)
+      -- | The decorrelation weight for channel B: @-1024 <= weight <=
+      --   1024@
+      , dpWeightB  ∷ !(STRef s Int16)
       -- | The decorrelation samples for channel A.
-      , dpSamplesA ∷ !(UV.Vector Int32)
+      , dpSamplesA ∷ !(v s Int32)
       -- | The decorrelation samples for channel B.
-      , dpSamplesB ∷ !(UV.Vector Int32)
+      , dpSamplesB ∷ !(v s Int32)
       }
-    deriving (Eq, Show)
+
+-- |General function to perform stereo decorrelation pass on specified
+-- buffer (although since this is the reverse function it might
+-- technically be called /correlation/ instead). This version handles
+-- all sample resolutions and weight deltas. The dpSamplesX data is
+-- *not* returned normalized for term values 1-8, so it should be
+-- normalized if it is going to be used to call this function again.
+decorrStereoPass ∷ ∀v s. MV.MVector v Int32
+                 ⇒ DecorrPass v s
+                 → v s 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
+                 MV.unsafeWrite (dpSamplesA dp) 1 a0
+                 tmpA ← MV.unsafeRead buffer (n ⋅ 2)
+                 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"
+
+applyWeight ∷ Int16 → Int32 → Int32
+{-# INLINE applyWeight #-}
+applyWeight weight sample
+    | sample `shiftR` 0xFFFF ≡ 0
+        = (fromIntegral weight ⋅ sample + 512) `shiftR` 10
+    | otherwise
+        = ( (((sample .&. 0xFFFF) ⋅ fromIntegral weight) `shiftR` 9) +
+            (((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