diff --git a/benchmark/Array.hs b/benchmark/Array.hs index c64908e5b..918bb1af4 100644 --- a/benchmark/Array.hs +++ b/benchmark/Array.hs @@ -86,7 +86,8 @@ main = -} , benchPureSink "<" Ops.ordInstance , benchPureSink "min" Ops.ordInstanceMin - , benchPureSink "IsList.toList" GHC.toList + -- length is used to check for foldr/build fusion + , benchPureSink "length . IsList.toList" (length . GHC.toList) , benchIOSink "foldl'" Ops.pureFoldl' , benchIOSink "read" (S.drain . S.unfold A.read) , benchIOSink "toStreamRev" (S.drain . IA.toStreamRev) diff --git a/benchmark/Linear.hs b/benchmark/Linear.hs index a50e62784..44dcf2574 100644 --- a/benchmark/Linear.hs +++ b/benchmark/Linear.hs @@ -124,7 +124,8 @@ main = , benchPureSink "<" Ops.ordInstance , benchPureSink "min" Ops.ordInstanceMin , benchPureSrc "IsList.fromList" Ops.sourceIsList - , benchPureSink "IsList.toList" GHC.toList + -- length is used to check for foldr/build fusion + , benchPureSink "length . IsList.toList" (length . GHC.toList) , benchPureSrc "IsString.fromString" Ops.sourceIsString , mkString `deepseq` (bench "readsPrec pure streams" $ nf Ops.readInstance mkString) diff --git a/src/Streamly/Internal/Memory/Array/Types.hs b/src/Streamly/Internal/Memory/Array/Types.hs index 64dd704cf..64d4a1c95 100644 --- a/src/Streamly/Internal/Memory/Array/Types.hs +++ b/src/Streamly/Internal/Memory/Array/Types.hs @@ -102,7 +102,7 @@ import Foreign.Storable (Storable(..)) import Prelude hiding (length, foldr, read, unlines, splitAt) import Text.Read (readPrec, readListPrec, readListPrecDefault) -import GHC.Base (Addr#, nullAddr#, realWorld#) +import GHC.Base (Addr#, nullAddr#, realWorld#, build) import GHC.Exts (IsList, IsString(..)) import GHC.ForeignPtr (ForeignPtr(..), newForeignPtr_) import GHC.IO (IO(IO), unsafePerformIO) @@ -738,12 +738,33 @@ fromStreamD m = runFold write m runFold (Fold step begin done) = D.foldlMx' step begin done -} +-- Use foldr/build fusion to fuse with list consumers +-- This can be useful when using the IsList instance +{-# INLINE_LATE toListFB #-} +toListFB :: forall a b. Storable a => (a -> b -> b) -> b -> Array a -> b +toListFB c n Array{..} = go (unsafeForeignPtrToPtr aStart) + where + + go p | p == aEnd = n + go p = + -- unsafeInlineIO allows us to run this in Identity monad for pure + -- toList/foldr case which makes them much faster due to not + -- accumulating the list and fusing better with the pure consumers. + -- + -- This should be safe as the array contents are guaranteed to be + -- evaluated/written to before we peek at them. + let !x = unsafeInlineIO $ do + r <- peek p + touchForeignPtr aStart + return r + in c x (go (p `plusPtr` (sizeOf (undefined :: a)))) + -- | Convert an 'Array' into a list. -- -- @since 0.7.0 -{-# INLINABLE toList #-} +{-# INLINE toList #-} toList :: Storable a => Array a -> [a] -toList = foldr (:) [] +toList s = build (\c n -> toListFB c n s) instance (Show a, Storable a) => Show (Array a) where {-# INLINE showsPrec #-} diff --git a/src/Streamly/Streams/StreamD/Type.hs b/src/Streamly/Streams/StreamD/Type.hs index f69f03824..68ebded33 100644 --- a/src/Streamly/Streams/StreamD/Type.hs +++ b/src/Streamly/Streams/StreamD/Type.hs @@ -68,6 +68,8 @@ where import Control.Applicative (liftA2) import Control.Monad (ap, when) import Control.Monad.Trans (lift, MonadTrans) +import Data.Functor.Identity (Identity(..)) +import GHC.Base (build) import GHC.Types (SPEC(..)) import Prelude hiding (map, mapM, foldr, take, concatMap) @@ -344,10 +346,26 @@ foldrT f final (Stream step state) = go SPEC state Skip s -> go SPEC s Stop -> final -{-# INLINE toList #-} +{-# INLINE_NORMAL toList #-} toList :: Monad m => Stream m a -> m [a] toList = foldr (:) [] +-- Use foldr/build fusion to fuse with list consumers +-- This can be useful when using the IsList instance +{-# INLINE_LATE toListFB #-} +toListFB :: (a -> b -> b) -> b -> Stream Identity a -> b +toListFB c n (Stream step state) = go state + where + go st = case runIdentity (step defState st) of + Yield x s -> x `c` go s + Skip s -> go s + Stop -> n + +{-# RULES "toList Identity" toList = toListId #-} +{-# INLINE_EARLY toListId #-} +toListId :: Stream Identity a -> Identity [a] +toListId s = Identity $ build (\c n -> toListFB c n s) + -- XXX run begin action only if the stream is not empty. {-# INLINE_NORMAL foldlMx' #-} foldlMx' :: Monad m => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> m b