refactor

fumieval · Dec 21, 2018 · 7292510 · 7292510
1 parent 08693c6
commit 7292510
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 32 deletions.
diff --git a/benchmarks/benchmark.hs b/benchmarks/benchmark.hs
@@ -4,9 +4,10 @@ import Gauge.Main
 import Codec.EliasFano
 import qualified Data.Binary as B
 import qualified Data.Vector.Unboxed as V
+import Data.Word (Word64)
 
-td :: V.Vector Int
-td = V.scanl (+) 0 $ V.fromList $ map fromEnum $ "Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum."
+td :: V.Vector Word64
+td = V.scanl (+) 0 $ V.fromList $ map (toEnum . fromEnum) $ "Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum."
 
 main = do
   let ef = unsafeFromVector td

diff --git a/src/Codec/EliasFano.hs b/src/Codec/EliasFano.hs
@@ -4,13 +4,11 @@
 {-# LANGUAGE LambdaCase #-}
 module Codec.EliasFano (EliasFano(..), unsafeFromVector, (!), prop_access) where
 
-import Control.Monad.Primitive
 import Control.Monad.ST
 import Data.Bits
 import qualified Data.Vector.Generic as V
 import qualified Data.Vector.Generic.Mutable as GM
 import qualified Data.Vector.Unboxed as UV
-import qualified Data.Vector.Unboxed.Mutable as MV
 import qualified Data.Vector.Fusion.Bundle.Monadic as B
 import qualified Data.Vector.Fusion.Bundle.Size as B
 import qualified Data.Vector.Fusion.Stream.Monadic as S
@@ -20,21 +18,21 @@ import Codec.EliasFano.Internal
 
 import qualified Test.QuickCheck as QC
 
-unsafeFromStreamNMax :: PrimMonad m => Int -> Int -> S.Stream m Int -> m EliasFano
-unsafeFromStreamNMax efLength maxValue (S.Stream upd s0) = do
-  let counterSize = 1 `unsafeShiftL` ceiling (logBase 2 $ fromIntegral efLength :: Double)
-  mcounter <- MV.replicate counterSize 0
+data EncoderState s = ESCont !Int !Word64 !Int | ESDone
 
-  let go s = upd s >>= \case
-        S.Done -> pure S.Done
-        S.Skip s' -> pure $ S.Skip s'
-        S.Yield v s' -> do
-          MV.unsafeModify mcounter (+1) (v `unsafeShiftR` efWidth)
-          return $ S.Yield (efWidth `B` fromIntegral v) s'
-  counter <- UV.unsafeFreeze mcounter
-  mefLower <- GM.munstream $ B.fromStream (chunk64 $ S.Stream go s0)
+unsafeFromVector :: V.Vector v Word64 => v Word64 -> EliasFano
+unsafeFromVector vec = runST $ do
+  let go ESDone = pure S.Done
+      go (ESCont i current n)
+        | current > maxValue `unsafeShiftR` efWidth = pure $ S.Yield n ESDone
+        | otherwise = pure $ case fromIntegral $ V.unsafeIndex vec i `unsafeShiftR` efWidth of
+          u | u == current -> S.Skip $ ESCont (i + 1) current (n + 1)
+            | otherwise -> S.Yield n (ESCont i (current + 1) 0)
+
+  mefLower <- GM.munstream $ B.fromStream
+    (chunk64 $ S.map (B efWidth . fromIntegral) $ B.elements $ B.fromVector vec)
     $ B.Exact $ (efWidth * efLength + 63) `div` 64
-  mefUpper <- GM.munstream $ B.fromStream (chunk64 $ S.Stream (upper counter) 0)
+  mefUpper <- GM.munstream $ B.fromStream (chunk64 $ unary $ S.Stream go $ ESCont 0 0 0)
     $ B.Exact $ (efLength + 3) `div` 4
   efUpper <- UV.unsafeFreeze mefUpper
   efLower <- UV.unsafeFreeze mefLower
@@ -43,18 +41,12 @@ unsafeFromStreamNMax efLength maxValue (S.Stream upd s0) = do
     , ..
     }
   where
+    efLength = V.length vec
+    maxValue
+      | V.null vec = 1
+      | otherwise = V.last vec + 1
     efWidth = max 1 $ ceiling $ logBase 2 (fromIntegral maxValue / fromIntegral efLength :: Double)
-    upper counter i
-      | i == V.length counter = pure S.Done
-      | otherwise = let !n = V.unsafeIndex counter i in pure $ S.Yield ((n + 1) `B` mask n) (i + 1)
-{-# INLINE unsafeFromStreamNMax #-}
-
-unsafeFromVector :: V.Vector v Int => v Int -> EliasFano
-unsafeFromVector vec
-  | V.null vec = runST $ unsafeFromStreamNMax 0 1 S.empty
-  | otherwise = runST $ unsafeFromStreamNMax (V.length vec) (V.last vec + 1)
-    $ B.elements $ B.fromVector vec
-{-# SPECIALISE unsafeFromVector :: UV.Vector Int -> EliasFano #-}
+{-# SPECIALISE unsafeFromVector :: UV.Vector Word64 -> EliasFano #-}
 
 data EliasFano = EliasFano
     { efLength :: !Int
@@ -65,14 +57,14 @@ data EliasFano = EliasFano
     }
     deriving Show
 
-(!) :: EliasFano -> Int -> Int
-(!) (EliasFano _ width upper ranks lower) i = unsafeShiftL (select ranks upper i - i) width
-  .|. fromIntegral (readBits lower width (i * width))
+(!) :: EliasFano -> Int -> Word64
+(!) (EliasFano _ width upper ranks lower) i = fromIntegral (unsafeShiftL (select ranks upper i - i) width)
+  .|. readBits lower width (i * width)
 
 prop_access :: [QC.NonNegative Int] -> QC.NonNegative Int -> QC.Property
 prop_access xs i_ = QC.counterexample (show (base, ef, i))
   $ ef ! i == base !! i
   where
     i = QC.getNonNegative i_ `mod` length base
-    base = scanl (+) 0 $ map QC.getNonNegative xs
+    base = scanl (+) 0 $ map (fromIntegral . QC.getNonNegative) xs
     ef = unsafeFromVector $ UV.fromList base
diff --git a/src/Codec/EliasFano/Internal.hs b/src/Codec/EliasFano/Internal.hs
@@ -3,6 +3,7 @@
 {-# LANGUAGE FlexibleContexts #-}
 {-# LANGUAGE LambdaCase #-}
 module Codec.EliasFano.Internal (B(..), chunk64
+  , unary
   , mask
   , readBits
   , select
@@ -32,6 +33,20 @@ chunk64 (S.Stream upd s0) = S.Stream go $ Chunker s0 zeroBits 0 where
          | otherwise -> S.Skip $ Chunker s' (acc .|. unsafeShiftL w' len) (len + width)
 {-# INLINE chunk64 #-}
 
+data Unary s = Unary s | UnaryCont !Int s
+
+unary :: Applicative m => S.Stream m Int -> S.Stream m B
+unary (S.Stream upd s0) = S.Stream go $ Unary s0 where
+  go (Unary s) = flip fmap (upd s) $ \case
+    S.Done -> S.Done
+    S.Skip s' -> S.Skip (Unary s')
+    S.Yield n s' -> step n s'
+  go (UnaryCont n s') = pure $ step n s'
+  step n s'
+    | n < 64 = S.Yield ((n + 1) `B` mask n) (Unary s')
+    | otherwise = S.Yield (B 64 (complement zeroBits)) (UnaryCont (n - 64) s')
+{-# INLINE unary #-}
+
 mask :: Int -> Word64
 mask n = unsafeShiftL 1 n - 1
 {-# INLINE mask #-}