added unzipVectorWith

2 files changed, 27 insertions, 3 deletions

diff --git a/lib/Data/Packed/Internal/Vector.hs b/lib/Data/Packed/Internal/Vector.hs
index d97a86d..06db806 100644
--- a/lib/Data/Packed/Internal/Vector.hs
+++ b/lib/Data/Packed/Internal/Vector.hs
@@ -1,4 +1,4 @@
-{-# LANGUAGE MagicHash, CPP, UnboxedTuples, BangPatterns #-}
+{-# LANGUAGE MagicHash, CPP, UnboxedTuples, BangPatterns, FlexibleContexts #-}
 -----------------------------------------------------------------------------
 -- |
 -- Module      :  Data.Packed.Internal.Vector
@@ -17,7 +17,7 @@ module Data.Packed.Internal.Vector (
     Vector, dim,
     fromList, toList, (|>),
     join, (@>), safe, at, at', subVector, takesV,
-    mapVector, zipVector,
+    mapVector, zipVector, unzipVectorWith,
     foldVector, foldVectorG, foldLoop,
     createVector, vec,
     asComplex, asReal,
@@ -318,6 +318,25 @@ zipVector f u v = unsafePerformIO $ do
     return w
 {-# INLINE zipVector #-}
 
+-- | unzipWith for Vectors
+unzipVectorWith :: (Storable (a,b), Storable c, Storable d) 
+                   => (a -> c) -> (b -> d) -> Vector (a,b) -> (Vector c,Vector d)
+unzipVectorWith f g u = unsafePerformIO $ do
+      let n = dim u
+      v <- createVector n
+      w <- createVector n
+      unsafeWith u $ \pu ->
+          unsafeWith v $ \pv ->
+              unsafeWith w $ \pw -> do
+                  let go (-1) = return ()
+                      go !k   = do (x,y) <- peekElemOff pu k
+                                   pokeElemOff          pv k (f x)
+                                   pokeElemOff          pw k (g y)
+                                   go (k-1)
+                  go (n-1)
+      return (v,w)
+{-# INLINE unzipVectorWith #-}
+
 foldVector f x v = unsafePerformIO $
     unsafeWith (v::Vector Double) $ \p -> do
         let go (-1) s = return s
diff --git a/lib/Data/Packed/Vector.hs b/lib/Data/Packed/Vector.hs
index c5b8ba7..ad2ae10 100644
--- a/lib/Data/Packed/Vector.hs
+++ b/lib/Data/Packed/Vector.hs
@@ -1,3 +1,4 @@
+{-# LANGUAGE FlexibleContexts #-}
 -----------------------------------------------------------------------------
 -- |
 -- Module      :  Data.Packed.Vector
@@ -23,7 +24,7 @@ module Data.Packed.Vector (
 --    vectorFMax, vectorFMin, vectorFMaxIndex, vectorFMinIndex,
 --    vectorMax, vectorMin,
     vectorMaxIndex, vectorMinIndex,
-    mapVector, zipVector,
+    mapVector, zipVector, unzipVector, unzipVectorWith,
     fscanfVector, fprintfVector, freadVector, fwriteVector,
     foldLoop, foldVector, foldVectorG
 ) where
@@ -149,3 +150,7 @@ vecdisp f v
     . unwords . lines .  tail . dropWhile (not . (`elem` " \n"))
     . f . trans . reshape 1
     $ v
+
+-- | unzip for Vectors
+unzipVector :: (Storable a, Storable b, Storable (a,b)) => Vector (a,b) -> (Vector a,Vector b)
+unzipVector = unzipVectorWith id id