git.iko.soy/streamly
on e692bad help sign in
streamly/src/Streamly/Memory/ArrayStream.hs
215 lines 7.0 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 , interposeSuffix
30 , intercalateSuffix
32 -- * Transformation
33 , splitOn
34 , splitOnSuffix
35 , compact -- compact
37 -- * Elimination
38 , toArray
39 )
40where
42import Control.Monad.IO.Class (MonadIO(..))
43-- import Data.Functor.Identity (Identity)
44import Data.Word (Word8)
45import Foreign.ForeignPtr (withForeignPtr)
46import Foreign.Ptr (minusPtr, plusPtr, castPtr)
47import Foreign.Storable (Storable(..))
48import Prelude hiding (length, null, last, map, (!!), read, concat)
50import Streamly.Memory.Array.Types (Array(..), length)
51import Streamly.Streams.Serial (SerialT)
52import Streamly.Streams.StreamK.Type (IsStream)
54import qualified Streamly.Unfold as UF
55import qualified Streamly.Memory.Array.Types as A
56import qualified Streamly.Prelude.Internal as S
57import qualified Streamly.Streams.StreamD as D
58import qualified Streamly.Streams.Prelude as P
60-- XXX efficiently compare two streams of arrays. Two streams can have chunks
61-- of different sizes, we can handle that in the stream comparison abstraction.
62-- This could be useful e.g. to fast compare whether two files differ.
64-- | Convert a stream of arrays into a stream of their elements.
65--
66-- Same as the following but more efficient:
67--
68-- > concat = S.concatMap A.read
69--
70-- @since 0.7.0
71{-# INLINE concat #-}
72concat :: (IsStream t, MonadIO m, Storable a) => t m (Array a) -> t m a
73-- concat m = D.fromStreamD $ A.flattenArrays (D.toStreamD m)
74-- concat m = D.fromStreamD $ D.concatMap A.toStreamD (D.toStreamD m)
75concat m = D.fromStreamD $ D.concatMapU A.readU (D.toStreamD m)
77-- XXX should we have a reverseArrays API to reverse the stream of arrays
78-- instead?
79--
80-- | Convert a stream of arrays into a stream of their elements reversing the
81-- contents of each array before flattening.
82--
83-- @since 0.7.0
84{-# INLINE concatRev #-}
85concatRev :: (IsStream t, MonadIO m, Storable a) => t m (Array a) -> t m a
86concatRev m = D.fromStreamD $ A.flattenArraysRev (D.toStreamD m)
88{-# INLINE intercalateSuffix #-}
89intercalateSuffix :: (MonadIO m, IsStream t, Storable a)
90 => Array a -> t m (Array a) -> t m a
91intercalateSuffix arr = S.intercalateSuffix A.readU arr A.readU
93-- | Flatten a stream of arrays appending the given element after each
94-- array.
95--
96-- @since 0.7.0
97{-# INLINE interposeSuffix #-}
98interposeSuffix :: (MonadIO m, IsStream t, Storable a)
99 => a -> t m (Array a) -> t m a
100-- interposeSuffix x = D.fromStreamD . A.unlines x . D.toStreamD
101interposeSuffix x = S.intercalateSuffix UF.singleton x A.readU
103-- | Split a stream of arrays on a given separator byte, dropping the separator
104-- and coalescing all the arrays between two separators into a single array.
105--
106-- @since 0.7.0
107{-# INLINE splitOn #-}
108splitOn
109 :: (IsStream t, MonadIO m)
110 => Word8
111 -> t m (Array Word8)
112 -> t m (Array Word8)
113splitOn byte s =
114 D.fromStreamD $ D.splitInnerBy (A.breakOn byte) A.spliceTwo $ D.toStreamD s
116{-# INLINE splitOnSuffix #-}
117splitOnSuffix
118 :: (IsStream t, MonadIO m)
119 => Word8
120 -> t m (Array Word8)
121 -> t m (Array Word8)
122-- splitOn byte s = D.fromStreamD $ A.splitOn byte $ D.toStreamD s
123splitOnSuffix byte s =
124 D.fromStreamD $ D.splitInnerBySuffix (A.breakOn byte) A.spliceTwo $ D.toStreamD s
126-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
127-- maximum specified size in bytes.
128--
129-- @since 0.7.0
130{-# INLINE compact #-}
131compact :: (MonadIO m, Storable a)
132 => Int -> SerialT m (Array a) -> SerialT m (Array a)
133compact n xs = D.fromStreamD $ A.packArraysChunksOf n (D.toStreamD xs)
135-- | @arraysOf n stream@ groups the elements in the input stream into arrays of
136-- @n@ elements each.
137--
138-- Same as the following but more efficient:
139--
140-- > arraysOf n = S.chunksOf n (A.writeN n)
141--
142-- @since 0.7.0
143{-# INLINE arraysOf #-}
144arraysOf :: (IsStream t, MonadIO m, Storable a)
145 => Int -> t m a -> t m (Array a)
146arraysOf n str =
147 D.fromStreamD $ A.fromStreamDArraysOf n (D.toStreamD str)
149-- XXX Both of these implementations of splicing seem to perform equally well.
150-- We need to perform benchmarks over a range of sizes though.
152-- CAUTION! length must more than equal to lengths of all the arrays in the
153-- stream.
154{-# INLINE spliceArraysLenUnsafe #-}
155spliceArraysLenUnsafe :: (MonadIO m, Storable a)
156 => Int -> SerialT m (Array a) -> m (Array a)
157spliceArraysLenUnsafe len buffered = do
158 arr <- liftIO $ A.newArray len
159 end <- S.foldlM' writeArr (aEnd arr) buffered
160 return $ arr {aEnd = end}
162 where
164 writeArr dst Array{..} =
165 liftIO $ withForeignPtr aStart $ \src -> do
166 let count = aEnd `minusPtr` src
167 A.memcpy (castPtr dst) (castPtr src) count
168 return $ dst `plusPtr` count
170{-# INLINE _spliceArraysBuffered #-}
171_spliceArraysBuffered :: (MonadIO m, Storable a)
172 => SerialT m (Array a) -> m (Array a)
173_spliceArraysBuffered s = do
174 buffered <- P.foldr S.cons S.nil s
175 len <- S.sum (S.map length buffered)
176 spliceArraysLenUnsafe len s
178{-# INLINE spliceArraysRealloced #-}
179spliceArraysRealloced :: forall m a. (MonadIO m, Storable a)
180 => SerialT m (Array a) -> m (Array a)
181spliceArraysRealloced s = do
182 idst <- liftIO $ A.newArray (A.bytesToElemCount (undefined :: a)
183 (A.mkChunkSizeKB 4))
185 arr <- S.foldlM' A.spliceWithDoubling idst s
186 liftIO $ A.shrinkToFit arr
188-- | Given a stream of arrays, splice them all together to generate a single
189-- array. The stream must be /finite/.
190--
191-- @since 0.7.0
192{-# INLINE toArray #-}
193toArray :: (MonadIO m, Storable a) => SerialT m (Array a) -> m (Array a)
194toArray = spliceArraysRealloced
195-- spliceArrays = _spliceArraysBuffered
197-- exponentially increasing sizes of the chunks upto the max limit.
198-- XXX this will be easier to implement with parsers/terminating folds
199-- With this we should be able to reduce the number of chunks/allocations.
200-- The reallocation/copy based toArray can also be implemented using this.
201--
202{-
203{-# INLINE toArraysInRange #-}
204toArraysInRange :: (IsStream t, MonadIO m, Storable a)
205 => Int -> Int -> Fold m (Array a) b -> Fold m a b
206toArraysInRange low high (Fold step initial extract) =
207-}
209{-
210-- | Fold the input to a pure buffered stream (List) of arrays.
211{-# INLINE _toArraysOf #-}
212_toArraysOf :: (MonadIO m, Storable a)
213 => Int -> Fold m a (SerialT Identity (Array a))
214_toArraysOf n = FL.lchunksOf n (A.writeNF n) FL.toStream
215-}