git.iko.soy/streamly
on cc072be help sign in
streamly/src/Streamly/Internal/Memory/ArrayStream.hs
236 lines 8.0 kB history raw
1#include "inline.hs"
3-- |
4-- Module : Streamly.Internal.Memory.ArrayStream
5-- Copyright : (c) 2020 Composewell Technologies
6--
7-- License : BSD3-3-Clause
8-- Maintainer : streamly@composewell.com
9-- Stability : experimental
10-- Portability : GHC
11--
12-- Combinators to efficiently manipulate streams of arrays.
13--
14module Streamly.Internal.Memory.ArrayStream
15 (
16 -- * Creation
17 arraysOf
19 -- * Flattening to elements
20 , concat
21 , concatRev
22 , interpose
23 , interposeSuffix
24 , intercalateSuffix
26 -- * Transformation
27 , splitOn
28 , splitOnSuffix
29 , compact -- compact
31 -- * Elimination
32 , toArray
33 )
34where
36import Control.Monad.IO.Class (MonadIO(..))
37-- import Data.Functor.Identity (Identity)
38import Data.Word (Word8)
39import Foreign.ForeignPtr (withForeignPtr)
40import Foreign.Ptr (minusPtr, plusPtr, castPtr)
41import Foreign.Storable (Storable(..))
42import Prelude hiding (length, null, last, map, (!!), read, concat)
44import Streamly.Internal.Memory.Array.Types (Array(..), length)
45import Streamly.Internal.Data.Stream.Serial (SerialT)
46import Streamly.Internal.Data.Stream.StreamK.Type (IsStream)
48import qualified Streamly.Internal.Memory.Array as A
49import qualified Streamly.Internal.Memory.Array.Types as A
50import qualified Streamly.Internal.Memory.Mutable.Array.Types as MA
51import qualified Streamly.Internal.Prelude as S
52import qualified Streamly.Internal.Data.Stream.StreamD as D
53import qualified Streamly.Internal.Data.Stream.Prelude as P
55-- XXX efficiently compare two streams of arrays. Two streams can have chunks
56-- of different sizes, we can handle that in the stream comparison abstraction.
57-- This could be useful e.g. to fast compare whether two files differ.
59-- | Convert a stream of arrays into a stream of their elements.
60--
61-- Same as the following but more efficient:
62--
63-- > concat = S.concatMap A.read
64--
65-- @since 0.7.0
66{-# INLINE concat #-}
67concat :: (IsStream t, MonadIO m, Storable a) => t m (Array a) -> t m a
68-- concat m = D.fromStreamD $ A.flattenArrays (D.toStreamD m)
69-- concat m = D.fromStreamD $ D.concatMap A.toStreamD (D.toStreamD m)
70concat m = D.fromStreamD $ D.concatMapU A.read (D.toStreamD m)
72-- XXX should we have a reverseArrays API to reverse the stream of arrays
73-- instead?
74--
75-- | Convert a stream of arrays into a stream of their elements reversing the
76-- contents of each array before flattening.
77--
78-- @since 0.7.0
79{-# INLINE concatRev #-}
80concatRev :: (IsStream t, MonadIO m, Storable a) => t m (Array a) -> t m a
81concatRev m = D.fromStreamD $ A.flattenArraysRev (D.toStreamD m)
83-- | Flatten a stream of arrays after inserting the given element between
84-- arrays.
85--
86-- /Internal/
87{-# INLINE interpose #-}
88interpose :: (MonadIO m, IsStream t, Storable a) => a -> t m (Array a) -> t m a
89interpose x = S.interpose x A.read
91{-# INLINE intercalateSuffix #-}
92intercalateSuffix :: (MonadIO m, IsStream t, Storable a)
93 => Array a -> t m (Array a) -> t m a
94intercalateSuffix arr = S.intercalateSuffix arr A.read
96-- | Flatten a stream of arrays appending the given element after each
97-- array.
98--
99-- @since 0.7.0
100{-# INLINE interposeSuffix #-}
101interposeSuffix :: (MonadIO m, IsStream t, Storable a)
102 => a -> t m (Array a) -> t m a
103-- interposeSuffix x = D.fromStreamD . A.unlines x . D.toStreamD
104interposeSuffix x = S.interposeSuffix x A.read
106-- | Split a stream of arrays on a given separator byte, dropping the separator
107-- and coalescing all the arrays between two separators into a single array.
108--
109-- @since 0.7.0
110{-# INLINE splitOn #-}
111splitOn
112 :: (IsStream t, MonadIO m)
113 => Word8
114 -> t m (Array Word8)
115 -> t m (Array Word8)
116splitOn byte s =
117 D.fromStreamD $ D.splitInnerBy (A.breakOn byte) A.spliceTwo $ D.toStreamD s
119{-# INLINE splitOnSuffix #-}
120splitOnSuffix
121 :: (IsStream t, MonadIO m)
122 => Word8
123 -> t m (Array Word8)
124 -> t m (Array Word8)
125-- splitOn byte s = D.fromStreamD $ A.splitOn byte $ D.toStreamD s
126splitOnSuffix byte s =
127 D.fromStreamD $ D.splitInnerBySuffix (A.breakOn byte) A.spliceTwo $ D.toStreamD s
129-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
130-- maximum specified size in bytes.
131--
132-- @since 0.7.0
133{-# INLINE compact #-}
134compact :: (MonadIO m, Storable a)
135 => Int -> SerialT m (Array a) -> SerialT m (Array a)
136compact n xs = D.fromStreamD $ A.packArraysChunksOf n (D.toStreamD xs)
138-- | @arraysOf n stream@ groups the elements in the input stream into arrays of
139-- @n@ elements each.
140--
141-- Same as the following but more efficient:
142--
143-- > arraysOf n = S.chunksOf n (A.writeN n)
144--
145-- @since 0.7.0
146{-# INLINE arraysOf #-}
147arraysOf :: (IsStream t, MonadIO m, Storable a)
148 => Int -> t m a -> t m (Array a)
149arraysOf n str =
150 D.fromStreamD $ A.fromStreamDArraysOf n (D.toStreamD str)
152-- XXX Both of these implementations of splicing seem to perform equally well.
153-- We need to perform benchmarks over a range of sizes though.
155-- CAUTION! length must more than equal to lengths of all the arrays in the
156-- stream.
157{-# INLINE spliceArraysLenUnsafe #-}
158spliceArraysLenUnsafe :: (MonadIO m, Storable a)
159 => Int -> SerialT m (MA.Array a) -> m (MA.Array a)
160spliceArraysLenUnsafe len buffered = do
161 arr <- liftIO $ MA.newArray len
162 end <- S.foldlM' writeArr (return $ MA.aEnd arr) buffered
163 return $ arr {MA.aEnd = end}
165 where
167 writeArr dst (MA.Array as ae _) =
168 liftIO $ withForeignPtr as $ \src -> do
169 let count = ae `minusPtr` src
170 A.memcpy (castPtr dst) (castPtr src) count
171 return $ dst `plusPtr` count
173{-# INLINE _spliceArrays #-}
174_spliceArrays :: (MonadIO m, Storable a)
175 => SerialT m (Array a) -> m (Array a)
176_spliceArrays s = do
177 buffered <- P.foldr S.cons S.nil s
178 len <- S.sum (S.map length buffered)
179 arr <- liftIO $ MA.newArray len
180 end <- S.foldlM' writeArr (return $ MA.aEnd arr) s
181 return $ A.unsafeFreeze $ arr {MA.aEnd = end}
183 where
185 writeArr dst (Array as ae) =
186 liftIO $ withForeignPtr as $ \src -> do
187 let count = ae `minusPtr` src
188 A.memcpy (castPtr dst) (castPtr src) count
189 return $ dst `plusPtr` count
191{-# INLINE _spliceArraysBuffered #-}
192_spliceArraysBuffered :: (MonadIO m, Storable a)
193 => SerialT m (Array a) -> m (Array a)
194_spliceArraysBuffered s = do
195 buffered <- P.foldr S.cons S.nil s
196 len <- S.sum (S.map length buffered)
197 A.unsafeFreeze <$> spliceArraysLenUnsafe len (S.map A.unsafeThaw s)
199{-# INLINE spliceArraysRealloced #-}
200spliceArraysRealloced :: forall m a. (MonadIO m, Storable a)
201 => SerialT m (Array a) -> m (Array a)
202spliceArraysRealloced s = do
203 let idst = liftIO $ MA.newArray (A.bytesToElemCount (undefined :: a)
204 (A.mkChunkSizeKB 4))
206 arr <- S.foldlM' MA.spliceWithDoubling idst (S.map A.unsafeThaw s)
207 liftIO $ A.unsafeFreeze <$> MA.shrinkToFit arr
209-- | Given a stream of arrays, splice them all together to generate a single
210-- array. The stream must be /finite/.
211--
212-- @since 0.7.0
213{-# INLINE toArray #-}
214toArray :: (MonadIO m, Storable a) => SerialT m (Array a) -> m (Array a)
215toArray = spliceArraysRealloced
216-- spliceArrays = _spliceArraysBuffered
218-- exponentially increasing sizes of the chunks upto the max limit.
219-- XXX this will be easier to implement with parsers/terminating folds
220-- With this we should be able to reduce the number of chunks/allocations.
221-- The reallocation/copy based toArray can also be implemented using this.
222--
223{-
224{-# INLINE toArraysInRange #-}
225toArraysInRange :: (IsStream t, MonadIO m, Storable a)
226 => Int -> Int -> Fold m (Array a) b -> Fold m a b
227toArraysInRange low high (Fold step initial extract) =
228-}
230{-
231-- | Fold the input to a pure buffered stream (List) of arrays.
232{-# INLINE _toArraysOf #-}
233_toArraysOf :: (MonadIO m, Storable a)
234 => Int -> Fold m a (SerialT Identity (Array a))
235_toArraysOf n = FL.lchunksOf n (A.writeNF n) FL.toStream
236-}