git.iko.soy/streamly
on defafa6 help sign in
streamly/core/src/Streamly/Internal/Data/MutArray.hs
190 lines 6.8 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 , compactOnByte
18 , compactOnByteSuffix
19 -- * Unboxed IORef
20 , module Streamly.Internal.Data.IORef.Unboxed
22 -- * Deprecated
23 , genSlicesFromLen
24 , getSlicesFromLen
25 )
26where
28#include "inline.hs"
30import Control.Monad.IO.Class (MonadIO(..))
31import Data.Word (Word8)
32import Streamly.Internal.Data.Stream.Type (Stream)
33import Streamly.Internal.Data.Unbox (Unbox)
34import Streamly.Internal.Data.Unfold.Type (Unfold(..))
36import qualified Streamly.Internal.Data.Stream.Type as Stream
37import qualified Streamly.Internal.Data.Stream.Nesting as Stream
38-- import qualified Streamly.Internal.Data.Stream.Transform as Stream
39import qualified Streamly.Internal.Data.Unfold as Unfold
41import Prelude hiding (foldr, length, read, splitAt)
42import Streamly.Internal.Data.MutArray.Type
43import Streamly.Internal.Data.IORef.Unboxed
45-- | Generate a stream of array slice descriptors ((index, len)) of specified
46-- length from an array, starting from the supplied array index. The last slice
47-- may be shorter than the requested length depending on the array length.
48--
49-- /Pre-release/
50{-# INLINE sliceIndexerFromLen #-}
51sliceIndexerFromLen :: forall m a. (Monad m, Unbox a)
52 => Int -- ^ from index
53 -> Int -- ^ length of the slice
54 -> Unfold m (MutArray a) (Int, Int)
55sliceIndexerFromLen from len =
56 let fromThenTo n = (from, from + len, n - 1)
57 mkSlice n i = return (i, min len (n - i))
58 in Unfold.lmap length
59 $ Unfold.mapM2 mkSlice
60 $ Unfold.lmap fromThenTo Unfold.enumerateFromThenTo
62{-# DEPRECATED genSlicesFromLen "Please use sliceIndexerFromLen instead." #-}
63genSlicesFromLen :: forall m a. (Monad m, Unbox a)
64 => Int -- ^ from index
65 -> Int -- ^ length of the slice
66 -> Unfold m (MutArray a) (Int, Int)
67genSlicesFromLen = sliceIndexerFromLen
69-- | Generate a stream of slices of specified length from an array, starting
70-- from the supplied array index. The last slice may be shorter than the
71-- requested length depending on the array length.
72--
73-- /Pre-release/
74{-# INLINE slicerFromLen #-}
75slicerFromLen :: forall m a. (Monad m, Unbox a)
76 => Int -- ^ from index
77 -> Int -- ^ length of the slice
78 -> Unfold m (MutArray a) (MutArray a)
79slicerFromLen from len =
80 let mkSlice arr (i, n) = return $ getSliceUnsafe i n arr
81 in Unfold.mapM2 mkSlice (sliceIndexerFromLen from len)
83{-# DEPRECATED getSlicesFromLen "Please use slicerFromLen instead." #-}
84getSlicesFromLen :: forall m a. (Monad m, Unbox a)
85 => Int -- ^ from index
86 -> Int -- ^ length of the slice
87 -> Unfold m (MutArray a) (MutArray a)
88getSlicesFromLen = slicerFromLen
90-------------------------------------------------------------------------------
91-- Compacting Streams of Arrays
92-------------------------------------------------------------------------------
94-- | Scan @compactLE n@ coalesces adjacent arrays in the input stream
95-- only if the combined size would be less than or equal to n.
96{-# INLINE compactLE #-}
97compactLE :: (MonadIO m, Unbox a) =>
98 Int -> Stream m (MutArray a) -> Stream m (MutArray a)
99-- XXX compactLE can be moved to MutArray/Type if we are not using the parser
100-- to implement it.
101-- compactLE n = Stream.catRights . Stream.parseManyD (pCompactLE n)
102compactLE = rCompactLE
104data SplitState s arr
105 = Initial s
106 | Buffering s arr
107 | Splitting s arr
108 | Yielding arr (SplitState s arr)
109 | Finishing
111-- | Split a stream of arrays on a given separator byte, dropping the separator
112-- and coalescing all the arrays between two separators into a single array.
113--
114{-# INLINE_NORMAL _compactOnByteCustom #-}
115_compactOnByteCustom
116 :: MonadIO m
117 => Word8
118 -> Stream m (MutArray Word8)
119 -> Stream m (MutArray Word8)
120_compactOnByteCustom byte (Stream.Stream step state) =
121 Stream.Stream step' (Initial state)
123 where
125 {-# INLINE_LATE step' #-}
126 step' gst (Initial st) = do
127 r <- step gst st
128 case r of
129 Stream.Yield arr s -> do
130 (arr1, marr2) <- breakOn byte arr
131 return $ case marr2 of
132 Nothing -> Stream.Skip (Buffering s arr1)
133 Just arr2 -> Stream.Skip (Yielding arr1 (Splitting s arr2))
134 Stream.Skip s -> return $ Stream.Skip (Initial s)
135 Stream.Stop -> return Stream.Stop
137 step' gst (Buffering st buf) = do
138 r <- step gst st
139 case r of
140 Stream.Yield arr s -> do
141 (arr1, marr2) <- breakOn byte arr
142 -- XXX Use spliceExp instead and then rightSize?
143 buf1 <- splice buf arr1
144 return $ case marr2 of
145 Nothing -> Stream.Skip (Buffering s buf1)
146 Just x -> Stream.Skip (Yielding buf1 (Splitting s x))
147 Stream.Skip s -> return $ Stream.Skip (Buffering s buf)
148 Stream.Stop -> return $
149 if byteLength buf == 0
150 then Stream.Stop
151 else Stream.Skip (Yielding buf Finishing)
153 step' _ (Splitting st buf) = do
154 (arr1, marr2) <- breakOn byte buf
155 return $ case marr2 of
156 Nothing -> Stream.Skip $ Buffering st arr1
157 Just arr2 -> Stream.Skip $ Yielding arr1 (Splitting st arr2)
159 step' _ (Yielding arr next) = return $ Stream.Yield arr next
160 step' _ Finishing = return Stream.Stop
162-- XXX implement predicate based version of this
163-- XXX Naming of predicate based vs custom version
165-- | Split a stream of arrays on a given separator byte, dropping the separator
166-- and coalescing all the arrays between two separators into a single array.
167--
168{-# INLINE compactOnByte #-}
169compactOnByte
170 :: (MonadIO m)
171 => Word8
172 -> Stream m (MutArray Word8)
173 -> Stream m (MutArray Word8)
174-- XXX compare perf of custom vs idiomatic version
175-- compactOnByte = _compactOnByteCustom
176-- XXX use spliceExp and rightSize?
177compactOnByte byte = Stream.splitInnerBy (breakOn byte) splice
179-- | Like 'compactOnByte' considers the separator in suffix position instead of
180-- infix position.
181{-# INLINE compactOnByteSuffix #-}
182compactOnByteSuffix
183 :: (MonadIO m)
184 => Word8
185 -> Stream m (MutArray Word8)
186 -> Stream m (MutArray Word8)
187compactOnByteSuffix byte =
188 -- XXX use spliceExp and rightSize?
189 Stream.splitInnerBySuffix
190 (\arr -> byteLength arr == 0) (breakOn byte) splice