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