on master help sign in
Use the SPEC argument on all recursive cases
commit00f5d9f42d149e6854c8f41e712b84e816f0a34a authorHarendra Kumar <harendra@composewell.com> dateSat 16 Jul 2022, 19:44 (4 years ago) parentbc49c3d
2 files changed+69−59
core/src/Streamly/Internal/Data/Stream/StreamD/Eliminate.hs+64−56 core/src/Streamly/Internal/Data/Stream/StreamD/Type.hs+5−3
core/src/Streamly/Internal/Data/Stream/StreamD/Eliminate.hs+64−56
@@ -228,6 +228,7 @@ parseBreak (PRD.Parser pstep initial extract) stream@(Stream step state) = do
228228-- Specialized Folds
229229------------------------------------------------------------------------------
230230
231+-- benchmark after dropping 1 item from stream or using unfolds
231232{-# INLINE_NORMAL null #-}
232233null :: Monad m => Stream m a -> m Bool
233234null = foldrM (\_ _ -> return False) (return True)
@@ -243,13 +244,13 @@ headElse a = foldrM (\x _ -> return x) (return a)
243244-- Does not fuse, has the same performance as the StreamK version.
244245{-# INLINE_NORMAL tail #-}
245246tail :: Monad m => Stream m a -> m (Maybe (Stream m a))
246−tail (UnStream step state) = go state
247+tail (UnStream step state) = go SPEC state
247248 where
248− go st = do
249+ go !_ st = do
249250 r <- step defState st
250251 case r of
251252 Yield _ s -> return (Just $ Stream step s)
252− Skip s -> go s
253+ Skip s -> go SPEC s
253254 Stop -> return Nothing
254255
255256-- XXX will it fuse? need custom impl?
@@ -257,18 +258,19 @@ tail (UnStream step state) = go state
257258last :: Monad m => Stream m a -> m (Maybe a)
258259last = foldl' (\_ y -> Just y) Nothing
259260
261+-- XXX Use the foldrM based impl instead
260262{-# INLINE_NORMAL elem #-}
261263elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
262264-- elem e m = foldrM (\x xs -> if x == e then return True else xs) (return False) m
263−elem e (Stream step state) = go state
265+elem e (Stream step state) = go SPEC state
264266 where
265− go st = do
267+ go !_ st = do
266268 r <- step defState st
267269 case r of
268270 Yield x s
269− | x == e -> return True
270− | otherwise -> go s
271− Skip s -> go s
271+ | x == e -> return True
272+ | otherwise -> go SPEC s
273+ Skip s -> go SPEC s
272274 Stop -> return False
273275
274276{-# INLINE_NORMAL notElem #-}
@@ -278,118 +280,124 @@ notElem e s = fmap not (elem e s)
278280{-# INLINE_NORMAL all #-}
279281all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
280282-- all p m = foldrM (\x xs -> if p x then xs else return False) (return True) m
281−all p (Stream step state) = go state
283+all p (Stream step state) = go SPEC state
282284 where
283− go st = do
285+ go !_ st = do
284286 r <- step defState st
285287 case r of
286288 Yield x s
287− | p x -> go s
289+ | p x -> go SPEC s
288290 | otherwise -> return False
289− Skip s -> go s
291+ Skip s -> go SPEC s
290292 Stop -> return True
291293
292294{-# INLINE_NORMAL any #-}
293295any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
294296-- any p m = foldrM (\x xs -> if p x then return True else xs) (return False) m
295−any p (Stream step state) = go state
297+any p (Stream step state) = go SPEC state
296298 where
297− go st = do
299+ go !_ st = do
298300 r <- step defState st
299301 case r of
300302 Yield x s
301− | p x -> return True
302− | otherwise -> go s
303− Skip s -> go s
303+ | p x -> return True
304+ | otherwise -> go SPEC s
305+ Skip s -> go SPEC s
304306 Stop -> return False
305307
306308{-# INLINE_NORMAL maximum #-}
307309maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
308−maximum (Stream step state) = go Nothing state
310+maximum (Stream step state) = go SPEC Nothing' state
309311 where
310− go Nothing st = do
312+ go !_ Nothing' st = do
311313 r <- step defState st
312314 case r of
313− Yield x s -> go (Just x) s
314− Skip s -> go Nothing s
315+ Yield x s -> go SPEC (Just' x) s
316+ Skip s -> go SPEC Nothing' s
315317 Stop -> return Nothing
316− go (Just acc) st = do
318+ go !_ (Just' acc) st = do
317319 r <- step defState st
318320 case r of
319321 Yield x s
320− | acc <= x -> go (Just x) s
321− | otherwise -> go (Just acc) s
322− Skip s -> go (Just acc) s
322+ | acc <= x -> go SPEC (Just' x) s
323+ | otherwise -> go SPEC (Just' acc) s
324+ Skip s -> go SPEC (Just' acc) s
323325 Stop -> return (Just acc)
324326
325327{-# INLINE_NORMAL maximumBy #-}
326328maximumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
327−maximumBy cmp (Stream step state) = go Nothing state
329+maximumBy cmp (Stream step state) = go SPEC Nothing' state
328330 where
329− go Nothing st = do
331+ go !_ Nothing' st = do
330332 r <- step defState st
331333 case r of
332− Yield x s -> go (Just x) s
333− Skip s -> go Nothing s
334+ Yield x s -> go SPEC (Just' x) s
335+ Skip s -> go SPEC Nothing' s
334336 Stop -> return Nothing
335− go (Just acc) st = do
337+ go !_ (Just' acc) st = do
336338 r <- step defState st
337339 case r of
338340 Yield x s -> case cmp acc x of
339− GT -> go (Just acc) s
340− _ -> go (Just x) s
341− Skip s -> go (Just acc) s
341+ GT -> go SPEC (Just' acc) s
342+ _ -> go SPEC (Just' x) s
343+ Skip s -> go SPEC (Just' acc) s
342344 Stop -> return (Just acc)
343345
344346{-# INLINE_NORMAL minimum #-}
345347minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
346−minimum (Stream step state) = go Nothing state
347− where
348− go Nothing st = do
348+minimum (Stream step state) = go SPEC Nothing' state
349+
350+ where
351+
352+ go !_ Nothing' st = do
349353 r <- step defState st
350354 case r of
351− Yield x s -> go (Just x) s
352− Skip s -> go Nothing s
355+ Yield x s -> go SPEC (Just' x) s
356+ Skip s -> go SPEC Nothing' s
353357 Stop -> return Nothing
354− go (Just acc) st = do
358+ go !_ (Just' acc) st = do
355359 r <- step defState st
356360 case r of
357361 Yield x s
358− | acc <= x -> go (Just acc) s
359− | otherwise -> go (Just x) s
360− Skip s -> go (Just acc) s
362+ | acc <= x -> go SPEC (Just' acc) s
363+ | otherwise -> go SPEC (Just' x) s
364+ Skip s -> go SPEC (Just' acc) s
361365 Stop -> return (Just acc)
362366
363367{-# INLINE_NORMAL minimumBy #-}
364368minimumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
365−minimumBy cmp (Stream step state) = go Nothing state
366− where
367− go Nothing st = do
369+minimumBy cmp (Stream step state) = go SPEC Nothing' state
370+
371+ where
372+
373+ go !_ Nothing' st = do
368374 r <- step defState st
369375 case r of
370− Yield x s -> go (Just x) s
371− Skip s -> go Nothing s
376+ Yield x s -> go SPEC (Just' x) s
377+ Skip s -> go SPEC Nothing' s
372378 Stop -> return Nothing
373− go (Just acc) st = do
379+ go !_ (Just' acc) st = do
374380 r <- step defState st
375381 case r of
376382 Yield x s -> case cmp acc x of
377− GT -> go (Just x) s
378− _ -> go (Just acc) s
379− Skip s -> go (Just acc) s
383+ GT -> go SPEC (Just' x) s
384+ _ -> go SPEC (Just' acc) s
385+ Skip s -> go SPEC (Just' acc) s
380386 Stop -> return (Just acc)
381387
382388{-# INLINE_NORMAL (!!) #-}
383389(!!) :: (Monad m) => Stream m a -> Int -> m (Maybe a)
384−(Stream step state) !! i = go i state
385− where
386− go n st = do
390+(Stream step state) !! i = go SPEC i state
391+
392+ where
393+
394+ go !_ !n st = do
387395 r <- step defState st
388396 case r of
389397 Yield x s | n < 0 -> return Nothing
390398 | n == 0 -> return $ Just x
391− | otherwise -> go (n - 1) s
392− Skip s -> go n s
399+ | otherwise -> go SPEC (n - 1) s
400+ Skip s -> go SPEC n s
393401 Stop -> return Nothing
394402
395403{-# INLINE_NORMAL lookup #-}
core/src/Streamly/Internal/Data/Stream/StreamD/Type.hs+5−3
@@ -159,13 +159,13 @@ consM m (Stream step state) = Stream step1 Nothing
159159-- | Does not fuse, has the same performance as the StreamK version.
160160{-# INLINE_NORMAL uncons #-}
161161uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
162−uncons (UnStream step state) = go state
162+uncons (UnStream step state) = go SPEC state
163163 where
164− go st = do
164+ go !_ st = do
165165 r <- step defState st
166166 case r of
167167 Yield x s -> return $ Just (x, Stream step s)
168− Skip s -> go s
168+ Skip s -> go SPEC s
169169 Stop -> return Nothing
170170
171171------------------------------------------------------------------------------
@@ -373,6 +373,8 @@ foldrMx fstep final convert (Stream step state) = convert $ go SPEC state
373373 Skip s -> go SPEC s
374374 Stop -> final
375375
376+-- XXX Should we make all argument strict wherever we use SPEC?
377+
376378-- Note that foldr works on pure values, therefore it becomes necessarily
377379-- strict when the monad m is strict. In that case it cannot terminate early,
378380-- it would evaluate all of its input. Though, this should work fine with lazy