git.iko.soy/streamly
on da62ada help sign in
streamly/src/Streamly/Internal/Memory/ArrayStream.hs
235 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.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
54-- XXX efficiently compare two streams of arrays. Two streams can have chunks
55-- of different sizes, we can handle that in the stream comparison abstraction.
56-- This could be useful e.g. to fast compare whether two files differ.
58-- | Convert a stream of arrays into a stream of their elements.
59--
60-- Same as the following but more efficient:
61--
62-- > concat = S.concatMap A.read
63--
64-- @since 0.7.0
65{-# INLINE concat #-}
66concat :: (IsStream t, MonadIO m, Storable a) => t m (Array a) -> t m a
67-- concat m = D.fromStreamD $ A.flattenArrays (D.toStreamD m)
68-- concat m = D.fromStreamD $ D.concatMap A.toStreamD (D.toStreamD m)
69concat m = D.fromStreamD $ D.concatMapU A.read (D.toStreamD m)
71-- XXX should we have a reverseArrays API to reverse the stream of arrays
72-- instead?
73--
74-- | Convert a stream of arrays into a stream of their elements reversing the
75-- contents of each array before flattening.
76--
77-- @since 0.7.0
78{-# INLINE concatRev #-}
79concatRev :: (IsStream t, MonadIO m, Storable a) => t m (Array a) -> t m a
80concatRev m = D.fromStreamD $ A.flattenArraysRev (D.toStreamD m)
82-- | Flatten a stream of arrays after inserting the given element between
83-- arrays.
84--
85-- /Internal/
86{-# INLINE interpose #-}
87interpose :: (MonadIO m, IsStream t, Storable a) => a -> t m (Array a) -> t m a
88interpose x = S.interpose x A.read
90{-# INLINE intercalateSuffix #-}
91intercalateSuffix :: (MonadIO m, IsStream t, Storable a)
92 => Array a -> t m (Array a) -> t m a
93intercalateSuffix arr = S.intercalateSuffix arr A.read
95-- | Flatten a stream of arrays appending the given element after each
96-- array.
97--
98-- @since 0.7.0
99{-# INLINE interposeSuffix #-}
100interposeSuffix :: (MonadIO m, IsStream t, Storable a)
101 => a -> t m (Array a) -> t m a
102-- interposeSuffix x = D.fromStreamD . A.unlines x . D.toStreamD
103interposeSuffix x = S.interposeSuffix x A.read
105-- | Split a stream of arrays on a given separator byte, dropping the separator
106-- and coalescing all the arrays between two separators into a single array.
107--
108-- @since 0.7.0
109{-# INLINE splitOn #-}
110splitOn
111 :: (IsStream t, MonadIO m)
112 => Word8
113 -> t m (Array Word8)
114 -> t m (Array Word8)
115splitOn byte s =
116 D.fromStreamD $ D.splitInnerBy (A.breakOn byte) A.spliceTwo $ D.toStreamD s
118{-# INLINE splitOnSuffix #-}
119splitOnSuffix
120 :: (IsStream t, MonadIO m)
121 => Word8
122 -> t m (Array Word8)
123 -> t m (Array Word8)
124-- splitOn byte s = D.fromStreamD $ A.splitOn byte $ D.toStreamD s
125splitOnSuffix byte s =
126 D.fromStreamD $ D.splitInnerBySuffix (A.breakOn byte) A.spliceTwo $ D.toStreamD s
128-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
129-- maximum specified size in bytes.
130--
131-- @since 0.7.0
132{-# INLINE compact #-}
133compact :: (MonadIO m, Storable a)
134 => Int -> SerialT m (Array a) -> SerialT m (Array a)
135compact n xs = D.fromStreamD $ A.packArraysChunksOf n (D.toStreamD xs)
137-- | @arraysOf n stream@ groups the elements in the input stream into arrays of
138-- @n@ elements each.
139--
140-- Same as the following but more efficient:
141--
142-- > arraysOf n = S.chunksOf n (A.writeN n)
143--
144-- @since 0.7.0
145{-# INLINE arraysOf #-}
146arraysOf :: (IsStream t, MonadIO m, Storable a)
147 => Int -> t m a -> t m (Array a)
148arraysOf n str =
149 D.fromStreamD $ A.fromStreamDArraysOf n (D.toStreamD str)
151-- XXX Both of these implementations of splicing seem to perform equally well.
152-- We need to perform benchmarks over a range of sizes though.
154-- CAUTION! length must more than equal to lengths of all the arrays in the
155-- stream.
156{-# INLINE spliceArraysLenUnsafe #-}
157spliceArraysLenUnsafe :: (MonadIO m, Storable a)
158 => Int -> SerialT m (MA.Array a) -> m (MA.Array a)
159spliceArraysLenUnsafe len buffered = do
160 arr <- liftIO $ MA.newArray len
161 end <- S.foldlM' writeArr (return $ MA.aEnd arr) buffered
162 return $ arr {MA.aEnd = end}
164 where
166 writeArr dst (MA.Array as ae _) =
167 liftIO $ withForeignPtr as $ \src -> do
168 let count = ae `minusPtr` src
169 A.memcpy (castPtr dst) (castPtr src) count
170 return $ dst `plusPtr` count
172{-# INLINE _spliceArrays #-}
173_spliceArrays :: (MonadIO m, Storable a)
174 => SerialT m (Array a) -> m (Array a)
175_spliceArrays s = do
176 buffered <- S.foldr S.cons S.nil s
177 len <- S.sum (S.map length buffered)
178 arr <- liftIO $ MA.newArray len
179 end <- S.foldlM' writeArr (return $ MA.aEnd arr) s
180 return $ A.unsafeFreeze $ arr {MA.aEnd = end}
182 where
184 writeArr dst (Array as ae) =
185 liftIO $ withForeignPtr as $ \src -> do
186 let count = ae `minusPtr` src
187 A.memcpy (castPtr dst) (castPtr src) count
188 return $ dst `plusPtr` count
190{-# INLINE _spliceArraysBuffered #-}
191_spliceArraysBuffered :: (MonadIO m, Storable a)
192 => SerialT m (Array a) -> m (Array a)
193_spliceArraysBuffered s = do
194 buffered <- S.foldr S.cons S.nil s
195 len <- S.sum (S.map length buffered)
196 A.unsafeFreeze <$> spliceArraysLenUnsafe len (S.map A.unsafeThaw s)
198{-# INLINE spliceArraysRealloced #-}
199spliceArraysRealloced :: forall m a. (MonadIO m, Storable a)
200 => SerialT m (Array a) -> m (Array a)
201spliceArraysRealloced s = do
202 let idst = liftIO $ MA.newArray (A.bytesToElemCount (undefined :: a)
203 (A.mkChunkSizeKB 4))
205 arr <- S.foldlM' MA.spliceWithDoubling idst (S.map A.unsafeThaw s)
206 liftIO $ A.unsafeFreeze <$> MA.shrinkToFit arr
208-- | Given a stream of arrays, splice them all together to generate a single
209-- array. The stream must be /finite/.
210--
211-- @since 0.7.0
212{-# INLINE toArray #-}
213toArray :: (MonadIO m, Storable a) => SerialT m (Array a) -> m (Array a)
214toArray = spliceArraysRealloced
215-- spliceArrays = _spliceArraysBuffered
217-- exponentially increasing sizes of the chunks upto the max limit.
218-- XXX this will be easier to implement with parsers/terminating folds
219-- With this we should be able to reduce the number of chunks/allocations.
220-- The reallocation/copy based toArray can also be implemented using this.
221--
222{-
223{-# INLINE toArraysInRange #-}
224toArraysInRange :: (IsStream t, MonadIO m, Storable a)
225 => Int -> Int -> Fold m (Array a) b -> Fold m a b
226toArraysInRange low high (Fold step initial extract) =
227-}
229{-
230-- | Fold the input to a pure buffered stream (List) of arrays.
231{-# INLINE _toArraysOf #-}
232_toArraysOf :: (MonadIO m, Storable a)
233 => Int -> Fold m a (SerialT Identity (Array a))
234_toArraysOf n = FL.chunksOf n (A.writeNF n) FL.toStream
235-}