git.iko.soy/streamly
on 0ef35be help sign in
streamly/src/Streamly/Memory/ArrayStream.hs
198 lines 6.4 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.Memory.ArrayStream
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.Memory.ArrayStream
22 (
23 -- * Creation
24 arraysOf
26 -- * Flattening to elements
27 , concat
28 , concatRev
29 , concatWithSuffix
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, concat)
48import Streamly.Memory.Array.Types (Array(..), length)
49import Streamly.Streams.Serial (SerialT)
50import Streamly.Streams.StreamK.Type (IsStream)
52import qualified Streamly.Memory.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-- XXX efficiently compare two streams of arrays. Two streams can have chunks
58-- of different sizes, we can handle that in the stream comparison abstraction.
59-- This could be useful e.g. to fast compare whether two files differ.
61-- | Convert a stream of arrays into a stream of their elements.
62--
63-- Same as the following but more efficient:
64--
65-- > concat = S.concatMap A.read
66--
67-- @since 0.7.0
68{-# INLINE concat #-}
69concat :: (IsStream t, MonadIO m, Storable a) => t m (Array a) -> t m a
70concat m = D.fromStreamD $ A.flattenArrays (D.toStreamD m)
71-- concat m = D.fromStreamD $ D.concatMap A.toStreamD (D.toStreamD m)
72-- concat m = D.fromStreamD $ D.concatMapU A.readU (D.toStreamD m)
74-- XXX should we have a reverseArrays API to reverse the stream of arrays
75-- instead?
76--
77-- | Convert a stream of arrays into a stream of their elements reversing the
78-- contents of each array before flattening.
79--
80-- @since 0.7.0
81{-# INLINE concatRev #-}
82concatRev :: (IsStream t, MonadIO m, Storable a) => t m (Array a) -> t m a
83concatRev m = D.fromStreamD $ A.flattenArraysRev (D.toStreamD m)
85-- XXX use an Array instead as separator? Or use a separate unlinesArraysBySeq
86-- API for that?
87--
88-- | Flatten a stream of arrays appending the given element after each
89-- array.
90--
91-- @since 0.7.0
92{-# INLINE concatWithSuffix #-}
93concatWithSuffix :: (MonadIO m, IsStream t, Storable a)
94 => a -> t m (Array a) -> t m a
95concatWithSuffix x = D.fromStreamD . A.unlines x . D.toStreamD
97-- | Split a stream of arrays on a given separator byte, dropping the separator
98-- and coalescing all the arrays between two separators into a single array.
99--
100-- @since 0.7.0
101{-# INLINE splitOn #-}
102splitOn
103 :: (IsStream t, MonadIO m)
104 => Word8
105 -> t m (Array Word8)
106 -> t m (Array Word8)
107splitOn byte s = D.fromStreamD $ A.splitOn byte $ D.toStreamD s
109-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
110-- maximum specified size in bytes.
111--
112-- @since 0.7.0
113{-# INLINE compact #-}
114compact :: (MonadIO m, Storable a)
115 => Int -> SerialT m (Array a) -> SerialT m (Array a)
116compact n xs = D.fromStreamD $ A.packArraysChunksOf n (D.toStreamD xs)
118-- | @arraysOf n stream@ groups the elements in the input stream into arrays of
119-- @n@ elements each.
120--
121-- Same as the following but more efficient:
122--
123-- > arraysOf n = S.chunksOf n (A.writeN n)
124--
125-- @since 0.7.0
126{-# INLINE arraysOf #-}
127arraysOf :: (IsStream t, MonadIO m, Storable a)
128 => Int -> t m a -> t m (Array a)
129arraysOf n str =
130 D.fromStreamD $ A.fromStreamDArraysOf n (D.toStreamD str)
132-- XXX Both of these implementations of splicing seem to perform equally well.
133-- We need to perform benchmarks over a range of sizes though.
135-- CAUTION! length must more than equal to lengths of all the arrays in the
136-- stream.
137{-# INLINE spliceArraysLenUnsafe #-}
138spliceArraysLenUnsafe :: (MonadIO m, Storable a)
139 => Int -> SerialT m (Array a) -> m (Array a)
140spliceArraysLenUnsafe len buffered = do
141 arr <- liftIO $ A.newArray len
142 end <- S.foldlM' writeArr (aEnd arr) buffered
143 return $ arr {aEnd = end}
145 where
147 writeArr dst Array{..} =
148 liftIO $ withForeignPtr aStart $ \src -> do
149 let count = aEnd `minusPtr` src
150 A.memcpy (castPtr dst) (castPtr src) count
151 return $ dst `plusPtr` count
153{-# INLINE _spliceArraysBuffered #-}
154_spliceArraysBuffered :: (MonadIO m, Storable a)
155 => SerialT m (Array a) -> m (Array a)
156_spliceArraysBuffered s = do
157 buffered <- P.foldr S.cons S.nil s
158 len <- S.sum (S.map length buffered)
159 spliceArraysLenUnsafe len s
161{-# INLINE spliceArraysRealloced #-}
162spliceArraysRealloced :: forall m a. (MonadIO m, Storable a)
163 => SerialT m (Array a) -> m (Array a)
164spliceArraysRealloced s = do
165 idst <- liftIO $ A.newArray (A.bytesToElemCount (undefined :: a)
166 (A.mkChunkSizeKB 4))
168 arr <- S.foldlM' A.spliceWithDoubling idst s
169 liftIO $ A.shrinkToFit arr
171-- | Given a stream of arrays, splice them all together to generate a single
172-- array. The stream must be /finite/.
173--
174-- @since 0.7.0
175{-# INLINE toArray #-}
176toArray :: (MonadIO m, Storable a) => SerialT m (Array a) -> m (Array a)
177toArray = spliceArraysRealloced
178-- spliceArrays = _spliceArraysBuffered
180-- exponentially increasing sizes of the chunks upto the max limit.
181-- XXX this will be easier to implement with parsers/terminating folds
182-- With this we should be able to reduce the number of chunks/allocations.
183-- The reallocation/copy based toArray can also be implemented using this.
184--
185{-
186{-# INLINE toArraysInRange #-}
187toArraysInRange :: (IsStream t, MonadIO m, Storable a)
188 => Int -> Int -> Fold m (Array a) b -> Fold m a b
189toArraysInRange low high (Fold step initial extract) =
190-}
192{-
193-- | Fold the input to a pure buffered stream (List) of arrays.
194{-# INLINE _toArraysOf #-}
195_toArraysOf :: (MonadIO m, Storable a)
196 => Int -> Fold m a (SerialT Identity (Array a))
197_toArraysOf n = FL.lchunksOf n (A.writeNF n) FL.toStream
198-}