git.iko.soy/streamly
on 00a2f7b help sign in
streamly/src/Streamly/Internal/Memory/ArrayStream.hs
236 lines 8.1 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.Data.Array.Storable.Foreign.Types (Array(..), length)
45import Streamly.Internal.Data.Stream.Serial (SerialT)
46import Streamly.Internal.Data.Stream.StreamK.Type (IsStream)
48import qualified Streamly.Internal.Data.Array.Storable.Foreign as A
49import qualified Streamly.Internal.Data.Array.Storable.Foreign.Types as A
50import qualified Streamly.Internal.Data.Array.Storable.Foreign.Mut.Types as MA
51import qualified Streamly.Internal.Data.Stream.IsStream 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-}