diff --git a/Data/Repa/Chain.hs b/Data/Repa/Chain.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Chain.hs @@ -0,0 +1,26 @@ + +-- | * See the "Data.Repa.Vector.Unboxed" module for examples of how these +-- functions can be used. +module Data.Repa.Chain + ( -- * Chain Fusion + Chain (..) + , Step (..) + , liftChain + , resumeChain + + -- * Weaves + , weaveC, Weave, Turn (..), Move(..), move + + -- * Folding + , foldsC, Folds(..) + + -- * Scanning + , scanMaybeC + + -- * Grouping + , groupsByC) +where +import Data.Repa.Chain.Base +import Data.Repa.Chain.Scan +import Data.Repa.Chain.Weave +import Data.Repa.Chain.Folds diff --git a/Data/Repa/Chain/Base.hs b/Data/Repa/Chain/Base.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Chain/Base.hs @@ -0,0 +1,57 @@ + +module Data.Repa.Chain.Base + ( Step (..) + , Chain (..) + , liftChain + , resumeChain) +where +import qualified Data.Vector.Fusion.Stream.Size as S +import Control.Monad.Identity +#include "repa-stream.h" + + +-- | A chain is an abstract, stateful producer of elements. It is similar +-- a stream as used in stream fusion, except that internal state is visible +-- in its type. This allows the computation to be paused and resumed at a +-- later point. +data Chain m s a + = Chain + { -- | Expected size of the output. + mchainSize :: S.Size + + -- | Starting state. + , mchainState :: s + + -- | Step the chain computation. + , mchainStep :: s -> m (Step s a) } + + +-- | Result of a chain computation step. +data Step s a + + -- | Yield an output value and a new seed. + = Yield !a !s + + -- | Provide just a new seed. + | Skip !s + + -- | Signal that the computation has finished. + | Done !s + deriving Show + + +-- | Lift a pure chain to a monadic chain. +liftChain :: Monad m => Chain Identity s a -> Chain m s a +liftChain (Chain sz s step) + = Chain sz s (return . runIdentity . step) +{-# INLINE_STREAM liftChain #-} + + +-- | Resume a chain computation from a previous state. +resumeChain + :: Monad m + => s -> Chain m s a -> Chain m s a +resumeChain s' (Chain sz _s step) + = Chain sz s' step +{-# INLINE_STREAM resumeChain #-} + diff --git a/Data/Repa/Chain/Folds.hs b/Data/Repa/Chain/Folds.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Chain/Folds.hs @@ -0,0 +1,148 @@ + +module Data.Repa.Chain.Folds + (foldsC, Folds (..)) +where +import Data.Repa.Option +import Data.Repa.Chain.Base +import Data.Vector.Fusion.Stream.Size as S +#include "repa-stream.h" + + +-- | 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. +-- +foldsC :: Monad m + => (a -> b -> m b) -- ^ Worker function. + -> b -- ^ Initial state when folding rest of segments. + -> Option3 n Int b -- ^ Name, length and initial state for first segment. + -> Chain m sLen (n, Int) -- ^ Segment names and lengths. + -> Chain m sVal a -- ^ Input data to fold. + -> Chain m (Folds sLen sVal n a b) (n, b) + +foldsC f zN s0 + (Chain _szLens sLens0 stepLens) + (Chain _szVals sVals0 stepVals) + = Chain S.Unknown (init_foldsC s0) step + where + init_foldsC s + = case s of + None3 -> Folds sLens0 sVals0 None 0 zN + Some3 n len acc -> Folds sLens0 sVals0 (Some n) len acc + {-# NOINLINE init_foldsC #-} + -- NOINLINE to hide the case match from the simplifier so it + -- doesn't unswitch it at top-level and duplicate the follow-on code. + + step ss@(Folds sLens sVals nameSeg lenSeg valSeg) + = case nameSeg of + -- If we don't have a segment length we need to load the next one. + None + -> stepLens sLens >>= \rLens + -> case rLens of + -- We got a segment length, so load it into the state and + -- initialise the accumulator. + Yield (name, xLen) sLens' + -> return $ Skip ss { _stateLens = sLens' + , _nameSeg = Some name + , _lenSeg = xLen + , _valSeg = zN } + + -- Lengths input takes a step. + Skip sLens' + -> return $ Skip ss { _stateLens = sLens' } + + -- We're not currently folding a segment, and no more segment + -- lengths are available, so we're done. + Done sLens' + -> return $ Done ss { _stateLens = sLens' } + + -- We're currently folding a segment. + Some name + -- We've reached the end of the segment, so emit the result. + | lenSeg == 0 + -> return $ Yield (name, valSeg) + ss { _nameSeg = None } + + -- We still need more values for this segment. + | otherwise + -> stepVals sVals >>= \rVals + -> case rVals of + -- We got a new value, so accumulate it into the state. + Yield xVal sVals' + -> f xVal valSeg >>= \rAcc + -> return $ Skip ss { _stateVals = sVals' + , _lenSeg = lenSeg - 1 + , _valSeg = rAcc } + + -- Vals input takes a step. + Skip sVals' + -> return $ Skip ss { _stateVals = sVals' } + + -- We're in a non-zero lengthed segment, but haven't got + -- all the values, so we're done for now. + Done sVals' + -> return $ Done ss { _stateVals = sVals' } + {-# INLINE_INNER step #-} +{-# INLINE_STREAM foldsC #-} + + +-- | Return state of a folds operation. +data Folds sLens sVals n a b + = Folds + { -- | State of lengths chain. + _stateLens :: !sLens + + -- | State of values chain. + , _stateVals :: !sVals + + -- | If we're currently in a segment, then hold its name, + , _nameSeg :: !(Option n) + + -- | Length of current segment. + , _lenSeg :: !Int + + -- | Accumulated value of current segment. + , _valSeg :: !b } + deriving Show + + +{- + + -- Defining folds in terms of weave doesn't work because if all the + -- segment lengths are 0 then we don't want to load any values at all. + + = weaveC work s0 cLens cVals + where + work !ms !mxLen !mxVal + = case ms of + -- If we haven't got a current state then load the next + -- segment length. + None2 + -> case mxLen of + None -> return $ Finish ms MoveNone + Some xLen -> return $ Next (Some2 xLen zN) MoveLeft + + Some2 len acc + | len == 0 -> return $ Give acc None2 MoveNone + | otherwise + -> case mxVal of + None -> return $ Finish ms MoveNone + Some xVal + -> do r <- f xVal acc + return $ Next (Some2 (len - 1) r) MoveRight + {-# INLINE [1] work #-} + + +-- | Pack the weave state of a folds operation into a `Folds` record, +-- which has better field names. +packFolds :: Weave sLens Int sVals a (Option2 Int b) + -> Folds sLens sVals a b + +packFolds (Weave stateL elemL _endL stateR elemR _endR mLenAcc) + = (Folds stateL elemL stateR elemR mLenAcc) +{-# INLINE packFolds #-} +-} diff --git a/Data/Repa/Chain/Scan.hs b/Data/Repa/Chain/Scan.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Chain/Scan.hs @@ -0,0 +1,65 @@ + +module Data.Repa.Chain.Scan + ( scanMaybeC + , groupsByC) +where +import Data.Repa.Chain.Base +import qualified Data.Vector.Fusion.Stream.Size as S +#include "repa-stream.h" + + +------------------------------------------------------------------------------- +-- | 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. +scanMaybeC + :: Monad m + => (k -> a -> m (k, Maybe b)) -- ^ Worker function. + -> k -- ^ Initial state for scan. + -> Chain m s a -- ^ Input elements. + -> Chain m (s, k) b -- ^ Output elements and final state. + +scanMaybeC f k0 (Chain sz s0 istep) + = Chain (S.toMax sz) (s0, k0) ostep + where + ostep (s1, k1) + = istep s1 >>= \rs + -> case rs of + Yield x s2 + -> f k1 x >>= \rk + -> case rk of + (k2, Nothing) -> return $ Skip (s2, k2) + (k2, Just y) -> return $ Yield y (s2, k2) + + Skip s2 -> return $ Skip (s2, k1) + Done s2 -> return $ Done (s2, k1) + {-# INLINE_INNER ostep #-} +{-# INLINE_STREAM scanMaybeC #-} + + +------------------------------------------------------------------------------- +-- | From a stream of values which has consecutive runs of idential values, +-- produce a stream of the lengths of these runs. +groupsByC + :: Monad m + => (a -> a -> m Bool) -- ^ Comparison function. + -> Maybe (a, Int) -- ^ Starting element and count. + -> Chain m s a -- ^ Input elements. + -> Chain m (s, Maybe (a, Int)) (a, Int) + +groupsByC f !s !vec + = scanMaybeC work_groupsByC s vec + where + work_groupsByC !acc !y + = case acc of + Nothing + -> return $ (Just (y, 1), Nothing) + + Just (x, n) + -> f x y >>= \rk + -> if rk + then return (Just (x, n + 1), Nothing) + else return (Just (y, 1), Just (x, n)) + {-# INLINE_INNER work_groupsByC #-} +{-# INLINE_STREAM groupsByC #-} + diff --git a/Data/Repa/Chain/Weave.hs b/Data/Repa/Chain/Weave.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Chain/Weave.hs @@ -0,0 +1,115 @@ + +module Data.Repa.Chain.Weave + ( weaveC + , Weave (..) + , Turn (..) + , Move (..), move + , Option (..)) +where +import Data.Repa.Option +import Data.Repa.Chain.Base +import qualified Data.Vector.Fusion.Stream.Size as S +#include "repa-stream.h" + + +-- | A weave is a generalized merge of two input chains. +-- +-- The worker function takes the current state, values from the +-- left and right input chains, and produces a `Turn` which +-- describes any output at that point, as well as how the input +-- chains should be advanced. +-- +weaveC :: Monad m + => (k -> Option aL -> Option aR -> m (Turn k aX)) + -- ^ Worker function. + -> k -- ^ Initial state. + -> Chain m sL aL -- ^ Left input chain. + -> Chain m sR aR -- ^ Right input chain. + -> Chain m (Weave sL aL sR aR k) aX -- ^ Result chain. + +weaveC f !ki (Chain _sz1 s1i step1) (Chain _sz2 s2i step2) + = Chain S.Unknown + (Weave s1i None False s2i None False ki) + step + where + step ss@(Weave s1 m1 e1 s2 m2 e2 k) + = case (m1, e1, m2, e2) of + (None, False, _, _) + -> step1 s1 >>= \r1 + -> return $ Skip + $ case r1 of + Yield x1 sL' -> ss { _stateL = sL', _elemL = Some x1 } + Skip sL' -> ss { _stateL = sL' } + Done sL' -> ss { _stateL = sL', _endL = True } + + (_, _, None, False) + -> step2 s2 >>= \r2 + -> return $ Skip + $ case r2 of + Yield x2 sR' -> ss { _stateR = sR', _elemR = Some x2 } + Skip sR' -> ss { _stateR = sR' } + Done sR' -> ss { _stateR = sR', _endR = True } + _ + -> f k m1 m2 >>= \t + -> case t of + Give x k' m -> return $ Yield x (move k' m ss) + Next k' m -> return $ Skip (move k' m ss) + Finish k' m -> return $ Done (move k' m ss) + {-# INLINE_INNER step #-} +{-# INLINE_STREAM weaveC #-} + + +-- | Internal state of a weave. +data Weave sL aL sR aR k + = Weave + { -- | State of the left input chain. + _stateL :: !sL + + -- | Current value loaded from the left input. + , _elemL :: !(Option aL) + + -- | Whether we've hit the end of the left input + , _endL :: Bool + + -- | State of the right input chain. + , _stateR :: !sR + + -- | Current value loaded from the right input. + , _elemR :: !(Option aR) + + -- | Whether we've hit the end of the right input. + , _endR :: Bool + + -- | Worker state at this point in the weave. + , _here :: !k } + deriving Show + + +-- | What to do after considering two input elements. +data Turn k a + = Give !a !k !Move -- ^ Give an element and a new state. + | Next !k !Move -- ^ Move to the next input. + | Finish !k !Move -- ^ Weave is finished for now. + deriving Show + + +-- | How to move the input chains after considering to input elements. +data Move + = MoveLeft + | MoveRight + | MoveBoth + | MoveNone + deriving Show + + +-- | Apply a `Move` instruction to a weave state. +move :: k -> Move + -> Weave s1 a1 s2 a2 k -> Weave s1 a1 s2 a2 k +move !k' !mm !ss + = case mm of + MoveLeft -> ss { _here = k', _elemL = None } + MoveRight -> ss { _here = k', _elemR = None } + MoveBoth -> ss { _here = k', _elemL = None, _elemR = None } + MoveNone -> ss { _here = k' } +{-# INLINE_INNER move #-} + diff --git a/Data/Repa/Option.hs b/Data/Repa/Option.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Option.hs @@ -0,0 +1,86 @@ + +-- | 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 new file mode 100644 --- /dev/null +++ b/Data/Repa/Stream.hs @@ -0,0 +1,22 @@ + +-- | * See the "Data.Repa.Vector.Unboxed" module for examples of how these +-- functions can be used. +module Data.Repa.Stream + ( extractS + , mergeS + , findSegmentsS + , diceSepS + , startLengthsOfSegsS + , padForwardS + + -- * Unsafe operators + , unsafeRatchetS) + +where +import Data.Repa.Stream.Extract +import Data.Repa.Stream.Merge +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/Dice.hs b/Data/Repa/Stream/Dice.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Stream/Dice.hs @@ -0,0 +1,117 @@ + +module Data.Repa.Stream.Dice + ( diceSepS ) +where +import Data.Vector.Fusion.Stream.Monadic (Stream(..), Step(..)) +import qualified Data.Vector.Fusion.Stream.Size as S +#include "repa-stream.h" + + +------------------------------------------------------------------------------- +-- | Given predicates that detect the begining and end of interesting segments +-- of information, scan through a vector looking for when these begin +-- and end. +-- +diceSepS + :: Monad m + => (a -> Bool) -- ^ Detect the end of a column. + -> (a -> Bool) -- ^ Detect the end of a row. + -> Stream m a + -> Stream m (Maybe (Int, Int), Maybe (Int, Int)) + -- ^ Segment starts and lengths + +diceSepS pEndCol pEndRow (Stream istep s0 sz) + = Stream ostep (s0, True, 0, 0, 0, False, Nothing) + (case sz of + S.Exact n -> S.Max (n + 1) + S.Max n -> S.Max (n + 1) + S.Unknown -> S.Unknown) + where + ostep (_, False, _, _, _, _, _) + = return Done + + -- We're not in an inner segment, so look for the next starting element. + ostep (si, f, iSrc, iRowStart, iSeps, inRow, mCol@Nothing) + = iSrc `seq` iRowStart `seq` iSeps `seq` inRow `seq` + istep si >>= \m + -> case m of + Yield x si' + -- Line ended outside a word. + | pEndRow x + -> let nRow = if inRow then iSeps + 1 else 0 + in return $ Yield + ( Just (iSrc, 0) + , Just (iRowStart, nRow)) + (si', f, iSrc + 1, iRowStart + nRow, 0, False, mCol) + + -- Inner segment started and ended on the same element. + | pEndCol x + -> return $ Yield + (Just (iSrc, 0), Nothing) + (si', f, iSrc + 1, iRowStart, iSeps + 1, True, mCol) + + -- Segment has started on this element. + | otherwise + -> return $ Skip + (si', f, iSrc + 1, iRowStart, iSeps, True, Just (iSrc, 1)) + + -- We didn't get an element this time. + Skip si' + -> return $ Skip + (si', f, iSrc, iRowStart, iSeps, inRow, mCol) + + -- Found end of input outside a segment. + Done + | inRow + -> return $ Yield + (Just (0, 0), Just (iRowStart, iSeps + 1)) + (si, False, iSrc, iRowStart, iSeps, False, mCol) + + | otherwise + -> return $ Yield + (Just (0, 0), Nothing) + (si, False, iSrc, iRowStart, iSeps, False, mCol) + + + -- We're in an inner segment, looking for the ending element. + ostep ( si, f, iSrc, iRowStart, iSeps, inRow + , mCol@(Just (iColStart, iColLen))) + + = iSrc `seq` iRowStart `seq` iSeps `seq` inRow `seq` + istep si >>= \m + -> case m of + Yield x si' + -- Both inner and outer ended at this point, + -- and now we're looking for a new segment start. + | pEndRow x + -> return $ Yield + ( Just (iColStart, iColLen) + , Just (iRowStart, iSeps + 1)) + ( si', f, iSrc + 1, iRowStart + iSeps + 1, 0, False, Nothing) + + -- Inner segment ended at this point, + -- but we're still in the outer segment. + | pEndCol x + -> return $ Yield + ( Just (iColStart, iColLen) + , Nothing) + ( si', f, iSrc + 1, iRowStart, iSeps + 1, inRow, Nothing) + + -- Another element of the inner segment. + | otherwise + -> return $ Skip + ( si', f, iSrc + 1, iRowStart, iSeps, inRow + , Just (iColStart, iColLen + 1)) + + -- We didn't get an element this time. + Skip si' + -> return $ Skip + ( si', f, iSrc, iRowStart, iSeps, inRow, mCol) + + -- Found end of input during a segment. + Done + -> return $ Yield + ( Just (iColStart, iColLen) + , Nothing) + ( si, True, iSrc, iRowStart, iSeps, inRow, Nothing) +{-# INLINE_STREAM diceSepS #-} diff --git a/Data/Repa/Stream/Extract.hs b/Data/Repa/Stream/Extract.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Stream/Extract.hs @@ -0,0 +1,35 @@ + +module Data.Repa.Stream.Extract + (extractS) +where +import Data.Vector.Fusion.Stream.Monadic (Stream(..), Step(..)) +import qualified Data.Vector.Fusion.Stream.Size as S +#include "repa-stream.h" + + +-- | Extract segments from some source array and concatenate them. +extractS + :: Monad m + => (Int -> a) -- ^ Function to get elements from the source. + -> Stream m (Int, Int) -- ^ Segment start positions and lengths. + -> Stream m a -- ^ Result elements. + +extractS get (Stream istep si0 _) + = Stream ostep (si0, Nothing) S.Unknown + where + -- Start a new segment. + ostep (si, Nothing) + = istep si >>= \m + -> case m of + Yield (iStart, iLen) si' + -> return $ Skip (si', Just (iStart, iStart + iLen)) + Skip si' -> return $ Skip (si', Nothing) + Done -> return $ Done + + -- Emit data from a segment. + ostep (si, Just (iPos, iTop)) + | iPos >= iTop = return $ Skip (si, Nothing) + | otherwise = return $ Yield (get iPos) + (si, Just (iPos + 1, iTop)) + {-# INLINE_INNER ostep #-} +{-# INLINE_STREAM extractS #-} diff --git a/Data/Repa/Stream/Merge.hs b/Data/Repa/Stream/Merge.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Stream/Merge.hs @@ -0,0 +1,79 @@ + +module Data.Repa.Stream.Merge + (mergeS) +where +import Data.Vector.Fusion.Stream.Monadic (Stream(..), Step(..)) +import Data.Repa.Option +import qualified Data.Vector.Fusion.Stream.Size as S +#include "repa-stream.h" + + +-- | Merge two key-value streams. +-- +-- The streams are assumed to be pre-sorted on the keys. +-- +mergeS :: (Monad m, Ord k) + => (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. + -> Stream m (k, a) -- ^ Stream of keys and left values. + -> Stream m (k, b) -- ^ Stream of keys and right values. + -> Stream m (k, c) -- ^ Stream of keys and results. + +mergeS fBoth fLeft fRight (Stream istepA sA0 _) (Stream istepB sB0 _) + = Stream ostep (sA0, sB0, None2, True, None2, True) S.Unknown + where + -- Merge where both streams match. + ostep (sA, sB, kxA@(Some2 kA xA), hasA + , kxB@(Some2 kB xB), hasB) + + = return $ Yield (if | kA == kB -> (kA, fBoth kA xA xB) + | kB < kA -> (kB, fRight kB xB) + | otherwise -> (kA, fLeft kA xA)) + + (if | kA == kB -> (sA, sB, None2, hasA, None2, hasB) + | kB < kA -> (sA, sB, kxA, hasA, None2, hasB) + | otherwise -> (sA, sB, None2, hasA, kxB, hasB)) + + -- Drain left stream. + ostep (sA, sB, Some2 kA xA, hasA, kxB@None2, hasB@False) + = return $ Yield (kA, fLeft kA xA) + (sA, sB, None2, hasA, kxB, hasB) + + -- Drain right stream. + ostep (sA, sB, kxA@None2, hasA@False, Some2 kB xB, hasB) + = return $ Yield (kB, fRight kB xB) + (sA, sB, kxA, hasA, None2, hasB) + + -- Advance left stream. + ostep (sA, sB, kxA@None2, hasA@True, kxB, hasB) + = istepA sA >>= \mA + -> case mA of + Yield (kA, xA) sA' + -> return $ Skip (sA', sB, Some2 kA xA, True, kxB, hasB) + + Skip sA' + -> return $ Skip (sA', sB, kxA, hasA, kxB, hasB) + + Done + -> return $ Skip (sA, sB, kxA, False, kxB, hasB) + + -- Advance the right stream. + ostep (sA, sB, kxA, hasA, kxB@None2, hasB@True) + = istepB sB >>= \mB + -> case mB of + Yield (kB, xB) sB' + -> return $ Skip (sA, sB', kxA, hasA, Some2 kB xB, True) + + Skip sB' + -> return $ Skip (sA, sB', kxA, hasA, kxB, hasB) + + Done + -> return $ Skip (sA, sB, kxA, hasA, kxB, False) + + -- Done + ostep (_sA, _sB, None2, False, None2, False) + = return $ Done + {-# INLINE_INNER ostep #-} +{-# INLINE_STREAM mergeS #-} + diff --git a/Data/Repa/Stream/Pad.hs b/Data/Repa/Stream/Pad.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Stream/Pad.hs @@ -0,0 +1,66 @@ + +module Data.Repa.Stream.Pad + (padForwardS) +where +import Data.Vector.Fusion.Stream.Monadic (Stream(..), Step(..)) +import Data.Repa.Option +import qualified Data.Vector.Fusion.Stream.Size as S +#include "repa-stream.h" + + +-- | 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. +-- +padForwardS + :: (Monad m, Ord k) + => (k -> k) -- ^ Successor functinon for keys. + -> Stream m (k, v) -- ^ Input stream. + -> Stream m (k, v) + +padForwardS ksucc (Stream istep si0 _) + = Stream ostep (si0, None2, None2) S.Unknown + where + -- Load the first element. + ostep (si, sPrev@None2, sBound) + = istep si >>= \m + -> case m of + Yield (k0, v0) si' + -> return $ Yield (k0, v0) (si', Some2 k0 v0, sBound) + + Skip si' + -> return $ Skip (si', sPrev, sBound) + + Done + -> return $ Done + + -- Load the next element element. + ostep (si, sPrev@(Some2 kPrev _vPrev), sTarget@None2) + = istep si >>= \m + -> case m of + Yield (kStep, vStep) si' + -- The next element from the input is more than the expected + -- one then there is a gap in the input. + | kExpect <- ksucc kPrev + , kStep > kExpect + -> return $ Skip (si', sPrev, Some2 kStep vStep) + + -- Otherwise there is no gap. + | otherwise + -> return $ Yield (kStep, vStep) (si', Some2 kStep vStep, None2) + + Skip si' -> return $ Skip (si', sPrev, sTarget) + Done -> return $ Done + + -- Fill in missing elements. + ostep (si, _sPrev@(Some2 kPrev vPrev), sTarget@(Some2 kTarget vTarget)) + = let kNext = ksucc kPrev + in if kNext >= kTarget + -- We've reached the target, so the gap is filled. + then return $ Yield (kTarget, vTarget) (si, Some2 kTarget vTarget, None2) + + -- We're still filling the gap. + else return $ Yield (kNext, vPrev) (si, Some2 kNext vPrev, sTarget) + {-# INLINE_INNER ostep #-} +{-# INLINE_STREAM padForwardS #-} + diff --git a/Data/Repa/Stream/Ratchet.hs b/Data/Repa/Stream/Ratchet.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Stream/Ratchet.hs @@ -0,0 +1,101 @@ + +module Data.Repa.Stream.Ratchet + ( unsafeRatchetS) +where +import Data.IORef +import Data.Vector.Fusion.Stream.Monadic (Stream(..), Step(..)) +import qualified Data.Vector.Generic as G +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" + + +-- | 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 +-- @ +-- +-- The function takes the starting values in a mutable vector and +-- updates it during computation. Computation proceeds by making passes +-- through the mutable vector and updating the starting values until +-- they match the stopping values. +-- +-- UNSAFE: Both input vectors must have the same length, +-- 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. + -> Stream IO Int + +unsafeRatchetS !mvStarts !vMax !rmvLens + = Stream ostep (0, Nothing, 0, 0) S.Unknown + where + !iSegMax = GM.length mvStarts - 1 + + ostep (iSeg, mvmLens, oSeg, oLen) + = ostep' iSeg mvmLens oSeg oLen + {-# INLINE ostep #-} + + ostep' !iSeg !mvmLens !oSeg !oLen + | iSeg <= iSegMax + = do !iVal <- GM.unsafeRead mvStarts iSeg + let !iNext = vMax `G.unsafeIndex` iSeg + if iVal >= iNext + then return $ Skip (iSeg + 1, mvmLens, oSeg, oLen) + else do + GM.unsafeWrite mvStarts iSeg (iVal + 1) + return $ Yield iVal (iSeg + 1, mvmLens, oSeg, oLen + 1) + + -- We're at the end of an output segment, + -- so write the output length into the lengths vector. + | oLen > 0 + = do -- Get the current output vector. + !vmLens <- case mvmLens of + Nothing -> readIORef rmvLens + Just vmLens -> return $ vmLens + + -- If the output vector is full then we need to grow it. + let !oSegLen = UM.length vmLens + if oSeg >= oSegLen + then do + !vmLens' <- UM.unsafeGrow vmLens (UM.length vmLens) + writeIORef rmvLens vmLens' + UM.unsafeWrite vmLens' oSeg oLen + return $ Skip (0, Just vmLens', oSeg + 1, 0) + + else do + UM.unsafeWrite vmLens oSeg oLen + return $ Skip (0, Just vmLens, oSeg + 1, 0) + + | otherwise + = do !vmLens <- case mvmLens of + Nothing -> readIORef rmvLens + Just vmLens -> return $ vmLens + + let !vmLens' = UM.unsafeSlice 0 oSeg vmLens + writeIORef rmvLens vmLens' + return Done + {-# INLINE_INNER ostep' #-} +{-# INLINE_STREAM unsafeRatchetS #-} + + diff --git a/Data/Repa/Stream/Segment.hs b/Data/Repa/Stream/Segment.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Stream/Segment.hs @@ -0,0 +1,111 @@ + +module Data.Repa.Stream.Segment + ( findSegmentsS + , startLengthsOfSegsS) +where +import Data.Vector.Fusion.Stream.Monadic (Stream(..), Step(..)) +import qualified Data.Vector.Fusion.Stream.Size as S +#include "repa-stream.h" + + +------------------------------------------------------------------------------- +-- | 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. +findSegmentsS + :: Monad m + => (a -> Bool) -- ^ Predicate to check for start of segment. + -> (a -> Bool) -- ^ Predicate to check for end of segment. + -> i -- ^ Index of final element in stream. + -> Stream m (i, a) -- ^ Stream of indices and elements. + -> Stream m (i, i) -- ^ Stream of segment start and end indices. + +findSegmentsS pStart pEnd iEnd (Stream istep s sz) + = Stream ostep (s, True, Nothing) (S.toMax sz) + where + -- We've hit the end of the stream + ostep (_, False, _) + = return Done + + -- We're not in a segment, so look for the next starting element. + ostep (si, f, n@Nothing) + = do m <- istep si + case m of + Yield (i, x) si' + | pStart x + -> if pEnd x + -- Segment started and ended on the same element. + then return $ Yield (i, i) (si', f, Nothing) + + -- Segment has started on this element. + else return $ Skip (si', f, Just i) + + -- Still looking for the starting element. + | otherwise -> return $ Skip (si', f, n) + + -- We didn't get an element this time. + Skip si' -> return $ Skip (si', f, n) + + -- Found end of imput. + Done -> return $ Done + + -- We're in a segment, so look for ending element. + ostep (si, f, j@(Just iStart)) + = do m <- istep si + case m of + Yield (i, x) si' + -- Segment ended here. + | pEnd x -> return $ Yield (iStart, i) + (si', f, Nothing) + + -- Still looking for the ending element. + | otherwise -> return $ Skip (si', f, j) + + -- We didn't get an element this time. + Skip si' -> return $ Skip (si', f, j) + + -- Found end of input during a segment. + Done -> return $ Yield (iStart, iEnd) + (si, False, Nothing) + {-# INLINE_INNER ostep #-} +{-# INLINE_STREAM findSegmentsS #-} + + +------------------------------------------------------------------------------- +-- | Given a stream of starting and ending indices for some segments, +-- convert it to a stream of starting indices and segment lengths. +-- +-- * The ending indices must be after the starting indices, +-- otherwise the result will contain negative lengths. +-- +startLengthsOfSegsS + :: Monad m + => Stream m (Int, Int) -- ^ Start and end indices. + -> Stream m (Int, Int) -- ^ Start indices and lengths of segments. + +startLengthsOfSegsS (Stream istep s sz) + = Stream ostep (s, True, Nothing) sz + where + ostep (_, False, _) + = return Done + + ostep (si, f, n@Nothing) + = do m <- istep si + case m of + Yield x si' -> return $ Skip (si', f, Just x) + Skip si' -> return $ Skip (si', f, n) + Done -> return $ Done + + ostep (si, f, j@(Just (iStart, iEnd))) + = do m <- istep si + case m of + Yield x si' -> return $ Yield (iStart, iEnd - iStart + 1) + (si', f, Just x) + + Skip si' -> return $ Skip (si', f, j) + + Done -> return $ Yield (iStart, iEnd - iStart + 1) + (si, False, Nothing) + {-# INLINE_INNER ostep #-} +{-# INLINE_STREAM startLengthsOfSegsS #-} + diff --git a/Data/Repa/Vector/Generic.hs b/Data/Repa/Vector/Generic.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Vector/Generic.hs @@ -0,0 +1,181 @@ + +-- | Converting `Stream`s and `Chain`s to and from generic `Vector`s. +-- +-- * NOTE: Support for streams of unknown length is not complete. +-- +module Data.Repa.Vector.Generic + ( -- * Stream functions + unstreamToVector2 + , unstreamToMVector2 + + -- * Chain functions + , chainOfVector + , unchainToVector + , unchainToMVector) +where +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.Size as S +import Control.Monad.Primitive +#include "repa-stream.h" + + +------------------------------------------------------------------------------- +-- | Unstream some elements to two separate vectors. +-- +-- `Nothing` values are ignored. +-- +unstreamToVector2 + :: (PrimMonad m, GV.Vector v a, GV.Vector v b) + => S.Stream m (Maybe a, Maybe b) + -- ^ Source data. + -> m (v a, v b) -- ^ Resulting vectors. + +unstreamToVector2 s + = do (mvec1, mvec2) <- unstreamToMVector2 s + vec1 <- GV.unsafeFreeze mvec1 + vec2 <- GV.unsafeFreeze mvec2 + return (vec1, vec2) +{-# INLINE_STREAM unstreamToVector2 #-} + + +-- | Unstream some elements to two separate mutable vectors. +-- +-- `Nothing` values are ignored. +-- +unstreamToMVector2 + :: (PrimMonad m, GM.MVector v a, GM.MVector v b) + => S.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) + = case sz of + S.Exact i -> unstreamToMVector2_max i s0 step + S.Max i -> unstreamToMVector2_max i s0 step + S.Unknown -> error "repa-stream: finish unstreamToMVector2" +{-# INLINE_STREAM unstreamToMVector2 #-} + +unstreamToMVector2_max nMax s0 step + = GM.unsafeNew nMax >>= \vecL + -> GM.unsafeNew nMax >>= \vecR + -> let + go !sPEC !iL !iR !s + = step s >>= \m + -> case m of + S.Yield (mL, mR) s' + -> do !iL' <- case mL of + Nothing -> return iL + Just xL -> do GM.unsafeWrite vecL iL xL + return (iL + 1) + + !iR' <- case mR of + Nothing -> return iR + Just xR -> do GM.unsafeWrite vecR iR xR + return (iR + 1) + + go sPEC iL' iR' s' + + S.Skip s' + -> go sPEC iL iR s' + + S.Done + -> return ( GM.unsafeSlice 0 iL vecL + , GM.unsafeSlice 0 iR vecR) + in go S.SPEC 0 0 s0 +{-# INLINE_STREAM unstreamToMVector2_max #-} + + +------------------------------------------------------------------------------- +-- | 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 #-} + + +------------------------------------------------------------------------------- +-- | 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 #-} + + +-- | 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) + +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 #-} + + +-- 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 #-} + + in go S.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 S.SPEC vec0 0 nStart s0 +{-# INLINE_STREAM unchainToMVector_unknown #-} + diff --git a/Data/Repa/Vector/Unboxed.hs b/Data/Repa/Vector/Unboxed.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Vector/Unboxed.hs @@ -0,0 +1,337 @@ + +module Data.Repa.Vector.Unboxed + ( -- * Conversion + chainOfVector + , unchainToVector + , unchainToMVector + + -- * Generators + , ratchet + + -- * Extract + , extract + + -- * Merging + , merge + + -- * Splitting + , findSegments + , findSegmentsFrom + , diceSep + + -- * Padding + , padForward + + + -- * Scanning + , scanMaybe + + -- * Grouping + , groupsBy + + -- * Folding + , 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" + + +------------------------------------------------------------------------------- +-- | Produce a chain from a generic vector. +chainOfVector + :: (Monad m, Unbox a) + => Vector a -> Chain m Int a +chainOfVector = G.chainOfVector +{-# INLINE chainOfVector #-} + + +-- | Compute a chain into a vector. +unchainToVector + :: (PrimMonad m, Unbox a) + => C.Chain m s a -> m (Vector a, s) +unchainToVector = G.unchainToVector +{-# INLINE unchainToVector #-} + + +-- | Compute a chain into a mutable vector. +unchainToMVector + :: (PrimMonad m, Unbox a) + => C.Chain m s a + -> m (MVector (PrimState m) a, s) +unchainToMVector = G.unchainToMVector +{-# INLINE unchainToMVector #-} + + + +------------------------------------------------------------------------------- +-- | 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 :: 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') +{-# 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 :: Unbox a + => (Int -> a) -- ^ Function to get elements from the source. + -> U.Vector (Int, Int) -- ^ Segment starts and lengths. + -> U.Vector a -- ^ Result elements. + +extract get vStartLen + = G.unstream $ extractS get $ G.stream vStartLen +{-# INLINE extract #-} + + +------------------------------------------------------------------------------- +-- | 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. + -> (k -> a -> c) -- ^ Handle a left value without a right value. + -> (k -> b -> 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. + +merge fBoth fLeft fRight vA vB + = G.unstream + $ mergeS fBoth fLeft fRight + (G.stream vA) + (G.stream vB) +{-# INLINE merge #-} + + +------------------------------------------------------------------------------- +-- | 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 + :: (Unbox a, Unbox b) + => (s -> a -> (s, Maybe b)) -- ^ Worker function. + -> s -- ^ Initial state for scan. + -> 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 +{-# 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 + :: Unbox a + => (a -> a -> Bool) -- ^ Comparison function. + -> Maybe (a, Int) -- ^ Starting element and count. + -> 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 +{-# 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 + :: U.Unbox a + => (a -> Bool) -- ^ Predicate to check for start of segment. + -> (a -> Bool) -- ^ Predicate to check for end of segment. + -> 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 +{-# 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 + :: (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. + -> (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 +{-# INLINE findSegmentsFrom #-} + + +------------------------------------------------------------------------------- +-- | Dice a vector stream into rows and columns. +-- +diceSep :: Unbox a + => (a -> Bool) -- ^ Detect the end of a column. + -> (a -> Bool) -- ^ Detect the end of a row. + -> U.Vector a + -> (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 +{-# INLINE diceSep #-} + + +------------------------------------------------------------------------------- +-- | 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 :: (Unbox n, Unbox a, Unbox 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. + -> U.Vector (n, Int) -- ^ Segment names and lengths. + -> 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) +{-# 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 + :: (Unbox k, Unbox v, Ord k) + => (k -> k) -- ^ Successor function. + -> U.Vector (k, v) -- ^ Input keys and values. + -> U.Vector (k, v) + +padForward ksucc vec + = G.unstream + $ padForwardS ksucc + $ G.stream vec +{-# INLINE padForward #-} + diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,25 @@ +Copyright (c) 2014-2015, The Repa Development Team + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +- Redistributions of source code must retain the above copyright notice, + this list of conditions and the following disclaimer. + +- Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +- The names of the copyright holders may not be used to endorse or promote + products derived from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, +INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND +FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE +COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, +OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE +OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF +ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/include/repa-stream.h b/include/repa-stream.h new file mode 100644 --- /dev/null +++ b/include/repa-stream.h @@ -0,0 +1,10 @@ + +#define PHASE_FLOW [3] +#define PHASE_ARRAY [2] +#define PHASE_STREAM [1] +#define PHASE_INNER [0] + +#define INLINE_FLOW INLINE PHASE_FLOW +#define INLINE_ARRAY INLINE PHASE_ARRAY +#define INLINE_STREAM INLINE PHASE_STREAM +#define INLINE_INNER INLINE PHASE_INNER diff --git a/repa-stream.cabal b/repa-stream.cabal new file mode 100644 --- /dev/null +++ b/repa-stream.cabal @@ -0,0 +1,67 @@ +Name: repa-stream +Version: 4.0.0.1 +License: BSD3 +License-file: LICENSE +Author: The Repa Development Team +Maintainer: Ben Lippmeier +Build-Type: Simple +Cabal-Version: >=1.6 +Stability: experimental +Category: Data Structures +Homepage: http://repa.ouroborus.net +Bug-reports: repa@ouroborus.net +Description: Stream functions not present in the vector library. +Synopsis: Stream functions not present in the vector library. + +source-repository head + type: git + location: https://github.com/DDCSF/repa.git + +Library + build-Depends: + base == 4.7.*, + vector == 0.10.*, + primitive == 0.5.4.*, + mtl == 2.2.* + + exposed-modules: + Data.Repa.Chain + Data.Repa.Option + Data.Repa.Stream + Data.Repa.Vector.Generic + Data.Repa.Vector.Unboxed + + + other-modules: + Data.Repa.Chain.Base + Data.Repa.Chain.Scan + Data.Repa.Chain.Weave + Data.Repa.Chain.Folds + + Data.Repa.Stream.Extract + Data.Repa.Stream.Ratchet + Data.Repa.Stream.Segment + Data.Repa.Stream.Dice + Data.Repa.Stream.Pad + Data.Repa.Stream.Merge + + include-dirs: + include + + install-includes: + repa-stream.h + + ghc-options: + -Wall -fno-warn-missing-signatures + -O2 + + extensions: + CPP + NoMonomorphismRestriction + ExistentialQuantification + BangPatterns + FlexibleContexts + PatternGuards + MultiWayIf + +