git.iko.soy/streamly
on 3cdecba help sign in
streamly/src/Streamly/Mem/Array/Stream.hs
173 lines 5.3 kB history raw
1{-# LANGUAGE BangPatterns #-}
2{-# LANGUAGE CPP #-}
3{-# LANGUAGE MagicHash #-}
4{-# LANGUAGE RecordWildCards #-}
5{-# LANGUAGE UnboxedTuples #-}
6{-# LANGUAGE ScopedTypeVariables #-}
8#include "inline.hs"
10-- |
11-- Module : Streamly.Mem.Array.Stream
12-- Copyright : (c) 2019 Composewell Technologies
13--
14-- License : BSD3
15-- Maintainer : harendra.kumar@gmail.com
16-- Stability : experimental
17-- Portability : GHC
18--
19-- Combinators to efficiently manipulate streams of arrays.
20--
21module Streamly.Mem.Array.Stream
22 (
23 -- * Creation
24 arraysOf
26 -- * Flattening to elements
27 , flatten
28 -- , _flattenArraysRev
29 , unlinesBy
31 -- * Transformation
32 , splitOn
33 , compact -- compact
35 -- * Elimination
36 , toArray
37 )
38where
40import Control.Monad.IO.Class (MonadIO(..))
41-- import Data.Functor.Identity (Identity)
42import Data.Word (Word8)
43import Foreign.ForeignPtr (withForeignPtr)
44import Foreign.Ptr (minusPtr, plusPtr, castPtr)
45import Foreign.Storable (Storable(..))
46import Prelude hiding (length, null, last, map, (!!), read)
48import Streamly.Mem.Array.Types (Array(..), length)
49import Streamly.Streams.Serial (SerialT)
50import Streamly.Streams.StreamK.Type (IsStream)
52import qualified Streamly.Mem.Array.Types as A
53import qualified Streamly.Prelude as S
54import qualified Streamly.Streams.StreamD as D
55import qualified Streamly.Streams.Prelude as P
57-- | Convert a stream of arrays into a stream of their elements.
58--
59-- Same as the following but more efficient:
60--
61-- > flatten = S.concatMap A.read
62--
63-- @since 0.7.0
64{-# INLINE flatten #-}
65flatten :: (IsStream t, MonadIO m, Storable a) => t m (Array a) -> t m a
66flatten m = D.fromStreamD $ A.flattenArrays (D.toStreamD m)
68-- XXX should we have a reverseArrays API to reverse the stream of arrays
69-- instead?
70--
71-- | Convert a stream of arrays into a stream of their elements reversing the
72-- contents of each array before flattening.
73--
74-- @since 0.7.0
75{-# INLINE _flattenArraysRev #-}
76_flattenArraysRev :: (IsStream t, MonadIO m, Storable a) => t m (Array a) -> t m a
77_flattenArraysRev m = D.fromStreamD $ A.flattenArraysRev (D.toStreamD m)
79-- XXX use an Array instead as separator? Or use a separate unlinesArraysBySeq
80-- API for that?
81--
82-- | Flatten a stream of arrays appending the given element after each
83-- array.
84--
85-- @since 0.7.0
86{-# INLINE unlinesBy #-}
87unlinesBy :: (MonadIO m, IsStream t, Storable a)
88 => a -> t m (Array a) -> t m a
89unlinesBy x = D.fromStreamD . A.unlines x . D.toStreamD
91-- | Split a stream of arrays on a given separator byte, dropping the separator
92-- and coalescing all the arrays between two separators into a single array.
93--
94-- @since 0.7.0
95{-# INLINE splitOn #-}
96splitOn
97 :: (IsStream t, MonadIO m)
98 => Word8
99 -> t m (Array Word8)
100 -> t m (Array Word8)
101splitOn byte s = D.fromStreamD $ A.splitOn byte $ D.toStreamD s
103-- | Coalesce adajcent arrays in incoming stream to form bigger arrays of a
104-- maximum specified size.
105--
106-- @since 0.7.0
107{-# INLINE compact #-}
108compact :: (MonadIO m, Storable a)
109 => Int -> SerialT m (Array a) -> SerialT m (Array a)
110compact n xs =
111 D.fromStreamD $ A.packArraysChunksOf n (D.toStreamD xs)
113-- | Groups the elements in an input stream into arrays of given size.
114--
115-- Same as the following but more efficient:
116--
117-- > arraysOf n = FL.groupsOf n (FL.toArrayN n)
118--
119-- @since 0.7.0
120{-# INLINE arraysOf #-}
121arraysOf :: (IsStream t, MonadIO m, Storable a)
122 => Int -> t m a -> t m (Array a)
123arraysOf n str =
124 D.fromStreamD $ A.fromStreamDArraysOf n (D.toStreamD str)
126-- XXX Both of these implementations of splicing seem to perform equally well.
127-- We need to perform benchmarks over a range of sizes though.
129-- CAUTION! length must more than equal to lengths of all the arrays in the
130-- stream.
131{-# INLINE spliceArraysLenUnsafe #-}
132spliceArraysLenUnsafe :: (MonadIO m, Storable a)
133 => Int -> SerialT m (Array a) -> m (Array a)
134spliceArraysLenUnsafe len buffered = do
135 arr <- liftIO $ A.newArray len
136 end <- S.foldlM' writeArr (aEnd arr) buffered
137 return $ arr {aEnd = end}
139 where
141 writeArr dst Array{..} =
142 liftIO $ withForeignPtr aStart $ \src -> do
143 let count = aEnd `minusPtr` src
144 A.memcpy (castPtr dst) (castPtr src) count
145 return $ dst `plusPtr` count
147{-# INLINE _spliceArraysBuffered #-}
148_spliceArraysBuffered :: (MonadIO m, Storable a)
149 => SerialT m (Array a) -> m (Array a)
150_spliceArraysBuffered s = do
151 buffered <- P.foldr S.cons S.nil s
152 len <- S.sum (S.map length buffered)
153 spliceArraysLenUnsafe len s
155{-# INLINE spliceArraysRealloced #-}
156spliceArraysRealloced :: forall m a. (MonadIO m, Storable a)
157 => SerialT m (Array a) -> m (Array a)
158spliceArraysRealloced s = do
159 idst <- liftIO $ A.newArray (A.bytesToCount (undefined :: a)
160 (A.mkChunkSizeKB 4))
162 arr <- S.foldlM' A.spliceWithDoubling idst s
163 liftIO $ A.shrinkToFit arr
165-- | Given a stream of arrays, splice them all together to generate a single
166-- array. The stream must be /finite/.
167--
168-- @since 0.7.0
169{-# INLINE toArray #-}
170toArray :: (MonadIO m, Storable a) => SerialT m (Array a) -> m (Array a)
171toArray = spliceArraysRealloced
172-- spliceArrays = _spliceArraysBuffered