Merge pull request #98 from haskell-works/new-unzipFromListN2-function

New unzipFromListN2 function

Author: newhoggy

src/HaskellWorks/Data/Vector/Storable.hs

@@ -10,6 +10,7 @@ module HaskellWorks.Data.Vector.Storable
   , constructSI
   , construct2N
   , construct64UnzipN
+  , unzipFromListN2
   ) where
 
 import Control.Monad.ST                   (ST, runST)
@@ -18,7 +19,7 @@ import Data.Vector.Storable               (Storable)
 import Data.Word
 import Foreign.ForeignPtr
 import HaskellWorks.Data.Vector.AsVector8
-import Prelude                            hiding (foldMap)
+import Prelude                            hiding (abs, foldMap)
 
 import qualified Data.ByteString              as BS
 import qualified Data.Vector.Generic          as DVG
@@ -126,3 +127,25 @@ construct64UnzipN nBytes xs = (DVS.unsafeCast ibv, DVS.unsafeCast bpv)
           DVSM.set (DVSM.take ibl ibmv) 0
           DVSM.set (DVSM.take bpl bpmv) 0
           return (DVSM.length (DVSM.drop ibl ibmv), DVSM.length (DVSM.drop bpl bpmv))
+
+unzipFromListN2 :: (Storable a, Storable b) => Int -> [(a, b)] -> (DVS.Vector a, DVS.Vector b)
+unzipFromListN2 n abs = runST $ do
+  mas <- DVSM.unsafeNew n
+  mbs <- DVSM.unsafeNew n
+  len <- go 0 mas mbs abs
+  as <- DVG.unsafeFreeze (DVSM.take len mas)
+  bs <- DVG.unsafeFreeze (DVSM.take len mbs)
+  return (as, bs)
+  where go :: (Storable c, Storable d)
+          => Int
+          -> DVSM.MVector s c
+          -> DVSM.MVector s d
+          -> [(c, d)]
+          -> ST s Int
+        go i _   _   []           = return i
+        go i mvc mvd ((c, d):cds) = if i < n
+          then do
+            DVSM.write mvc i c
+            DVSM.write mvd i d
+            go (i + 1) mvc mvd cds
+          else return i

test/HaskellWorks/Data/Vector/StorableSpec.hs

@@ -6,11 +6,11 @@ module HaskellWorks.Data.Vector.StorableSpec
   ( spec
   ) where
 
-import Control.Monad.ST                  (ST)
-import Data.Vector.Storable              (Storable)
-import HaskellWorks.Data.Vector.Storable (construct64UnzipN)
+import Control.Monad.ST            (ST)
+import Data.Vector.Storable        (Storable)
 import HaskellWorks.Hspec.Hedgehog
 import Hedgehog
+import Prelude                     hiding (abs)
 import Test.Hspec
 
 import qualified Data.ByteString                   as BS
@@ -51,11 +51,23 @@ spec = describe "HaskellWorks.Data.Vector.StorableSpec" $ do
       al    <- forAll $ pure $ BS.length (mconcat ass)
       bl    <- forAll $ pure $ BS.length (mconcat bss)
       len   <- forAll $ pure $ al + bl
-      let res = construct64UnzipN len (zip ass bss)
+      let res = DVS.construct64UnzipN len (zip ass bss)
       let ra = BS.toByteString (fst res)
       let rb = BS.toByteString (snd res)
       ra === BS.padded ((BS.length as + 7) `div` 8 * 8) as
       rb === BS.padded ((BS.length bs + 7) `div` 8 * 8) bs
+  it "unzipFromListN2" $ requireProperty $ do
+    abs       <- forAll $ G.list (R.linear 0 8) $ (,)
+      <$> G.word8 R.constantBounded
+      <*> G.word8 R.constantBounded
+    len       <- forAll $ G.int (R.linear 0 8)
+    as        <- forAll $ pure $ fmap fst abs
+    bs        <- forAll $ pure $ fmap snd abs
+    (va, vb)  <- forAll $ pure $ DVS.unzipFromListN2 len abs
+    let eva = DVS.fromList (take len as)
+    let evb = DVS.fromList (take len bs)
+    va === eva
+    vb === evb
 
 dupList :: [a] -> [a]
 dupList (a:as) = (a:a:dupList as)