on master help sign in
Simplify fromParser and take
fb1ee2eHarendra Kumar
src/Streamly/Internal/Data/Array/Stream/Fold/Foreign.hs+23−28
@@ -136,18 +136,10 @@ fromFold (Fold.Fold fstep finitial fextract) =
136fromParser :: forall m a b. (MonadIO m, Storable a) =>
136fromParser :: forall m a b. (MonadIO m, Storable a) =>
137 ParserD.Parser m a b -> Fold m a b
137 ParserD.Parser m a b -> Fold m a b
138fromParser (ParserD.Parser step1 initial1 extract1) =
138fromParser (ParserD.Parser step1 initial1 extract1) =
139− Fold (ParserD.Parser step initial extract1)
139+ Fold (ParserD.Parser step initial1 extract1)
140
140
141 where
141 where
142
142
143− initial = do
144− res <- initial1
145− return
146− $ case res of
147− IPartial s1 -> IPartial s1
148− IDone b -> IDone b
149− IError err -> IError err
150−
151 step s (Array fp@(ForeignPtr start _) end) = do
143 step s (Array fp@(ForeignPtr start _) end) = do
152 if Ptr start >= end
144 if Ptr start >= end
153 then return $ Continue 0 s
145 then return $ Continue 0 s
@@ -155,6 +147,13 @@ fromParser (ParserD.Parser step1 initial1 extract1) =
155
147
156 where
148 where
157
149
150+ {-# INLINE partial #-}
151+ partial arrRem cur next elemSize st n fs1 = do
152+ let next1 = next `plusPtr` negate (n * elemSize)
153+ if next1 >= Ptr start && cur < end
154+ then goArray SPEC next1 fs1
155+ else return $ st (arrRem + n) fs1
156+
158 goArray !_ !cur !fs = do
157 goArray !_ !cur !fs = do
159 x <- liftIO $ peek cur
158 x <- liftIO $ peek cur
160 liftIO $ touchForeignPtr fp
159 liftIO $ touchForeignPtr fp
@@ -165,16 +164,10 @@ fromParser (ParserD.Parser step1 initial1 extract1) =
165 case res of
164 case res of
166 ParserD.Done n b -> do
165 ParserD.Done n b -> do
167 return $ Done (arrRem + n) b
166 return $ Done (arrRem + n) b
168− ParserD.Partial n fs1 -> do
167+ ParserD.Partial n fs1 ->
169− let next1 = next `plusPtr` negate (n * elemSize)
168+ partial arrRem cur next elemSize Partial n fs1
170− if next1 >= Ptr start && cur < end
171− then goArray SPEC next1 fs1
172− else return $ Partial (arrRem + n) fs1
173 ParserD.Continue n fs1 -> do
169 ParserD.Continue n fs1 -> do
174− let next1 = next `plusPtr` negate (n * elemSize)
170+ partial arrRem cur next elemSize Continue n fs1
175− if next1 >= Ptr start && cur < end
176− then goArray SPEC next1 fs1
177− else return $ Continue (arrRem + n) fs1
178 Error err -> return $ Error err
171 Error err -> return $ Error err
179
172
180-- | Adapt an array stream fold.
173-- | Adapt an array stream fold.
@@ -309,20 +302,22 @@ take n (Fold (ParserD.Parser step1 initial1 extract1)) =
309 IDone b -> return $ IDone b
302 IDone b -> return $ IDone b
310 IError err -> return $ IError err
303 IError err -> return $ IError err
311
304
305+ {-# INLINE partial #-}
306+ partial i1 st j s =
307+ let i2 = i1 + j
308+ in if i2 > 0
309+ then return $ st j (Tuple' i2 s)
310+ else Done 0 <$> extract1 s -- i2 == i1 == j == 0
311+
312 step (Tuple' i r) arr = do
312 step (Tuple' i r) arr = do
313 let len = Array.length arr
313 let len = Array.length arr
314− if len <= i
314+ i1 = i - len
315+ if i1 >= 0
315 then do
316 then do
316 res <- step1 r arr
317 res <- step1 r arr
317 case res of
318 case res of
318− Partial j s ->
319+ Partial j s -> partial i1 Partial j s
319− if (i + j - len) > 0
320− then return $ Partial j (Tuple' (i + j - len) s)
321− else Done j <$> extract1 s
322− Continue j s ->
320+ Continue j s -> partial i1 Continue j s
323− if (i + j - len) > 0
324− then return $ Continue j (Tuple' (i + j - len) s)
325− else Done j <$> extract1 s
326 Done j b -> return $ Done j b
321 Done j b -> return $ Done j b
327 Error err -> return $ Error err
322 Error err -> return $ Error err
328 else do
323 else do
@@ -330,7 +325,7 @@ take n (Fold (ParserD.Parser step1 initial1 extract1)) =
330 sz = sizeOf (undefined :: a)
325 sz = sizeOf (undefined :: a)
331 end = Ptr start `plusPtr` (i * sz)
326 end = Ptr start `plusPtr` (i * sz)
332 arr1 = Array (ForeignPtr start contents) end
327 arr1 = Array (ForeignPtr start contents) end
333− remaining = len - i
328+ remaining = negate i1
334 res <- step1 r arr1
329 res <- step1 r arr1
335 case res of
330 case res of
336 Partial 0 s -> Done remaining <$> extract1 s
331 Partial 0 s -> Done remaining <$> extract1 s