Merge pull request #90 from haskell-works/fix-construct64UnzipN-with-tests

Fix construct64UnzipN with tests

Author: newhoggy

.vscode/tasks.json

@@ -37,7 +37,7 @@
       "label": "Test",
       "type": "shell",
       "command": "bash",
-      "args": ["-lc", "cabal new-test --enable-tests && echo 'Done'"],
+      "args": ["-lc", "cabal new-test --enable-tests --test-show-details=direct && echo 'Done'"],
       "group": {
         "kind": "test",
         "isDefault": true

hw-prim.cabal

@@ -35,7 +35,7 @@ common deepseq              { build-depends: deepseq              >= 1.4
 common directory            { build-depends: directory            >= 1.2        && < 1.4    }
 common exceptions           { build-depends: exceptions           >= 0.8        && < 0.11   }
 common ghc-prim             { build-depends: ghc-prim             >= 0.5        && < 0.6    }
-common hedgehog             { build-depends: hedgehog             >= 0.5        && < 1.1    }
+common hedgehog             { build-depends: hedgehog             >= 1.0        && < 1.1    }
 common hspec                { build-depends: hspec                >= 2.4        && < 2.8    }
 common hw-hspec-hedgehog    { build-depends: hw-hspec-hedgehog    >= 0.1        && < 0.2    }
 common mmap                 { build-depends: mmap                 >= 0.5        && < 0.6    }

src/HaskellWorks/Data/ByteString.hs

@@ -8,6 +8,7 @@ module HaskellWorks.Data.ByteString
   , ToByteString(..)
   , ToByteStrings(..)
   , mmap
+  , padded
   , rechunk
   , rechunkPadded
   , resegment
@@ -219,3 +220,7 @@ mmap filepath = do
   (fptr :: ForeignPtr Word8, offset, size) <- IO.mmapFileForeignPtr filepath IO.ReadOnly Nothing
   let !bs = BSI.fromForeignPtr (castForeignPtr fptr) offset size
   return bs
+
+padded :: Int -> BS.ByteString -> BS.ByteString
+padded n v = v <> BS.replicate ((n - BS.length v) `max` 0) 0
+{-# INLINE padded #-}

src/HaskellWorks/Data/Vector/Storable.hs

@@ -108,20 +108,21 @@ construct2N nb fb nc fc as = runST $ do
 construct64UnzipN :: Int -> [(BS.ByteString, BS.ByteString)] -> (DVS.Vector Word64, DVS.Vector Word64)
 construct64UnzipN nBytes xs = (DVS.unsafeCast ibv, DVS.unsafeCast bpv)
   where [ibv, bpv] = DVS.createT $ do
-          let nW64s     = (nBytes + 7) `div` 8
-          ibmv <- DVSM.new nW64s
-          bpmv <- DVSM.new nW64s
+          let nW8s     = (nBytes + 7) `div` 8 * 8
+          ibmv <- DVSM.new nW8s
+          bpmv <- DVSM.new nW8s
           (ibmvRemaining, bpmvRemaining) <- go ibmv bpmv xs
-          return
-            [ DVSM.take (((DVSM.length ibmv - ibmvRemaining) `div` 8) * 8) ibmv
-            , DVSM.take (((DVSM.length bpmv - bpmvRemaining) `div` 8) * 8) bpmv
-            ]
+          let ibl = ((DVSM.length ibmv - ibmvRemaining + 7) `div` 8) * 8
+          let bpl = ((DVSM.length bpmv - bpmvRemaining + 7) `div` 8) * 8
+          return [ DVSM.take ibl ibmv, DVSM.take bpl bpmv]
         go :: DVSM.MVector s Word8 -> DVSM.MVector s Word8 -> [(BS.ByteString, BS.ByteString)] -> ST s (Int, Int)
         go ibmv bpmv ((ib, bp):ys) = do
           DVS.copy (DVSM.take (BS.length ib) ibmv) (asVector8 ib)
           DVS.copy (DVSM.take (BS.length bp) bpmv) (asVector8 bp)
           go (DVSM.drop (BS.length ib) ibmv) (DVSM.drop (BS.length bp) bpmv) ys
         go ibmv bpmv [] = do
-          DVSM.set (DVSM.take 8 ibmv) 0
-          DVSM.set (DVSM.take 8 bpmv) 0
-          return (DVSM.length (DVSM.drop 8 ibmv), DVSM.length (DVSM.drop 8 bpmv))
+          let ibl = DVSM.length ibmv `mod` 8
+          let bpl = DVSM.length bpmv `mod` 8
+          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))

test/HaskellWorks/Data/Vector/StorableSpec.hs

@@ -6,15 +6,18 @@ module HaskellWorks.Data.Vector.StorableSpec
   ( spec
   ) where
 
-import Control.Monad.ST            (ST)
-import Data.Vector.Storable        (Storable)
+import Control.Monad.ST                  (ST)
+import Data.Vector.Storable              (Storable)
+import HaskellWorks.Data.Vector.Storable (construct64UnzipN)
 import HaskellWorks.Hspec.Hedgehog
 import Hedgehog
 import Test.Hspec
 
+import qualified Data.ByteString                   as BS
 import qualified Data.List                         as L
 import qualified Data.Vector.Storable              as DVS
 import qualified Data.Vector.Storable.Mutable      as DVSM
+import qualified HaskellWorks.Data.ByteString      as BS
 import qualified HaskellWorks.Data.Vector.Storable as DVS
 import qualified Hedgehog.Gen                      as G
 import qualified Hedgehog.Range                    as R
@@ -36,6 +39,23 @@ spec = describe "HaskellWorks.Data.Vector.StorableSpec" $ do
     let (bs, cs) = DVS.construct2N (length as) stepb (length as * 2) stepc as
     DVS.fromList as           === bs
     DVS.fromList (dupList as) === cs
+  describe "construct64UnzipN" $ do
+    it "property" $ requireProperty $ do
+      abss  <- forAll $ G.list (R.linear 1 8) $ (,)
+        <$> (fmap BS.pack (G.list (R.linear 1 8) (G.word8 R.constantBounded)))
+        <*> (fmap BS.pack (G.list (R.linear 1 8) (G.word8 R.constantBounded)))
+      ass   <- forAll $ pure $ fmap fst abss
+      bss   <- forAll $ pure $ fmap snd abss
+      as    <- forAll $ pure $ mconcat ass
+      bs    <- forAll $ pure $ mconcat bss
+      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 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
 
 dupList :: [a] -> [a]
 dupList (a:as) = (a:a:dupList as)