diff --git a/Data/Repa/Stream.hs b/Data/Repa/Stream.hs --- a/Data/Repa/Stream.hs +++ b/Data/Repa/Stream.hs @@ -3,7 +3,11 @@ -- functions can be used. module Data.Repa.Stream ( extractS + , insertS , mergeS + , compactS + , compactInS + , catMaybesS , findSegmentsS , diceSepS , startLengthsOfSegsS @@ -13,10 +17,13 @@ , unsafeRatchetS) where +import Data.Repa.Stream.Concat +import Data.Repa.Stream.Compact +import Data.Repa.Stream.Dice import Data.Repa.Stream.Extract +import Data.Repa.Stream.Insert import Data.Repa.Stream.Merge +import Data.Repa.Stream.Pad import Data.Repa.Stream.Ratchet import Data.Repa.Stream.Segment -import Data.Repa.Stream.Dice -import Data.Repa.Stream.Pad diff --git a/Data/Repa/Stream/Compact.hs b/Data/Repa/Stream/Compact.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Stream/Compact.hs @@ -0,0 +1,69 @@ + +module Data.Repa.Stream.Compact + ( compactS + , compactInS ) +where +import Data.Vector.Fusion.Stream.Monadic (Stream(..), Step(..)) +import qualified Data.Vector.Fusion.Stream.Size as S +#include "repa-stream.h" + + +-- | Combination of `fold` and `filter`. +-- +-- We walk over the stream front to back, maintaining an accumulator. +-- At each point we can chose to emit an element (or not) +-- +compactS + :: Monad m + => (s -> a -> (Maybe b, s)) -- ^ Worker function. + -> s -- ^ Starting state + -> Stream m a -- ^ Input elements. + -> Stream m b + +compactS f s0 (Stream istep si0 sz) + = Stream ostep (si0, s0) (S.toMax sz) + where + ostep (si, s) + = istep si >>= \m + -> case m of + Yield x si' + -> case f s x of + (Nothing, s') -> return $ Skip (si', s') + (Just y, s') -> return $ Yield y (si', s') + + Skip si' -> return $ Skip (si', s) + Done -> return $ Done + {-# INLINE_INNER ostep #-} +{-# INLINE_STREAM compactS #-} + + +-- | Like `compact` but use the first value of the stream as the +-- initial state, and add the final state to the end of the output. +compactInS + :: Monad m + => (a -> a -> (Maybe a, a)) -- ^ Worker function. + -> Stream m a -- ^ Input elements. + -> Stream m a + +compactInS f (Stream istep si0 sz) + = Stream ostep (si0, Nothing) (S.toMax sz) + where + ostep (si, ms@Nothing) + = istep si >>= \m + -> case m of + Yield x si' -> return $ Skip (si', Just x) + Skip si' -> return $ Skip (si', ms) + Done -> return $ Done + + ostep (si, ms@(Just s)) + = istep si >>= \m + -> case m of + Yield x si' + -> case f s x of + (Nothing, s') -> return $ Skip (si', Just s') + (Just y, s') -> return $ Yield y (si', Just s') + + Skip si' -> return $ Skip (si', ms) + Done -> return $ Yield s (si, Nothing) + {-# INLINE_INNER ostep #-} +{-# INLINE_STREAM compactInS #-} diff --git a/Data/Repa/Stream/Concat.hs b/Data/Repa/Stream/Concat.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Stream/Concat.hs @@ -0,0 +1,28 @@ + +module Data.Repa.Stream.Concat + (catMaybesS) +where +import Data.Vector.Fusion.Stream.Monadic (Stream(..), Step(..)) +import qualified Data.Vector.Fusion.Stream.Size as S +#include "repa-stream.h" + + +-- | Return the `Just` elements from a stream, dropping the `Nothing`s. +catMaybesS + :: Monad m + => Stream m (Maybe a) + -> Stream m a + +catMaybesS (Stream istep si0 sz) + = Stream ostep si0 (S.toMax sz) + where + ostep si + = istep si >>= \m + -> case m of + Yield Nothing si' -> return $ Skip si' + Yield (Just x) si' -> return $ Yield x si' + Skip si' -> return $ Skip si' + Done -> return $ Done + {-# INLINE_INNER ostep #-} +{-# INLINE_STREAM catMaybesS #-} + diff --git a/Data/Repa/Stream/Insert.hs b/Data/Repa/Stream/Insert.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Stream/Insert.hs @@ -0,0 +1,38 @@ + +module Data.Repa.Stream.Insert + (insertS) +where +import Data.Vector.Fusion.Stream.Monadic (Stream(..), Step(..)) +import qualified Data.Vector.Fusion.Stream.Size as S +#include "repa-stream.h" + + +-- | Insert elements produced by the given function in to a stream. +insertS :: Monad m + => (Int -> Maybe a) -- ^ Produce a new element for this index. + -> Stream m a -- ^ Source stream. + -> Stream m a + +insertS fNew (Stream istep si0 _) + = Stream ostep (si0, 0, True, False) S.Unknown + where + ostep (si, ix, tryNew, srcDone) + | tryNew + = case fNew ix of + Just x -> return $ Yield x (si, ix + 1, True, srcDone) + Nothing + | srcDone -> return Done + | otherwise -> ostep_next si ix + + | otherwise + = ostep_next si ix + {-# INLINE_INNER ostep #-} + + ostep_next !si !ix + = istep si >>= \m + -> case m of + Yield x si' -> return $ Yield x (si', ix + 1, True, False) + Skip si' -> return $ Skip (si', ix, False, False) + Done -> return $ Skip (si, ix, True, True) + {-# INLINE_INNER ostep_next #-} +{-# INLINE_STREAM insertS #-} diff --git a/Data/Repa/Stream/Ratchet.hs b/Data/Repa/Stream/Ratchet.hs --- a/Data/Repa/Stream/Ratchet.hs +++ b/Data/Repa/Stream/Ratchet.hs @@ -2,12 +2,11 @@ module Data.Repa.Stream.Ratchet ( unsafeRatchetS) where +import Control.Monad.Primitive import Data.IORef import Data.Vector.Fusion.Stream.Monadic (Stream(..), Step(..)) -import qualified Data.Vector.Generic as G +import qualified Data.Vector.Generic as GV import qualified Data.Vector.Generic.Mutable as GM -import qualified Data.Vector.Unboxed as U -import qualified Data.Vector.Unboxed.Mutable as UM import qualified Data.Vector.Fusion.Stream.Size as S #include "repa-stream.h" @@ -42,9 +41,10 @@ -- but this is not checked. -- unsafeRatchetS - :: UM.IOVector Int -- ^ Starting values. Overwritten duing computation. - -> U.Vector Int -- ^ Ending values - -> IORef (UM.IOVector Int) -- ^ Vector holding segment lengths. + :: (GM.MVector vm Int, GV.Vector vv Int) + => vm (PrimState IO) Int -- ^ Starting values. Overwritten duing computation. + -> vv Int -- ^ Ending values + -> IORef (vm (PrimState IO) Int) -- ^ Vector holding segment lengths. -> Stream IO Int unsafeRatchetS !mvStarts !vMax !rmvLens @@ -59,7 +59,7 @@ ostep' !iSeg !mvmLens !oSeg !oLen | iSeg <= iSegMax = do !iVal <- GM.unsafeRead mvStarts iSeg - let !iNext = vMax `G.unsafeIndex` iSeg + let !iNext = vMax `GV.unsafeIndex` iSeg if iVal >= iNext then return $ Skip (iSeg + 1, mvmLens, oSeg, oLen) else do @@ -75,16 +75,16 @@ Just vmLens -> return $ vmLens -- If the output vector is full then we need to grow it. - let !oSegLen = UM.length vmLens + let !oSegLen = GM.length vmLens if oSeg >= oSegLen then do - !vmLens' <- UM.unsafeGrow vmLens (UM.length vmLens) + !vmLens' <- GM.unsafeGrow vmLens (GM.length vmLens) writeIORef rmvLens vmLens' - UM.unsafeWrite vmLens' oSeg oLen + GM.unsafeWrite vmLens' oSeg oLen return $ Skip (0, Just vmLens', oSeg + 1, 0) else do - UM.unsafeWrite vmLens oSeg oLen + GM.unsafeWrite vmLens oSeg oLen return $ Skip (0, Just vmLens, oSeg + 1, 0) | otherwise @@ -92,7 +92,7 @@ Nothing -> readIORef rmvLens Just vmLens -> return $ vmLens - let !vmLens' = UM.unsafeSlice 0 oSeg vmLens + let !vmLens' = GM.unsafeSlice 0 oSeg vmLens writeIORef rmvLens vmLens' return Done {-# INLINE_INNER ostep' #-} diff --git a/Data/Repa/Vector/Generic.hs b/Data/Repa/Vector/Generic.hs --- a/Data/Repa/Vector/Generic.hs +++ b/Data/Repa/Vector/Generic.hs @@ -11,14 +11,62 @@ -- * Chain functions , chainOfVector , unchainToVector - , unchainToMVector) + , unchainToMVector + + -- * Generating + , ratchet + + -- * Extracting + , extract + + -- * Inserting + , insert + + -- * Merging + , merge + , mergeMaybe + + -- * Splitting + , findSegments + , findSegmentsFrom + , diceSep + + -- * Compacting + , compact + , compactIn + + -- * Padding + , padForward + + -- * Scanning + , scanMaybe + + -- * Grouping + , groupsBy + + -- * Folding + , folds, C.Folds(..)) where +import Data.Repa.Stream.Compact +import Data.Repa.Stream.Concat +import Data.Repa.Stream.Dice +import Data.Repa.Stream.Extract +import Data.Repa.Stream.Insert +import Data.Repa.Stream.Merge +import Data.Repa.Stream.Pad +import Data.Repa.Stream.Ratchet +import Data.Repa.Stream.Segment +import Data.Repa.Option +import Data.IORef +import System.IO.Unsafe +import Control.Monad.ST +import Control.Monad.Primitive import Data.Repa.Chain as C import qualified Data.Vector.Generic as GV import qualified Data.Vector.Generic.Mutable as GM -import qualified Data.Vector.Fusion.Stream.Monadic as S +import qualified Data.Vector.Fusion.Stream.Monadic as SM import qualified Data.Vector.Fusion.Stream.Size as S -import Control.Monad.Primitive +import qualified Data.Vector.Fusion.Stream as S #include "repa-stream.h" @@ -29,7 +77,7 @@ -- unstreamToVector2 :: (PrimMonad m, GV.Vector v a, GV.Vector v b) - => S.Stream m (Maybe a, Maybe b) + => SM.Stream m (Maybe a, Maybe b) -- ^ Source data. -> m (v a, v b) -- ^ Resulting vectors. @@ -47,12 +95,12 @@ -- unstreamToMVector2 :: (PrimMonad m, GM.MVector v a, GM.MVector v b) - => S.Stream m (Maybe a, Maybe b) + => SM.Stream m (Maybe a, Maybe b) -- ^ Source data. -> m (v (PrimState m) a, v (PrimState m) b) -- ^ Resulting vectors. -unstreamToMVector2 (S.Stream step s0 sz) +unstreamToMVector2 (SM.Stream step s0 sz) = case sz of S.Exact i -> unstreamToMVector2_max i s0 step S.Max i -> unstreamToMVector2_max i s0 step @@ -66,7 +114,7 @@ go !sPEC !iL !iR !s = step s >>= \m -> case m of - S.Yield (mL, mR) s' + SM.Yield (mL, mR) s' -> do !iL' <- case mL of Nothing -> return iL Just xL -> do GM.unsafeWrite vecL iL xL @@ -79,13 +127,13 @@ go sPEC iL' iR' s' - S.Skip s' + SM.Skip s' -> go sPEC iL iR s' - S.Done + SM.Done -> return ( GM.unsafeSlice 0 iL vecL , GM.unsafeSlice 0 iR vecR) - in go S.SPEC 0 0 s0 + in go SM.SPEC 0 0 s0 {-# INLINE_STREAM unstreamToMVector2_max #-} @@ -151,7 +199,7 @@ Done s' -> return (GM.unsafeSlice 0 i vec, s') {-# INLINE_INNER go #-} - in go S.SPEC 0 s0 + in go SM.SPEC 0 s0 {-# INLINE_STREAM unchainToMVector_max #-} @@ -176,6 +224,339 @@ Done s' -> return (GM.unsafeSlice 0 i vec, s') {-# INLINE_INNER go #-} - in go S.SPEC vec0 0 nStart s0 + in go SM.SPEC vec0 0 nStart s0 {-# INLINE_STREAM unchainToMVector_unknown #-} + + +------------------------------------------------------------------------------- +-- | Interleaved `enumFromTo`. +-- +-- Given a vector of starting values, and a vector of stopping values, +-- produce an stream of elements where we increase each of the starting +-- values to the stopping values in a round-robin order. Also produce a +-- vector of result segment lengths. +-- +-- @ +-- unsafeRatchetS [10,20,30,40] [15,26,33,47] +-- = [10,20,30,40 -- 4 +-- ,11,21,31,41 -- 4 +-- ,12,22,32,42 -- 4 +-- ,13,23 ,43 -- 3 +-- ,14,24 ,44 -- 3 +-- ,25 ,45 -- 2 +-- ,46] -- 1 +-- +-- ^^^^ ^^^ +-- Elements Lengths +-- @ +-- +ratchet :: (GV.Vector v Int, GV.Vector v (Int, Int)) + => v (Int, Int) -- ^ Starting and ending values. + -> (v Int, v Int) -- ^ Elements and Lengths vectors. +ratchet vStartsMax + = unsafePerformIO + $ do + -- Make buffers for the start values and unpack the max values. + let (vStarts, vMax) = GV.unzip vStartsMax + mvStarts <- GV.thaw vStarts + + -- Make a vector for the output lengths. + mvLens <- GM.unsafeNew (GV.length vStartsMax) + rmvLens <- newIORef mvLens + + -- Run the computation + mvStarts' <- GM.munstream $ unsafeRatchetS mvStarts vMax rmvLens + + -- Read back the output segment lengths and freeze everything. + mvLens' <- readIORef rmvLens + vStarts' <- GV.unsafeFreeze mvStarts' + vLens' <- GV.unsafeFreeze mvLens' + return (vStarts', vLens') +{-# INLINE ratchet #-} + + +------------------------------------------------------------------------------- +-- | Extract segments from some source array and concatenate them. +-- +-- @ +-- let arr = [10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20] +-- in extractS (index arr) [(0, 1), (3, 3), (2, 6)] +-- +-- => [10, 13, 14, 15, 12, 13, 14, 15, 16, 17] +-- @ +-- +extract :: (GV.Vector v (Int, Int), GV.Vector v a) + => (Int -> a) -- ^ Function to get elements from the source. + -> v (Int, Int) -- ^ Segment starts and lengths. + -> v a -- ^ Result elements. + +extract get vStartLen + = GV.unstream $ extractS get $ GV.stream vStartLen +{-# INLINE extract #-} + + +------------------------------------------------------------------------------- +-- | Insert elements produced by the given function into a vector. +insert :: GV.Vector v a + => (Int -> Maybe a) -- ^ Produce a new element for this index. + -> v a -- ^ Source vector. + -> v a + +insert fNew vec + = GV.unstream $ insertS fNew $ GV.stream vec +{-# INLINE insert #-} + + +------------------------------------------------------------------------------- +-- | Merge two pre-sorted key-value streams. +merge :: ( Ord k + , GV.Vector v k + , GV.Vector v (k, a) + , GV.Vector v (k, b) + , GV.Vector v (k, c)) + => (k -> a -> b -> c) -- ^ Combine two values with the same key. + -> (k -> a -> c) -- ^ Handle a left value without a right value. + -> (k -> b -> c) -- ^ Handle a right value without a left value. + -> v (k, a) -- ^ Vector of keys and left values. + -> v (k, b) -- ^ Vector of keys and right values. + -> v (k, c) -- ^ Vector of keys and results. + +merge fBoth fLeft fRight vA vB + = GV.unstream + $ mergeS fBoth fLeft fRight + (GV.stream vA) + (GV.stream vB) +{-# INLINE merge #-} + + +-- | Like `merge`, but only produce the elements where the worker functions +-- return `Just`. +mergeMaybe + :: ( Ord k + , GV.Vector v k + , GV.Vector v (k, a) + , GV.Vector v (k, b) + , GV.Vector v (k, c)) + => (k -> a -> b -> Maybe c) -- ^ Combine two values with the same key. + -> (k -> a -> Maybe c) -- ^ Handle a left value without a right value. + -> (k -> b -> Maybe c) -- ^ Handle a right value without a left value. + -> v (k, a) -- ^ Vector of keys and left values. + -> v (k, b) -- ^ Vector of keys and right values. + -> v (k, c) -- ^ Vector of keys and results. + +mergeMaybe fBoth fLeft fRight vA vB + = GV.unstream + $ catMaybesS + $ SM.map munge_mergeMaybe + $ mergeS fBoth fLeft fRight + (GV.stream vA) + (GV.stream vB) + + where munge_mergeMaybe (_k, Nothing) = Nothing + munge_mergeMaybe (k, Just x) = Just (k, x) + {-# INLINE munge_mergeMaybe #-} +{-# INLINE mergeMaybe #-} + + +------------------------------------------------------------------------------- +-- | Perform a left-to-right scan through an input vector, maintaining a state +-- value between each element. For each element of input we may or may not +-- produce an element of output. +scanMaybe + :: forall v1 v2 a b s + . (GV.Vector v1 a, GV.Vector v2 b) + => (s -> a -> (s, Maybe b)) -- ^ Worker function. + -> s -- ^ Initial state for scan. + -> v1 a -- ^ Input elements. + -> (v2 b, s) -- ^ Output elements. + +scanMaybe f k0 vec0 + = (vec1, snd k1) + where + f' s x = return $ f s x + + (vec1 :: v2 b, k1 :: (Int, s)) + = runST $ unchainToVector $ C.liftChain + $ C.scanMaybeC f' k0 $ chainOfVector vec0 +{-# INLINE scanMaybe #-} + + +------------------------------------------------------------------------------- +-- | From a stream of values which has consecutive runs of idential values, +-- produce a stream of the lengths of these runs. +-- +-- @ +-- groupsBy (==) (Just ('a', 4)) +-- [\'a\', \'a\', \'a\', \'b\', \'b\', \'c\', \'d\', \'d\'] +-- => ([('a', 7), ('b', 2), ('c', 1)], Just (\'d\', 2)) +-- @ +-- +groupsBy + :: forall v1 v2 a + . (GV.Vector v1 a, GV.Vector v2 (a, Int)) + => (a -> a -> Bool) -- ^ Comparison function. + -> Maybe (a, Int) -- ^ Starting element and count. + -> v1 a -- ^ Input elements. + -> (v2 (a, Int), Maybe (a, Int)) + +groupsBy f !c !vec0 + = (vec1, snd k1) + where + f' x y = return $ f x y + + (vec1 :: v2 (a, Int), k1) + = runST $ unchainToVector $ C.liftChain + $ C.groupsByC f' c $ chainOfVector vec0 +{-# INLINE groupsBy #-} + + +------------------------------------------------------------------------------- +-- | Given predicates that detect the beginning and end of some interesting +-- segment of information, scan through a vector looking for when these +-- segments begin and end. Return vectors of the segment starting positions +-- and lengths. +-- +-- * As each segment must end on a element where the ending predicate returns +-- True, the miniumum segment length returned is 1. +-- +findSegments + :: (GV.Vector v a, GV.Vector v Int, GV.Vector v (Int, Int)) + => (a -> Bool) -- ^ Predicate to check for start of segment. + -> (a -> Bool) -- ^ Predicate to check for end of segment. + -> v a -- ^ Input vector. + -> (v Int, v Int) + +findSegments pStart pEnd src + = GV.unzip + $ GV.unstream + $ startLengthsOfSegsS + $ findSegmentsS pStart pEnd (GV.length src - 1) + $ SM.indexed + $ GV.stream src +{-# INLINE findSegments #-} + + +------------------------------------------------------------------------------- +-- | Given predicates that detect the beginning and end of some interesting +-- segment of information, scan through a vector looking for when these +-- segments begin and end. Return vectors of the segment starting positions +-- and lengths. +findSegmentsFrom + :: (GV.Vector v Int, GV.Vector v (Int, Int)) + => (a -> Bool) -- ^ Predicate to check for start of segment. + -> (a -> Bool) -- ^ Predicate to check for end of segment. + -> Int -- ^ Input length. + -> (Int -> a) -- ^ Get an element from the input. + -> (v Int, v Int) + +findSegmentsFrom pStart pEnd len get + = GV.unzip + $ GV.unstream + $ startLengthsOfSegsS + $ findSegmentsS pStart pEnd (len - 1) + $ SM.map (\ix -> (ix, get ix)) + $ SM.enumFromStepN 0 1 len +{-# INLINE findSegmentsFrom #-} + + +------------------------------------------------------------------------------- +-- | Dice a vector stream into rows and columns. +-- +diceSep :: (GV.Vector v a, GV.Vector v (Int, Int)) + => (a -> Bool) -- ^ Detect the end of a column. + -> (a -> Bool) -- ^ Detect the end of a row. + -> v a + -> (v (Int, Int), v (Int, Int)) -- ^ Segment starts and lengths + +diceSep pEndInner pEndBoth vec + = runST + $ unstreamToVector2 + $ diceSepS pEndInner pEndBoth + $ S.liftStream + $ GV.stream vec +{-# INLINE diceSep #-} + + +------------------------------------------------------------------------------- +-- | Combination of `fold` and `filter`. +-- +-- We walk over the stream front to back, maintaining an accumulator. +-- At each point we can chose to emit an element (or not) +-- +compact + :: (GV.Vector v a, GV.Vector v b) + => (s -> a -> (Maybe b, s)) -- ^ Worker function + -> s -- ^ Starting state + -> v a -- ^ Input vector + -> v b + +compact f s0 vec + = GV.unstream $ compactS f s0 $ GV.stream vec +{-# INLINE compact #-} + + +-- | Like `compact` but use the first value of the stream as the +-- initial state, and add the final state to the end of the output. +compactIn + :: GV.Vector v a + => (a -> a -> (Maybe a, a)) -- ^ Worker function. + -> v a -- ^ Input elements. + -> v a + +compactIn f vec + = GV.unstream $ compactInS f $ GV.stream vec +{-# INLINE compactIn #-} + + +------------------------------------------------------------------------------- +-- | Segmented fold over vectors of segment lengths and input values. +-- +-- The total lengths of all segments need not match the length of the +-- input elements vector. The returned `C.Folds` state can be inspected +-- to determine whether all segments were completely folded, or the +-- vector of segment lengths or elements was too short relative to the +-- other. In the resulting state, `C.foldLensState` is the index into +-- the lengths vector *after* the last one that was consumed. If this +-- equals the length of the lengths vector then all segment lengths were +-- consumed. Similarly for the elements vector. +-- +folds :: forall v n a b + . (GV.Vector v (n, Int), GV.Vector v a, GV.Vector v (n, b)) + => (a -> b -> b) -- ^ Worker function to fold each segment. + -> b -- ^ Initial state when folding segments. + -> Option3 n Int b -- ^ Length and initial state for first segment. + -> v (n, Int) -- ^ Segment names and lengths. + -> v a -- ^ Elements. + -> (v (n, b), C.Folds Int Int n a b) + +folds f zN s0 vLens vVals + = let + f' x y = return $ f x y + {-# INLINE f' #-} + + (vResults :: v (n, b), state) + = runST $ unchainToVector + $ C.foldsC f' zN s0 + (chainOfVector vLens) + (chainOfVector vVals) + + in (vResults, state) +{-# INLINE folds #-} + + +------------------------------------------------------------------------------- +-- | Given a stream of keys and values, and a successor function for keys, +-- if the stream is has keys missing in the sequence then insert +-- the missing key, copying forward the the previous value. +padForward + :: (Ord k, GV.Vector v (k, a)) + => (k -> k) -- ^ Successor function. + -> v (k, a) -- ^ Input keys and values. + -> v (k, a) + +padForward ksucc vec + = GV.unstream + $ padForwardS ksucc + $ GV.stream vec +{-# INLINE padForward #-} diff --git a/Data/Repa/Vector/Unboxed.hs b/Data/Repa/Vector/Unboxed.hs --- a/Data/Repa/Vector/Unboxed.hs +++ b/Data/Repa/Vector/Unboxed.hs @@ -5,24 +5,31 @@ , unchainToVector , unchainToMVector - -- * Generators + -- * Generating , ratchet - -- * Extract + -- * Extracting , extract + -- * Inserting + , insert + -- * Merging , merge + , mergeMaybe -- * Splitting , findSegments , findSegmentsFrom , diceSep + -- * Compacting + , compact + , compactIn + -- * Padding , padForward - -- * Scanning , scanMaybe @@ -33,26 +40,13 @@ , folds, C.Folds(..)) where import Data.Repa.Option -import Data.Repa.Stream.Extract -import Data.Repa.Stream.Ratchet -import Data.Repa.Stream.Segment -import Data.Repa.Stream.Dice -import Data.Repa.Stream.Merge -import Data.Repa.Stream.Pad import Data.Vector.Unboxed (Unbox, Vector) import Data.Vector.Unboxed.Mutable (MVector) import Data.Repa.Chain (Chain) import qualified Data.Repa.Vector.Generic as G import qualified Data.Repa.Chain as C import qualified Data.Vector.Unboxed as U -import qualified Data.Vector.Unboxed.Mutable as UM -import qualified Data.Vector.Generic as G -import qualified Data.Vector.Generic.Mutable as GM -import qualified Data.Vector.Fusion.Stream as S -import Control.Monad.ST import Control.Monad.Primitive -import System.IO.Unsafe -import Data.IORef #include "repa-stream.h" @@ -82,7 +76,6 @@ {-# INLINE unchainToMVector #-} - ------------------------------------------------------------------------------- -- | Interleaved `enumFromTo`. -- @@ -107,25 +100,7 @@ -- ratchet :: U.Vector (Int, Int) -- ^ Starting and ending values. -> (U.Vector Int, U.Vector Int) -- ^ Elements and Lengths vectors. -ratchet vStartsMax - = unsafePerformIO - $ do - -- Make buffers for the start values and unpack the max values. - let (vStarts, vMax) = U.unzip vStartsMax - mvStarts <- U.thaw vStarts - - -- Make a vector for the output lengths. - mvLens <- UM.unsafeNew (U.length vStartsMax) - rmvLens <- newIORef mvLens - - -- Run the computation - mvStarts' <- GM.munstream $ unsafeRatchetS mvStarts vMax rmvLens - - -- Read back the output segment lengths and freeze everything. - mvLens' <- readIORef rmvLens - vStarts' <- G.unsafeFreeze mvStarts' - vLens' <- G.unsafeFreeze mvLens' - return (vStarts', vLens') +ratchet = G.ratchet {-# INLINE ratchet #-} @@ -144,12 +119,22 @@ -> U.Vector (Int, Int) -- ^ Segment starts and lengths. -> U.Vector a -- ^ Result elements. -extract get vStartLen - = G.unstream $ extractS get $ G.stream vStartLen +extract = G.extract {-# INLINE extract #-} ------------------------------------------------------------------------------- +-- | Insert elements produced by the given function into a vector. +insert :: Unbox a + => (Int -> Maybe a) -- ^ Produce a new element for this index. + -> U.Vector a -- ^ Source vector. + -> U.Vector a + +insert = G.insert +{-# INLINE insert #-} + + +------------------------------------------------------------------------------- -- | Merge two pre-sorted key-value streams. merge :: (Ord k, Unbox k, Unbox a, Unbox b, Unbox c) => (k -> a -> b -> c) -- ^ Combine two values with the same key. @@ -159,14 +144,25 @@ -> U.Vector (k, b) -- ^ Vector of keys and right values. -> U.Vector (k, c) -- ^ Vector of keys and results. -merge fBoth fLeft fRight vA vB - = G.unstream - $ mergeS fBoth fLeft fRight - (G.stream vA) - (G.stream vB) +merge = G.merge {-# INLINE merge #-} +-- | Like `merge`, but only produce the elements where the worker functions +-- return `Just`. +mergeMaybe + :: (Ord k, Unbox k, Unbox a, Unbox b, Unbox c) + => (k -> a -> b -> Maybe c) -- ^ Combine two values with the same key. + -> (k -> a -> Maybe c) -- ^ Handle a left value without a right value. + -> (k -> b -> Maybe c) -- ^ Handle a right value without a left value. + -> U.Vector (k, a) -- ^ Vector of keys and left values. + -> U.Vector (k, b) -- ^ Vector of keys and right values. + -> U.Vector (k, c) -- ^ Vector of keys and results. + +mergeMaybe = G.mergeMaybe +{-# INLINE mergeMaybe #-} + + ------------------------------------------------------------------------------- -- | Perform a left-to-right scan through an input vector, maintaining a state -- value between each element. For each element of input we may or may not @@ -178,14 +174,7 @@ -> U.Vector a -- ^ Input elements. -> (U.Vector b, s) -- ^ Output elements. -scanMaybe f k0 vec0 - = (vec1, snd k1) - where - f' s x = return $ f s x - - (vec1, k1) - = runST $ unchainToVector $ C.liftChain - $ C.scanMaybeC f' k0 $ chainOfVector vec0 +scanMaybe = G.scanMaybe {-# INLINE scanMaybe #-} @@ -206,14 +195,7 @@ -> U.Vector a -- ^ Input elements. -> (U.Vector (a, Int), Maybe (a, Int)) -groupsBy f !c !vec0 - = (vec1, snd k1) - where - f' x y = return $ f x y - - (vec1, k1) - = runST $ unchainToVector $ C.liftChain - $ C.groupsByC f' c $ chainOfVector vec0 +groupsBy = G.groupsBy {-# INLINE groupsBy #-} @@ -233,13 +215,7 @@ -> U.Vector a -- ^ Input vector. -> (U.Vector Int, U.Vector Int) -findSegments pStart pEnd src - = U.unzip - $ G.unstream - $ startLengthsOfSegsS - $ findSegmentsS pStart pEnd (U.length src - 1) - $ S.indexed - $ G.stream src +findSegments = G.findSegments {-# INLINE findSegments #-} @@ -255,13 +231,7 @@ -> (Int -> a) -- ^ Get an element from the input. -> (U.Vector Int, U.Vector Int) -findSegmentsFrom pStart pEnd len get - = U.unzip - $ G.unstream - $ startLengthsOfSegsS - $ findSegmentsS pStart pEnd (len - 1) - $ S.map (\ix -> (ix, get ix)) - $ S.enumFromStepN 0 1 len +findSegmentsFrom = G.findSegmentsFrom {-# INLINE findSegmentsFrom #-} @@ -275,16 +245,40 @@ -> (U.Vector (Int, Int), U.Vector (Int, Int)) -- ^ Segment starts and lengths -diceSep pEndInner pEndBoth vec - = runST - $ G.unstreamToVector2 - $ diceSepS pEndInner pEndBoth - $ S.liftStream - $ G.stream vec +diceSep = G.diceSep {-# INLINE diceSep #-} ------------------------------------------------------------------------------- +-- | Combination of `fold` and `filter`. +-- +-- We walk over the stream front to back, maintaining an accumulator. +-- At each point we can chose to emit an element (or not) +-- +compact + :: (Unbox a, Unbox b) + => (s -> a -> (Maybe b, s)) -- ^ Worker function + -> s -- ^ Starting state + -> Vector a -- ^ Input vector + -> Vector b + +compact = G.compact +{-# INLINE compact #-} + + +-- | Like `compact` but use the first value of the stream as the +-- initial state, and add the final state to the end of the output. +compactIn + :: Unbox a + => (a -> a -> (Maybe a, a)) -- ^ Worker function. + -> Vector a -- ^ Input elements. + -> Vector a + +compactIn = G.compactIn +{-# INLINE compactIn #-} + + +------------------------------------------------------------------------------- -- | Segmented fold over vectors of segment lengths and input values. -- -- The total lengths of all segments need not match the length of the @@ -304,18 +298,7 @@ -> U.Vector a -- ^ Elements. -> (U.Vector (n, b), C.Folds Int Int n a b) -folds f zN s0 vLens vVals - = let - f' x y = return $ f x y - {-# INLINE f' #-} - - (vResults, state) - = runST $ unchainToVector - $ C.foldsC f' zN s0 - (chainOfVector vLens) - (chainOfVector vVals) - - in (vResults, state) +folds = G.folds {-# INLINE folds #-} @@ -329,9 +312,6 @@ -> U.Vector (k, v) -- ^ Input keys and values. -> U.Vector (k, v) -padForward ksucc vec - = G.unstream - $ padForwardS ksucc - $ G.stream vec +padForward = G.padForward {-# INLINE padForward #-} diff --git a/repa-stream.cabal b/repa-stream.cabal --- a/repa-stream.cabal +++ b/repa-stream.cabal @@ -1,5 +1,5 @@ Name: repa-stream -Version: 4.0.0.1 +Version: 4.1.0.1 License: BSD3 License-file: LICENSE Author: The Repa Development Team @@ -38,12 +38,15 @@ Data.Repa.Chain.Weave Data.Repa.Chain.Folds + Data.Repa.Stream.Compact + Data.Repa.Stream.Concat + Data.Repa.Stream.Dice Data.Repa.Stream.Extract + Data.Repa.Stream.Insert + Data.Repa.Stream.Merge + Data.Repa.Stream.Pad Data.Repa.Stream.Ratchet Data.Repa.Stream.Segment - Data.Repa.Stream.Dice - Data.Repa.Stream.Pad - Data.Repa.Stream.Merge include-dirs: include @@ -63,5 +66,6 @@ FlexibleContexts PatternGuards MultiWayIf + ScopedTypeVariables