diff --git a/Data/Conduit/Combinators.hs b/Data/Conduit/Combinators.hs --- a/Data/Conduit/Combinators.hs +++ b/Data/Conduit/Combinators.hs @@ -172,6 +172,9 @@ , unlinesAscii , linesUnbounded , linesUnboundedAscii + + -- * Special + , vectorBuilder ) where -- BEGIN IMPORTS @@ -193,6 +196,7 @@ import Control.Monad.Trans.Class (lift) import Control.Monad.Trans.Resource (MonadResource, MonadThrow) import Data.Conduit +import Data.Conduit.Internal (ConduitM (..), Pipe (..)) import qualified Data.Conduit.List as CL import Data.IOData import Data.Monoid (Monoid (..)) @@ -201,6 +205,7 @@ import Data.Sequences.Lazy import qualified Data.Vector.Generic as V import qualified Data.Vector.Generic.Mutable as VM +import Data.Void (absurd) import qualified Filesystem as F import Filesystem.Path (FilePath, ()) import Filesystem.Path.CurrentOS (encodeString, decodeString) @@ -208,7 +213,8 @@ Maybe (..), Monad (..), Num (..), Ord (..), fromIntegral, maybe, ($), Functor (..), Enum, seq, Show, Char, (||), - mod, otherwise, Either (..)) + mod, otherwise, Either (..), + ($!), succ) import Data.Word (Word8) import qualified Prelude import System.IO (Handle) @@ -220,6 +226,8 @@ import qualified System.Random.MWC as MWC import Data.Conduit.Combinators.Internal import qualified System.PosixCompat.Files as PosixC +import Data.Primitive.MutVar (MutVar, newMutVar, readMutVar, + writeMutVar) #ifndef WINDOWS import qualified System.Posix.Directory as Dir @@ -1749,3 +1757,89 @@ else yield x >> loop (Seq.drop 1 y) where (x, y) = Seq.break (== 10) t + +-- | Generally speaking, yielding values from inside a Conduit requires +-- some allocation for constructors. This can introduce an overhead, +-- similar to the overhead needed to represent a list of values instead of +-- a vector. This overhead is even more severe when talking about unboxed +-- values. +-- +-- This combinator allows you to overcome this overhead, and efficiently +-- fill up vectors. It takes two parameters. The first is the size of each +-- mutable vector to be allocated. The second is a function. The function +-- takes an argument which will yield the next value into a mutable +-- vector. +-- +-- Under the surface, this function uses a number of tricks to get high +-- performance. For more information on both usage and implementation, +-- please see: +-- . +-- +-- Since 1.0.0 +vectorBuilder :: (PrimMonad base, MonadBase base m, V.Vector v e, MonadBase base n) + => Int -- ^ size + -> ((e -> n ()) -> Sink i m r) + -> ConduitM i (v e) m r +vectorBuilder size inner = do + ref <- liftBase $ do + mv <- VM.new size + newMutVar $! S 0 mv id + res <- onAwait (yieldS ref) (inner (liftBase . addE ref)) + vs <- liftBase $ do + S idx mv front <- readMutVar ref + end <- + if idx == 0 + then return [] + else do + v <- V.unsafeFreeze mv + return [V.unsafeTake idx v] + return $ front end + Prelude.mapM_ yield vs + return res +{-# INLINE vectorBuilder #-} + +data S s v e = S + {-# UNPACK #-} !Int -- ^ index + !(V.Mutable v s e) + ([v e] -> [v e]) + +onAwait :: Monad m + => ConduitM i o m () + -> Sink i m r + -> ConduitM i o m r +onAwait (ConduitM callback) = + ConduitM . go . unConduitM + where + go (Done r) = Done r + go (HaveOutput _ _ o) = absurd o + go (NeedInput f g) = callback >> NeedInput (go . f) (go . g) + go (PipeM mp) = PipeM (liftM go mp) + go (Leftover f i) = Leftover (go f) i +{-# INLINE onAwait #-} + +yieldS :: (PrimMonad base, MonadBase base m) + => MutVar (PrimState base) (S (PrimState base) v e) + -> Producer m (v e) +yieldS ref = do + S idx mv front <- liftBase $ readMutVar ref + Prelude.mapM_ yield (front []) + liftBase $ writeMutVar ref $! S idx mv id +{-# INLINE yieldS #-} + +addE :: (PrimMonad m, V.Vector v e) + => MutVar (PrimState m) (S (PrimState m) v e) + -> e + -> m () +addE ref e = do + S idx mv front <- readMutVar ref + VM.write mv idx e + let idx' = succ idx + size = VM.length mv + if idx' >= size + then do + v <- V.unsafeFreeze mv + let front' = front . (v:) + mv' <- VM.new size + writeMutVar ref $! S 0 mv' front' + else writeMutVar ref $! S idx' mv front +{-# INLINE addE #-} diff --git a/Data/Conduit/Combinators/Unqualified.hs b/Data/Conduit/Combinators/Unqualified.hs --- a/Data/Conduit/Combinators/Unqualified.hs +++ b/Data/Conduit/Combinators/Unqualified.hs @@ -161,6 +161,9 @@ , unlinesAsciiC , linesUnboundedC , linesUnboundedAsciiC + + -- ** Special + , vectorBuilderC ) where -- BEGIN IMPORTS @@ -185,6 +188,7 @@ import Control.Monad.Trans.Class (lift) import Control.Monad.Trans.Resource (MonadResource, MonadThrow) import Data.Conduit +import Data.Conduit.Internal (ConduitM (..), Pipe (..)) import qualified Data.Conduit.List as CL import Data.IOData import Data.Monoid (Monoid (..)) @@ -193,6 +197,7 @@ import Data.Sequences.Lazy import qualified Data.Vector.Generic as V import qualified Data.Vector.Generic.Mutable as VM +import Data.Void (absurd) import qualified Filesystem as F import Filesystem.Path (FilePath, ()) import Filesystem.Path.CurrentOS (encodeString, decodeString) @@ -200,7 +205,8 @@ Maybe (..), Monad (..), Num (..), Ord (..), fromIntegral, maybe, ($), Functor (..), Enum, seq, Show, Char, (||), - mod, otherwise, Either (..)) + mod, otherwise, Either (..), + ($!), succ) import Data.Word (Word8) import qualified Prelude import System.IO (Handle) @@ -212,6 +218,8 @@ import qualified System.Random.MWC as MWC import Data.Conduit.Combinators.Internal import qualified System.PosixCompat.Files as PosixC +import Data.Primitive.MutVar (MutVar, newMutVar, readMutVar, + writeMutVar) #ifndef WINDOWS import qualified System.Posix.Directory as Dir @@ -1373,3 +1381,28 @@ => Conduit seq m seq linesUnboundedAsciiC = CC.linesUnboundedAscii {-# INLINE linesUnboundedAsciiC #-} + +-- | Generally speaking, yielding values from inside a Conduit requires +-- some allocation for constructors. This can introduce an overhead, +-- similar to the overhead needed to represent a list of values instead of +-- a vector. This overhead is even more severe when talking about unboxed +-- values. +-- +-- This combinator allows you to overcome this overhead, and efficiently +-- fill up vectors. It takes two parameters. The first is the size of each +-- mutable vector to be allocated. The second is a function. The function +-- takes an argument which will yield the next value into a mutable +-- vector. +-- +-- Under the surface, this function uses a number of tricks to get high +-- performance. For more information on both usage and implementation, +-- please see: +-- . +-- +-- Since 1.0.0 +vectorBuilderC :: (PrimMonad base, MonadBase base m, V.Vector v e, MonadBase base n) + => Int -- ^ size + -> ((e -> n ()) -> Sink i m r) + -> ConduitM i (v e) m r +vectorBuilderC = CC.vectorBuilder +{-# INLINE vectorBuilderC #-} diff --git a/conduit-combinators.cabal b/conduit-combinators.cabal --- a/conduit-combinators.cabal +++ b/conduit-combinators.cabal @@ -1,5 +1,5 @@ name: conduit-combinators -version: 0.2.7 +version: 0.2.8 synopsis: Commonly used conduit functions, for both chunked and unchunked data description: Provides a replacement for Data.Conduit.List, as well as a convenient Conduit module. homepage: https://github.com/fpco/conduit-combinators diff --git a/test/Spec.hs b/test/Spec.hs --- a/test/Spec.hs +++ b/test/Spec.hs @@ -1,5 +1,6 @@ {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ViewPatterns #-} +{-# LANGUAGE ScopedTypeVariables #-} {-# OPTIONS_GHC -fno-warn-type-defaults #-} import Conduit import Prelude hiding (FilePath) @@ -37,6 +38,7 @@ import qualified Data.ByteString.Base64.Lazy as B64L import qualified Data.ByteString.Base64.URL.Lazy as B64LU import qualified Data.ByteString.Base64.URL as B64U +import Control.Monad.ST (runST) main :: IO () main = hspec $ do @@ -609,6 +611,20 @@ it "slidingWindow 6" $ let res = runIdentity $ yieldMany [1..5] $= slidingWindow 6 $$ sinkList in res `shouldBe` [[1,2,3,4,5]] + prop "vectorBuilder" $ \(values :: [[Int]]) ((+1) . (`mod` 30) . abs -> size) -> do + let res = runST + $ yieldMany values + $$ vectorBuilderC size mapM_CE + =$ sinkList + expected = + loop $ concat values + where + loop [] = [] + loop x = + VU.fromList y : loop z + where + (y, z) = splitAt size x + res `shouldBe` expected evenInt :: Int -> Bool evenInt = even