diff --git a/Data/Repa/Flow.hs b/Data/Repa/Flow.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow.hs @@ -0,0 +1,209 @@ +{-# OPTIONS -fno-warn-unused-imports #-} +-- | +-- +-- = Getting Started +-- +-- A flow consists of a bundle of individual streams. Here we create +-- a bundle of two streams, using different files for each. Data will +-- be read in chunks, using the default chunk size of 64kBytes. +-- +-- @ +-- > import Data.Repa.Flow +-- > import Data.Repa.Flow.Default.Debug +-- > ws <- fromFiles [\"\/usr\/share\/dict\/words\", \"\/usr\/share\/dict\/cracklib-small\"] sourceLines +-- @ +-- +-- Show the first few elements of the first chunk of the first file. +-- +-- @ +-- > more 0 ws +-- Just [\"A\",\"A's\",\"AA's\",\"AB's\",\"ABM's\",\"AC's\",\"ACTH's\",\"AI's\" ...] +-- @ +-- +-- The `more` function is helpful for debugging. It pulls a whole chunk from a +-- source, displays the requested number of elements from the front of it, then +-- discards the rest. In production code you could use `head_i` to split a few +-- elements from a stream while retaining the rest. +-- +-- Use `more'` to show more elements at a time. We've already pulled the first chunk, +-- so here are the first 100 elements from the second chunk: +-- +-- @ +-- > more' 0 100 ws +-- Just [\"Jubal\",\"Judah\",\"Judaic\",\"Judaism\",\"Judaism's\",\"Judaisms\",\"Judas\" ...] +-- @ +-- +-- Use `moret` to display elements in tabular form. Here are the first few elements of +-- the second stream in the bundle: +-- +-- @ +-- > moret 1 ws +-- "10th" +-- "1st" +-- "2nd" +-- "3rd" +-- "4th" +-- "5th" +-- ... +-- @ +-- +-- Lets convert the characters to upper-case. +-- +-- @ +-- > import Data.Char +-- > up <- map_i B (mapS U toUpper) ws +-- > more 0 up +-- Just [\"UTOPIAN\",\"UTOPIAN'S\",\"UTOPIANS\",\"UTOPIAS\",\"UTRECHT\" ...] +-- @ +-- +-- The `B` and `U` are `Layout` names that indicate how the chunks for the +-- result streams should be arranged in memory. In this case the chunks +-- are `B`-oxed arrays of `U`-nboxed arrays of characters. Other useful +-- layouts are `F` which stores data in foreign memory, and `N` for nested +-- arrays. +-- +-- Flows are data-parallel, which means operators like `map_i` apply to all +-- streams in the bundle. The second stream has been converted to upper-case +-- as well: +-- +-- @ +-- > more 1 up +-- Just [\"BROWNER\",\"BROWNEST\",\"BROWNIAN\",\"BROWNIE\",\"BROWNIE'S\" ...] +-- @ +-- +-- Lets write out the data to some files. There are two streams in the bundle, +-- so open a file for each stream: +-- +-- @ +-- > out <- toFiles ["out1.txt", "out2.txt"] $ sinkLines B U +-- @ +-- +-- Note that the @ws@ and @up@ we used before were bundles of stream +-- `Sources` whereas @out@ is a bundle of stream `Sinks`. When we used +-- the `map_i` operator before the @_i@ (input) suffix indicates that +-- this is a transformer of `Sources`. There is a related `map_o` +-- (output) operator for `Sinks`. +-- +-- Now that we have a bundle of `Sources`, and some matching `Sinks`, +-- we can `drainS` all of the data from the former into the latter. +-- +-- @ +-- > drainS up out +-- @ +-- +-- At this point we can run an external shell command to check the output. +-- +-- @ +-- > :! head out1.txt +-- BEARSKIN'S +-- BEARSKINS +-- BEAST +-- BEAST'S +-- BEASTLIER +-- BEASTLIEST +-- BEASTLINESS +-- BEASTLINESS'S +-- BEASTLY +-- BEASTLY'S +-- @ +-- +-- = Performance +-- +-- Althogh @repa-flow@ can be used productively in the ghci REPL, +-- performance won't be great because you will be running unspecialised, +-- polymorphic code. For best results you should write a complete +-- program and compile it with @ghc -fllvm -O2 Main.hs@. +-- +module Data.Repa.Flow + ( -- * Flow types + Sources + , Sinks + , Flow + , sourcesArity + , sinksArity + + -- * States and Arrays + , module Data.Repa.Flow.States + , module Data.Repa.Eval.Array + , module Data.Repa.Array + , module Data.Repa.Array.Material + + -- * Evaluation + , drainS + , drainP + + -- * Conversion + , fromList, fromLists + , toList1, toLists1 + + -- * Finalizers + , finalize_i, finalize_o + + -- * Flow Operators + -- ** Mapping + -- | If you want to work on a chunk at a time then use + -- `Data.Repa.Flow.Generic.map_i` and + -- `Data.Repa.Flow.Generic.map_o` from "Data.Repa.Flow.Generic". + , map_i, map_o + + -- ** Connecting + , dup_oo + , dup_io + , dup_oi + , connect_i + + -- ** Watching + , watch_i, watch_o + , trigger_o + + -- ** Ignorance + , discard_o + , ignore_o + + -- ** Splitting + , head_i + + -- ** Grouping + , groups_i + , groupsBy_i + , GroupsDict + + -- ** Folding + , foldlS, foldlAllS + , folds_i, FoldsDict + , foldGroupsBy_i, FoldGroupsDict + + -- * Flow I/O + , defaultChunkSize + + -- ** Buckets + , module Data.Repa.Flow.IO.Bucket + + -- ** Sourcing + , sourceCSV + , sourceTSV + , sourceRecords + , sourceLines + , sourceChars + , sourceBytes + + -- ** Sinking + , sinkChars + , sinkLines + , sinkBytes) +where +import Data.Repa.Flow.Default +import Data.Repa.Flow.Default.Debug +import Data.Repa.Flow.Default.IO +import Data.Repa.Flow.IO.Bucket +import Data.Repa.Flow.States + +import Data.Repa.Eval.Array + +import Data.Repa.Array + hiding (fromList, Index, GroupsDict, FoldsDict) + +import Data.Repa.Array.Material + hiding (fromLists) + + diff --git a/Data/Repa/Flow/Chunked.hs b/Data/Repa/Flow/Chunked.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow/Chunked.hs @@ -0,0 +1,61 @@ + +module Data.Repa.Flow.Chunked + ( module Data.Repa.Flow.States + + , Sources, Sinks + , Flow + + -- * Evaluation + , drainS + + -- * Conversion + , fromList + , fromLists + , toList1 + , toLists1 + + -- * Finalizers + , finalize_i, finalize_o + + -- * Flow Operators + -- ** Mapping + -- | If you want to work on a chunk at a time then use + -- `Data.Repa.Flow.Generic.map_i` and + -- `Data.Repa.Flow.Generic.map_o` from "Data.Repa.Flow.Generic". + , smap_i, smap_o + , szipWith_ii + + -- ** Splitting + , head_i + + -- ** Grouping + , groupsBy_i, GroupsDict + + -- ** Folding + , foldlS, foldlAllS + , folds_i, FoldsDict + + -- ** Watching + , watch_i, watch_o + , trigger_o + + -- ** Ignorance + , discard_o + , ignore_o) +where +import Data.Repa.Flow.Chunked.Base +import Data.Repa.Flow.Chunked.Map +import Data.Repa.Flow.Chunked.Fold +import Data.Repa.Flow.Chunked.Folds +import Data.Repa.Flow.Chunked.Groups +import Data.Repa.Flow.Chunked.Operator +import Data.Repa.Flow.States +import qualified Data.Repa.Flow.Generic as G +#include "repa-flow.h" + + +-- | Pull all available values from the sources and push them to the sinks. +drainS :: (Next i, Monad m) + => Sources i m r a -> Sinks i m r a -> m () +drainS = G.drainS +{-# INLINE drainS #-} diff --git a/Data/Repa/Flow/Chunked/Base.hs b/Data/Repa/Flow/Chunked/Base.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow/Chunked/Base.hs @@ -0,0 +1,169 @@ + +module Data.Repa.Flow.Chunked.Base + ( Sources, Sinks + , Flow + , Data.Repa.Flow.Chunked.Base.fromList + , fromLists + , toList1 + , toLists1 + , head_i + , finalize_i, finalize_o) +where +import qualified Data.Sequence as Q +import qualified Data.Foldable as Q +import Data.Repa.Flow.States +import Data.Repa.Array as A +import Data.Repa.Eval.Array as A +import qualified Data.Repa.Flow.Generic as G +import Control.Monad +import Prelude as P +#include "repa-flow.h" + + +-- | A bundle of sources, where the elements are chunked into arrays. +type Sources i m l e + = G.Sources i m (A.Array l e) + + +-- | A bundle of sinks, where the elements are chunked into arrays. +type Sinks i m l e + = G.Sinks i m (A.Array l e) + + +-- | Shorthand for common type classes. +type Flow i m l a + = (Ord i, Monad m, BulkI l a, States i m) + + +-- Conversion ----------------------------------------------------------------- +-- | Given an arity and a list of elements, yield sources that each produce all +-- the elements. +-- +-- * All elements are stuffed into a single chunk, and each stream is given +-- the same chunk. +fromList :: (States i m, A.TargetI l a) + => Name l -> i -> [a] -> m (Sources i m l a) +fromList nDst n xs + = G.fromList n [A.fromList nDst xs] +{-# INLINE fromList #-} + + +-- | Like `fromLists` but take a list of lists, where each of the inner +-- lists is packed into a single chunk. +fromLists :: (States i m, A.TargetI l a) + => Name l -> i -> [[a]] -> m (Sources i m l a) + +fromLists nDst n xs + = G.fromList n $ P.map (A.fromList nDst) xs +{-# INLINE fromLists #-} + + +-- | Drain a single source into a list of elements. +toList1 :: (States i m, A.BulkI l a) + => i -> Sources i m l a -> m [a] +toList1 i sources + = do chunks <- G.toList1 i sources + return $ P.concat $ P.map A.toList chunks +{-# INLINE toList1 #-} + + +-- | Drain a single source into a list of chunks. +toLists1 :: (States i m, A.BulkI l a) + => i -> Sources i m l a -> m [[a]] +toLists1 i sources + = do chunks <- G.toList1 i sources + return $ P.map A.toList chunks +{-# INLINE toLists1 #-} + + +-- | Split the given number of elements from the head of a source, +-- retrurning those elements in a list, and yielding a new source +-- for the rest. +-- +-- * We pull /whole chunks/ from the source stream until we have +-- at least the desired number of elements. The leftover elements +-- in the final chunk are visible in the result `Sources`. +-- +head_i :: (States i m, A.Windowable l a, A.Index l ~ Int) + => Int -> Sources i m l a -> i -> m ([a], Sources i m l a) + +head_i len s0 i + = do + (s1, s2) <- G.connect_i s0 + + let G.Sources n pull_chunk = s1 + + -- Pull chunks from the source until we have enough elements to return. + refsList <- newRefs n Q.empty + refsChunk <- newRefs n Nothing + + let loop_takeList1 !has !acc !mchunk + | has >= len + = do writeRefs refsList i acc + writeRefs refsChunk i mchunk + + | otherwise + = pull_chunk i eat_toList eject_toList + where + eat_toList x + = loop_takeList1 + (has + A.length x) + (acc Q.>< (Q.fromList $ A.toList x)) + (Just x) + + eject_toList + = do writeRefs refsList i acc + writeRefs refsChunk i mchunk + + {-# INLINE loop_takeList1 #-} + + loop_takeList1 0 Q.empty Nothing + + -- Split off the required number of elements. + has <- readRefs refsList i + mFinal <- readRefs refsChunk i + let (here, rest) = Q.splitAt len has + + -- As we've pulled whole chunks from the input stream, + -- we now prepend the remaining ones back on. + let start = Q.length has - Q.length rest + let stash = case mFinal of + Nothing -> [] + Just c -> [A.window start (Q.length rest) c] + + s2' <- G.prependOn_i (\i' -> i' == i) stash s2 + return (Q.toList here, s2') +{-# INLINE_FLOW head_i #-} + + +-- Finalizers ----------------------------------------------------------------- +-- | Attach a finalizer to a bundle of sources. +-- +-- For each stream in the bundle, the finalizer will be called the first +-- time a consumer of that stream tries to pull an element when no more +-- are available. +-- +-- The provided finalizer will be run after any finalizers already +-- attached to the source. +-- +finalize_i + :: States i m + => (i -> m ()) + -> Sources i m l a -> m (Sources i m l a) +finalize_i = G.finalize_i +{-# INLINE finalize_i #-} + + +-- | Attach a finalizer to a bundle of sinks. +-- +-- The finalizer will be called the first time the stream is ejected. +-- +-- The provided finalizer will be run after any finalizers already +-- attached to the sink. +-- +finalize_o + :: States i m + => (i -> m ()) + -> Sinks i m l a -> m (Sinks i m l a) +finalize_o = G.finalize_o +{-# INLINE finalize_o #-} diff --git a/Data/Repa/Flow/Chunked/Folds.hs b/Data/Repa/Flow/Chunked/Folds.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow/Chunked/Folds.hs @@ -0,0 +1,210 @@ + +module Data.Repa.Flow.Chunked.Folds + ( folds_i + , FoldsDict) +where +import Data.Repa.Flow.Chunked.Base +import Data.Repa.Flow.States +import Data.Repa.Fusion.Unpack +import Data.Repa.Option +import Data.Repa.Array as A hiding (FoldsDict) +import Data.Repa.Eval.Array as A +import qualified Data.Repa.Flow.Generic as G +#include "repa-flow.h" + + +-- | Dictionaries needed to perform a segmented fold. +type FoldsDict i m lSeg tSeg lElt tElt lGrp tGrp lRes tRes n a b + = ( States i m + , Windowable lSeg (n, Int), Windowable lElt a + , BulkI lSeg (n, Int) + , BulkI lElt a + , BulkI lGrp n + , BulkI lRes b + , TargetI lElt a + , TargetI lGrp n + , TargetI lRes b + , Unpack (IOBuffer lGrp n) tGrp + , Unpack (IOBuffer lRes b) tRes) + + +-- Folds ---------------------------------------------------------------------- +-- | Segmented fold over vectors of segment lengths and input values. +folds_i :: FoldsDict i m lSeg tSeg lElt tElt lGrp tGrp lRes tRes n a b + => Name lGrp -- ^ Layout for group names. + -> Name lRes -- ^ Layout for fold results. + -> (a -> b -> b) -- ^ Worker function. + -> b -- ^ Initial state when folding each segment. + -> Sources i m lSeg (n, Int) -- ^ Segment lengths. + -> Sources i m lElt a -- ^ Input elements to fold. + -> m (Sources i m (T2 lGrp lRes) (n, b)) -- ^ Result elements. + +folds_i _ _ f z sLens@(G.Sources nLens _) + sVals@(G.Sources nVals _) + = do + -- Arity of the result bundle is the minimum of the two inputs. + let nFolds = min nLens nVals + + -- Refs to hold partial fold states between chunks. + refsState <- newRefs nFolds None3 + + -- Refs to hold the current chunk of lengths data for each stream. + refsNameLens <- newRefs nFolds Nothing + + -- Refs to hold the current chunk of vals data for each stream. + refsVals <- newRefs nFolds Nothing + refsValsDone <- newRefs nFolds False + + let pull_folds i eat eject + = do mNameLens <- folds_loadChunkNameLens sLens refsNameLens i + mVals <- folds_loadChunkVals sVals refsVals refsValsDone i + + case (mNameLens, mVals) of + -- If we couldn't get a chunk for both sides then we can't + -- produce anymore results, and the merge is done. + (Nothing, _) -> eject + (_, Nothing) -> eject + + -- We've got a chunk for both sides, time to do some work. + (Just cNameLens, Just cVals) + -> cNameLens `seq` cVals `seq` + do + mState <- readRefs refsState i + + let (cResults, sFolds) + = A.foldsWith name name f z + (fromOption3 mState) cNameLens cVals + + folds_update + refsState refsNameLens refsVals i + cNameLens cVals sFolds + + valsDone <- readRefs refsValsDone i + + -- If we're not producing output while we still + -- have segment lengths then we're done. + if A.length cResults == 0 + && A.length cNameLens >= 0 + && valsDone + then eject + else eat cResults + {-# INLINE pull_folds #-} + + return $ G.Sources nFolds pull_folds +{-# INLINE_FLOW folds_i #-} + + +-- Load the current chunk of lengths data. +-- If we already have one in the state then use that, +-- otherwise try to pull a new chunk from the source. +folds_loadChunkNameLens + :: States i m + => Sources i m l1 (n, Int) + -> Refs i m (Maybe (Array l1 (n, Int))) + -> i + -> m (Maybe (Array l1 (n, Int))) + +folds_loadChunkNameLens (G.Sources _ pullLens) refsLens i + = do mChunkLens <- readRefs refsLens i + case mChunkLens of + Nothing + -> let eatLens_folds chunk + = writeRefs refsLens i (Just chunk) + {-# INLINE eatLens_folds #-} + + ejectLens_folds = return () + {-# INLINE ejectLens_folds #-} + + in do + pullLens i eatLens_folds ejectLens_folds + readRefs refsLens i + + jc@(Just _) + -> return jc +{-# NOINLINE folds_loadChunkNameLens #-} +-- NOINLINE as this doesn't need to be specialized, +--- and we want to hide the case from the simplifier. + + +-- Grab the current chunk of values data. +-- If we already have one in the state then use that, +-- otherwise try to pull a new chunk from the source. +folds_loadChunkVals + :: (States i m, TargetI l2 a) + => Sources i m l2 a + -> Refs i m (Maybe (Array l2 a)) + -> Refs i m Bool + -> i + -> m (Maybe (Array l2 a)) + +folds_loadChunkVals (G.Sources _ pullVals) refsVals refsValsDone i + = do mChunkVals <- readRefs refsVals i + case mChunkVals of + Nothing + -> let eatVals_folds chunk + = writeRefs refsVals i (Just chunk) + {-# INLINE eatVals_folds #-} + + -- When there are no more values then shim in an + -- empty chunk so that we can keep calling A.folds + -- this is needed when there are zero lengthed + -- segments on the end of the stream + ejectVals_folds + = do writeRefs refsVals i (Just $ A.fromList name []) + writeRefs refsValsDone i True + {-# INLINE ejectVals_folds #-} + + in do + pullVals i eatVals_folds ejectVals_folds + readRefs refsVals i + + jc@(Just _) + -> return jc +{-# NOINLINE folds_loadChunkVals #-} +-- NOINLINE as this doesn't need to be specialized, +-- and we want to hide the case from the simplifier. + + +folds_update + :: ( States i m + , Windowable l1 (n, Int), Windowable l2 a + , A.Index l1 ~ Int, A.Index l2 ~ Int) + => Refs i m (Option3 n Int b) + -> Refs i m (Maybe (Array l1 (n, Int))) + -> Refs i m (Maybe (Array l2 a)) + -> i + -> Array l1 (n, Int) + -> Array l2 a + -> Folds Int Int n a b + -> m () + +folds_update refsState refsLens refsVals i cLens cVals sFolds + = do + -- Remember state for the final segment. + writeRefs refsState i + $ case _nameSeg sFolds of + Some n -> Some3 n (_lenSeg sFolds) (_valSeg sFolds) + None -> None3 + + -- Slice down the lengths chunk to just the elements + -- that we haven't already consumed. If we've consumed + -- them all then clear the chunk reference so a new one + -- will be loaded the next time around. + let !posLens = _stateLens sFolds + let !nLensRemain = A.length cLens - posLens + writeRefs refsLens i + $ if nLensRemain <= 0 + then Nothing + else Just $ A.window posLens nLensRemain cLens + + -- Likewise for the values chunk. + let !posVals = _stateVals sFolds + let !nValsRemain = A.length cVals - posVals + writeRefs refsVals i + $ if nValsRemain <= 0 + then Nothing + else Just $ A.window posVals nValsRemain cVals +{-# NOINLINE folds_update #-} +-- NOINLINE because it is only called once per chunk +-- and does not need to be specialised. + diff --git a/Data/Repa/Flow/Chunked/Groups.hs b/Data/Repa/Flow/Chunked/Groups.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow/Chunked/Groups.hs @@ -0,0 +1,74 @@ + +module Data.Repa.Flow.Chunked.Groups + ( groupsBy_i + , GroupsDict) +where +import Data.Repa.Flow.Chunked.Base +import Data.Repa.Flow.States +import Data.Repa.Fusion.Unpack +import Data.Repa.Array as A hiding (GroupsDict) +import Data.Repa.Eval.Array as A +import qualified Data.Repa.Flow.Generic as G +#include "repa-flow.h" + + +-- | Dictionaries needed to perform a grouping. +type GroupsDict i m lVal lGrp tGrp lLen tLen a + = ( Flow i m lVal a, A.Index lVal ~ Int + , TargetI lGrp a + , TargetI lLen Int + , Unpack (IOBuffer lGrp a) tGrp + , Unpack (IOBuffer lLen Int) tLen) + + +-- Grouping ------------------------------------------------------------------- +-- | From a stream of values which has consecutive runs of idential values, +-- produce a stream of the lengths of these runs. +-- +-- @ +-- groupsBy (==) [4, 4, 4, 3, 3, 1, 1, 1, 4] +-- => [3, 2, 3, 1] +-- @ +-- +groupsBy_i + :: GroupsDict i m lVal lGrp tGrp lLen tLen a + => Name lGrp -- ^ Layout for group names. + -> Name lLen -- ^ Layout for group lengths. + -> (a -> a -> Bool) -- ^ Whether successive elements should be grouped. + -> Sources i m lVal a -- ^ Source values. + -> m (Sources i m (T2 lGrp lLen) (a, Int)) + +groupsBy_i _ _ f (G.Sources n pull_chunk) + = do + -- Refs to hold partial counts between chunks. + refs <- newRefs n Nothing + + let pull_groupsBy i eat eject + = pull_chunk i eat_groupsBy eject_groupsBy + where + -- Process a chunk from the a source stream, + -- using the current state we have for that stream. + eat_groupsBy chunk + = do state <- readRefs refs i + let (segs, state') = A.groupsWith name name f state chunk + writeRefs refs i state' + eat segs + {-# INLINE eat_groupsBy #-} + + -- When there are no more chunks of source data we still + -- need to pass on the last count we have stored in the + -- state. + eject_groupsBy + = do state <- readRefs refs i + case state of + Nothing -> eject + Just seg + -> do writeRefs refs i Nothing + eat (A.fromList name [seg]) + {-# INLINE eject_groupsBy #-} + {-# INLINE pull_groupsBy #-} + + return $ G.Sources n pull_groupsBy +{-# INLINE_FLOW groupsBy_i #-} + + diff --git a/Data/Repa/Flow/Chunked/IO.hs b/Data/Repa/Flow/Chunked/IO.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow/Chunked/IO.hs @@ -0,0 +1,67 @@ + +-- | Input and Output for Chunked Flows. +-- +-- Most functions in this module are re-exports of the ones from +-- "Data.Repa.Flow.Generic.IO", but using the `Sources` and `Sinks` +-- type synonyms for chunked flows. +-- +module Data.Repa.Flow.Chunked.IO + ( -- * Sourcing + sourceRecords + , sourceChars + , sourceBytes + + -- * Sinking + , sinkChars + , sinkBytes) +where +import Data.Repa.Flow.IO.Bucket +import Data.Repa.Flow.Chunked.Base +import Data.Repa.Array as A +import Data.Repa.Array.Material as A +import qualified Data.Repa.Flow.Generic.IO as G +import Data.Word +#include "repa-flow.h" + + +-- | Like `fileSourceRecords`, but taking an existing file handle. +sourceRecords + :: BulkI l Bucket + => Integer -- ^ Size of chunk to read in bytes. + -> (Word8 -> Bool) -- ^ Detect the end of a record. + -> IO () -- ^ Action to perform if we can't get a whole record. + -> Array l Bucket -- ^ File handles. + -> IO (Sources Int IO N (Array F Word8)) +sourceRecords = G.sourceRecords +{-# INLINE sourceRecords #-} + + +-- | Read 8-bit ASCII characters from some files, using the given chunk length. +sourceChars + :: BulkI l Bucket + => Integer -> Array l Bucket -> IO (Sources Int IO F Char) +sourceChars = G.sourceChars +{-# INLINE sourceChars #-} + + +-- | Read data from some files, using the given chunk length. +sourceBytes + :: BulkI l Bucket + => Integer -> Array l Bucket -> IO (Sources Int IO F Word8) +sourceBytes = G.sourceBytes +{-# INLINE sourceBytes #-} + + +-- | Write 8-bit ASCII characters to the given file handles. +sinkChars :: BulkI l Bucket + => Array l Bucket -> IO (Sinks Int IO F Char) +sinkChars = G.sinkChars +{-# INLINE sinkChars #-} + + +-- | Write chunks of data to the given file handles. +sinkBytes :: BulkI l Bucket + => Array l Bucket -> IO (Sinks Int IO F Word8) +sinkBytes = G.sinkBytes +{-# INLINE sinkBytes #-} + diff --git a/Data/Repa/Flow/Chunked/Map.hs b/Data/Repa/Flow/Chunked/Map.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow/Chunked/Map.hs @@ -0,0 +1,97 @@ + +module Data.Repa.Flow.Chunked.Map + ( smap_i, smap_o + , szipWith_ii) +where +import Data.Repa.Flow.Chunked.Base +import Data.Repa.Flow.States +import Data.Repa.Array as A +import Data.Repa.Eval.Array as A +import qualified Data.Repa.Flow.Generic as G +#include "repa-flow.h" + + +-- | Map a function over elements pulled from a source. +smap_i :: (Flow i m l1 a, A.TargetI l2 b) + => (i -> a -> b) -> Sources i m l1 a -> m (Sources i m l2 b) +smap_i f s0 = G.smap_i (\i c -> A.computeS name $ A.map (f i) c) s0 +{-# INLINE smap_i #-} + + +-- | Map a function over elements pushed into a sink. +smap_o :: (Flow i m l1 a, A.TargetI l2 b) + => (i -> a -> b) -> Sinks i m l2 b -> m (Sinks i m l1 a) +smap_o f s0 = G.smap_o (\i c -> A.computeS name $ A.map (f i) c) s0 +{-# INLINE smap_o #-} + + +-- | Combine the elements of two flows with the given function. +szipWith_ii + :: ( Ord i, States i m + , BulkI lSrc1 a, BulkI lSrc2 b + , TargetI lDst c + , Windowable lSrc1 a, Windowable lSrc2 b) + => Name lDst + -> (i -> a -> b -> c) + -> Sources i m lSrc1 a -> Sources i m lSrc2 b + -> m (Sources i m lDst c) + +szipWith_ii nDst f (G.Sources nA pullA) (G.Sources nB pullB) + = do + let nC = min nA nB + + -- Refs to hold leftover pieces of chunks. + bitsA <- newRefs nC Nothing + bitsB <- newRefs nC Nothing + + let pullC i eatC ejectC + | not $ check i nC = ejectC + | otherwise = getA + where + getA + = do mA <- readRefs bitsA i + case mA of + Just chunkA -> getB chunkA + Nothing -> pullA i getB ejectC + {-# INLINE getA #-} + + getB chunkA + = do mB <- readRefs bitsB i + case mB of + Just chunkB -> zipAB chunkA chunkB + Nothing -> pullB i (zipAB chunkA) ejectC + {-# INLINE getB #-} + + zipAB chunkA chunkB + = do + let !lenA = A.length chunkA + let !lenB = A.length chunkB + let !lenC = min lenA lenB + + -- Split the chunks into the bits that we will zip + -- in this round, and the bits that we will leave + -- until later. + let !hereA = A.window 0 lenC chunkA + let !restA = A.window lenC (lenA - lenC) chunkA + + let !hereB = A.window 0 lenC chunkB + let !restB = A.window lenC (lenB - lenC) chunkB + + -- Zip the common parts we have now. + let Just !hereC = A.map2S nDst (f i) hereA hereB + + (if A.length restA > 0 + then writeRefs bitsA i (Just restA) + else writeRefs bitsA i Nothing) + + (if A.length restB > 0 + then writeRefs bitsB i (Just restB) + else writeRefs bitsB i Nothing) + + eatC hereC + {-# INLINE zipAB #-} + {-# INLINE pullC #-} + + return $ G.Sources nC pullC +{-# INLINE_FLOW szipWith_ii #-} + diff --git a/Data/Repa/Flow/Chunked/Operator.hs b/Data/Repa/Flow/Chunked/Operator.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow/Chunked/Operator.hs @@ -0,0 +1,73 @@ + +-- | Operators for chunked flows. +-- +-- Most functions in this module are re-exports of the ones from +-- "Data.Repa.Flow.Generic.IO", but using the `Sources` and `Sinks` +-- type synonyms for chunked flows. +-- +module Data.Repa.Flow.Chunked.Operator + ( -- * Watching + watch_i, watch_o + , trigger_o + + -- * Ignorance + , discard_o + , ignore_o) +where +import Data.Repa.Flow.Chunked.Base +import Data.Repa.Array as A +import qualified Data.Repa.Flow.Generic as G +#include "repa-flow.h" + + +-- Watch ---------------------------------------------------------------------- +-- | Hook a monadic function to some sources, which will be passed every +-- chunk that is pulled from the result. +watch_i :: Monad m + => (i -> Array l a -> m ()) + -> Sources i m l a -> m (Sources i m l a) +watch_i = G.watch_i +{-# INLINE watch_i #-} + + +-- | Hook a monadic function to some sinks, which will be passed every +-- chunk that is pushed to the result. +watch_o :: Monad m + => (i -> Array l a -> m ()) + -> Sinks i m l a -> m (Sinks i m l a) + +watch_o = G.watch_o +{-# INLINE watch_o #-} + + +-- | Like `watch_o` but discard the incoming chunks after they are passed +-- to the function. +trigger_o :: Monad m + => i -> (i -> Array l a -> m ()) -> m (Sinks i m l a) +trigger_o = G.trigger_o +{-# INLINE trigger_o #-} + + +-- Ignorance ------------------------------------------------------------------ +-- | A sink that ignores all incoming data. +-- +-- This sink is non-strict in the chunks. +-- Haskell tracing thunks attached to the chunks will *not* be demanded. +-- +ignore_o :: Monad m => i -> m (Sinks i m l a) +ignore_o = G.ignore_o +{-# INLINE ignore_o #-} + + +-- | Yield a bundle of sinks of the given arity that drops all data on the +-- floor. +-- +-- * The sinks is strict in the *chunks*, so they are demanded before being +-- discarded. Haskell debugging thunks attached to the chunks will be +-- demanded, but thunks attached to elements may not be -- depending on +-- whether the chunk representation is strict in the elements. +-- +discard_o :: Monad m => i -> m (Sinks i m l a) +discard_o = G.discard_o +{-# INLINE discard_o #-} + diff --git a/Data/Repa/Flow/Default.hs b/Data/Repa/Flow/Default.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow/Default.hs @@ -0,0 +1,516 @@ + +-- | This module defines the default specialisation of flows that +-- appears in "Data.Repa.Flow". Each stream in the bundle is indexed +-- by a single integer, and stream state is stored using the IO monad. +-- +module Data.Repa.Flow.Default + ( -- * Flow types + Sources + , Sinks + , Flow + , sourcesArity + , sinksArity + + -- * States and Arrays + , module Data.Repa.Flow.States + , module Data.Repa.Eval.Array + , module Data.Repa.Array + , module Data.Repa.Array.Material + + -- * Evaluation + , drainS + , drainP + + -- * Conversion + , fromList, fromLists + , toList1, toLists1 + + -- * Finalizers + , finalize_i, finalize_o + + -- * Flow Operators + -- ** Mapping + -- | If you want to work on a chunk at a time then use + -- `Data.Repa.Flow.Generic.map_i` and + -- `Data.Repa.Flow.Generic.map_o` from "Data.Repa.Flow.Generic". + , map_i, map_o + + -- ** Connecting + , dup_oo + , dup_io + , dup_oi + , connect_i + + -- ** Watching + , watch_i, watch_o + , trigger_o + + -- ** Ignorance + , discard_o + , ignore_o + + -- ** Splitting + , head_i + + -- ** Grouping + , groups_i + , groupsBy_i + , GroupsDict + + -- ** Folding + , foldlS, foldlAllS + , folds_i, FoldsDict + , foldGroupsBy_i, FoldGroupsDict) +where +import Data.Repa.Flow.States +import Data.Repa.Eval.Array +import Data.Repa.Eval.Array as A + +import Data.Repa.Array + hiding (FoldsDict, GroupsDict, Index, fromList) + +import Data.Repa.Array as A + hiding (FoldsDict, GroupsDict, fromList) + +import Data.Repa.Array.Material hiding (fromLists) +import Data.Repa.Fusion.Unpack as A +import qualified Data.Repa.Flow.Chunked as C hiding (next) +import qualified Data.Repa.Flow.Generic as G hiding (next) +import Control.Monad +#include "repa-flow.h" + + +-- | A bundle of stream sources, where the elements of the stream +-- are chunked into arrays. +-- +-- The chunks have some `Layout` @l@ and contain elements of type @a@. +-- See "Data.Repa.Array" for the available layouts. +type Sources l a = C.Sources Int IO l a + + +-- | A bundle of stream sinks, where the elements of the stream +-- are chunked into arrays. +-- +type Sinks l a = C.Sinks Int IO l a + + +-- | Yield the number of streams in the bundle. +sourcesArity :: Sources l a -> Int +sourcesArity = G.sourcesArity + + +-- | Yield the number of streams in the bundle. +sinksArity :: Sinks l a -> Int +sinksArity = G.sinksArity + + +-- | Shorthand for common type classes. +type Flow l a = C.Flow Int IO l a + + +-- Evaluation ----------------------------------------------------------------- +-- | Pull all available values from the sources and push them to the sinks. +-- Streams in the bundle are processed sequentially, from first to last. +-- +-- * If the `Sources` and `Sinks` have different numbers of streams then +-- we only evaluate the common subset. +-- +drainS :: Sources l a -> Sinks l a -> IO () +drainS = G.drainS +{-# INLINE drainS #-} + + +-- | Pull all available values from the sources and push them to the sinks, +-- in parallel. We fork a thread for each of the streams and evaluate +-- them all in parallel. +-- +-- * If the `Sources` and `Sinks` have different numbers of streams then +-- we only evaluate the common subset. +-- +drainP :: Sources l a -> Sinks l a -> IO () +drainP = G.drainP +{-# INLINE drainP #-} + + +-- Conversion ----------------------------------------------------------------- +-- | Given an arity and a list of elements, yield sources that each produce all +-- the elements. +-- +-- * All elements are stuffed into a single chunk, and each stream is given +-- the same chunk. +-- +fromList :: A.TargetI l a + => Name l -> Int -> [a] -> IO (Sources l a) +fromList l xs = C.fromList l xs +{-# INLINE fromList #-} + + +-- | Like `fromLists_i` but take a list of lists. Each each of the inner +-- lists is packed into a single chunk. +fromLists :: A.TargetI l a + => Name l -> Int -> [[a]] -> IO (Sources l a) +fromLists nDst xss = C.fromLists nDst xss +{-# INLINE fromLists #-} + + +-- | Drain a single source from a bundle into a list of elements. +toList1 :: A.BulkI l a + => Int -> Sources l a -> IO [a] +toList1 ix s + | ix >= G.sourcesArity s = return [] + | otherwise = C.toList1 ix s +{-# INLINE toList1 #-} + + +-- | Drain a single source from a bundle into a list of chunks. +toLists1 :: A.BulkI l a + => Int -> Sources l a -> IO [[a]] +toLists1 ix s + | ix >= G.sourcesArity s = return [] + | otherwise = C.toLists1 ix s +{-# INLINE toLists1 #-} + + +-- Finalizers ----------------------------------------------------------------- +-- | Attach a finalizer to some sources. +-- +-- * For a given source, the finalizer will be called the first time a +-- consumer of that source tries to pull an element when no more +-- are available. +-- +-- * The finalizer is given the index of the source that ended. +-- +-- * The finalizer will be run after any finalizers already attached +-- to the source. +-- +--- +-- TODO: make the finalizer run just the first time. +-- +finalize_i + :: (Int -> IO ()) + -> Sources l a -> IO (Sources l a) +finalize_i f s + = G.finalize_i f s +{-# INLINE finalize_i #-} + + +-- | Attach a finalizer to some sinks. +-- +-- * For a given sink, the finalizer will be called the first time +-- that sink is ejected. +-- +-- * The finalizer is given the index of the sink that was ejected. +-- +-- * The finalizer will be run after any finalizers already attached +-- to the sink. +-- +--- +-- TODO: make the finalizer run just the first time. +-- +finalize_o + :: (Int -> IO ()) + -> Sinks l a -> IO (Sinks l a) +finalize_o f k + = G.finalize_o f k +{-# INLINE finalize_o #-} + + +-- Mapping -------------------------------------------------------------------- +-- | Apply a function to all elements pulled from some sources. +map_i :: (Flow l1 a, A.TargetI l2 b) + => Name l2 -> (a -> b) -> Sources l1 a -> IO (Sources l2 b) +map_i _ f s = C.smap_i (\_ x -> f x) s +{-# INLINE map_i #-} + + +-- | Apply a function to all elements pushed to some sinks. +map_o :: (Flow l1 a, A.TargetI l2 b) + => Name l1 -> (a -> b) -> Sinks l2 b -> IO (Sinks l1 a) +map_o _ f s = C.smap_o (\_ x -> f x) s +{-# INLINE map_o #-} + + +-- Connecting ----------------------------------------------------------------- +-- | Send the same data to two consumers. +-- +-- Given two argument sinks, yield a result sink. +-- Pushing to the result sink causes the same element to be pushed to both +-- argument sinks. +dup_oo :: Sinks l a -> Sinks l a -> IO (Sinks l a) +dup_oo = G.dup_oo +{-# INLINE dup_oo #-} + + +-- | Send the same data to two consumers. +-- +-- Given an argument source and argument sink, yield a result source. +-- Pulling an element from the result source pulls from the argument source, +-- and pushes that element to the sink, as well as returning it via the +-- result source. +-- +dup_io :: Sources l a -> Sinks l a -> IO (Sources l a) +dup_io = G.dup_io +{-# INLINE dup_io #-} + + +-- | Send the same data to two consumers. +-- +-- Like `dup_io` but with the arguments flipped. +-- +dup_oi :: Sinks l a -> Sources l a -> IO (Sources l a) +dup_oi = G.dup_oi +{-# INLINE dup_oi #-} + + +-- | Connect an argument source to two result sources. +-- +-- Pulling from either result source pulls from the argument source. +-- Each result source only gets the elements pulled at the time, +-- so if one side pulls all the elements the other side won't get any. +-- +connect_i :: Sources l a -> IO (Sources l a, Sources l a) +connect_i = G.connect_i +{-# INLINE connect_i #-} + + +-- Watching ------------------------------------------------------------------- +-- | Hook a worker function to some sources, which will be passed every +-- chunk that is pulled from each source. +-- +-- * The worker is also passed the source index of the chunk that was pulled. +-- +watch_i :: (Int -> Array l a -> IO ()) + -> Sources l a -> IO (Sources l a) +watch_i f s = G.watch_i f s +{-# INLINE watch_i #-} + + +-- | Hook a worker function to some sinks, which will be passed every +-- chunk that is pushed to each sink. +-- +-- * The worker is also passed the source index of the chunk that was pushed. +-- +watch_o :: (Int -> Array l a -> IO ()) + -> Sinks l a -> IO (Sinks l a) +watch_o f k = G.watch_o f k +{-# INLINE watch_o #-} + + +-- | Create a bundle of sinks of the given arity that pass incoming chunks +-- to a worker function. +-- +-- * This is like `watch_o`, except that the incoming chunks are discarded +-- after they are passed to the worker function +-- +trigger_o :: Int -> (Int -> Array l a -> IO ()) + -> IO (Sinks l a) +trigger_o arity f + = G.trigger_o arity f +{-# INLINE trigger_o #-} + + +-- Ignorance ------------------------------------------------------------------ +-- | Create a bundle of sinks of the given arity that drop all data on the +-- floor. +-- +-- * The sinks is strict in the *chunks*, so they are demanded before being +-- discarded. +-- * Haskell debugging thunks attached to the chunks will be +-- demanded, but thunks attached to elements may not be -- depending on +-- whether the chunk representation is strict in the elements. +-- +discard_o :: Int -> IO (Sinks l a) +discard_o = G.discard_o +{-# INLINE discard_o #-} + + +-- | Create a bundle of sinks of the given arity that drop all data on the +-- floor. +-- +-- * As opposed to `discard_o` the sinks are non-strict in the chunks. +-- * Haskell debugging thunks attached to the chunks will *not* be +-- demanded. +-- +ignore_o :: Int -> IO (Sinks l a) +ignore_o = G.ignore_o +{-# INLINE ignore_o #-} + + +-- Splitting ------------------------------------------------------------------ +-- | Given a source index and a length, split the a list of that +-- length from the front of the source. Yields a new source for the +-- remaining elements. +-- +-- * We pull /whole chunks/ from the source stream until we have +-- at least the desired number of elements. The leftover elements +-- in the final chunk are visible in the result `Sources`. +-- +head_i :: (A.Windowable l a, A.Index l ~ Int) + => Int -> Int -> Sources l a -> IO (Maybe ([a], Sources l a)) +head_i ix len s + | ix >= G.sourcesArity s = return Nothing + | otherwise + = liftM Just $ C.head_i len s ix +{-# INLINE head_i #-} + + +-- Grouping ------------------------------------------------------------------- +-- | Scan through some sources to find runs of matching elements, +-- and count the lengths of those runs. +-- +-- @ +-- > import Data.Repa.Flow +-- > toList1 0 =<< groups_i U U =<< fromList U 1 "waabbbblle" +-- Just [(\'w\',1),(\'a\',2),(\'b\',4),(\'l\',2),(\'e\',1)] +-- @ +-- +groups_i + :: (GroupsDict lVal lGrp tGrp lLen tLen a, Eq a) + => Name lGrp -- ^ Layout of result groups. + -> Name lLen -- ^ Layout of result lengths. + -> Sources lVal a -- ^ Input elements. + -> IO (Sources (T2 lGrp lLen) (a, Int)) + -- ^ Starting element and length of groups. +groups_i nGrp nLen s + = groupsBy_i nGrp nLen (==) s +{-# INLINE groups_i #-} + + +-- | Like `groupsBy`, but take a function to determine whether two consecutive +-- values should be in the same group. +groupsBy_i + :: GroupsDict lVal lGrp tGrp lLen tLen a + => Name lGrp -- ^ Layout of result groups. + -> Name lLen -- ^ Layout of result lengths. + -> (a -> a -> Bool) -- ^ Fn to check if consecutive elements + -- are in the same group. + -> Sources lVal a -- ^ Input elements. + -> IO (Sources (T2 lGrp lLen) (a, Int)) + -- ^ Starting element and length of groups. +groupsBy_i nGrp nLen f s + = C.groupsBy_i nGrp nLen f s +{-# INLINE groupsBy_i #-} + + +-- | Dictionaries needed to perform a grouping. +type GroupsDict lVal lGrp tGrp lLen tLen a + = C.GroupsDict Int IO lVal lGrp tGrp lLen tLen a + + +-- Folding -------------------------------------------------------------------- +-- | Fold all the elements of each stream in a bundle, one stream after the +-- other, returning an array of fold results. +foldlS + :: ( A.Target lDst a, A.Index lDst ~ Int + , A.BulkI lSrc b) + => A.Name lDst -- ^ Layout for result. + -> (a -> b -> a) -- ^ Combining funtion. + -> a -- ^ Starting value. + -> Sources lSrc b -- ^ Input elements to fold. + -> IO (A.Array lDst a) + +foldlS n f z ss + = C.foldlS n f z ss +{-# INLINE foldlS #-} + + +-- | Fold all the elements of each stream in a bundle, one stream after the +-- other, returning an array of fold results. +foldlAllS + :: A.BulkI lSrc b + => (a -> b -> a) -- ^ Combining funtion. + -> a -- ^ Starting value. + -> Sources lSrc b -- ^ Input elements to fold. + -> IO a + +foldlAllS f z ss + = C.foldlAllS f z ss +{-# INLINE foldlAllS #-} + + +-- | Given streams of lengths and values, perform a segmented fold where +-- fold segments of values of the corresponding lengths are folded +-- together. +-- +-- @ +-- > import Data.Repa.Flow +-- > sSegs <- fromList U 1 [(\'a\', 1), (\'b\', 2), (\'c\', 4), (\'d\', 0), (\'e\', 1), (\'f\', 5 :: Int)] +-- > sVals <- fromList U 1 [10, 20, 30, 40, 50, 60, 70, 80, 90 :: Int] +-- > toList1 0 =<< folds_i U U (+) 0 sSegs sVals +-- Just [(\'a\',10),(\'b\',50),(\'c\',220),(\'d\',0),(\'e\',80)] +-- @ +-- +-- If not enough input elements are available to fold a complete segment +-- then no output is produced for that segment. However, trailing zero +-- length segments still produce the initial value for the fold. +-- +-- @ +-- > import Data.Repa.Flow +-- > sSegs <- fromList U 1 [(\'a\', 1), (\'b\', 2), (\'c\', 0), (\'d\', 0), (\'e\', 0 :: Int)] +-- > sVals <- fromList U 1 [10, 20, 30 :: Int] +-- > toList1 0 =<< folds_i U U (*) 1 sSegs sVals +-- Just [(\'a\',10),(\'b\',600),(\'c\',1),(\'d\',1),(\'e\',1)] +-- @ +-- +folds_i :: (FoldsDict lSeg tSeg lElt tElt lGrp tGrp lRes tRes n a b) + => Name lGrp -- ^ Layout for group names. + -> Name lRes -- ^ Layout for fold results. + -> (a -> b -> b) -- ^ Worker function. + -> b -- ^ Initial state when folding each segment. + -> Sources lSeg (n, Int) -- ^ Segment lengths. + -> Sources lElt a -- ^ Input elements to fold. + -> IO (Sources (T2 lGrp lRes) (n, b)) -- ^ Result elements. + +folds_i nGrp nRes f z sLen sVal + = C.folds_i nGrp nRes f z sLen sVal +{-# INLINE folds_i #-} + +-- | Dictionaries needed to perform a segmented fold. +type FoldsDict lSeg tSeg lElt tElt lGrp tGrp lRes tRes n a b + = C.FoldsDict Int IO lSeg tSeg lElt tElt lGrp tGrp lRes tRes n a b + + +-- | Combination of `groupsBy_i` and `folds_i`. We determine the the segment +-- lengths while performing the folds. +-- +-- Note that a SQL-like groupby aggregations can be performed using this +-- function, provided the data is pre-sorted on the group key. For example, +-- we can take the average of some groups of values: +-- +-- @ +-- > import Data.Repa.Flow +-- > sKeys <- fromList U 1 "waaaabllle" +-- > sVals <- fromList U 1 [10, 20, 30, 40, 50, 60, 70, 80, 90, 100 :: Double] +-- +-- > sResult \<- map_i U (\\(key, (acc, n)) -\> (key, acc / n)) +-- =\<\< foldGroupsBy_i U U (==) (\\x (acc, n) -> (acc + x, n + 1)) (0, 0) sKeys sVals +-- +-- > toList1 0 sResult +-- Just [10.0,35.0,60.0,80.0,100.0] +-- @ +-- +foldGroupsBy_i + :: ( FoldGroupsDict lSeg tSeg lVal tVal lGrp tGrp lRes tRes n a b) + => Name lGrp -- ^ Layout for group names. + -> Name lRes -- ^ Layout for fold results. + -> (n -> n -> Bool) -- ^ Fn to check if consecutive elements + -- are in the same group. + -> (a -> b -> b) -- ^ Worker function for the fold. + -> b -- ^ Initial when folding each segment. + -> Sources lSeg n -- ^ Names that determine groups. + -> Sources lVal a -- ^ Values to fold. + -> IO (Sources (T2 lGrp lRes) (n, b)) + +foldGroupsBy_i nGrp nRes pGroup f z sNames sVals + = do segLens <- groupsBy_i nGrp U pGroup sNames + folds_i nGrp nRes f z segLens sVals +{-# INLINE foldGroupsBy_i #-} + + +type FoldGroupsDict lSeg tSeg lElt tElt lGrp tGrp lRes tRes n a b + = ( A.BulkI lSeg n + , A.Material lElt a, A.Index lElt ~ Int + , A.Material lGrp n, A.Index lGrp ~ Int + , A.Material lRes b, A.Index lRes ~ Int + , Unpack (IOBuffer lGrp n) tGrp + , Unpack (IOBuffer lRes b) tRes) diff --git a/Data/Repa/Flow/Default/Debug.hs b/Data/Repa/Flow/Default/Debug.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow/Default/Debug.hs @@ -0,0 +1,94 @@ +{-# LANGUAGE OverlappingInstances, TypeSynonymInstances, FlexibleInstances #-} +module Data.Repa.Flow.Default.Debug + (-- * More + more, more' + + -- * More (tabular) + , moret, moret' + + -- * More (raw) + , morer, morer' + + -- * Nicer + , Nicer (..) + , Presentable (..)) +where +import Data.Repa.Nice.Present +import Data.Repa.Nice.Tabulate +import Data.Repa.Nice +import Data.Repa.Flow.Default hiding (next) +import qualified Data.Repa.Array as A +import Control.Monad +import Data.List as L +import Data.Text as T +import Prelude as P +#include "repa-flow.h" + + +------------------------------------------------------------------------------- +-- | Given a source index and a length, pull enough chunks from the source +-- to build a list of the requested length, and discard the remaining +-- elements in the final chunk. +-- +-- * This function is intended for interactive debugging. +-- If you want to retain the rest of the final chunk then use `head_i`. +-- +more :: (Windowable l a, A.Index l ~ Int, Nicer a) + => Int -- ^ Index of source in bundle. + -> Sources l a -- ^ Bundle of sources. + -> IO (Maybe [Nice a]) +more i ss = more' i 20 ss +{-# INLINE more #-} + + +-- | Like `more` but also specify now many elements you want. +more' :: (Windowable l a, A.Index l ~ Int, Nicer a) + => Int -> Int -> Sources l a -> IO (Maybe [Nice a]) +more' ix len s + = liftM (liftM (L.map nice . fst)) $ head_i ix len s +{-# INLINE_FLOW more' #-} + + +------------------------------------------------------------------------------- +-- | Like `more`, but print results in a tabular form to the console. +moret :: ( A.Windowable l a, A.Index l ~ Int + , Nicer [a], Presentable (Nice [a])) + => Int -- ^ Index of source in bundle. + -> Sources l a -- ^ Bundle of sources. + -> IO () + +moret i ss = moret' i 20 ss +{-# INLINE moret #-} + + +-- | Like `more'`, but print results in tabular form to the console. +moret' :: ( A.Windowable l a, A.Index l ~ Int + , Nicer [a], Presentable (Nice [a])) + => Int -> Int -> Sources l a -> IO () + +moret' ix len s + = do Just (vals, _) <- head_i ix len s + putStrLn $ T.unpack $ tabulate $ nice vals +{-# INLINE_FLOW moret' #-} + + +------------------------------------------------------------------------------- +-- | Like `more`, but show elements in their raw format. +morer :: (A.Windowable l a, A.Index l ~ Int) + => Int -- ^ Index of source in bundle. + -> Sources l a -- ^ Bundle of sources. + -> IO (Maybe [a]) + +morer i ss = morer' i 20 ss +{-# INLINE morer #-} + + +-- | Like `more'`, but show elements in their raw format. +morer' :: (A.Windowable l a, A.Index l ~ Int) + => Int -> Int -> Sources l a -> IO (Maybe [a]) +morer' ix len s + = liftM (liftM fst) $ head_i ix len s +{-# INLINE_FLOW morer' #-} + + + diff --git a/Data/Repa/Flow/Default/IO.hs b/Data/Repa/Flow/Default/IO.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow/Default/IO.hs @@ -0,0 +1,159 @@ + +-- | Read and write files. +-- +-- The functions in this module are wrappers for the ones in +-- "Data.Repa.Flow.Default.SizedIO" that use a default chunk size of +-- 64kBytes and just call `error` if the source file appears corruped. +module Data.Repa.Flow.Default.IO + ( defaultChunkSize + + -- * Buckets + , module Data.Repa.Flow.IO.Bucket + + -- * Sourcing + , sourceCSV + , sourceTSV + , sourceRecords + , sourceLines + , sourceChars + , sourceBytes + + -- * Sinking + , sinkChars + , sinkLines + , sinkBytes) +where +import Data.Repa.Flow.Default +import Data.Repa.Flow.IO.Bucket +import Data.Repa.Fusion.Unpack +import Data.Word +import qualified Data.Repa.Flow.Default.SizedIO as F +#include "repa-flow.h" + + +-- | The default chunk size of 64kBytes. +defaultChunkSize :: Integer +defaultChunkSize = 64 * 1024 + + +-- | Read a file containing Comma-Separated-Values. +sourceCSV + :: BulkI l Bucket + => Array l Bucket -> IO (Sources N (Array N (Array F Char))) +sourceCSV + = F.sourceCSV defaultChunkSize + $ error $ "Line exceeds chunk size of " + ++ show defaultChunkSize ++ "bytes." +{-# INLINE sourceCSV #-} + + +-- | Read a file containing Tab-Separated-Values. +sourceTSV + :: BulkI l Bucket + => Array l Bucket -> IO (Sources N (Array N (Array F Char))) +sourceTSV + = F.sourceTSV defaultChunkSize + $ error $ "Line exceeds chunk size of " + ++ show defaultChunkSize ++ "bytes." +{-# INLINE sourceTSV #-} + + +-- | Read complete records of data form a file, into chunks of the given length. +-- We read as many complete records as will fit into each chunk. +-- +-- The records are separated by a special terminating character, which the +-- given predicate detects. After reading a chunk of data we seek the file to +-- just after the last complete record that was read, so we can continue to +-- read more complete records next time. +-- +-- If we cannot fit at least one complete record in the chunk then perform +-- the given failure action. Limiting the chunk length guards against the +-- case where a large input file is malformed, as we won't try to read the +-- whole file into memory. +-- +-- +-- * Data is read into foreign memory without copying it through the GHC heap. +-- * The provided file handle must support seeking, else you'll get an +-- exception. +-- * Each file is closed the first time the consumer tries to pull a +-- record from the associated stream when no more are available. +-- +sourceRecords + :: BulkI l Bucket + => (Word8 -> Bool) -- ^ Detect the end of a record. + -> Array l Bucket -- ^ Buckets. + -> IO (Sources N (Array F Word8)) +sourceRecords pSep + = F.sourceRecords defaultChunkSize pSep + $ error $ "Record exceeds chunk size of " + ++ show defaultChunkSize ++ "bytes." +{-# INLINE sourceRecords #-} + + +-- | Read complete lines of data from a text file, using the given chunk length. +-- We read as many complete lines as will fit into each chunk. +-- +-- * The trailing new-line characters are discarded. +-- * Data is read into foreign memory without copying it through the GHC heap. +-- * The provided file handle must support seeking, else you'll get an +-- exception. +-- * Each file is closed the first time the consumer tries to pull a line +-- from the associated stream when no more are available. +-- +sourceLines + :: BulkI l Bucket + => Array l Bucket -> IO (Sources N (Array F Char)) +sourceLines + = F.sourceLines defaultChunkSize + $ error $ "Line exceeds chunk size of " + ++ show defaultChunkSize ++ "bytes." +{-# INLINE sourceLines #-} + + +-- | Read 8-bit ASCII characters from some files, using the given chunk length. +sourceChars + :: BulkI l Bucket + => Array l Bucket -> IO (Sources F Char) +sourceChars = F.sourceChars defaultChunkSize +{-# INLINE sourceChars #-} + + +-- | Read data from some files, using the given chunk length. +sourceBytes + :: BulkI l Bucket + => Array l Bucket -> IO (Sources F Word8) +sourceBytes = F.sourceBytes defaultChunkSize +{-# INLINE sourceBytes #-} + + +-- | Write vectors of text lines to the given files handles. +-- +-- * Data is copied into a new buffer to insert newlines before being +-- written out. +-- +sinkLines + :: ( BulkI l Bucket + , BulkI l1 (Array l2 Char) + , BulkI l2 Char, Unpack (Array l2 Char) t2) + => Name l1 -- ^ Layout of chunks. + -> Name l2 -- ^ Layout of lines in chunks. + -> Array l Bucket -- ^ Buckets + -> IO (Sinks l1 (Array l2 Char)) +sinkLines = F.sinkLines +{-# INLINE sinkLines #-} + + +-- | Write 8-bit ASCII characters to some files. +sinkChars + :: BulkI l Bucket + => Array l Bucket -> IO (Sinks F Char) +sinkChars = F.sinkChars +{-# INLINE sinkChars #-} + + +-- | Write bytes to some file. +sinkBytes + :: BulkI l Bucket + => Array l Bucket -> IO (Sinks F Word8) +sinkBytes = F.sinkBytes +{-# INLINE sinkBytes #-} diff --git a/Data/Repa/Flow/Default/IO/CSV.hs b/Data/Repa/Flow/Default/IO/CSV.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow/Default/IO/CSV.hs @@ -0,0 +1,52 @@ + +module Data.Repa.Flow.Default.IO.CSV + (sourceCSV) +where +import Data.Repa.Flow.Default +import Data.Repa.Flow.IO.Bucket +import Data.Repa.Array as A +import Data.Repa.Array.Material as A +import Data.Char +import qualified Data.Repa.Flow.Generic.IO as G +import qualified Data.Repa.Flow.Generic as G +#include "repa-flow.h" + + +-- | Read a file containing Tab-Separated-Values. +-- +-- TODO: handle escaped commas. +-- TODO: check CSV file standard. +-- +sourceCSV + :: BulkI l Bucket + => Integer -- Chunk length. + -> IO () -- Action to perform if we find line longer + -- than the chunk length. + -> Array l Bucket -- File paths. + -> IO (Sources N (Array N (Array F Char))) + +sourceCSV nChunk aFail bs + = do + -- Rows are separated by new lines, + -- fields are separated by commas. + let !nl = fromIntegral $ ord '\n' + let !nr = fromIntegral $ ord '\r' + let !nt = fromIntegral $ ord ',' + + -- Stream chunks of data from the input file, where the chunks end + -- cleanly at line boundaries. + sChunk <- G.sourceChunks nChunk (== nl) aFail bs + sRows8 <- G.map_i (A.diceSep nt nl . A.filter U (/= nr)) sChunk + + -- Convert element data from Word8 to Char. + -- Chars take 4 bytes each, but are standard Haskell and pretty + -- print properly. We've done the dicing on the smaller Word8 + -- version, and now map across the elements vector in the array + -- to do the conversion. + sRows <- G.map_i + (A.mapElems (A.mapElems + (A.computeS F . A.map (chr . fromIntegral)))) + sRows8 + + return sRows +{-# INLINE sourceCSV #-} diff --git a/Data/Repa/Flow/Default/IO/TSV.hs b/Data/Repa/Flow/Default/IO/TSV.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow/Default/IO/TSV.hs @@ -0,0 +1,48 @@ + +module Data.Repa.Flow.Default.IO.TSV + (sourceTSV) +where +import Data.Repa.Flow.Default +import Data.Repa.Flow.IO.Bucket +import Data.Repa.Array as A +import Data.Repa.Array.Material as A +import Data.Char +import qualified Data.Repa.Flow.Generic.IO as G +import qualified Data.Repa.Flow.Generic as G +#include "repa-flow.h" + + +-- | Read a file containing Tab-Separated-Values. +sourceTSV + :: BulkI l Bucket + => Integer -- Chunk length. + -> IO () -- Action to perform if we find line longer + -- than the chunk length. + -> Array l Bucket -- File paths. + -> IO (Sources N (Array N (Array F Char))) + +sourceTSV nChunk aFail bs + = do + -- Rows are separated by new lines, + -- fields are separated by tabs. + let !nl = fromIntegral $ ord '\n' + let !nr = fromIntegral $ ord '\r' + let !nt = fromIntegral $ ord '\t' + + -- Stream chunks of data from the input file, where the chunks end + -- cleanly at line boundaries. + sChunk <- G.sourceChunks nChunk (== nl) aFail bs + sRows8 <- G.map_i (A.diceSep nt nl . A.filter U (/= nr)) sChunk + + -- Convert element data from Word8 to Char. + -- Chars take 4 bytes each, but are standard Haskell and pretty + -- print properly. We've done the dicing on the smaller Word8 + -- version, and now map across the elements vector in the array + -- to do the conversion. + sRows <- G.map_i + (A.mapElems (A.mapElems + (A.computeS F . A.map (chr . fromIntegral)))) + sRows8 + + return sRows +{-# INLINE sourceTSV #-} diff --git a/Data/Repa/Flow/Default/SizedIO.hs b/Data/Repa/Flow/Default/SizedIO.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow/Default/SizedIO.hs @@ -0,0 +1,120 @@ + +-- | Read and write files. +module Data.Repa.Flow.Default.SizedIO + ( module Data.Repa.Flow.IO.Bucket + + -- * Sourcing + , sourceBytes + , sourceChars + , sourceLines + , sourceRecords + , sourceTSV + , sourceCSV + + -- * Sinking + , sinkBytes + , sinkChars + , sinkLines) +where +import Data.Repa.Flow.Default +import Data.Repa.Flow.IO.Bucket +import Data.Repa.Flow.Default.IO.TSV as F +import Data.Repa.Flow.Default.IO.CSV as F +import Data.Repa.Eval.Array as A +import Data.Repa.Array.Material as A +import Data.Repa.Fusion.Unpack as F +import Data.Repa.Array as A +import qualified Data.Repa.Flow.Generic as G +import qualified Data.Repa.Flow.Generic.IO as G +import Data.Word +import Data.Char +#include "repa-flow.h" + + +-- Sourcing --------------------------------------------------------------------------------------- +-- | Like `F.sourceBytes`, but with the default chunk size. +sourceBytes + :: BulkI l Bucket + => Integer -> Array l Bucket -> IO (Sources F Word8) +sourceBytes i bs = G.sourceBytes i bs +{-# INLINE sourceBytes #-} + + +-- | Like `F.sourceChars`, but with the default chunk size. +sourceChars + :: BulkI l Bucket + => Integer -> Array l Bucket -> IO (Sources F Char) +sourceChars i bs = G.sourceChars i bs +{-# INLINE sourceChars #-} + + +-- | Like `F.sourceLines`, but with the default chunk size and error action. +sourceLines + :: BulkI l Bucket + => Integer -- ^ Size of chunk to read in bytes. + -> IO () -- ^ Action to perform if we can't get a + -- whole record. + -> Array l Bucket -- ^ Buckets. + -> IO (Sources N (Array F Char)) +sourceLines nChunk fails bs + = G.map_i chopChunk + =<< G.sourceRecords nChunk isNewLine fails bs + where + isNewLine :: Word8 -> Bool + isNewLine x = x == nl + {-# INLINE isNewLine #-} + + chopChunk chunk + = A.mapElems (A.computeS name . A.map (chr . fromIntegral)) + $ A.trimEnds (== nl) chunk + {-# INLINE chopChunk #-} + + nl :: Word8 + !nl = fromIntegral $ ord '\n' +{-# NOINLINE sourceLines #-} + + +-- | Like `F.sourceRecords`, but with the default chunk size and error action. +sourceRecords + :: BulkI l Bucket + => Integer -- ^ Size of chunk to read in bytes. + -> (Word8 -> Bool) -- ^ Detect the end of a record. + -> IO () -- ^ Action to perform if we can't get a + -- whole record. + -> Array l Bucket -- ^ File handles. + -> IO (Sources N (Array F Word8)) +sourceRecords i pSep aFail bs + = G.sourceRecords i pSep aFail bs +{-# INLINE sourceRecords #-} + + +-- Sinking ---------------------------------------------------------------------------------------- +-- | An alias for `F.sinkBytes`. +sinkBytes + :: BulkI l Bucket + => Array l Bucket -> IO (Sinks F Word8) +sinkBytes bs = G.sinkBytes bs +{-# INLINE sinkBytes #-} + + +-- | An alias for `F.sinkChars`. +sinkChars + :: BulkI l Bucket + => Array l Bucket -> IO (Sinks F Char) +sinkChars bs = G.sinkChars bs +{-# INLINE sinkChars #-} + + +-- | An alias for `F.sinkLines`. +sinkLines + :: ( BulkI l Bucket + , BulkI l1 (Array l2 Char) + , BulkI l2 Char, Unpack (Array l2 Char) t2) + => Name l1 -- ^ Layout for chunks of lines. + -> Name l2 -- ^ Layout for lines. + -> Array l Bucket -- ^ Buckets + -> IO (Sinks l1 (Array l2 Char)) +sinkLines n1 n2 bs + = G.sinkLines n1 n2 bs +{-# INLINE sinkLines #-} + diff --git a/Data/Repa/Flow/Generic.hs b/Data/Repa/Flow/Generic.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow/Generic.hs @@ -0,0 +1,142 @@ + +-- | Everything flows. +-- +-- This module defines generic flows. The other flow types defined +-- in "Data.Repa.Flow.Chunked" and "Data.Repa.Flow.Simple" are +-- specialisations of this generic one. +-- +module Data.Repa.Flow.Generic + ( Sources (..) + , Sinks (..) + + -- * Stream State and Thread Safety + + -- $threadsafety + + , module Data.Repa.Flow.States + + -- * Evaluation + , drainS + , drainP + + -- * Conversion + , fromList + , toList1 + , takeList1 + + , pushList + , pushList1 + + -- * Stream Indices + , mapIndex_i + , mapIndex_o + , flipIndex2_i + , flipIndex2_o + + -- * Finalizers + , finalize_i + , finalize_o + + -- * Flow Operators + -- ** Projection + , project_i + , project_o + + -- ** Constructors + , repeat_i + , replicate_i + , prepend_i, prependOn_i + + -- ** Mapping + , map_i, map_o + , smap_i, smap_o + , szipWith_ii, szipWith_io, szipWith_oi + + -- ** Connecting + , dup_oo, dup_io, dup_oi + , connect_i + , funnel_i + , funnel_o + + -- ** Splitting + , head_i + + -- ** Grouping + , groups_i + + -- ** Packing + , pack_ii + + -- ** Folding + , folds_ii + + -- ** Watching + , watch_i + , watch_o + , trigger_o + + -- ** Capturing + , capture_o + , rcapture_o + + -- ** Ignorance + , discard_o + , ignore_o + + -- ** Tracing + , trace_o + + -- * Vector Flow Operators + -- ** 1-dimensional distribution + , distribute_o + , ddistribute_o + + -- ** 2-dimensional distribution + , distribute2_o + , ddistribute2_o + + -- ** Shuffling + , shuffle_o + , dshuffle_o + , dshuffleBy_o + + -- ** Chunking + , chunk_i + , unchunk_i) +where +import Data.Repa.Flow.States +import Data.Repa.Flow.Generic.Base +import Data.Repa.Flow.Generic.Connect +import Data.Repa.Flow.Generic.List +import Data.Repa.Flow.Generic.Map +import Data.Repa.Flow.Generic.Operator +import Data.Repa.Flow.Generic.Eval +import Data.Repa.Flow.Generic.Array.Distribute +import Data.Repa.Flow.Generic.Array.Shuffle +import Data.Repa.Flow.Generic.Array.Chunk +import Data.Repa.Flow.Generic.Array.Unchunk + + +-- $threadsafety +-- As most functions in this library produce `IO` actions, thread safety is not +-- guaranteed by their types. +-- +-- It is /not safe/ to concurrently pull from the same stream of a `Sources` +-- bundle, or concurrently push to the same stream of a `Sinks` bundle. +-- Both `Sources` and `Sinks` may hold per-stream state information, and +-- accessing the same stream concurrently may cause a race condition. +-- +-- It is safe to concurrently push or pull from /different/ streams of a bundle, +-- as the state information for each stream is guaranteed to be separate. +-- Any inter-stream communication is protected by appropriate locks. +-- +-- Unless stated otherwise, any worker function passed to a flow operator may +-- be invoked concurrently. For example, if you pass an `IO` action to +-- `trigger_o` then that action may be invoked concurrently. +-- +-- In practice, if you use just the bulk operators provided by this library +-- then you won't have a problem. However, if you construct your own +-- `Sources` or `Sinks` by providing raw @push@, @pull@ and @eject@ functions +-- then you must obey the above rules. +-- + diff --git a/Data/Repa/Flow/Generic/Array/Chunk.hs b/Data/Repa/Flow/Generic/Array/Chunk.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow/Generic/Array/Chunk.hs @@ -0,0 +1,73 @@ + +module Data.Repa.Flow.Generic.Array.Chunk + (chunk_i) +where +import Data.Repa.Flow.Generic.Base +import Data.Repa.Eval.Array as A +import Data.Repa.Array as A +#include "repa-flow.h" + + +-- | Take elements from a flow and pack them into chunks of the given +-- maximum length. +chunk_i :: (Target lDst a, Index lDst ~ Int, States i IO) + => Name lDst -- ^ Layout for result chunks. + -> Int -- ^ Maximum chunk length. + -> Sources i IO a -- ^ Element sources. + -> IO (Sources i IO (Array lDst a)) -- ^ Chunk sources. + +chunk_i nDst !maxLen (Sources n pullX) + = do + -- Refs for signalling how many elements we managed to read for + -- each chunk. + final <- newRefs n Nothing + + let pull_chunk i eat eject + = do + -- New buffer to hold elements we read from the source. + chunk <- unsafeNewBuffer (A.create nDst maxLen) + + let loop_chunk !ix + -- The chunk is already full. + | ix >= maxLen + = writeRefs final i (Just ix) + + | otherwise + = pullX i eat_chunk eject_chunk + where + -- Write the next element to the chunk. + eat_chunk x + = do unsafeWriteBuffer chunk ix x + loop_chunk (ix + 1) + + -- There are no more elements available from the soruce. + eject_chunk + -- We don't have a current chunk so we're done. + | ix == 0 = writeRefs final i Nothing + + -- We've got a current chunk, so signal + -- that it needs to be passed on downstream. + | otherwise = writeRefs final i (Just ix) + {-# INLINE loop_chunk #-} + + -- Pull as many elements as we can into a chunk. + loop_chunk 0 + + -- See what happened. + mlen <- readRefs final i + + case mlen of + -- We couldn't read any more elements to start a new + -- chunk, so the source is empty. + Nothing -> eject + + -- Pass this chunk downstream. + Just len + -> do chunk' <- unsafeSliceBuffer 0 len chunk + arr <- unsafeFreezeBuffer chunk' + eat arr + {-# INLINE pull_chunk #-} + + return $ Sources n pull_chunk +{-# INLINE_FLOW chunk_i #-} + diff --git a/Data/Repa/Flow/Generic/Array/Distribute.hs b/Data/Repa/Flow/Generic/Array/Distribute.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow/Generic/Array/Distribute.hs @@ -0,0 +1,161 @@ + +module Data.Repa.Flow.Generic.Array.Distribute + ( -- | 1-dimensional distribution. + distribute_o + , ddistribute_o + + -- | 2-dimensional distribution. + , distribute2_o + , ddistribute2_o) +where +import Data.Repa.Flow.Generic.Base +import Data.Repa.Array +import Prelude hiding (length) +#include "repa-flow.h" + + +------------------------------------------------------------------------------- +-- | Given a bundle of sinks indexed by an `Int`, +-- produce a result sink for arrays. +-- +-- Each time an array is pushed to the result sink, its elements are +-- pushed to the corresponding streams of the argument sink. If there +-- are more elements than sinks then then give them to the spill action. +-- +-- @ +-- +-- | .. | +-- | [w0, x0, y0, z0] | :: Sinks () IO (Array l a) +-- | [w1, x1, y1, z1, u1] | (sink for a single stream of arrays) +-- | .. | +-- +-- | | | | | +-- v v v v .------> spilled +-- +-- | .. | .. | .. | .. | +-- | w0 | x0 | y0 | z0 | :: Sinks Int IO a +-- | w1 | x1 | y1 | z1 | (sink for several streams of elements) +-- | .. | .. | .. | .. | +-- @ +-- +distribute_o + :: BulkI l a + => (Int -> a -> IO ()) -- ^ Spill action, given the spilled element + -- along with its index in the array. + -> Sinks Int IO a -- ^ Sinks to push elements into. + -> IO (Sinks () IO (Array l a)) + +distribute_o aSpill (Sinks nSinks push eject) + = do + let push_distribute _ !xs + = loop_distribute 0 + where !nx = length xs + + loop_distribute !ix + | ix >= nx + = return () + + | ix >= nSinks + = do aSpill ix (index xs ix) + loop_distribute (ix + 1) + + | otherwise + = do push ix (index xs ix) + loop_distribute (ix + 1) + {-# INLINE loop_distribute #-} + {-# INLINE push_distribute #-} + + let eject_distribute _ + = loop_distribute 0 + where + loop_distribute !ix + | ix >= nSinks + = return () + + | otherwise + = do eject ix + loop_distribute (ix + 1) + {-# INLINE loop_distribute #-} + {-# INLINE eject_distribute #-} + + return $ Sinks () push_distribute eject_distribute +{-# INLINE_FLOW distribute_o #-} + + +-- | Like `distribute_o`, but drop spilled elements on the floor. +ddistribute_o + :: BulkI l a + => Sinks Int IO a + -> IO (Sinks () IO (Array l a)) + +ddistribute_o sinks + = distribute_o (\_ _ -> return ()) sinks +{-# INLINE ddistribute_o #-} + + +------------------------------------------------------------------------------- +-- | Like `distribute_o`, but with 2-d stream indexes. +-- +-- Given the argument and result sinks, when pushing to the result the +-- stream index is used as the first component for the argument sink, +-- and the index of the element in its array is used as the second +-- component. +-- +-- If you want to the components of stream index the other way around +-- then apply `flipIndex2_o` to the argument sinks. +-- +distribute2_o + :: BulkI l a + => (SH2 -> a -> IO ()) -- ^ Spill action, given the spilled element + -- along with its index in the array. + -> Sinks SH2 IO a -- ^ Sinks to push elements into. + -> IO (Sinks Int IO (Array l a)) + +distribute2_o aSpill (Sinks (Z :. a1 :. a0) push eject) + = do + let push_distribute i1 !xs + = loop_distribute 0 + where !nx = length xs + + loop_distribute !ix + | ix >= nx + = return () + + | ix >= a0 + = do aSpill (ish2 i1 ix) (index xs ix) + loop_distribute (ix + 1) + + | otherwise + = do push (ish2 i1 ix) (index xs ix) + loop_distribute (ix + 1) + {-# INLINE loop_distribute #-} + {-# INLINE push_distribute #-} + + let eject_distribute i1 + = loop_distribute 0 + where + loop_distribute !ix + | ix >= a0 + = return () + + | otherwise + = do eject (ish2 i1 ix) + loop_distribute (ix + 1) + {-# INLINE loop_distribute #-} + {-# INLINE eject_distribute #-} + + return $ Sinks a1 push_distribute eject_distribute +{-# INLINE_FLOW distribute2_o #-} + + +-- | Like `distribute2_o`, but drop spilled elements on the floor. +ddistribute2_o + :: BulkI l a + => Sinks SH2 IO a + -> IO (Sinks Int IO (Array l a)) + +ddistribute2_o sinks + = distribute2_o (\_ _ -> return ()) sinks +{-# INLINE ddistribute2_o #-} + + diff --git a/Data/Repa/Flow/Generic/Array/Shuffle.hs b/Data/Repa/Flow/Generic/Array/Shuffle.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow/Generic/Array/Shuffle.hs @@ -0,0 +1,168 @@ +{-# OPTIONS -fno-warn-unused-imports #-} +module Data.Repa.Flow.Generic.Array.Shuffle + ( shuffle_o + , dshuffle_o + , dshuffleBy_o) +where +import Data.Repa.Flow.Generic.Base as F +import Data.Repa.Flow.Generic.Map as F +import Data.Repa.Flow.Generic.Operator as F +import Data.Repa.Array as A +import Data.Repa.Eval.Elt +import Control.Monad +#include "repa-flow.h" + + +-- | Given a bundle of argument sinks, produce a result sink. +-- Arrays of indices and elements are pushed to the result sink. +-- On doing so, the elements are pushed into the corresponding streams +-- of the argument sinks. +-- +-- If the index associated with an element does not have a corresponding +-- stream in the argument sinks, then pass it to the provided spill +-- function. +-- +-- +-- @ +-- | .. | +-- | [(0, v0), (1, v1), (0, v2), (0, v3), (2, v4)] | :: Sources Int IO (Array l (Int, a)) +-- | .. | +-- \\ \\ | +-- \\ .------------. | +-- v v .---------> spilled +-- +-- | .. | .. | +-- | [v0, v2, v3] | [v1] | :: Sinks Int IO (Array l a) +-- | .. | .. | +-- @ +-- +-- +-- The following example uses `capture_o` to demonstrate how the +-- `shuffle_o` operator can be used as one step of a bucket-sort. We start +-- with two arrays of key-value pairs. In the result, the values from each +-- block that had the same key are packed into the same tuple (bucket). +-- +-- @ +-- > import Data.Repa.Flow.Generic as G +-- > import Data.Repa.Array as A +-- > import Data.Repa.Array.Material as A +-- > import Data.Repa.Nice +-- +-- > let arr1 = A.fromList B [(0, \'a\'), (1, \'b\'), (2, \'c\'), (0, \'d\'), (0, \'c\')] +-- > let arr2 = A.fromList B [(0, \'A\'), (3, \'B\'), (3, \'C\')] +-- > result :: Array B (Int, Array U Char) +-- > \<- capture_o B 4 (\\k -> shuffle_o B (error \"spilled\") k +-- > >>= pushList1 () [arr1, arr2]) +-- +-- > nice result +-- [(0,\"adc\"),(1,\"b\"),(2,\"c\"),(0,\"A\"),(3,\"BC\")] +-- @ +-- +shuffle_o + :: ( BulkI lDst a, BulkI lSrc (Int, a) + , Windowable lDst a + , Target lDst a + , Elt a) + => Name lSrc -- ^ Name of source layout. + -> (Int -> Array lDst a -> IO ()) -- ^ Handle spilled elements. + -> Sinks Int IO (Array lDst a) -- ^ Sinks to push results to. + -> IO (Sinks () IO (Array lSrc (Int, a))) + +shuffle_o _ aSpill (Sinks nSinks opush oeject) + = return $ Sinks () shuffle_push shuffle_eject + where + shuffle_push _ !arr + = do -- Partition the elements by segment number. + let !parts = A.partition name nSinks arr + + -- Push the individual segments into the argument sinks. + let loop_shuffle_push !i + | i >= A.length parts + = return () + + | i >= nSinks + = do let !part = parts `index` i + when (A.length part > 0) + $ aSpill i part + + loop_shuffle_push (i + 1) + + | otherwise + = do let !part = parts `index` i + when (A.length part > 0) + $ opush i part + + loop_shuffle_push (i + 1) + + loop_shuffle_push 0 + {-# INLINE shuffle_push #-} + + shuffle_eject _ + = do + let loop_shuffle_eject !i + | i >= nSinks + = return () + + | otherwise + = do oeject i + loop_shuffle_eject (i + 1) + + loop_shuffle_eject 0 + {-# INLINE shuffle_eject #-} + +{-# INLINE_FLOW shuffle_o #-} + + +-- | Like `shuffle_o`, but drop spilled elements on the floor. +dshuffle_o + :: ( BulkI lDst a, BulkI lSrc (Int, a) + , Windowable lDst a + , Target lDst a + , Elt a) + => Name lSrc -- ^ Name of source layout. + -> Sinks Int IO (Array lDst a) -- ^ Sinks to push results to. + -> IO (Sinks () IO (Array lSrc (Int, a))) + +dshuffle_o nSrc sinks + = shuffle_o nSrc (\_ _ -> return ()) sinks +{-# INLINE dshuffle_o #-} + + +-- | Like `dshuffle_o`, but use the given function to decide which stream of +-- the argument bundle each element should be pushed into. +-- +-- @ +-- > import Data.Repa.Flow.Generic as G +-- > import Data.Repa.Array as A +-- > import Data.Repa.Array.Material as A +-- > import Data.Repa.Nice +-- > import Data.Char +-- +-- > let arr1 = A.fromList B \"FooBAr\" +-- > let arr2 = A.fromList B \"BazLIKE\" +-- > result :: Array B (Int, Array U Char) +-- \<- capture_o B 2 (\\k -> dshuffleBy_o B (\\x -> if isUpper x then 0 else 1) k +-- >>= pushList1 () [arr1, arr2]) +-- > nice result +-- [(0,\"FBA\"),(1,\"oor\"),(0,\"BLIKE\"),(1,\"az\")] +-- @ +-- +dshuffleBy_o + :: ( BulkI lDst a, BulkI lSrc a + , Windowable lDst a + , Target lDst a + , Elt a) + => Name lSrc -- ^ Name of source layout. + -> (a -> Int) -- ^ Get the stream number for an element. + -> Sinks Int IO (Array lDst a) -- ^ Sinks to push results to. + -> IO (Sinks () IO (Array lSrc a)) + +dshuffleBy_o _ fBucket sinks + = do kShuf <- dshuffle_o name sinks + + let chunk _ arr = A.tup2 (A.map fBucket arr) arr + {-# INLINE chunk #-} + + smap_o chunk kShuf +{-# INLINE dshuffleBy_o #-} + diff --git a/Data/Repa/Flow/Generic/Array/Unchunk.hs b/Data/Repa/Flow/Generic/Array/Unchunk.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow/Generic/Array/Unchunk.hs @@ -0,0 +1,110 @@ + +module Data.Repa.Flow.Generic.Array.Unchunk + (unchunk_i) +where +import Data.Repa.Flow.Generic.Base +import Data.Repa.Array as A +#include "repa-flow.h" + + +-- Unchunk -------------------------------------------------------------------- +-- | Take a flow of chunks and flatten it into a flow of the individual +-- elements. +unchunk_i :: (BulkI l a, States i IO) + => Sources i IO (Array l a) -- ^ Chunk sources. + -> IO (Sources i IO a) -- ^ Element sources. + +unchunk_i (Sources n pullC) + = do + -- States to hold the current chunk. + -- INVARIANT: if this holds a chunk then it is non-empty. + rChunks <- newRefs n Nothing + + -- States to hold the current index in each chunk. + rIxs <- newRefs n 0 + + let pullX i eat eject + = pullStart + where + -- If we already have a non-empty chunk then we can return + -- the next element from that. + pullStart + = do mchunk <- readRefs rChunks i + case mchunk of + Just chunk -> pullElem chunk + Nothing -> pullSource + {-# INLINE pullStart #-} + + -- Try to pull a non-empty chunk from the source, + -- and then pass on to 'pullElem' which will take the next + -- element from it. + pullSource + = pullC i eat_source eject_source + {-# INLINE pullSource #-} + + eat_source !chunk + | A.length chunk == 0 + = pullSource + + | otherwise + = do writeRefs rChunks i (Just chunk) + writeRefs rIxs i 0 + pullElem chunk + {-# INLINE eat_source #-} + + eject_source + = eject + {-# INLINE eject_source #-} + + -- We've got a chunk containing some elements + pullElem !chunk + = do !ix <- readRefs rIxs i + + _ <- if (ix + 1) >= A.length chunk + -- This was the last element of the chunk. + -- We need to pull a new one from the source + -- the next time around. + then do writeRefs rChunks i Nothing + writeRefs rIxs i 0 + + -- There are still more elements to read + -- from the current chunk. + else do writeRefs rIxs i (ix + 1) + + let !x = index chunk ix + eat x + {-# INLINE pullElem #-} + {-# INLINE pullX #-} + + return $ Sources n pullX +{-# INLINE_FLOW unchunk_i #-} + +{- +-- | Take an argument sink for individual elements, and produce a result sink +-- for chunks. +-- +-- When a chunk it pushed to the result sink then all its elements are +-- pushed to the argument sink. +-- +unchunk_o :: Monad m + => Bulk r DIM1 e + => Sink m e -> m (Sink m (Vector r e)) + +unchunk_o (Sink pushX ejectX) + = return $ Sink push_unchunk eject_unchunk + where + push_unchunk !chunk + = loop_unchunk 0 + where !len = size (extent chunk) + loop_unchunk !i + | i >= len = return () + | otherwise + = do pushX (chunk `index` (Z :. i)) + loop_unchunk (i + 1) + {-# INLINE push_unchunk #-} + + eject_unchunk + = ejectX + {-# INLINE eject_unchunk #-} +{-# INLINE_FLOW unchunk_o #-} +-} diff --git a/Data/Repa/Flow/Generic/Base.hs b/Data/Repa/Flow/Generic/Base.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow/Generic/Base.hs @@ -0,0 +1,186 @@ + +module Data.Repa.Flow.Generic.Base + ( module Data.Repa.Flow.States + , Sources (..) + , Sinks (..) + , mapIndex_i + , mapIndex_o + , flipIndex2_i + , flipIndex2_o + , finalize_i + , finalize_o) +where +import Data.Repa.Flow.States +import Data.Repa.Array as A +import Control.Monad +#include "repa-flow.h" + +-- | A bundle of stream sources, indexed by a value of type @i@, +-- in some monad @m@, returning elements of type @e@. +-- +-- Elements can be pulled from each stream in the bundle individually. +-- +data Sources i m e + = Sources + { -- | Number of sources in this bundle. + sourcesArity :: i + + -- | Function to pull data from a bundle. + -- Give it the index of the desired stream, a continuation that + -- accepts an element, and a continuation to invoke when no more + -- elements will ever be available. + , sourcesPull :: i -> (e -> m ()) -> m () -> m () } + + +-- | A bundle of stream sinks, indexed by a value of type @i@, +-- in some monad @m@, returning elements of type @e@. +-- +-- Elements can be pushed to each stream in the bundle individually. +-- +data Sinks i m e + = Sinks + { -- | Number of sources in the bundle. + sinksArity :: i + + -- | Push an element to one of the streams in the bundle. + , sinksPush :: i -> e -> m () + + -- | Signal that no more elements will ever be available for this + -- sink. It is ok to eject the same stream multiple times. + , sinksEject :: i -> m () } + + +------------------------------------------------------------------------------- +-- | Transform the stream indexes of a bundle of sources. +-- +-- The given transform functions should be inverses of each other, +-- else you'll get a confusing result. +mapIndex_i + :: Monad m + => (i1 -> i2) -> (i2 -> i1) + -> Sources i1 m a + -> m (Sources i2 m a) + +mapIndex_i to from (Sources n pullX) + = return $ Sources (to n) pull_mapIndex + where + pull_mapIndex i eat eject + = pullX (from i) eat eject + {-# INLINE pull_mapIndex #-} +{-# INLINE_FLOW mapIndex_i #-} + + +-- | Transform the stream indexes of a bundle of sinks. +-- +-- The given transform functions should be inverses of each other, +-- else you'll get a confusing result. +mapIndex_o + :: Monad m + => (i1 -> i2) -> (i2 -> i1) + -> Sinks i1 m a + -> m (Sinks i2 m a) + +mapIndex_o to from (Sinks n pushX ejectX) + = return $ Sinks (to n) push_mapIndex eject_mapIndex + where + push_mapIndex i x = pushX (from i) x + {-# INLINE push_mapIndex #-} + + eject_mapIndex i = ejectX (from i) + {-# INLINE eject_mapIndex #-} +{-# INLINE_FLOW mapIndex_o #-} + + +-- | For a bundle of sources with a 2-d stream index, +-- flip the components of the index. +flipIndex2_i + :: Monad m + => Sources SH2 m a + -> m (Sources SH2 m a) + +flipIndex2_i ss + = mapIndex_i + (\(Z :. y :. x) -> (Z :. x :. y)) + (\(Z :. y :. x) -> (Z :. x :. y)) + ss +{-# INLINE flipIndex2_i #-} + + +-- | For a bundle of sinks with a 2-d stream index, +-- flip the components of the index. +flipIndex2_o + :: Monad m + => Sinks SH2 m a + -> m (Sinks SH2 m a) + +flipIndex2_o ss + = mapIndex_o + (\(Z :. y :. x) -> (Z :. x :. y)) + (\(Z :. y :. x) -> (Z :. x :. y)) + ss +{-# INLINE flipIndex2_o #-} + + +------------------------------------------------------------------------------- +-- | Attach a finalizer to bundle of sources. +-- +-- For each stream in the bundle, the finalizer will be called the first +-- time a consumer of that stream tries to pull an element when no more +-- are available. +-- +-- The provided finalizer will be run after any finalizers already +-- attached to the source. +-- +finalize_i + :: States i m + => (i -> m ()) + -> Sources i m a -> m (Sources i m a) + +finalize_i f (Sources n pull) + = do + refs <- newRefs n False + + let pull_finalize i eat eject + = pull i eat eject_finalize + where + eject_finalize + = do eject + done <- readRefs refs i + when (not done) + $ do f i + writeRefs refs i False + {-# INLINE eject_finalize #-} + {-# INLINE pull_finalize #-} + + return $ Sources n pull_finalize +{-# INLINE_FLOW finalize_i #-} + + +-- | Attach a finalizer to a bundle of sinks. +-- +-- For each stream in the bundle, the finalizer will be called the first +-- time that stream is ejected. +-- +-- The provided finalizer will be run after any finalizers already +-- attached to the sink. +-- +finalize_o + :: States i m + => (i -> m ()) + -> Sinks i m a -> m (Sinks i m a) + +finalize_o f (Sinks n push eject) + = do + refs <- newRefs n False + + let eject_finalize i + = do eject i + done <- readRefs refs i + when (not done) + $ do f i + writeRefs refs i False + {-# INLINE eject_finalize #-} + + return $ Sinks n push eject_finalize +{-# INLINE_FLOW finalize_o #-} + diff --git a/Data/Repa/Flow/Generic/Connect.hs b/Data/Repa/Flow/Generic/Connect.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow/Generic/Connect.hs @@ -0,0 +1,213 @@ +module Data.Repa.Flow.Generic.Connect + ( -- * Dup + dup_oo + , dup_io + , dup_oi + + -- * Connect + , connect_i + + -- * Funnel + , funnel_i + , funnel_o) +where +import Data.Repa.Flow.Generic.Base +import Control.Monad +import Prelude as P +#include "repa-flow.h" + + +-- Dup ------------------------------------------------------------------------ +-- | Send the same data to two consumers. +-- +-- Given two argument sinks, yield a result sink. +-- Pushing to the result sink causes the same element to be pushed to both +-- argument sinks. +dup_oo :: (Ord i, States i m) + => Sinks i m a -> Sinks i m a -> m (Sinks i m a) +dup_oo (Sinks n1 push1 eject1) (Sinks n2 push2 eject2) + = return $ Sinks (min n1 n2) push_dup eject_dup + where + push_dup i x = push1 i x >> push2 i x + {-# INLINE push_dup #-} + + eject_dup i = eject1 i >> eject2 i + {-# INLINE eject_dup #-} +{-# INLINE_FLOW dup_oo #-} + + +-- | Send the same data to two consumers. +-- +-- Given an argument source and argument sink, yield a result source. +-- Pulling an element from the result source pulls from the argument +-- source, and pushes that element to the sink, as well as returning it +-- via the result source. +-- +dup_io :: (Ord i, Monad m) + => Sources i m a -> Sinks i m a -> m (Sources i m a) +dup_io (Sources n1 pull1) (Sinks n2 push2 eject2) + = return $ Sources (min n1 n2) pull_dup + where + pull_dup i eat eject + = pull1 i eat_x eject_x + where + eat_x x = eat x >> push2 i x + {-# INLINE eat_x #-} + + eject_x = eject >> eject2 i + {-# INLINE eject_x #-} + {-# INLINE pull_dup #-} +{-# INLINE_FLOW dup_io #-} + + +-- | Send the same data to two consumers. +-- +-- Like `dup_io` but with the arguments flipped. +-- +dup_oi :: (Ord i, Monad m) + => Sinks i m a -> Sources i m a -> m (Sources i m a) +dup_oi sink1 source2 = dup_io source2 sink1 +{-# INLINE_FLOW dup_oi #-} + + +-- Connect -------------------------------------------------------------------- +-- | Connect an argument source to two result sources. +-- +-- Pulling from either result source pulls from the argument source. +-- Each result source only gets the elements pulled at the time, +-- so if one side pulls all the elements the other side won't get any. +-- +connect_i + :: States i m + => Sources i m a -> m (Sources i m a, Sources i m a) + +connect_i (Sources n pullX) + = do + refs <- newRefs n Nothing + + -- IMPORTANT: the pump function is set to NOINLINE so that pullX + -- will not be inlined into both consumers. We do not want to + -- duplicate that code for both result sources. Instead, calling + -- pump writes its element into a ref, and then only the code + -- that reads the ref is duplicated. + let pump_connect i + = pullX i pump_eat pump_eject + where + pump_eat !x = writeRefs refs i (Just x) + {-# INLINE pump_eat #-} + + pump_eject + = writeRefs refs i Nothing + {-# INLINE pump_eject #-} + {-# NOINLINE pump_connect #-} + + let pull_splitAt i eat eject + = do pump_connect i + mx <- readRefs refs i + case mx of + Just x -> eat x + Nothing -> eject + {-# INLINE pull_splitAt #-} + + return ( Sources n pull_splitAt + , Sources n pull_splitAt ) + +{-# INLINE_FLOW connect_i #-} + + +-- Funneling ------------------------------------------------------------------ +-- | Given a bundle of sources containing several streams, produce a new +-- bundle containing a single stream that gets data from the former. +-- +-- Streams from the source are consumed in their natural order, +-- and a complete stream is consumed before moving onto the next one. +-- +-- @ +-- > import Data.Repa.Flow.Generic +-- > toList1 () =<< funnel_i =<< fromList (3 :: Int) [11, 22, 33] +-- [11,22,33,11,22,33,11,22,33] +-- @ +funnel_i :: (States i m, States () m) + => Sources i m a -> m (Sources () m a) + +funnel_i (Sources n pullX) + = do + -- Ref to hold the current stream index. + refCur <- newRefs () first + + let pull_funnel _ eat eject + = do i <- readRefs refCur () + pullX i (eat_funnel i) (eject_funnel i) + + where + eat_funnel _ x = eat x + {-# INLINE eat_funnel #-} + + eject_funnel i + = case next i n of + Nothing -> eject + Just i' + -> do writeRefs refCur () i' + pullX i' (eat_funnel i') (eject_funnel i') + {-# INLINE eject_funnel #-} + + return $ Sources () pull_funnel +{-# INLINE funnel_i #-} + + +-- | Given a bundle of sinks consisting of a single stream, produce a new +-- bundle of the given arity that sends all data to the former, ignoring +-- the stream index. +-- +-- The argument stream is ejected only when all of the streams in the +-- result bundle have been ejected. +-- +-- * Using this function in conjunction with parallel operators like +-- `drainP` introduces non-determinism. Elements pushed to different +-- streams in the result bundle could enter the single stream in the +-- argument bundle in any order. +-- +-- @ +-- > import Data.Repa.Flow.Generic +-- > import Data.Repa.Array.Material +-- > import Data.Repa.Nice +-- +-- > let things = [(0 :: Int, \"foo\"), (1, \"bar\"), (2, \"baz\")] +-- > result \<- capture_o B () (\\k -> funnel_o 4 k >>= pushList things) +-- > nice result +-- [((),\"foo\"),((),\"bar\"),((),\"baz\")] +-- @ +-- +funnel_o :: States i m + => i -> Sinks () m a -> m (Sinks i m a) +funnel_o nSinks (Sinks _ pushX ejectX) + = do + -- Refs to track which streams have been ejected. + refs <- newRefs nSinks False + + -- Push all received data into the single stream of the + -- argument bundle. + let push_funnel _ x + = pushX () x + {-# INLINE push_funnel #-} + + -- When all the result streams have been ejected, + -- eject the argument stream. + let eject_funnel i + = do + -- RACE: If two concurrent processes eject the final two + -- streams then they will both think they were the last + -- one, and eject the single argument stream. This is ok + -- as we allow the argument sink to be ejected multiple + -- times. + -- + -- See docs of `Sinks` type in "Data.Repa.Flow.Generic.Base". + -- + writeRefs refs i True + done <- foldRefsM (&&) True refs + when done $ ejectX () + {-# INLINE eject_funnel #-} + + return $ Sinks nSinks push_funnel eject_funnel +{-# INLINE_FLOW funnel_o #-} + diff --git a/Data/Repa/Flow/Generic/Debug.hs b/Data/Repa/Flow/Generic/Debug.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow/Generic/Debug.hs @@ -0,0 +1,91 @@ + +module Data.Repa.Flow.Generic.Debug + (-- * More + more, more' + + -- * More (tabular) + , moret, moret' + + -- * More (raw) + , morer, morer' + + -- * Nicer + , Nicer (..) + , Presentable (..)) +where +import Data.Repa.Nice.Present +import Data.Repa.Nice.Tabulate +import Data.Repa.Nice +import Data.Repa.Flow.Generic hiding (next) +import Control.Monad +import Data.List as L +import Data.Text as T +import Prelude as P +#include "repa-flow.h" + + +------------------------------------------------------------------------------- +-- | Given a source index and a length, pull enough chunks from the source +-- to build a list of the requested length, and discard the remaining +-- elements in the final chunk. +-- +-- * This function is intended for interactive debugging. +-- If you want to retain the rest of the final chunk then use `head_i`. +-- +more :: (States i IO, Nicer a) + => i -- ^ Index of source in bundle. + -> Sources i IO a -- ^ Bundle of sources. + -> IO [Nice a] +more i ss = more' i 20 ss +{-# INLINE more #-} + + +-- | Like `more` but also specify now many elements you want. +more' :: (States i IO, Nicer a) + => i -> Int -> Sources i IO a -> IO [Nice a] +more' ix len s + = liftM (L.map nice . fst) $ head_i len s ix +{-# INLINE_FLOW more' #-} + + +------------------------------------------------------------------------------- +-- | Like `more`, but print results in a tabular form to the console. +moret :: (States i IO, Nicer [a], Presentable (Nice [a])) + => i -- ^ Index of source in bundle. + -> Sources i IO a -- ^ Bundle of sources. + -> IO () + +moret i ss = moret' i 20 ss +{-# INLINE moret #-} + + +-- | Like `more'`, but print results in tabular form to the console. +moret' :: (States i IO, Nicer [a], Presentable (Nice [a])) + => i -> Int -> Sources i IO a -> IO () + +moret' i len s + = do (vals, _) <- head_i len s i + putStrLn $ T.unpack $ tabulate $ nice vals +{-# INLINE_FLOW moret' #-} + + +------------------------------------------------------------------------------- +-- | Like `more`, but show elements in their raw format. +morer :: States i IO + => i -- ^ Index of source in bundle. + -> Sources i IO a -- ^ Bundle of sources. + -> IO [a] + +morer i ss = morer' i 20 ss +{-# INLINE morer #-} + + +-- | Like `more'`, but show elements in their raw format. +morer' :: States i IO + => i -> Int -> Sources i IO a -> IO [a] +morer' i len s + = liftM fst $ head_i len s i +{-# INLINE_FLOW morer' #-} + + + diff --git a/Data/Repa/Flow/Generic/Eval.hs b/Data/Repa/Flow/Generic/Eval.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow/Generic/Eval.hs @@ -0,0 +1,72 @@ + +module Data.Repa.Flow.Generic.Eval + ( drainS + , drainP) +where +import Data.Repa.Flow.Generic.Base +import Data.Repa.Eval.Gang as Eval +import GHC.Exts +#include "repa-flow.h" + + +-- | Pull all available values from the sources and push them to the sinks. +-- Streams in the bundle are processed sequentially, from first to last. +-- +-- * If the `Sources` and `Sinks` have different numbers of streams then +-- we only evaluate the common subset. +-- +drainS :: (Next i, Monad m) + => Sources i m a -> Sinks i m a -> m () + +drainS (Sources nSources ipull) (Sinks nSinks opush oeject) + = loop_drain first + where + n = min nSources nSinks + + loop_drain !ix + = ipull ix eat_drain eject_drain + where eat_drain v + = do opush ix v + loop_drain ix + {-# INLINE eat_drain #-} + + eject_drain + = do oeject ix + case next ix n of + Nothing -> return () + Just ix' -> loop_drain ix' + {-# INLINE eject_drain #-} + {-# INLINE loop_drain #-} +{-# INLINE_FLOW drainS #-} + + +-- | Pull all available values from the sources and push them to the sinks, +-- in parallel. We fork a thread for each of the streams and evaluate +-- them all in parallel. +-- +-- * If the `Sources` and `Sinks` have different numbers of streams then +-- we only evaluate the common subset. +-- +drainP :: Sources Int IO a -> Sinks Int IO a -> IO () +drainP (Sources nSources ipull) (Sinks nSinks opush oeject) + = do + -- Create a new gang. + gang <- Eval.forkGang n + + -- Evalaute all the streams in different threads. + Eval.gangIO gang drainMe + + where + !n = min nSources nSinks + + drainMe !ix + = ipull (I# ix) eat_drain eject_drain + where eat_drain v + = do opush (I# ix) v + drainMe ix + {-# INLINE eat_drain #-} + + eject_drain = oeject (I# ix) + {-# INLINE eject_drain #-} + {-# INLINE drainMe #-} +{-# INLINE_FLOW drainP #-} diff --git a/Data/Repa/Flow/Generic/IO.hs b/Data/Repa/Flow/Generic/IO.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow/Generic/IO.hs @@ -0,0 +1,209 @@ + +module Data.Repa.Flow.Generic.IO + ( -- * Buckets + module Data.Repa.Flow.IO.Bucket + + -- * Sourcing + , sourceBytes + , sourceChars + , sourceChunks + , sourceRecords + + -- * Sinking + , sinkBytes + , sinkChars + , sinkLines + + -- * Sieving + , sieve_o) +where +import Data.Repa.Flow.IO.Bucket +import Data.Repa.Flow.Generic.IO.Sieve as F +import Data.Repa.Flow.Generic.Map as F +import Data.Repa.Flow.Generic.Base as F +import Data.Repa.Fusion.Unpack as F +import Data.Repa.Array.Material as A +import Data.Repa.Array as A +import Data.Char +import System.IO +import Data.Word +import Prelude as P +#include "repa-flow.h" + + +-- Sourcing --------------------------------------------------------------------------------------- +-- | Read complete records of data form a bucket, into chunks of the given length. +-- We read as many complete records as will fit into each chunk. +-- +-- The records are separated by a special terminating character, which the +-- given predicate detects. After reading a chunk of data we seek the bucket to +-- just after the last complete record that was read, so we can continue to +-- read more complete records next time. +-- +-- If we cannot fit at least one complete record in the chunk then perform +-- the given failure action. Limiting the chunk length guards against the +-- case where a large input file is malformed, as we won't try to read the +-- whole file into memory. +-- +-- * Data is read into foreign memory without copying it through the GHC heap. +-- * The provided file handle must support seeking, else you'll get an exception. +-- +sourceRecords + :: BulkI l Bucket + => Integer -- ^ Chunk length in bytes. + -> (Word8 -> Bool) -- ^ Detect the end of a record. + -> IO () -- ^ Action to perform if we can't get a whole record. + -> Array l Bucket -- ^ Source buckets. + -> IO (Sources Int IO (Array N (Array F Word8))) + +sourceRecords len pSep aFail hs + = smap_i (\_ !c -> A.segmentOn pSep c) + =<< sourceChunks len pSep aFail hs +{-# INLINE sourceRecords #-} + + +-- | Like `sourceRecords`, but produce all records in a single vector. +sourceChunks + :: BulkI l Bucket + => Integer -- ^ Chunk length in bytes. + -> (Word8 -> Bool) -- ^ Detect the end of a record. + -> IO () -- ^ Action to perform if we can't get a whole record. + -> Array l Bucket -- ^ Source buckets. + -> IO (Sources (Index l) IO (Array F Word8)) + +sourceChunks len pSep aFail bs + = return $ Sources (A.extent $ A.layout bs) pull_sourceChunks + where + pull_sourceChunks i eat eject + = let b = bs `index` i + in bAtEnd b >>= \eof -> + if eof + -- We're at the end of the file. + then eject + + else do + -- Read a new chunk from the file. + arr <- bGetArray b len + + -- Find the end of the last record in the file. + let !mLenSlack = findIndex pSep (A.reverse arr) + + case mLenSlack of + -- If we couldn't find the end of record then apply the failure action. + Nothing -> aFail + + -- Work out how long the record is. + Just lenSlack + -> do let !lenArr = A.length arr + let !ixSplit = lenArr - lenSlack + + -- Seek the file to just after the last complete record. + bSeek b RelativeSeek (fromIntegral $ negate lenSlack) + + -- Eat complete records at the start of the chunk. + eat $ window 0 ixSplit arr + {-# INLINE pull_sourceChunks #-} +{-# INLINE_FLOW sourceChunks #-} + + +-- | Read 8-byte ASCII characters from some files, using the given chunk length. +-- +-- * Data is read into foreign memory without copying it through the GHC heap. +-- * All chunks have the same size, except possibly the last one. +---- +sourceChars + :: Bulk l Bucket + => Integer -- ^ Chunk length in bytes. + -> Array l Bucket -- ^ Buckets. + -> IO (Sources (Index l) IO (Array F Char)) +sourceChars len bs + = smap_i (\_ !c -> A.computeS F $ A.map (chr . fromIntegral) c) + =<< sourceBytes len bs +{-# INLINE sourceChars #-} + + +-- | Read data from some files, using the given chunk length. +-- +-- * Data is read into foreign memory without copying it through the GHC heap. +-- * All chunks have the same size, except possibly the last one. +-- +sourceBytes + :: Bulk l Bucket + => Integer -- ^ Chunk length in bytes. + -> Array l Bucket -- ^ Buckets. + -> IO (Sources (Index l) IO (Array F Word8)) +sourceBytes len bs + = return $ Sources (A.extent $ A.layout bs) pull_sourceBytes + where + pull_sourceBytes i eat eject + = do let b = A.index bs i + op <- bIsOpen b + if not op + then eject + else do + eof <- bAtEnd b + if eof + then eject + else do + !chunk <- bGetArray b len + eat chunk + {-# INLINE pull_sourceBytes #-} +{-# INLINE_FLOW sourceBytes #-} + + +-- Sinking ---------------------------------------------------------------------------------------- +-- | Write vectors of text lines to the given files handles. +-- +-- * Data is copied into a new buffer to insert newlines before being +-- written out. +-- +sinkLines + :: ( Bulk l Bucket + , BulkI l1 (Array l2 Char) + , BulkI l2 Char, Unpack (Array l2 Char) t2) + => Name l1 -- ^ Layout of chunks of lines. + -> Name l2 -- ^ Layout of lines. + -> Array l Bucket -- ^ Buckets. + -> IO (Sinks (Index l) IO (Array l1 (Array l2 Char))) +sinkLines _ _ !bs + = smap_o (\_ !c -> computeS F $ A.map (fromIntegral . ord) $ concatWith F fl c) + =<< sinkBytes bs + where !fl = A.fromList F ['\n'] +{-# INLINE sinkLines #-} + + +-- | Write chunks of 8-byte ASCII characters to the given file handles. +-- +-- * Data is copied into a foreign buffer to truncate the characters +-- to 8-bits each before being written out. +-- +sinkChars + :: (Bulk l Bucket, BulkI r Char) + => Array l Bucket -- ^ Buckets. + -> IO (Sinks (Index l) IO (Array r Char)) +sinkChars !bs + = smap_o (\_ !c -> computeS F $ A.map (fromIntegral . ord) c) + =<< sinkBytes bs +{-# INLINE sinkChars #-} + + +-- | Write chunks of bytes to the given file handles. +-- +-- * Data is written out directly from the provided buffer. +-- +sinkBytes + :: Bulk l Bucket + => Array l Bucket -- ^ Buckets. + -> IO (Sinks (Index l) IO (Array F Word8)) +sinkBytes !bs + = do let push_sinkBytes i !chunk + = bPutArray (bs `index` i) chunk + {-# NOINLINE push_sinkBytes #-} + + let eject_sinkBytes i + = bClose (bs `index` i) + {-# INLINE eject_sinkBytes #-} + + return $ Sinks (A.extent $ A.layout bs) push_sinkBytes eject_sinkBytes +{-# INLINE_FLOW sinkBytes #-} + diff --git a/Data/Repa/Flow/Generic/IO/Sieve.hs b/Data/Repa/Flow/Generic/IO/Sieve.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow/Generic/IO/Sieve.hs @@ -0,0 +1,50 @@ + +module Data.Repa.Flow.Generic.IO.Sieve + (sieve_o) +where +import Data.Repa.Flow.Generic.Base +import Data.Repa.Array as A +import Data.Repa.Array.Material as A +import Data.Repa.IO.Array as A +import System.IO +import Data.Word +#include "repa-flow.h" + + +-- | Create an output sieve that writes data to an indeterminate number of +-- output files. Each new element is appended to its associated file. +-- +-- * TODO: +-- This function keeps a maximum of 8 files open at once, closing +-- and re-opening them in a least-recently-used order. +-- Due to this behaviour it's fine to create thousands of separate +-- output files without risking overflowing the process limit on +-- the maximum number of useable file handles. +-- +sieve_o :: (a -> Maybe (FilePath, Array F Word8)) + -- ^ Produce the desired file path and output + -- record for this element, or `Nothing` if + -- it should be discarded. + -> IO (Sinks () IO a) + +sieve_o diag + = do + + let push_sieve _ e + = case diag e of + Nothing + -> return () + + Just (path, arr) + -> do -- TODO: repeatededly opening and closing the file + -- will be very slow. + h <- openBinaryFile path AppendMode + hPutArray h arr + hClose h + + -- TODO: ignore any more incoming data after being ejected. + let eject_sieve _ + = return () + + return $ Sinks () push_sieve eject_sieve +{-# INLINE sieve_o #-} diff --git a/Data/Repa/Flow/Generic/List.hs b/Data/Repa/Flow/Generic/List.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow/Generic/List.hs @@ -0,0 +1,100 @@ + +module Data.Repa.Flow.Generic.List + ( fromList + , toList1 + , takeList1 + + , pushList + , pushList1) +where +import Data.Repa.Flow.Generic.Base +#include "repa-flow.h" + + +------------------------------------------------------------------------------- +-- | Given an arity and a list of elements, yield sources that each produce +-- all the elements. +fromList :: States i m + => i -> [a] -> m (Sources i m a) +fromList n xx0 + = do + refs <- newRefs n xx0 + + let pull_fromList i eat eject + = do xx <- readRefs refs i + case xx of + [] -> eject + x : xs -> do writeRefs refs i xs + eat x + {-# INLINE pull_fromList #-} + + return $ Sources n pull_fromList +{-# INLINE_FLOW fromList #-} + + +-- | Drain a single source into a list. +toList1 :: States i m + => i -> Sources i m a -> m [a] +toList1 i (Sources n pullX) + = do + refs <- newRefs n [] + + let loop_toList !acc = pullX i eat_toList eject_toList + where eat_toList x = loop_toList (x : acc) + eject_toList = writeRefs refs i (reverse acc) + {-# INLINE loop_toList #-} + + loop_toList [] + xx <- readRefs refs i + return xx +{-# INLINE_FLOW toList1 #-} + + +-- | Drain the given number of elements from a single source into a list. +takeList1 :: States i m + => Int -> i -> Sources i m a -> m [a] +takeList1 len i (Sources n pullX) + = do + refs <- newRefs n [] + + let loop_toList !ix !acc + | ix >= len = writeRefs refs i (reverse acc) + | otherwise = pullX i eat_toList eject_toList + where eat_toList x = loop_toList (ix + 1) (x : acc) + eject_toList = writeRefs refs i (reverse acc) + {-# INLINE loop_toList #-} + + loop_toList 0 [] + xx <- readRefs refs i + return xx +{-# INLINE_FLOW takeList1 #-} + + +------------------------------------------------------------------------------- +-- | Push elements into the associated streams of a bundle of sinks. +pushList :: Monad m => [(i, a)] -> Sinks i m a -> m () +pushList xx (Sinks _nSinks eat _eject) + = loop_pushList xx + where + loop_pushList [] + = return () + + loop_pushList ((i, x) : ixs) + = do eat i x + loop_pushList ixs +{-# INLINE_FLOW pushList #-} + + +-- | Push the elements of a list into the given stream of a +-- bundle of sinks. +pushList1 :: Monad m => i -> [a] -> Sinks i m a -> m () +pushList1 i xx (Sinks _nSinks eat _eject) + = loop_pushList1 xx + where + loop_pushList1 [] + = return () + + loop_pushList1 (x : xs) + = do eat i x + loop_pushList1 xs +{-# INLINE_FLOW pushList1 #-} diff --git a/Data/Repa/Flow/Generic/Map.hs b/Data/Repa/Flow/Generic/Map.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow/Generic/Map.hs @@ -0,0 +1,130 @@ + +module Data.Repa.Flow.Generic.Map + ( map_i, map_o + , smap_i, smap_o + + , szipWith_ii, szipWith_io, szipWith_oi) +where +import Data.Repa.Flow.Generic.Base +import Control.Monad +import Prelude as P +#include "repa-flow.h" + + +-- | Apply a function to every element pulled from some sources, +-- producing some new sources. +map_i :: Monad m + => (a -> b) -> Sources i m a -> m (Sources i m b) +map_i f s = smap_i (\_ x -> f x) s +{-# INLINE map_i #-} + + +-- | Like `map_i`, but the worker function is also given the stream index. +smap_i :: Monad m + => (i -> a -> b) -> Sources i m a -> m (Sources i m b) +smap_i f (Sources n pullsA) + = return $ Sources n pullsB_map + where + pullsB_map i eat eject + = pullsA i eat_a eject_a + where + eat_a v = eat (f i v) + {-# INLINE eat_a #-} + + eject_a = eject + {-# INLINE eject_a #-} + + {-# INLINE [1] pullsB_map #-} +{-# INLINE_FLOW smap_i #-} + + +-- | Apply a function to every element pulled from some sources, +-- producing some new sources. +map_o :: Monad m + => (a -> b) -> Sinks i m b -> m (Sinks i m a) +map_o f k = smap_o (\_ x -> f x) k +{-# INLINE map_o #-} + + +-- | Like `map_o`, but the worker function is also given the stream index. +smap_o :: Monad m + => (i -> a -> b) -> Sinks i m b -> m (Sinks i m a) +smap_o f (Sinks n pushB ejectB) + = return $ Sinks n pushA_map ejectA_map + where + pushA_map i a = pushB i (f i a) + {-# INLINE pushA_map #-} + + ejectA_map i = ejectB i + {-# INLINE ejectA_map #-} +{-# INLINE_FLOW smap_o #-} + + +-- | Combine the elements of two flows with the given function. +-- The worker function is also given the stream index. +szipWith_ii + :: (Ord i, Monad m) + => (i -> a -> b -> c) + -> Sources i m a -> Sources i m b + -> m (Sources i m c) + +szipWith_ii f (Sources nA pullA) (Sources nB pullB) + = return $ Sources (min nA nB) pull_szipWith + where + pull_szipWith i eat eject + = pullA i eatA eject + where + eatA xA = pullB i eatB eject + where + eatB xB = eat (f i xA xB) + {-# INLINE eatB #-} + {-# INLINE eatA #-} + {-# INLINE pull_szipWith #-} +{-# INLINE_FLOW szipWith_ii #-} + + +-- | Like `szipWith_ii`, but take a bundle of `Sinks` for the result +-- elements, and yield a bundle of `Sinks` to accept the @b@ elements. +szipWith_io + :: (Ord i, Monad m) + => (i -> a -> b -> c) + -> Sinks i m c -> Sources i m a + -> m (Sinks i m b) + +szipWith_io f (Sinks nC pushC ejectC) (Sources nA pullA) + = return $ Sinks nB pushB ejectC + where + !nB = min nC nA + + pushB i xB + | i > nB = return () + | otherwise = pullA i eatA (ejectC i) + where + eatA xA = pushC i (f i xA xB) + {-# INLINE eatA #-} + {-# INLINE pushB #-} +{-# INLINE_FLOW szipWith_io #-} + + +-- | Like `szipWith_ii`, but take a bundle of `Sinks` for the result +-- elements, and yield a bundle of `Sinks` to accept the @a@ elements. +szipWith_oi + :: (Ord i, Monad m) + => (i -> a -> b -> c) + -> Sinks i m c -> Sources i m b + -> m (Sinks i m a) + +szipWith_oi f (Sinks nC pushC ejectC) (Sources nB pullB) + = return $ Sinks nA pushA ejectC + where + !nA = min nC nB + + pushA i xA + | i > nA = return () + | otherwise = pullB i eatB (ejectC i) + where + eatB xB = pushC i (f i xA xB) + {-# INLINE eatB #-} + {-# INLINE pushA #-} +{-# INLINE_FLOW szipWith_oi #-} + diff --git a/Data/Repa/Flow/Generic/Operator.hs b/Data/Repa/Flow/Generic/Operator.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow/Generic/Operator.hs @@ -0,0 +1,428 @@ +{-# OPTIONS -fno-warn-unused-imports #-} +module Data.Repa.Flow.Generic.Operator + ( -- * Projection + project_i + , project_o + + -- * Funneling + , funnel_o + + -- * Constructors + , repeat_i + , replicate_i + , prepend_i, prependOn_i + + -- * Splitting + , head_i + + -- * Grouping + , groups_i + + -- * Packing + , pack_ii + + -- * Folding + , folds_ii + + -- * Watching + , watch_i + , watch_o + , trigger_o + + -- * Capturing + , capture_o + , rcapture_o + + -- * Ignorance + , discard_o + , ignore_o + + -- * Tracing + , trace_o) +where +import Data.Repa.Flow.Generic.Eval +import Data.Repa.Flow.Generic.List +import Data.Repa.Flow.Generic.Connect +import Data.Repa.Flow.Generic.Base +import Data.Repa.Array as A +import Data.IORef +import Control.Monad +import Debug.Trace +import GHC.Exts +import Prelude as P +#include "repa-flow.h" + + +-- Projection ----------------------------------------------------------------- +-- | Project out a single stream source from a bundle. +project_i :: Monad m + => i -> Sources i m a -> m (Sources () m a) +project_i ix (Sources _ pull) + = return $ Sources () pull_project + where pull_project _ eat eject + = pull ix eat eject +{-# INLINE_FLOW project_i #-} + + +-- | Project out a single stream sink from a bundle. +project_o :: Monad m + => i -> Sinks i m a -> m (Sinks () m a) +project_o ix (Sinks _ push eject) + = return $ Sinks () push_project eject_project + where + push_project _ v = push ix v + eject_project _ = eject ix +{-# INLINE_FLOW project_o #-} + + +-- Constructors --------------------------------------------------------------- +-- | Yield sources that always produce the same value. +repeat_i :: Monad m + => i -> (i -> a) + -> m (Sources i m a) +repeat_i n f + = return $ Sources n pull_repeat + where pull_repeat i eat _eject + = eat (f i) + {-# INLINE pull_repeat #-} +{-# INLINE_FLOW repeat_i #-} + + +-- | Yield sources of the given length that always produce the same value. +replicate_i + :: States i m + => i -> Int -> (i -> a) + -> m (Sources i m a) + +replicate_i n len f + = do + refs <- newRefs n 0 + let pull_replicate i eat eject + = do !n' <- readRefs refs i + if n' >= len + then eject + else eat (f i) + {-# INLINE pull_replicate #-} + + return $ Sources n pull_replicate +{-# INLINE_FLOW replicate_i #-} + + +-- | Prepend some more elements into the front of some sources. +prepend_i :: States i m + => [a] -> Sources i m a -> m (Sources i m a) +prepend_i xs (Sources n pullX) + = do + refs <- newRefs n xs + + let pull_prepend i eat eject + = do xs' <- readRefs refs i + case xs' of + x : xs'' -> do + writeRefs refs i xs'' + eat x + + [] -> pullX i eat eject + {-# INLINE pull_prepend #-} + + return (Sources n pull_prepend) +{-# INLINE_FLOW prepend_i #-} + + +-- | Like `prepend_i` but only prepend the elements to the streams +-- that match the given predicate. +prependOn_i + :: States i m + => (i -> Bool) -> [a] -> Sources i m a -> m (Sources i m a) +prependOn_i p xs (Sources n pullX) + = do + refs <- newRefs n xs + + let pull_prependOn i eat eject + | p i + = do xs' <- readRefs refs i + case xs' of + x : xs'' -> do + writeRefs refs i xs'' + eat x + + [] -> pullX i eat eject + + | otherwise + = pullX i eat eject + {-# INLINE pull_prependOn #-} + + return (Sources n pull_prependOn) +{-# INLINE_FLOW prependOn_i #-} + + +-- Splitting ------------------------------------------------------------------ +-- | Split the given number of elements from the head of a source +-- returning those elements in a list, and yielding a new source +-- for the rest. +head_i :: States i m + => Int -> Sources i m a -> i -> m ([a], Sources i m a) +head_i len s0 i + = do + (s1, s2) <- connect_i s0 + xs <- takeList1 len i s1 + return (xs, s2) +{-# INLINE head_i #-} + + +-- Groups --------------------------------------------------------------------- +-- | From a stream of values which has consecutive runs of idential values, +-- produce a stream of the lengths of these runs. +-- +-- Example: groups [4, 4, 4, 3, 3, 1, 1, 1, 4] = [3, 2, 3, 1] +-- +groups_i + :: (Ord i, Monad m, Eq a) + => Sources i m a -> m (Sources i m Int) + +groups_i (Sources n pullV) + = return $ Sources n pull_n + where + -- Pull a whole run from the source, so that we can produce. + -- the output element. + pull_n i eat eject + = loop_groups Nothing 1# + where + loop_groups !mx !count + = pullV i eat_v eject_v + where eat_v v + = case mx of + -- This is the first element that we've read from + -- the source. + Nothing -> loop_groups (Just v) count + + -- See if the next element is the same as the one + -- we read previously + Just x -> if x == v + then loop_groups (Just x) (count +# 1#) + else eat (I# count) + -- TODO: ** STORE PULLED VALUE FOR LATER + {-# INLINE eat_v #-} + + eject_v + = case mx of + -- We've already written our last count, + -- and there are no more elements in the source. + Nothing -> eject + + -- There are no more elements in the source, + -- so emit the final count + Just _ -> eat (I# count) + {-# INLINE eject_v #-} + {-# INLINE loop_groups #-} + {-# INLINE pull_n #-} +{-# INLINE_FLOW groups_i #-} + + +-- Pack ----------------------------------------------------------------------- +-- | Given a stream of flags and a stream of values, produce a new stream +-- of values where the corresponding flag was True. The length of the result +-- is the length of the shorter of the two inputs. +pack_ii :: (Ord i, Monad m) + => Sources i m Bool -> Sources i m a -> m (Sources i m a) + +pack_ii (Sources nF pullF) (Sources nX pullX) + = return $ Sources (min nF nX) pull_pack + where + pull_pack i eat eject + = pullF i eat_f eject_f + where eat_f f = pack_x f + eject_f = eject + + pack_x f + = pullX i eat_x eject_x + where eat_x x = if f then eat x + else pull_pack i eat eject + + eject_x = eject + {-# INLINE pack_x #-} + {-# INLINE pull_pack #-} +{-# INLINE_FLOW pack_ii #-} + + +-- Folds ---------------------------------------------------------------------- +-- | Segmented fold. +folds_ii + :: (Ord i, Monad m) + => (a -> a -> a) -> a + -> Sources i m Int + -> Sources i m a + -> m (Sources i m a) + +folds_ii f z (Sources nL pullLen) + (Sources nX pullX) + = return $ Sources (min nL nX) pull_folds + where + pull_folds i eat eject + = pullLen i eat_len eject_len + where + eat_len (I# len) = loop_folds len z + eject_len = eject + + loop_folds !c !acc + | tagToEnum# (c ==# 0#) = eat acc + | otherwise + = pullX i eat_x eject_x + where + eat_x x = loop_folds (c -# 1#) (f acc x) + eject_x = eject + {-# INLINE loop_folds #-} + {-# INLINE pull_folds #-} +{-# INLINE_FLOW folds_ii #-} + + +-- Watch ---------------------------------------------------------------------- +-- | Apply a monadic function to every element pulled from some sources, +-- producing some new sources. +watch_i :: Monad m + => (i -> a -> m ()) + -> Sources i m a -> m (Sources i m a) + +watch_i f (Sources n pullX) + = return $ Sources n pull_watch + where + pull_watch i eat eject + = pullX i eat_watch eject_watch + where + eat_watch x = f i x >> eat x + eject_watch = eject + {-# INLINE pull_watch #-} +{-# INLINE_FLOW watch_i #-} + + +-- | Pass elements to the provided action as they are pushed into the sink. +watch_o :: Monad m + => (i -> a -> m ()) + -> Sinks i m a -> m (Sinks i m a) + +watch_o f (Sinks n push eject) + = return $ Sinks n push_watch eject_watch + where + push_watch !i !x = f i x >> push i x + eject_watch !i = eject i +{-# INLINE_FLOW watch_o #-} + + +-- | Create a bundle of sinks of the given arity and capture any +-- data pushed to it. +-- +-- @ +-- > import Data.Repa.Flow.Generic +-- > import Data.Repa.Array.Material +-- > import Data.Repa.Nice +-- > import Control.Monad +-- > liftM nice $ capture_o B 4 (\k -> pushList [(0 :: Int, "foo"), (1, "bar"), (0, "baz")] k) +-- > [(0,"foo"),(1,"bar"),(0,"baz")] +-- @ +-- +--- +-- TODO: avoid going via lists when accumulating. +-- +capture_o + :: (Target lDst (i, a), Index lDst ~ Int) + => Name lDst -- ^ Name of desination layout. + -> i -- ^ Arity of result bundle. + -> (Sinks i IO a -> IO ()) -- ^ Function to push data into the sinks. + -> IO (Array lDst (i, a)) + +capture_o nDst n use + = liftM fst $ rcapture_o nDst n use +{-# INLINE capture_o #-} + + +-- | Like `capture_o` but also return the @r@-esult of the push function. +rcapture_o + :: (Target lDst (i, a), Index lDst ~ Int) + => Name lDst -- ^ Name of desination layout. + -> i -- ^ Arity of result bundle. + -> (Sinks i IO a -> IO b) -- ^ Function to push data into the sinks. + -> IO (Array lDst (i, a), b) + +rcapture_o nDst n use + = do + ref <- newIORef [] + + let capture_eat i x + = atomicModifyIORef ref (\old -> ((i, x) : old, ())) + {-# INLINE capture_eat #-} + + k0 <- discard_o n + k1 <- watch_o capture_eat k0 + + x <- use k1 + result <- readIORef ref + let !arr = A.fromList nDst $ P.reverse result + + return (arr, x) +{-# INLINE_FLOW rcapture_o #-} + + +-- | Like `watch_o` but doesn't pass elements to another sink. +trigger_o :: Monad m + => i -> (i -> a -> m ()) -> m (Sinks i m a) +trigger_o i f + = discard_o i >>= watch_o f +{-# INLINE trigger_o #-} + + +-- Ignorance------------------------------------------------------------------- +-- | A sink that drops all data on the floor. +-- +-- This sink is strict in the elements, so they are demanded before being +-- discarded. Haskell debugging thunks attached to the elements will be +-- demanded. +-- +discard_o :: Monad m + => i -> m (Sinks i m a) +discard_o n + = return $ Sinks n push_discard eject_discard + where + -- IMPORTANT: push_discard should be strict in the element so that + -- and Haskell tracing thunks attached to it are evaluated. + -- We *discard* the elements, but don't completely ignore them. + push_discard !_ !_ = return () + eject_discard !_ = return () +{-# INLINE_FLOW discard_o #-} + + +-- | A sink that ignores all incoming data. +-- +-- This sink is non-strict in the elements. +-- Haskell tracing thunks attached to the elements will *not* be demanded. +-- +ignore_o :: Monad m + => i -> m (Sinks i m a) +ignore_o n + = return $ Sinks n push_ignore eject_ignore + where + push_ignore _ _ = return () + eject_ignore _ = return () +{-# INLINE_FLOW ignore_o #-} + + +-- Trace ---------------------------------------------------------------------- +-- | Use the `trace` function from "Debug.Trace" to print each element +-- that is pushed to a sink. +-- +-- * This function is intended for debugging only, and is not intended +-- for production code. +-- +trace_o :: (Show i, Show a, Monad m) + => i -> m (Sinks i m a) + +trace_o nSinks + = trigger_o nSinks eat + where + eat i x + = trace ("repa-flow trace_o: " ++ show i ++ "; " ++ show x) + (return ()) + +{-# NOINLINE trace_o #-} + + + diff --git a/Data/Repa/Flow/IO/Bucket.hs b/Data/Repa/Flow/IO/Bucket.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow/IO/Bucket.hs @@ -0,0 +1,504 @@ + +module Data.Repa.Flow.IO.Bucket + ( Bucket + , bucketLength + , openBucket + , hBucket + + -- * Reading + , fromFiles, fromFiles' + , fromDir + , fromSplitFile + , fromSplitFileAt + + -- * Writing + , toFiles, toFiles' + , toDir + , toDirs' + + -- * Bucket IO + , bClose + , bIsOpen + , bAtEnd + , bSeek + , bGetArray + , bPutArray) +where +import Data.Repa.Array as A +import Data.Repa.Array.Material as A +import Data.Repa.Array.RowWise as A +import Data.Repa.IO.Array as A +import qualified Foreign.Storable as Foreign +import qualified Foreign.Marshal.Alloc as Foreign +import Control.Monad +import Data.Word +import System.IO +import System.FilePath +import System.Directory +import Prelude as P + + +-- | A bucket represents portion of a whole data-set on disk, +-- and contains a file handle that points to the next piece of +-- data to be read or written. +-- +-- The bucket could be created from a portion of a single flat file, +-- or be one file of a pre-split data set. The main advantage over a +-- plain `Handle` is that a `Bucket` can represent a small portion +-- of a single large file. +-- +data Bucket + = Bucket + { -- | Physical location of the file, if known. + bucketFilePath :: Maybe FilePath + + -- | Starting position of the bucket in the file, in bytes. + , bucketStartPos :: Integer + + -- | Maximum length of the bucket, in bytes. + -- + -- If `Nothing` then the length is indeterminate, which is used + -- when writing to files. + , bucketLength :: Maybe Integer + + -- | File handle for the bucket. + -- + -- If several buckets have been created from a single file, + -- then all buckets will have handles bound to that file, + -- but they will be at different positions. + , bucketHandle :: Handle } + + +-- | Open a file as a single bucket. +openBucket :: FilePath -> IOMode -> IO Bucket +openBucket path mode + = do h <- openBinaryFile path mode + hSeek h SeekFromEnd 0 + lenTotal <- hTell h + hSeek h AbsoluteSeek 0 + + return $ Bucket + { bucketFilePath = Just path + , bucketStartPos = 0 + , bucketLength = Just lenTotal + , bucketHandle = h } +{-# NOINLINE openBucket #-} + + +-- | Wrap an existing file handle as a bucket. +hBucket :: Handle -> IO Bucket +hBucket h + = do pos <- hTell h + return $ Bucket + { bucketFilePath = Nothing + , bucketStartPos = pos + , bucketLength = Nothing + , bucketHandle = h } +{-# NOINLINE hBucket #-} + + +-- From Files ----------------------------------------------------------------- +-- | Open some files as buckets and use them as `Sources`. +fromFiles + :: (Bulk l FilePath, Target l Bucket) + => Array l FilePath -- ^ Files to open. + -> (Array l Bucket -> IO b) + -- ^ Consumer. + -> IO b + +fromFiles paths use + = do + -- Open all the files, ending up with a list of buckets. + bs <- mapM (flip openBucket ReadMode) $ A.toList paths + + -- Pack buckets back into an array with the same layout as + -- the original. + let Just bsArr = A.fromListInto (A.layout paths) bs + + use bsArr +{-# NOINLINE fromFiles #-} + + +-- | Like `fromFiles`, but take a list of file paths. +fromFiles' + :: [FilePath] + -> (Array B Bucket -> IO b) + -> IO b +fromFiles' files use + = fromFiles (A.fromList B files) use +{-# INLINE fromFiles' #-} + + +-- | Open all the files in a directory as separate buckets. +-- +-- This operation may fail with the same exceptions as `getDirectoryContents`. +-- +fromDir :: FilePath + -> (Array B Bucket -> IO b) + -> IO b + +fromDir dir use + = do fs <- getDirectoryContents dir + let fsRel + = P.map (dir ) + $ P.filter (\f -> f /= "." && f /= "..") fs + fromFiles' fsRel use +{-# INLINE fromDir #-} + + +-- | Open a file containing atomic records and split it into the given number +-- of evenly sized buckets. +-- +-- The records are separated by a special terminating charater, which the +-- given predicate detects. The file is split cleanly on record boundaries, +-- so we get a whole number of records in each bucket. As the records can be +-- of varying size the buckets are not guaranteed to have exactly the same +-- length, in either records or buckets, though we try to give them the +-- approximatly the same number of bytes. +-- +fromSplitFile + :: Int -- ^ Number of buckets. + -> (Word8 -> Bool) -- ^ Detect the end of a record. + -> FilePath -- ^ File to open. + -> (Array B Bucket -> IO b) -- ^ Consumer. + -> IO b + +fromSplitFile n pEnd path use + = fromSplitFileAt n pEnd path 0 use +{-# INLINE fromSplitFile #-} + + +-- | Like `fromSplitFile` but start at the given offset. +fromSplitFileAt + :: Int -- ^ Number of buckets. + -> (Word8 -> Bool) -- ^ Detect the end of a record. + -> FilePath -- ^ File to open. + -> Integer -- ^ Starting offset. + -> (Array B Bucket -> IO b) -- ^ Consumer. + -> IO b + +fromSplitFileAt n pEnd path offsetStart use + = do + -- Open the file first to check its length. + h0 <- openBinaryFile path ReadMode + hSeek h0 SeekFromEnd 0 + lenTotal <- hTell h0 + hClose h0 + + -- Open a file handle for each of the buckets. + -- The handles start at the begining of the file and still need + -- to be advanced. + -- + -- TODO: check at least one elem in list + hh@(h1_ : _) <- mapM (flip openBinaryFile ReadMode) (replicate n path) + + hSeek h1_ AbsoluteSeek offsetStart + + -- Advance all the handles to the start of their part of the file. + let loop_advances _ _ [] + = return [] + + loop_advances _ pos1 (_h1 : []) + = return [pos1] + + loop_advances remain pos1 (h1 : h2 : hs) + = do + -- Push the next handle an even distance into the + -- remaining part of the file. + let lenWanted + = remain `div` (fromIntegral $ P.length (h1 : h2 : hs)) + let posWanted = pos1 + lenWanted + hSeek h2 AbsoluteSeek posWanted + + -- Now advance it until we get to the end of a record. + pos2 <- advance h2 pEnd + let remain' = lenTotal - pos2 + + poss <- loop_advances remain' pos2 (h2 : hs) + return $ pos1 : poss + + starts <- loop_advances lenTotal offsetStart hh + + -- Ending positions and lengths for each bucket. + let ends = tail (starts ++ [lenTotal]) + let lens = P.map (\(start, end) -> end - start) + $ P.zip starts ends + + let bs = [ Bucket + { bucketFilePath = Just path + , bucketStartPos = start + , bucketLength = Just len + , bucketHandle = h } + | start <- starts + | len <- lens + | h <- hh ] + + use $ A.fromList B bs +{-# NOINLINE fromSplitFileAt #-} +-- NOINLINE to avoid polluting the core code of the consumer. +-- This prevents it from being specialised for the pEnd predicate, +-- but we're expecting pEnd to be applied a small number of times, +-- so it shouldn't matter. + + +-- | Advance a file handle until we reach a byte that, matches the given +-- predicate, then return the final file position. +advance :: Handle -> (Word8 -> Bool) -> IO Integer +advance h pEnd + = do buf <- Foreign.mallocBytes 1 + + let loop_advance + = do c <- hGetBuf h buf 1 + if c == 0 + then return () + else do + x <- Foreign.peek buf + if pEnd x + then return () + else loop_advance + loop_advance + Foreign.free buf + hTell h +{-# NOINLINE advance #-} + + +-- Writing -------------------------------------------------------------------- +-- | Open some files for writing as individual buckets and pass +-- them to the given consumer. +-- +toFiles :: (Bulk l FilePath, Target l Bucket) + => Array l FilePath -- ^ File paths. + -> (Array l Bucket -> IO b) + -- ^ Consumer. + -> IO b + +toFiles paths use + = do -- Open all the files, ending up with a list of buckets. + bs <- mapM (flip openBucket WriteMode) $ A.toList paths + + -- Pack buckets back into an array with the same layout as + -- the original. + let Just bsArr = A.fromListInto (A.layout paths) bs + + use bsArr +{-# NOINLINE toFiles #-} +--- +-- TODO: Attached finalizers to the sinks so that file assocated with +-- each stream is closed when that stream is ejected. + + +-- | Like `toFiles`, but take a list of file paths. +toFiles' :: [FilePath] + -> (Array B Bucket -> IO b) + -> IO b + +toFiles' paths use + = toFiles (A.fromList B paths) use +{-# INLINE toFiles' #-} + + +-- | Create a new directory of the given name, containing the given number +-- of buckets. +-- +-- If the directory is named @somedir@ then the files are named +-- @somedir/000000@, @somedir/000001@, @somedir/000002@ and so on. +toDir :: Int -- ^ Number of buckets to create. + -> FilePath -- ^ Path to directory. + -> (Array B Bucket -> IO b) -- ^ Consumer. + -> IO b + +toDir nBuckets path use + | nBuckets <= 0 + = use (A.fromList B []) + + | otherwise + = do + createDirectory path + + let makeName i = path ((replicate (6 - (P.length $ show i)) '0') ++ show i) + let names = [makeName i | i <- [0 .. nBuckets - 1]] + + let newBucket file + = do h <- openBinaryFile file WriteMode + return $ Bucket + { bucketFilePath = Just file + , bucketStartPos = 0 + , bucketLength = Nothing + , bucketHandle = h } + + bs <- mapM newBucket names + use (A.fromList B bs) +{-# NOINLINE toDir #-} + + +-- | Given a list of directories, create those directories and open +-- the given number of output files per directory. +-- +-- In the resulting array of buckets, the outer dimension indexes +-- each directory, and the inner one indexes each file in its +-- directory. +-- +-- For each directory @somedir@ the files are named +-- @somedir/000000@, @somedir/000001@, @somedir/000002@ and so on. +-- +toDirs' :: Int -- ^ Number of buckets to create per directory. + -> [FilePath] -- ^ Paths to directories. + -> (Array (E B DIM2) Bucket -> IO b) + -- ^ Consumer. + -> IO b + +toDirs' nBucketsPerDir paths use + | nBucketsPerDir <= 0 + = do let Just bsArr + = A.fromListInto (A.matrix B 0 0) [] + use bsArr + + | otherwise + = do + let makeName path i + = path ((replicate (6 - (P.length $ show i)) '0') ++ show i) + + let newBucket file + = do h <- openBinaryFile file WriteMode + return $ Bucket + { bucketFilePath = Just file + , bucketStartPos = 0 + , bucketLength = Nothing + , bucketHandle = h } + + let newDir path + = do createDirectory path + bs <- mapM newBucket + $ [makeName path i | i <- [0 .. nBucketsPerDir - 1]] + + return bs + + -- Make all the buckets, then pack them into a matrix. + bs <- liftM P.concat $ P.mapM newDir paths + + -- Pack the buckets into an array + let Just bsArr + = A.fromListInto + (A.matrix B (P.length paths) nBucketsPerDir) + bs + + use bsArr +{-# NOINLINE toDirs' #-} + + + +-- Bucket IO ------------------------------------------------------------------ +-- | Close a bucket, releasing the contained file handle. +bClose :: Bucket -> IO () +bClose bucket + = hClose $ bucketHandle bucket +{-# NOINLINE bClose #-} + + +-- | Check if the bucket is currently open. +bIsOpen :: Bucket -> IO Bool +bIsOpen bucket + = hIsOpen $ bucketHandle bucket +{-# NOINLINE bIsOpen #-} + + +-- | Check if the contained file handle is at the end of the bucket. +bAtEnd :: Bucket -> IO Bool +bAtEnd bucket + = do eof <- hIsEOF $ bucketHandle bucket + + -- Position in the file. + posFile <- hTell $ bucketHandle bucket + + -- Check for bogus position before we subtract the startPos. + -- If this happenes then something has messed with our handle. + when (posFile < bucketStartPos bucket) + $ error $ P.unlines + [ "repa-flow.bAtEnd: handle position is outside bucket." + , " bucket file path = " ++ show (bucketFilePath bucket) + , " bucket start pos = " ++ show (bucketStartPos bucket) + , " pos in file = " ++ show posFile ] + + -- Position in the bucket. + let posBucket = posFile - bucketStartPos bucket + + return $ eof || (case bucketLength bucket of + Nothing -> False + Just len -> posBucket >= len) +{-# NOINLINE bAtEnd #-} + + +-- | Seek to a position with a bucket. +bSeek :: Bucket -> SeekMode -> Integer -> IO () +bSeek bucket mode offset + = do + -- The current position in the underlying file. + posFile <- hTell $ bucketHandle bucket + + -- Apply the seek mode to get the wanted position in the file, + -- which might be outside the bucket. + -- If Nothing it means the end of the file. + let posWanted + = case mode of + AbsoluteSeek + -> Just $ bucketStartPos bucket + max 0 offset + + RelativeSeek + -> Just $ posFile + offset + + SeekFromEnd + -> case bucketLength bucket of + Nothing -> Nothing + Just len -> Just $ bucketStartPos bucket + + len - max 0 offset + + -- Clip the wanted position so that it's inside the bucket. + let posActual + = case posWanted of + Nothing + -> Nothing + + Just wanted + | Just len <- bucketLength bucket + -> if wanted < bucketStartPos bucket + then Just $ bucketStartPos bucket + else if wanted > bucketStartPos bucket + len + then Just $ bucketStartPos bucket + len + else Just wanted + + | otherwise + -> if wanted < bucketStartPos bucket + then Just $ bucketStartPos bucket + else Just wanted + + case posActual of + Nothing -> hSeek (bucketHandle bucket) SeekFromEnd 0 + Just pos -> hSeek (bucketHandle bucket) AbsoluteSeek pos +{-# NOINLINE bSeek #-} + + +-- | Get some data from a bucket. +bGetArray :: Bucket -> Integer -> IO (Array F Word8) +bGetArray bucket lenWanted + = do + -- Curent position in the file. + posFile <- hTell $ bucketHandle bucket + let posBucket = posFile - bucketStartPos bucket + + let len = case bucketLength bucket of + Nothing -> lenWanted + Just lenMax + -> let lenRemain = lenMax - posBucket + in min lenWanted lenRemain + + hGetArray (bucketHandle bucket) + $ fromIntegral len +{-# NOINLINE bGetArray #-} + + +-- | Put some data in a bucket. +bPutArray :: Bucket -> Array F Word8 -> IO () +bPutArray bucket arr + = hPutArray (bucketHandle bucket) arr +{-# NOINLINE bPutArray #-} + diff --git a/Data/Repa/Flow/Simple.hs b/Data/Repa/Flow/Simple.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow/Simple.hs @@ -0,0 +1,77 @@ + +module Data.Repa.Flow.Simple + ( module Data.Repa.Flow.States + , Source + , Sink + + -- * Evaluation + , drainS + + -- * Conversions + , fromList + , toList + , takeList + + -- * Finalizers + , finalize_i + , finalize_o + + -- * Flow Operators + -- ** Constructors + , repeat_i + , replicate_i + , prepend_i + + -- ** Mapping + , map_i, map_o + + -- ** Connecting + , dup_oo, dup_io, dup_oi + , connect_i + + -- ** Splitting + , head_i + , peek_i + + -- ** Grouping + , groups_i + + -- ** Packing + , pack_ii + + -- ** Folding + , folds_ii + + -- ** Watching + , watch_i + , watch_o + , trigger_o + + -- ** Ignorance + , discard_o + , ignore_o + + -- * Flow IO + -- ** Sourcing + , fromFiles + , sourceBytes + , sourceRecords + + -- ** Sinking + , toFiles + , sinkBytes) +where +import Data.Repa.Flow.States +import Data.Repa.Flow.Simple.Base +import Data.Repa.Flow.Simple.List +import Data.Repa.Flow.Simple.Operator +import Data.Repa.Flow.Simple.IO +import qualified Data.Repa.Flow.Generic.Eval as G +#include "repa-flow.h" + + +-- | Pull all available values from the source and push them to the sink. +drainS :: Monad m + => Source m a -> Sink m a -> m () +drainS = G.drainS +{-# INLINE drainS #-} diff --git a/Data/Repa/Flow/Simple/Base.hs b/Data/Repa/Flow/Simple/Base.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow/Simple/Base.hs @@ -0,0 +1,75 @@ + +module Data.Repa.Flow.Simple.Base + ( Source, Sink + , finalize_i + , finalize_o + , wrapI_i + , wrapI_o) +where +import Data.Repa.Flow.States +import qualified Data.Repa.Flow.Generic as G +#include "repa-flow.h" + + +-- | Source consisting of a single stream. +type Source m e = G.Sources () m e + +-- | Sink consisting of a single stream. +type Sink m e = G.Sinks () m e + + +-- Finalizers ----------------------------------------------------------------- +-- | Attach a finalizer to a source. +-- +-- The finalizer will be called the first time a consumer of that stream +-- tries to pull an element when no more are available. +-- +-- The provided finalizer will be run after any finalizers already +-- attached to the source. +-- +finalize_i + :: States () m + => m () + -> Source m a -> m (Source m a) + +finalize_i f s0 = G.finalize_i (\_ -> f) s0 +{-# INLINE finalize_i #-} + + +-- | Attach a finalizer to a sink. +-- +-- The finalizer will be called the first time the stream is ejected. +-- +-- The provided finalizer will be run after any finalizers already +-- attached to the sink. +-- +finalize_o + :: States () m + => m () + -> Sink m a -> m (Sink m a) + +finalize_o f s0 = G.finalize_o (\_ -> f) s0 +{-# INLINE finalize_o #-} + + +-- Wrapping ------------------------------------------------------------------- +wrapI_i :: G.Sources Int m e -> Maybe (Source m e) +wrapI_i (G.Sources n pullX) + | n /= 1 = Nothing + | otherwise + = let pullX' _ eat eject + = pullX 0 eat eject + {-# INLINE pullX' #-} + in Just $ G.Sources () pullX' +{-# INLINE_FLOW wrapI_i #-} + + +wrapI_o :: G.Sinks Int m e -> Maybe (Sink m e) +wrapI_o (G.Sinks n eatX ejectX) + | n /= 1 = Nothing + | otherwise + = let eatX' _ x = eatX 0 x + ejectX' _ = ejectX 0 + in Just $ G.Sinks () eatX' ejectX' +{-# INLINE_FLOW wrapI_o #-} + diff --git a/Data/Repa/Flow/Simple/IO.hs b/Data/Repa/Flow/Simple/IO.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow/Simple/IO.hs @@ -0,0 +1,84 @@ + +module Data.Repa.Flow.Simple.IO + ( G.fromFiles + , sourceBytes + , sourceRecords + , G.toFiles + , sinkBytes) +where +import Data.Repa.Flow.IO.Bucket +import Data.Repa.Flow.Simple.Base +import Data.Word +import Data.Repa.Array as A +import Data.Repa.Array.Material as A +import qualified Data.Repa.Flow.Generic.IO as G +#include "repa-flow.h" + + +-- Source Records --------------------------------------------------------------------------------- +-- | Read complete records of data from a file, using the given chunk length +-- +-- The records are separated by a special terminating character, which the +-- given predicate detects. After reading a chunk of data we seek to just after the +-- last complete record that was read, so we can continue to read more complete +-- records next time. +-- +-- If we cannot find an end-of-record terminator in the chunk then apply the given +-- failure action. The records can be no longer than the chunk length. This fact +-- guards against the case where a large input file is malformed and contains no +-- end-of-record terminators, as we won't try to read the whole file into memory. +-- +-- * Data is read into foreign memory without copying it through the GHC heap. +-- * All chunks have the same size, except possibly the last one. +-- * The provided file handle must support seeking, else you'll get an exception. +-- +-- The file will be closed the first time the consumer tries to pull an element +-- from the associated stream when no more are available. +-- +sourceRecords + :: Integer -- ^ Size of chunk to read in bytes. + -> (Word8 -> Bool) -- ^ Detect the end of a record. + -> IO () -- ^ Action to perform if we can't get a whole record. + -> Bucket -- ^ File handle. + -> IO (Source IO (Array N (Array F Word8))) + +sourceRecords len pSep aFail b + = do s0 <- G.sourceRecords len pSep aFail (A.fromList B [b]) + let Just s1 = wrapI_i s0 + return s1 +{-# INLINE sourceRecords #-} + + +-- Source Bytes ----------------------------------------------------------------------------------- +-- | Read data from a file, using the given chunk length. +-- +-- * Data is read into foreign memory without copying it through the GHC heap. +-- * All chunks have the same size, except possibly the last one. +-- +-- The file will be closed the first time the consumer tries to pull an element +-- from the associated stream when no more are available. +-- +sourceBytes + :: Integer + -> Bucket + -> IO (Source IO (Array F Word8)) + +sourceBytes len b + = do s0 <- G.sourceBytes len (A.fromList B [b]) + let Just s1 = wrapI_i s0 + return s1 +{-# INLINE sourceBytes #-} + + +-- Sinking Bytes ---------------------------------------------------------------------------------- +-- | Write chunks of data to the given files. +-- +-- The file will be closed when the associated stream is ejected. +-- +sinkBytes :: Bucket -> IO (Sink IO (Array F Word8)) +sinkBytes b + = do s0 <- G.sinkBytes (A.fromList B [b]) + let Just s1 = wrapI_o s0 + return s1 +{-# INLINE sinkBytes #-} + diff --git a/Data/Repa/Flow/Simple/List.hs b/Data/Repa/Flow/Simple/List.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow/Simple/List.hs @@ -0,0 +1,32 @@ + +module Data.Repa.Flow.Simple.List + ( fromList + , toList + , takeList) +where +import Data.Repa.Flow.Simple.Base +import Data.Repa.Flow.States (States) +import qualified Data.Repa.Flow.Generic as G +#include "repa-flow.h" + + +-- | Given an arity and a list of elements, yield a source that produces +-- all the elements. +fromList :: States () m + => [a] -> m (Source m a) +fromList xx = G.fromList () xx +{-# INLINE fromList #-} + + +-- | Drain a source into a list. +toList :: States () m + => Source m a -> m [a] +toList s = G.toList1 () s +{-# INLINE toList #-} + + +-- | Drain the given number of elements from a single source into a list. +takeList :: States () m + => Int -> Source m a -> m [a] +takeList len s = G.takeList1 len () s +{-# INLINE takeList #-} diff --git a/Data/Repa/Flow/Simple/Operator.hs b/Data/Repa/Flow/Simple/Operator.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow/Simple/Operator.hs @@ -0,0 +1,225 @@ + +module Data.Repa.Flow.Simple.Operator + ( -- * Constructors + repeat_i + , replicate_i + , prepend_i + + -- * Mapping + , map_i, map_o + + -- * Connecting + , dup_oo, dup_io, dup_oi + , connect_i + + -- * Splitting + , head_i + , peek_i + + -- * Grouping + , groups_i + + -- * Packing + , pack_ii + + -- * Folding + , folds_ii + + -- * Watching + , watch_i + , watch_o + , trigger_o + + -- * Ignorance + , discard_o + , ignore_o) +where +import Data.Repa.Flow.Simple.Base +import Data.Repa.Flow.States (States (..)) +import qualified Data.Repa.Flow.Generic as G +#include "repa-flow.h" + + +-- Constructors --------------------------------------------------------------- +-- | Yield a source that always produces the same value. +repeat_i :: States () m + => a -> m (Source m a) +repeat_i x + = G.repeat_i () (const x) +{-# INLINE repeat_i #-} + + +-- | Yield a source of the given length that always produces the same value. +replicate_i + :: States () m + => Int -> a -> m (Source m a) +replicate_i n x + = G.replicate_i () n (const x) +{-# INLINE replicate_i #-} + + +-- | Prepend some more elements to the front of a source. +prepend_i :: States () m + => [a] -> Source m a -> m (Source m a) +prepend_i = G.prepend_i +{-# INLINE prepend_i #-} + + +-- Mapping -------------------------------------------------------------------- +-- | Apply a function to every element pulled from some source, +-- producing a new source. +map_i :: States () m => (a -> b) -> Source m a -> m (Source m b) +map_i f s = G.smap_i (\_ x -> f x) s +{-# INLINE map_i #-} + + +-- | Apply a function to every element pushed to some sink, +-- producing a new sink. +map_o :: States () m => (a -> b) -> Sink m b -> m (Sink m a) +map_o f s = G.smap_o (\_ x -> f x) s +{-# INLINE map_o #-} + + +-- Connecting ----------------------------------------------------------------- +-- | Send the same data to two consumers. +-- +-- Given two argument sinks, yield a result sink. +-- Pushing to the result sink causes the same element to be pushed to both +-- argument sinks. +dup_oo :: States () m => Sink m a -> Sink m a -> m (Sink m a) +dup_oo = G.dup_oo +{-# INLINE dup_oo #-} + + +-- | Send the same data to two consumers. +-- +-- Given an argument source and argument sink, yield a result source. +-- Pulling an element from the result source pulls from the argument source, +-- and pushes that element to the sink, as well as returning it via the +-- result source. +dup_io :: States () m => Source m a -> Sink m a -> m (Source m a) +dup_io = G.dup_io +{-# INLINE dup_io #-} + + +-- | Send the same data to two consumers. +-- +-- Like `dup_io` but with the arguments flipped. +-- +dup_oi :: States () m => Sink m a -> Source m a -> m (Source m a) +dup_oi = G.dup_oi +{-# INLINE dup_oi #-} + + +-- | Connect an argument source to two result sources. +-- +-- Pulling from either result source pulls from the argument source. +-- Each result source only gets the elements pulled at the time, +-- so if one side pulls all the elements the other side won't get any. +connect_i :: States () m + => Source m a -> m (Source m a, Source m a) +connect_i = G.connect_i +{-# INLINE connect_i #-} + + +-- Splitting ------------------------------------------------------------------ +-- | Split the given number of elements from the head of a source, +-- returning those elements in a list, and yielding a new source +-- for the rest. +head_i :: States () m + => Int -> Source m a -> m ([a], Source m a) + +head_i len s0 + = G.head_i len s0 () +{-# INLINE head_i #-} + + +-- | Peek at the given number of elements in the stream, +-- returning a result stream that still produces them all. +peek_i :: States () m + => Int -> Source m a -> m ([a], Source m a) +peek_i n s0 + = do (s1, s2) <- G.connect_i s0 + xs <- G.takeList1 n () s1 + s3 <- G.prepend_i xs s2 + return (xs, s3) +{-# INLINE peek_i #-} + + +-- Grouping ------------------------------------------------------------------- +-- | From a stream of values which has consecutive runs of idential values, +-- produce a stream of the lengths of these runs. +-- +-- Example: groups [4, 4, 4, 3, 3, 1, 1, 1, 4] = [3, 2, 3, 1] +-- +groups_i :: (Monad m, Eq a) + => Source m a -> m (Source m Int) +groups_i = G.groups_i +{-# INLINE groups_i #-} + + +-- Packing -------------------------------------------------------------------- +-- | Given a stream of flags and a stream of values, produce a new stream +-- of values where the corresponding flag was True. The length of the result +-- is the length of the shorter of the two inputs. +pack_ii :: Monad m + => Source m Bool -> Source m a -> m (Source m a) +pack_ii s0 s1 = G.pack_ii s0 s1 +{-# INLINE pack_ii #-} + + +-- Folding -------------------------------------------------------------------- +-- | Segmented fold. +folds_ii :: Monad m + => (a -> a -> a) -> a + -> Source m Int -> Source m a + -> m (Source m a) +folds_ii f z s0 s1 = G.folds_ii f z s0 s1 +{-# INLINE folds_ii #-} + + +-- Watching ------------------------------------------------------------------- +-- | Apply a monadic function to every element pulled from a source +-- producing a new source. +watch_i :: Monad m + => (a -> m ()) + -> Source m a -> m (Source m a) +watch_i f s0 = G.watch_i (\_ x -> f x) s0 +{-# INLINE watch_i #-} + + +-- | Pass elements to the provided action as they are pushed to the sink. +watch_o :: Monad m + => (a -> m ()) + -> Sink m a -> m (Sink m a) +watch_o f s0 = G.watch_o (\_ x -> f x) s0 +{-# INLINE watch_o #-} + + +-- | Like `watch` but doesn't pass elements to another sink. +trigger_o :: Monad m + => (a -> m ()) -> m (Sink m a) +trigger_o f = G.trigger_o () (\_ x -> f x) +{-# INLINE trigger_o #-} + + +-- Ignorance ------------------------------------------------------------------ +-- | A sink that drops all data on the floor. +-- +-- This sink is strict in the elements, so they are demanded before being +-- discarded. Haskell debugging thunks attached to the elements will be demanded. +discard_o :: Monad m + => m (Sink m a) +discard_o = G.discard_o () +{-# INLINE discard_o #-} + + +-- | A sink that ignores all incoming elements. +-- +-- This sink is non-strict in the elements. +-- Haskell tracing thinks attached to the elements will *not* be demanded. +ignore_o :: Monad m + => m (Sink m a) +ignore_o = G.ignore_o () +{-# INLINE ignore_o #-} + diff --git a/Data/Repa/Flow/States.hs b/Data/Repa/Flow/States.hs new file mode 100644 --- /dev/null +++ b/Data/Repa/Flow/States.hs @@ -0,0 +1,163 @@ +{-# LANGUAGE UndecidableInstances #-} +module Data.Repa.Flow.States + ( Next (..) + , States (..) + , Refs (..) + , foldRefsM + , toListM) +where +import Control.Monad +import qualified Data.Vector.Mutable as VM +#include "repa-flow.h" + + +------------------------------------------------------------------------------- +class (Ord i, Eq i) => Next i where + + -- | Get the zero for this index type. + first :: i + + -- | Given an index an arity, get the next index after this one, + -- or `Nothing` if there aren't any more. + next :: i -> i -> Maybe i + + -- | Check if an index is valid for this arity. + check :: i -> i -> Bool + + +-- | Unit indices. +instance Next () where + first = () + next _ _ = Nothing + check _ _ = True + {-# INLINE first #-} + {-# INLINE next #-} + + +-- | Integer indices. +instance Next Int where + + first = 0 + {-# INLINE first #-} + + next i len + | i + 1 >= len = Nothing + | otherwise = Just (i + 1) + {-# INLINE next #-} + + check i len + = i >= 0 && len >= 0 && i < len + {-# INLINE check #-} + + +-- | Tuple indices. +instance Next (Int, Int) where + + first = (0, 0) + {-# INLINE first #-} + + next (ix1, ix0) (a1, a0) + | ix0 + 1 >= a0 + = if ix1 + 1 >= a1 + then Nothing + else Just (ix1 + 1, 0) + + | otherwise + = Just (ix1, ix0 + 1) + {-# INLINE next #-} + + check (ix1, ix2) (len1, len2) + = check ix1 len1 && check ix2 len2 + {-# INLINE check #-} + + +------------------------------------------------------------------------------- +class (Ord i, Next i, Monad m) => States i m where + + -- | A collection of mutable references. + data Refs i m a + + -- | Get the extent of the collection. + extentRefs :: Refs i m a -> i + + -- | Allocate a new state of the given arity, also returning an index to the + -- first element of the collection. + newRefs :: i -> a -> m (Refs i m a) + + -- | Write an element of the state. + readRefs :: Refs i m a -> i -> m a + + -- | Read an element of the state. + writeRefs :: Refs i m a -> i -> a -> m () + + +-- | Fold all the elements in a collection of refs. +foldRefsM + :: States i m + => (a -> b -> b) -> b -> Refs i m a -> m b + +foldRefsM f z refs + = loop_foldsRefsM first z + where + loop_foldsRefsM i acc + = do x <- readRefs refs i + let acc' = f x acc + case next i (extentRefs refs) of + Nothing -> return acc' + Just i' -> loop_foldsRefsM i' acc' + {-# INLINE loop_foldsRefsM #-} +{-# INLINE foldRefsM #-} + + +toListM :: States i m + => Refs i m a -> m [a] +toListM refs + = foldRefsM (:) [] refs +{-# NOINLINE toListM #-} + + +instance States Int IO where + data Refs Int IO a = Refs !(VM.IOVector a) + extentRefs (Refs !refs) = VM.length refs + newRefs !n !x = liftM Refs $ unsafeNewWithVM n x + readRefs (Refs !refs) !i = VM.unsafeRead refs i + writeRefs (Refs !refs) !i !x = VM.unsafeWrite refs i x + {-# NOINLINE newRefs #-} + {-# INLINE readRefs #-} + {-# INLINE writeRefs #-} + + +instance States Int m => States () m where + + data Refs () m a = URefs !(Refs Int m a) + + extentRefs _ = () + {-# INLINE extentRefs #-} + + newRefs _ !x + = do refs <- newRefs (1 :: Int) x + return $ URefs refs + {-# NOINLINE newRefs #-} + + readRefs (URefs !refs) _ = readRefs refs 0 + writeRefs (URefs !refs) _ !x = writeRefs refs 0 x + {-# INLINE readRefs #-} + {-# INLINE writeRefs #-} + + +------------------------------------------------------------------------------- +unsafeNewWithVM :: Int -> a -> IO (VM.IOVector a) +unsafeNewWithVM n x + = do vec <- VM.unsafeNew n + + let loop_newRefs !i + | i >= n = return () + | otherwise + = do VM.unsafeWrite vec i x + loop_newRefs (i + 1) + {-# INLINE loop_newRefs #-} + + loop_newRefs 0 + return vec +{-# INLINE unsafeNewWithVM #-} + 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-flow.h b/include/repa-flow.h new file mode 100644 --- /dev/null +++ b/include/repa-flow.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-flow.cabal b/repa-flow.cabal new file mode 100644 --- /dev/null +++ b/repa-flow.cabal @@ -0,0 +1,108 @@ +Name: repa-flow +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: Data-parallel data flows. +Synopsis: Data-parallel data flows. + +Library + build-Depends: + base == 4.7.*, + directory == 1.2.*, + filepath == 1.3.*, + vector == 0.10.*, + bytestring == 0.10.*, + primitive == 0.5.4.*, + containers == 0.5.*, + text == 1.2.*, + repa-stream == 4.0.0.*, + repa-eval == 4.0.0.*, + repa-array == 4.0.0.* + + exposed-modules: + Data.Repa.Flow.Chunked + Data.Repa.Flow.Chunked.IO + + Data.Repa.Flow.Default + Data.Repa.Flow.Default.Debug + Data.Repa.Flow.Default.IO + Data.Repa.Flow.Default.SizedIO + + Data.Repa.Flow.Generic + Data.Repa.Flow.Generic.Debug + Data.Repa.Flow.Generic.IO + + Data.Repa.Flow.IO.Bucket + + Data.Repa.Flow.Simple + + Data.Repa.Flow.States + + Data.Repa.Flow + + other-modules: + Data.Repa.Flow.Chunked.Base + Data.Repa.Flow.Chunked.Map + Data.Repa.Flow.Chunked.Folds + Data.Repa.Flow.Chunked.Groups + Data.Repa.Flow.Chunked.Operator + + Data.Repa.Flow.Default.IO.TSV + Data.Repa.Flow.Default.IO.CSV + + Data.Repa.Flow.Generic.Base + Data.Repa.Flow.Generic.Connect + Data.Repa.Flow.Generic.List + Data.Repa.Flow.Generic.Map + Data.Repa.Flow.Generic.Operator + Data.Repa.Flow.Generic.Eval + Data.Repa.Flow.Generic.Array.Distribute + Data.Repa.Flow.Generic.Array.Shuffle + Data.Repa.Flow.Generic.Array.Chunk + Data.Repa.Flow.Generic.Array.Unchunk + Data.Repa.Flow.Generic.IO.Sieve + + Data.Repa.Flow.Simple.Base + Data.Repa.Flow.Simple.List + Data.Repa.Flow.Simple.Operator + Data.Repa.Flow.Simple.IO + + include-dirs: + include + + install-includes: + repa-flow.h + + ghc-options: + -threaded + -Wall -fno-warn-missing-signatures + -O2 + -fcpr-off + + extensions: + CPP + BangPatterns + NoMonomorphismRestriction + RankNTypes + MagicHash + FlexibleContexts + FlexibleInstances + PatternGuards + TypeFamilies + MultiParamTypeClasses + ScopedTypeVariables + FunctionalDependencies + ConstraintKinds + ForeignFunctionInterface + StandaloneDeriving + ParallelListComp + +