git.iko.soy/streamly
on 5d096df help sign in
streamly/core/src/Streamly/Internal/Data/MutArray.hs
207 lines 7.5 kB history raw
1-- |
2-- Module : Streamly.Internal.Data.MutArray
3-- Copyright : (c) 2020 Composewell Technologies
4-- License : BSD-3-Clause
5-- Maintainer : streamly@composewell.com
6-- Stability : experimental
7-- Portability : GHC
8--
9module Streamly.Internal.Data.MutArray
10 (
11 -- * MutArray.Type module
12 module Streamly.Internal.Data.MutArray.Type
13 -- * MutArray module
14 , sliceIndexerFromLen
15 , slicerFromLen
16 , compactLE
17 , pinnedCompactLE
18 , compactOnByte
19 , compactOnByteSuffix
20 -- * Unboxed IORef
21 , module Streamly.Internal.Data.IORef.Unboxed
23 -- * Deprecated
24 , genSlicesFromLen
25 , getSlicesFromLen
26 )
27where
29#include "inline.hs"
31import Control.Monad.IO.Class (MonadIO(..))
32import Data.Word (Word8)
33import Streamly.Internal.Data.MutByteArray.Type (PinnedState(..))
34import Streamly.Internal.Data.Stream.Type (Stream)
35import Streamly.Internal.Data.Unbox (Unbox)
36import Streamly.Internal.Data.Unfold.Type (Unfold(..))
38import qualified Streamly.Internal.Data.Stream.Type as Stream
39import qualified Streamly.Internal.Data.Stream.Nesting as Stream
40-- import qualified Streamly.Internal.Data.Stream.Transform as Stream
41import qualified Streamly.Internal.Data.Unfold as Unfold
43import Prelude hiding (foldr, length, read, splitAt)
44import Streamly.Internal.Data.MutArray.Type
45import Streamly.Internal.Data.IORef.Unboxed
47-- | Generate a stream of array slice descriptors ((index, len)) of specified
48-- length from an array, starting from the supplied array index. The last slice
49-- may be shorter than the requested length depending on the array length.
50--
51-- /Pre-release/
52{-# INLINE sliceIndexerFromLen #-}
53sliceIndexerFromLen :: forall m a. (Monad m, Unbox a)
54 => Int -- ^ from index
55 -> Int -- ^ length of the slice
56 -> Unfold m (MutArray a) (Int, Int)
57sliceIndexerFromLen from len =
58 let fromThenTo n = (from, from + len, n - 1)
59 mkSlice n i = return (i, min len (n - i))
60 in Unfold.lmap length
61 $ Unfold.mapM2 mkSlice
62 $ Unfold.lmap fromThenTo Unfold.enumerateFromThenTo
64{-# DEPRECATED genSlicesFromLen "Please use sliceIndexerFromLen instead." #-}
65genSlicesFromLen :: forall m a. (Monad m, Unbox a)
66 => Int -- ^ from index
67 -> Int -- ^ length of the slice
68 -> Unfold m (MutArray a) (Int, Int)
69genSlicesFromLen = sliceIndexerFromLen
71-- | Generate a stream of slices of specified length from an array, starting
72-- from the supplied array index. The last slice may be shorter than the
73-- requested length depending on the array length.
74--
75-- /Pre-release/
76{-# INLINE slicerFromLen #-}
77slicerFromLen :: forall m a. (Monad m, Unbox a)
78 => Int -- ^ from index
79 -> Int -- ^ length of the slice
80 -> Unfold m (MutArray a) (MutArray a)
81slicerFromLen from len =
82 let mkSlice arr (i, n) = return $ getSliceUnsafe i n arr
83 in Unfold.mapM2 mkSlice (sliceIndexerFromLen from len)
85{-# DEPRECATED getSlicesFromLen "Please use slicerFromLen instead." #-}
86getSlicesFromLen :: forall m a. (Monad m, Unbox a)
87 => Int -- ^ from index
88 -> Int -- ^ length of the slice
89 -> Unfold m (MutArray a) (MutArray a)
90getSlicesFromLen = slicerFromLen
92-------------------------------------------------------------------------------
93-- Compacting Streams of Arrays
94-------------------------------------------------------------------------------
96-- | @compactLE maxElems@ coalesces adjacent arrays in the input stream
97-- only if the combined size would be less than or equal to @maxElems@
98-- elements. Note that it won't split an array if the original array is already
99-- larger than maxElems.
100--
101-- @maxElems@ must be greater than 0.
102--
103-- Generates unpinned arrays irrespective of the pinning status of input
104-- arrays.
105{-# INLINE compactLE #-}
106compactLE :: (MonadIO m, Unbox a) =>
107 Int -> Stream m (MutArray a) -> Stream m (MutArray a)
108-- XXX compactLE can be moved to MutArray/Type if we are not using the parser
109-- to implement it.
110compactLE = compactLeAs Unpinned
111-- The parser version turns out to be a little bit slower.
112-- compactLE n = Stream.catRights . Stream.parseManyD (pCompactLE n)
114-- | Pinned version of 'compactLE'.
115{-# INLINE pinnedCompactLE #-}
116pinnedCompactLE :: forall m a. (MonadIO m, Unbox a)
117 => Int -> Stream m (MutArray a) -> Stream m (MutArray a)
118pinnedCompactLE = compactLeAs Pinned
119-- pinnedCompactLE n = Stream.catRights . Stream.parseManyD (pPinnedCompactLE n)
121data SplitState s arr
122 = Initial s
123 | Buffering s arr
124 | Splitting s arr
125 | Yielding arr (SplitState s arr)
126 | Finishing
128-- | Split a stream of arrays on a given separator byte, dropping the separator
129-- and coalescing all the arrays between two separators into a single array.
130--
131{-# INLINE_NORMAL _compactOnByteCustom #-}
132_compactOnByteCustom
133 :: MonadIO m
134 => Word8
135 -> Stream m (MutArray Word8)
136 -> Stream m (MutArray Word8)
137_compactOnByteCustom byte (Stream.Stream step state) =
138 Stream.Stream step' (Initial state)
140 where
142 {-# INLINE_LATE step' #-}
143 step' gst (Initial st) = do
144 r <- step gst st
145 case r of
146 Stream.Yield arr s -> do
147 (arr1, marr2) <- breakOn byte arr
148 return $ case marr2 of
149 Nothing -> Stream.Skip (Buffering s arr1)
150 Just arr2 -> Stream.Skip (Yielding arr1 (Splitting s arr2))
151 Stream.Skip s -> return $ Stream.Skip (Initial s)
152 Stream.Stop -> return Stream.Stop
154 step' gst (Buffering st buf) = do
155 r <- step gst st
156 case r of
157 Stream.Yield arr s -> do
158 (arr1, marr2) <- breakOn byte arr
159 -- XXX Use spliceExp instead and then rightSize?
160 buf1 <- splice buf arr1
161 return $ case marr2 of
162 Nothing -> Stream.Skip (Buffering s buf1)
163 Just x -> Stream.Skip (Yielding buf1 (Splitting s x))
164 Stream.Skip s -> return $ Stream.Skip (Buffering s buf)
165 Stream.Stop -> return $
166 if byteLength buf == 0
167 then Stream.Stop
168 else Stream.Skip (Yielding buf Finishing)
170 step' _ (Splitting st buf) = do
171 (arr1, marr2) <- breakOn byte buf
172 return $ case marr2 of
173 Nothing -> Stream.Skip $ Buffering st arr1
174 Just arr2 -> Stream.Skip $ Yielding arr1 (Splitting st arr2)
176 step' _ (Yielding arr next) = return $ Stream.Yield arr next
177 step' _ Finishing = return Stream.Stop
179-- XXX implement predicate based version of this
180-- XXX Naming of predicate based vs custom version
182-- | Split a stream of arrays on a given separator byte, dropping the separator
183-- and coalescing all the arrays between two separators into a single array.
184--
185{-# INLINE compactOnByte #-}
186compactOnByte
187 :: (MonadIO m)
188 => Word8
189 -> Stream m (MutArray Word8)
190 -> Stream m (MutArray Word8)
191-- XXX compare perf of custom vs idiomatic version
192-- compactOnByte = _compactOnByteCustom
193-- XXX use spliceExp and rightSize?
194compactOnByte byte = Stream.splitInnerBy (breakOn byte) splice
196-- | Like 'compactOnByte' considers the separator in suffix position instead of
197-- infix position.
198{-# INLINE compactOnByteSuffix #-}
199compactOnByteSuffix
200 :: (MonadIO m)
201 => Word8
202 -> Stream m (MutArray Word8)
203 -> Stream m (MutArray Word8)
204compactOnByteSuffix byte =
205 -- XXX use spliceExp and rightSize?
206 Stream.splitInnerBySuffix
207 (\arr -> byteLength arr == 0) (breakOn byte) splice