diff --git a/Data/Repa/Chain.hs b/Data/Repa/Chain.hs --- a/Data/Repa/Chain.hs +++ b/Data/Repa/Chain.hs @@ -14,6 +14,10 @@ -- * Folding , foldsC, Folds(..) + -- * Unfolding + , unfoldsC + , StepUnfold (..) + -- * Scanning , scanMaybeC diff --git a/Data/Repa/Chain/Folds.hs b/Data/Repa/Chain/Folds.hs --- a/Data/Repa/Chain/Folds.hs +++ b/Data/Repa/Chain/Folds.hs @@ -2,8 +2,8 @@ module Data.Repa.Chain.Folds (foldsC, Folds (..)) where -import Data.Repa.Option import Data.Repa.Chain.Base +import Data.Repa.Scalar.Option import Data.Vector.Fusion.Stream.Size as S #include "repa-stream.h" diff --git a/Data/Repa/Chain/Scan.hs b/Data/Repa/Chain/Scan.hs --- a/Data/Repa/Chain/Scan.hs +++ b/Data/Repa/Chain/Scan.hs @@ -1,12 +1,63 @@ module Data.Repa.Chain.Scan - ( scanMaybeC + ( StepUnfold (..) + , unfoldsC + , scanMaybeC , groupsByC) where import Data.Repa.Chain.Base +import Data.Repa.Scalar.Option import qualified Data.Vector.Fusion.Stream.Size as S #include "repa-stream.h" + +------------------------------------------------------------------------------- +-- | Segmented unfold. +-- +-- The worker function takes the current element from the input stream +-- and current state. +-- +-- If the worker returns Just x then that output element will be placed +-- in the output stream, and it will be called again with the same +-- input elemenent and next state. +-- +-- If the worker returns Nothing then we advance to the next element +-- of the input stream. +-- +unfoldsC + :: Monad m + => (a -> k -> m (StepUnfold k b)) -- ^ Worker function. + -> k -- ^ Initial state for the unfold. + -> Chain m s a -- ^ Input elements. + -> Chain m (s, k, Option a) b -- ^ Output elements. + +unfoldsC f k0 (Chain _ s0 istep) + = Chain S.Unknown (s0, k0, None) ostep + where + ostep (s1, k1, None) + = istep s1 >>= \rs + -> case rs of + Yield xa s2 -> return $ Skip (s2, k1, Some xa) + Skip s2 -> return $ Skip (s2, k1, None) + Done s2 -> return $ Done (s2, k1, None) + + ostep (s1, k1, Some xa) + = f xa k1 >>= \kmb + -> case kmb of + StepUnfoldGive xb k2 -> return $ Yield xb (s1, k2, Some xa) + StepUnfoldNext xb k2 -> return $ Yield xb (s1, k2, None) + StepUnfoldBump k2 -> return $ Skip (s1, k2, Some xa) + StepUnfoldFinish k2 -> return $ Skip (s1, k2, None) + {-# INLINE_INNER ostep #-} +{-# INLINE_STREAM unfoldsC #-} + + +data StepUnfold s a + = StepUnfoldGive a s + | StepUnfoldNext a s + | StepUnfoldBump s + | StepUnfoldFinish s + deriving Show ------------------------------------------------------------------------------- -- | Perform a left-to-right scan through an input vector, maintaining a state diff --git a/Data/Repa/Chain/Weave.hs b/Data/Repa/Chain/Weave.hs --- a/Data/Repa/Chain/Weave.hs +++ b/Data/Repa/Chain/Weave.hs @@ -6,8 +6,8 @@ , Move (..), move , Option (..)) where -import Data.Repa.Option import Data.Repa.Chain.Base +import Data.Repa.Scalar.Option import qualified Data.Vector.Fusion.Stream.Size as S #include "repa-stream.h" diff --git a/Data/Repa/Option.hs b/Data/Repa/Option.hs deleted file mode 100644 --- a/Data/Repa/Option.hs +++ /dev/null @@ -1,86 +0,0 @@ - --- | Data types used during low-level fusion optimisations. --- --- These types are synonyms for @Maybe (a, b)@, which are strict in the --- components. They can be used to ensure that we do not suspend the --- computation that produces these components in fused code. --- -module Data.Repa.Option - ( -- * Single component - Option (..) - , fromOption, toOption - - -- * Two components - , Option2 (..) - , fromOption2, toOption2 - - -- * Three components - , Option3 (..) - , fromOption3, toOption3) -where - - -------------------------------------------------------------------------------- --- | A strict `Maybe` type. -data Option a - = Some !a - | None - deriving Show - --- | Convert a `Maybe` to an `Option`. -toOption :: Maybe a -> Option a -toOption Nothing = None -toOption (Just x) = Some x -{-# INLINE toOption #-} - - --- | Convert an `Option` to a `Maybe`. -fromOption :: Option a -> Maybe a -fromOption None = Nothing -fromOption (Some x) = Just x -{-# INLINE fromOption #-} - - -------------------------------------------------------------------------------- --- | A strict `Maybe` type, with two parameters. -data Option2 a b - = Some2 !a !b - | None2 - deriving Show - - --- | Convert a `Maybe` to an `Option2`. -toOption2 :: Maybe (a, b) -> Option2 a b -toOption2 Nothing = None2 -toOption2 (Just (x, y)) = Some2 x y -{-# INLINE toOption2 #-} - - --- | Convert an `Option2` to a `Maybe`. -fromOption2 :: Option2 a b -> Maybe (a, b) -fromOption2 None2 = Nothing -fromOption2 (Some2 x y) = Just (x, y) -{-# INLINE fromOption2 #-} - - -------------------------------------------------------------------------------- --- | A strict `Maybe` type with three parameters. -data Option3 a b c - = Some3 !a !b !c - | None3 - deriving Show - - --- | Convert a `Maybe` to an `Option3`. -toOption3 :: Maybe (a, b, c) -> Option3 a b c -toOption3 Nothing = None3 -toOption3 (Just (x, y, z)) = Some3 x y z -{-# INLINE toOption3 #-} - - --- | Convert an `Option2` to a `Maybe`. -fromOption3 :: Option3 a b c -> Maybe (a, b, c) -fromOption3 None3 = Nothing -fromOption3 (Some3 x y z) = Just (x, y, z) -{-# INLINE fromOption3 #-} - diff --git a/Data/Repa/Stream.hs b/Data/Repa/Stream.hs --- a/Data/Repa/Stream.hs +++ b/Data/Repa/Stream.hs @@ -2,20 +2,37 @@ -- | * See the "Data.Repa.Vector.Unboxed" module for examples of how these -- functions can be used. module Data.Repa.Stream - ( extractS - , insertS - , mergeS - , compactS + ( -- * Compacting + compactS , compactInS + + -- * Concatenating , catMaybesS - , findSegmentsS + + -- * Dicing , diceSepS - , startLengthsOfSegsS + + -- * Extracting + , extractS + + -- * Inserting + , insertS + + -- * Merging + , mergeS + + -- * Padding , padForwardS - -- * Unsafe operators - , unsafeRatchetS) + -- * Ratcheting + , unsafeRatchetS + -- * Replicating + , replicatesS + + -- * Segmenting + , findSegmentsS + , startLengthsOfSegsS) where import Data.Repa.Stream.Concat import Data.Repa.Stream.Compact @@ -25,5 +42,6 @@ import Data.Repa.Stream.Merge import Data.Repa.Stream.Pad import Data.Repa.Stream.Ratchet +import Data.Repa.Stream.Replicate import Data.Repa.Stream.Segment diff --git a/Data/Repa/Stream/Compact.hs b/Data/Repa/Stream/Compact.hs --- a/Data/Repa/Stream/Compact.hs +++ b/Data/Repa/Stream/Compact.hs @@ -15,7 +15,7 @@ -- compactS :: Monad m - => (s -> a -> (Maybe b, s)) -- ^ Worker function. + => (s -> a -> (s, Maybe b)) -- ^ Worker function. -> s -- ^ Starting state -> Stream m a -- ^ Input elements. -> Stream m b @@ -28,9 +28,8 @@ -> 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') - + (s', Nothing) -> return $ Skip (si', s') + (s', Just y) -> return $ Yield y (si', s') Skip si' -> return $ Skip (si', s) Done -> return $ Done {-# INLINE_INNER ostep #-} @@ -41,7 +40,7 @@ -- initial state, and add the final state to the end of the output. compactInS :: Monad m - => (a -> a -> (Maybe a, a)) -- ^ Worker function. + => (a -> a -> (a, Maybe a)) -- ^ Worker function. -> Stream m a -- ^ Input elements. -> Stream m a @@ -60,9 +59,8 @@ -> 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') - + (s', Nothing) -> return $ Skip (si', Just s') + (s', Just y) -> return $ Yield y (si', Just s') Skip si' -> return $ Skip (si', ms) Done -> return $ Yield s (si, Nothing) {-# INLINE_INNER ostep #-} diff --git a/Data/Repa/Stream/Merge.hs b/Data/Repa/Stream/Merge.hs --- a/Data/Repa/Stream/Merge.hs +++ b/Data/Repa/Stream/Merge.hs @@ -3,7 +3,7 @@ (mergeS) where import Data.Vector.Fusion.Stream.Monadic (Stream(..), Step(..)) -import Data.Repa.Option +import Data.Repa.Scalar.Option import qualified Data.Vector.Fusion.Stream.Size as S #include "repa-stream.h" diff --git a/Data/Repa/Stream/Pad.hs b/Data/Repa/Stream/Pad.hs --- a/Data/Repa/Stream/Pad.hs +++ b/Data/Repa/Stream/Pad.hs @@ -3,7 +3,7 @@ (padForwardS) where import Data.Vector.Fusion.Stream.Monadic (Stream(..), Step(..)) -import Data.Repa.Option +import Data.Repa.Scalar.Option import qualified Data.Vector.Fusion.Stream.Size as S #include "repa-stream.h" diff --git a/Data/Repa/Stream/Replicate.hs b/Data/Repa/Stream/Replicate.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Stream/Replicate.hs @@ -0,0 +1,40 @@ + +module Data.Repa.Stream.Replicate + (replicatesS) +where +import Data.Vector.Fusion.Stream.Monadic (Stream(..), Step(..)) +import qualified Data.Vector.Fusion.Stream.Size as S +#include "repa-stream.h" + + +-- | Segmented replicate. +-- +-- Given a stream of counts and values, produce a result stream where +-- each value is replciated the associated number of times. +-- +replicatesS + :: Monad m + => Stream m (Int, a) + -> Stream m a + +replicatesS (Stream istepA sA0 _) + = Stream ostep (sA0, 0, Nothing) S.Unknown + where + -- Pull the next element from the source stream. + ostep (sA, 0, mv) + = istepA sA >>= \rA + -> case rA of + Done -> return Done + Skip sA' -> return $ Skip (sA', 0, mv) + Yield (len, val) sA' -> return $ Skip (sA', len, Just val) + + -- Should never be entered, as we'll always have a value + -- when the count is non-zero. + ostep (_, _, Nothing) + = return Done + + ostep (sA, len, mv@(Just val)) + = return $ Yield val (sA, len - 1, mv) + {-# INLINE_INNER ostep #-} +{-# INLINE_STREAM replicatesS #-} + 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 @@ -4,48 +4,49 @@ -- * NOTE: Support for streams of unknown length is not complete. -- module Data.Repa.Vector.Generic - ( -- * Stream functions + ( -- * Stream operators unstreamToVector2 , unstreamToMVector2 - -- * Chain functions - , chainOfVector - , unchainToVector - , unchainToMVector + -- ** Compacting + , compact + , compactIn - -- * Generating - , ratchet + -- ** Dicing + , findSegments + , findSegmentsFrom + , diceSep - -- * Extracting + -- ** Extracting , extract - -- * Inserting + -- ** Inserting , insert - -- * Merging + -- ** Merging , merge , mergeMaybe - -- * Splitting - , findSegments - , findSegmentsFrom - , diceSep - - -- * Compacting - , compact - , compactIn - - -- * Padding + -- ** Padding , padForward - -- * Scanning - , scanMaybe + -- ** Ratcheting + , ratchet - -- * Grouping - , groupsBy + -- ** Replicating + , replicates + -- * Chain operators + , chainOfVector + , unchainToVector + , unchainToMVector + -- * Folding - , folds, C.Folds(..)) + , folds, C.Folds(..) + + -- * Scanning + , scanMaybe + , groupsBy) where import Data.Repa.Stream.Compact import Data.Repa.Stream.Concat @@ -55,8 +56,9 @@ import Data.Repa.Stream.Merge import Data.Repa.Stream.Pad import Data.Repa.Stream.Ratchet +import Data.Repa.Stream.Replicate import Data.Repa.Stream.Segment -import Data.Repa.Option +import Data.Repa.Scalar.Option import Data.IORef import System.IO.Unsafe import Control.Monad.ST @@ -138,141 +140,100 @@ ------------------------------------------------------------------------------- --- | Produce a chain from a generic vector. -chainOfVector - :: (Monad m, GV.Vector v a) - => v a -> Chain m Int a - -chainOfVector vec - = Chain (S.Exact len) 0 step - where - !len = GV.length vec - - step !i - | i >= len - = return $ Done i - - | otherwise - = return $ Yield (GV.unsafeIndex vec i) (i + 1) - {-# INLINE_INNER step #-} -{-# INLINE_STREAM chainOfVector #-} - +-- | 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 -> (s, Maybe b)) -- ^ Worker function + -> s -- ^ Starting state + -> v a -- ^ Input vector + -> v b -------------------------------------------------------------------------------- --- | Compute a chain into a generic vector. -unchainToVector - :: (PrimMonad m, GV.Vector v a) - => C.Chain m s a -> m (v a, s) -unchainToVector chain - = do (mvec, c') <- unchainToMVector chain - vec <- GV.unsafeFreeze mvec - return (vec, c') -{-# INLINE_STREAM unchainToVector #-} +compact f s0 vec + = GV.unstream $ compactS f s0 $ GV.stream vec +{-# INLINE compact #-} --- | Compute a chain into a generic mutable vector. -unchainToMVector - :: (PrimMonad m, GM.MVector v a) - => Chain m s a - -> m (v (PrimState m) a, s) +-- | 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 -> (a, Maybe a)) -- ^ Worker function. + -> v a -- ^ Input elements. + -> v a -unchainToMVector (Chain sz s0 step) - = case sz of - S.Exact i -> unchainToMVector_max i s0 step - S.Max i -> unchainToMVector_max i s0 step - S.Unknown -> unchainToMVector_unknown 32 s0 step -{-# INLINE_STREAM unchainToMVector #-} +compactIn f vec + = GV.unstream $ compactInS f $ GV.stream vec +{-# INLINE compactIn #-} --- unchain when we known the maximum size of the vector. -unchainToMVector_max nMax s0 step - = GM.unsafeNew nMax >>= \vec - -> let - go !sPEC !i !s - = step s >>= \m - -> case m of - Yield e s' - -> do GM.unsafeWrite vec i e - go sPEC (i + 1) s' - - Skip s' -> go sPEC i s' - Done s' -> return (GM.unsafeSlice 0 i vec, s') - {-# INLINE_INNER go #-} +------------------------------------------------------------------------------- +-- | 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) - in go SM.SPEC 0 s0 -{-# INLINE_STREAM unchainToMVector_max #-} +findSegments pStart pEnd src + = GV.unzip + $ GV.unstream + $ startLengthsOfSegsS + $ findSegmentsS pStart pEnd (GV.length src - 1) + $ SM.indexed + $ GV.stream src +{-# INLINE findSegments #-} --- unchain when we don't know the maximum size of the vector. -unchainToMVector_unknown nStart s0 step - = GM.unsafeNew nStart >>= \vec0 - -> let - go !sPEC !vec !i !n !s - = step s >>= \m - -> case m of - Yield e s' - | i >= n - -> do vec' <- GM.unsafeGrow vec n - GM.unsafeWrite vec' i e - go sPEC vec' (i + 1) (n + n) s' - - | otherwise - -> do GM.unsafeWrite vec i e - go sPEC vec (i + 1) n s' - - Skip s' -> go sPEC vec i n s' - Done s' -> return (GM.unsafeSlice 0 i vec, s') - {-# INLINE_INNER go #-} +------------------------------------------------------------------------------- +-- | 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) - in go SM.SPEC vec0 0 nStart s0 -{-# INLINE_STREAM unchainToMVector_unknown #-} +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 #-} ------------------------------------------------------------------------------- --- | 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 --- @ +-- | Dice a vector stream into rows and columns. -- -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 +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 - -- 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 #-} +diceSep pEndInner pEndBoth vec + = runST + $ unstreamToVector2 + $ diceSepS pEndInner pEndBoth + $ S.liftStream + $ GV.stream vec +{-# INLINE diceSep #-} ------------------------------------------------------------------------------- @@ -359,155 +320,174 @@ ------------------------------------------------------------------------------- --- | 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 +-- | 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) - (vec1 :: v2 b, k1 :: (Int, s)) - = runST $ unchainToVector $ C.liftChain - $ C.scanMaybeC f' k0 $ chainOfVector vec0 -{-# INLINE scanMaybe #-} +padForward ksucc vec + = GV.unstream + $ padForwardS ksucc + $ GV.stream vec +{-# INLINE padForward #-} ------------------------------------------------------------------------------- --- | From a stream of values which has consecutive runs of idential values, --- produce a stream of the lengths of these runs. --- +-- | 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. +-- -- @ --- groupsBy (==) (Just ('a', 4)) --- [\'a\', \'a\', \'a\', \'b\', \'b\', \'c\', \'d\', \'d\'] --- => ([('a', 7), ('b', 2), ('c', 1)], Just (\'d\', 2)) +-- 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 -- @ -- -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)) +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 -groupsBy f !c !vec0 - = (vec1, snd k1) - where - f' x y = return $ f x y + -- Make a vector for the output lengths. + mvLens <- GM.unsafeNew (GV.length vStartsMax) + rmvLens <- newIORef mvLens - (vec1 :: v2 (a, Int), k1) - = runST $ unchainToVector $ C.liftChain - $ C.groupsByC f' c $ chainOfVector vec0 -{-# INLINE groupsBy #-} + -- 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 #-} + ------------------------------------------------------------------------------- --- | 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) +-- | Segmented replicate. +replicates + :: (GV.Vector v (Int, a), GV.Vector v a) + => v (Int, a) + -> v a -findSegments pStart pEnd src - = GV.unzip - $ GV.unstream - $ startLengthsOfSegsS - $ findSegmentsS pStart pEnd (GV.length src - 1) - $ SM.indexed - $ GV.stream src -{-# INLINE findSegments #-} +replicates vec + = GV.unstream + $ replicatesS + $ GV.stream vec +{-# INLINE replicates #-} ------------------------------------------------------------------------------- --- | 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) +-- | Produce a chain from a generic vector. +chainOfVector + :: (Monad m, GV.Vector v a) + => v a -> Chain m Int a -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 #-} +chainOfVector vec + = Chain (S.Exact len) 0 step + where + !len = GV.length vec + step !i + | i >= len + = return $ Done i + | otherwise + = return $ Yield (GV.unsafeIndex vec i) (i + 1) + {-# INLINE_INNER step #-} +{-# INLINE_STREAM chainOfVector #-} + + ------------------------------------------------------------------------------- --- | 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 +-- | Compute a chain into a generic vector. +unchainToVector + :: (PrimMonad m, GV.Vector v a) + => C.Chain m s a -> m (v a, s) +unchainToVector chain + = do (mvec, c') <- unchainToMVector chain + vec <- GV.unsafeFreeze mvec + return (vec, c') +{-# INLINE_STREAM unchainToVector #-} -diceSep pEndInner pEndBoth vec - = runST - $ unstreamToVector2 - $ diceSepS pEndInner pEndBoth - $ S.liftStream - $ GV.stream vec -{-# INLINE diceSep #-} +-- | Compute a chain into a generic mutable vector. +unchainToMVector + :: (PrimMonad m, GM.MVector v a) + => Chain m s a + -> m (v (PrimState m) a, s) -------------------------------------------------------------------------------- --- | 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 +unchainToMVector (Chain sz s0 step) + = case sz of + S.Exact i -> unchainToMVector_max i s0 step + S.Max i -> unchainToMVector_max i s0 step + S.Unknown -> unchainToMVector_unknown 32 s0 step +{-# INLINE_STREAM unchainToMVector #-} -compact f s0 vec - = GV.unstream $ compactS f s0 $ GV.stream vec -{-# INLINE compact #-} +-- unchain when we known the maximum size of the vector. +unchainToMVector_max nMax s0 step + = GM.unsafeNew nMax >>= \vec + -> let + go !sPEC !i !s + = step s >>= \m + -> case m of + Yield e s' + -> do GM.unsafeWrite vec i e + go sPEC (i + 1) s' --- | 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 + Skip s' -> go sPEC i s' + Done s' -> return (GM.unsafeSlice 0 i vec, s') + {-# INLINE_INNER go #-} -compactIn f vec - = GV.unstream $ compactInS f $ GV.stream vec -{-# INLINE compactIn #-} + in go SM.SPEC 0 s0 +{-# INLINE_STREAM unchainToMVector_max #-} +-- unchain when we don't know the maximum size of the vector. +unchainToMVector_unknown nStart s0 step + = GM.unsafeNew nStart >>= \vec0 + -> let + go !sPEC !vec !i !n !s + = step s >>= \m + -> case m of + Yield e s' + | i >= n + -> do vec' <- GM.unsafeGrow vec n + GM.unsafeWrite vec' i e + go sPEC vec' (i + 1) (n + n) s' + + | otherwise + -> do GM.unsafeWrite vec i e + go sPEC vec (i + 1) n s' + + Skip s' -> go sPEC vec i n s' + Done s' -> return (GM.unsafeSlice 0 i vec, s') + {-# INLINE_INNER go #-} + + in go SM.SPEC vec0 0 nStart s0 +{-# INLINE_STREAM unchainToMVector_unknown #-} + + ------------------------------------------------------------------------------- -- | Segmented fold over vectors of segment lengths and input values. -- @@ -545,18 +525,53 @@ ------------------------------------------------------------------------------- --- | 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) +-- | 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. -padForward ksucc vec - = GV.unstream - $ padForwardS ksucc - $ GV.stream vec -{-# INLINE padForward #-} +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 #-} 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 @@ -1,48 +1,50 @@ module Data.Repa.Vector.Unboxed - ( -- * Conversion - chainOfVector - , unchainToVector - , unchainToMVector + ( -- * Stream operators + -- ** Compacting + compact + , compactIn - -- * Generating - , ratchet + -- ** Dicing + , findSegments + , findSegmentsFrom + , diceSep - -- * Extracting + -- ** Extracting , extract - -- * Inserting + -- ** Inserting , insert - -- * Merging + -- ** Merging , merge , mergeMaybe - -- * Splitting - , findSegments - , findSegmentsFrom - , diceSep + -- ** Padding + , padForward - -- * Compacting - , compact - , compactIn + -- ** Ratcheting + , ratchet - -- * Padding - , padForward + -- * Chain operators + , unchainToVector + , unchainToMVector - -- * Scanning - , scanMaybe + -- ** Folding + , folds, C.Folds(..) - -- * Grouping + -- ** Scanning + , scanMaybe , groupsBy - -- * Folding - , folds, C.Folds(..)) + -- ** Conversion + , chainOfVector + ) where -import Data.Repa.Option import Data.Vector.Unboxed (Unbox, Vector) import Data.Vector.Unboxed.Mutable (MVector) import Data.Repa.Chain (Chain) +import Data.Repa.Scalar.Option import qualified Data.Repa.Vector.Generic as G import qualified Data.Repa.Chain as C import qualified Data.Vector.Unboxed as U @@ -257,7 +259,7 @@ -- compact :: (Unbox a, Unbox b) - => (s -> a -> (Maybe b, s)) -- ^ Worker function + => (s -> a -> (s, Maybe b)) -- ^ Worker function -> s -- ^ Starting state -> Vector a -- ^ Input vector -> Vector b @@ -270,7 +272,7 @@ -- initial state, and add the final state to the end of the output. compactIn :: Unbox a - => (a -> a -> (Maybe a, a)) -- ^ Worker function. + => (a -> a -> (a, Maybe a)) -- ^ Worker function. -> Vector a -- ^ Input elements. -> Vector a 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.1.0.1 +Version: 4.2.2.1 License: BSD3 License-file: LICENSE Author: The Repa Development Team @@ -19,14 +19,14 @@ Library build-Depends: - base == 4.7.*, + base == 4.8.*, vector == 0.10.*, - primitive == 0.5.4.*, - mtl == 2.2.* + primitive == 0.6.*, + mtl == 2.2.*, + repa-scalar == 4.2.2.* exposed-modules: Data.Repa.Chain - Data.Repa.Option Data.Repa.Stream Data.Repa.Vector.Generic Data.Repa.Vector.Unboxed @@ -46,6 +46,7 @@ Data.Repa.Stream.Merge Data.Repa.Stream.Pad Data.Repa.Stream.Ratchet + Data.Repa.Stream.Replicate Data.Repa.Stream.Segment include-dirs: