bd13edb · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Reuse mutable array producer code in immutable arrays

bd13edbAdithya Kumar4 years ago

core/src/Streamly/Internal/Data/Array/Foreign.hs+4−22

compare with master

unifiedsplit

@@ -140,6 +140,7 @@ import qualified Streamly.Internal.Data.Array.Foreign.Mut.Type as MA

140import qualified Streamly.Internal.Data.Array.Foreign.Mut as MA

141import qualified Streamly.Internal.Data.Array.Foreign.Type as A

142import qualified Streamly.Internal.Data.Fold as FL

143+import qualified Streamly.Internal.Data.Producer.Type as Producer

143import qualified Streamly.Internal.Data.Producer as Producer

144import qualified Streamly.Internal.Data.Producer as Producer

144import qualified Streamly.Internal.Data.Stream.Common as P

145import qualified Streamly.Internal.Data.Stream.Common as P

145import qualified Streamly.Internal.Data.Stream.StreamD as D

146import qualified Streamly.Internal.Data.Stream.StreamD as D

@@ -181,28 +182,9 @@ fromStream (SerialT m) = P.fold A.write m

181

182

182{-# INLINE_NORMAL producer #-}

183{-# INLINE_NORMAL producer #-}

183producer :: forall m a. (Monad m, Storable a) => Producer m (Array a) a

184producer :: forall m a. (Monad m, Storable a) => Producer m (Array a) a

184−producer = Producer step inject extract

185+producer =

185− where

186+ Producer.lmap2 A.unsafeThaw A.unsafeFreeze

186−

187− {-# INLINE inject #-}

188− inject (Array contents start end) = return $ ReadUState contents end start

189−

190− {-# INLINE_LATE step #-}

191− step (ReadUState contents end cur)

192− | assert (cur <= end) (cur == end) =

193− let x = unsafeInlineIO $ touch contents

194− in x `seq` return D.Stop

195− step (ReadUState contents end cur) = do

196− -- unsafeInlineIO allows us to run this in Identity monad for pure

197− -- toList/foldr case which makes them much faster due to not

198− -- accumulating the list and fusing better with the pure consumers.

199− --

200− -- This should be safe as the array contents are guaranteed to be

201− -- evaluated/written to before we peek at them.

202− let !x = unsafeInlineIO $ peek cur

187+ $ MA.producerWith (return . unsafeInlineIO)

203− return $ D.Yield x (ReadUState contents end (PTR_NEXT(cur,a)))

204−

205− extract (ReadUState contents end cur) = return $ Array contents cur end

206

188

207-- | Unfold an array into a stream.

189-- | Unfold an array into a stream.

208--

190--

core/src/Streamly/Internal/Data/Array/Foreign/Mut/Type.hs+53−41

@@ -98,14 +98,14 @@ module Streamly.Internal.Data.Array.Foreign.Mut.Type

98 -- ** To streams

99 , ReadUState(..)

100 , read

101− , readRev_

101+ , readRevWith

102 , readRev

103

104 -- ** To containers

105− , toStreamD_

105+ , toStreamDWith

106− , toStreamDRev_

106+ , toStreamDRevWith

107− , toStreamK_

107+ , toStreamKWith

108− , toStreamKRev_

108+ , toStreamKRevWith

109 , toStreamD

110 , toStreamDRev

111 , toStreamK

@@ -113,6 +113,7 @@ module Streamly.Internal.Data.Array.Foreign.Mut.Type

113 , toList

114

115 -- experimental

116+ , producerWith

116 , producer

117 , producer

117

118

118 -- ** Random reads

119 -- ** Random reads

@@ -203,6 +204,10 @@ module Streamly.Internal.Data.Array.Foreign.Mut.Type

203 )

204 )

204where

205where

205

206

207+-- Notes:

208+-- When we use a purely lazy Monad like Identity, we need to force ordering of

209+-- some actions for correctness.

210+

206#include "inline.hs"

211#include "inline.hs"

207#include "ArrayMacros.h"

212#include "ArrayMacros.h"

208#include "MachDeps.h"

213#include "MachDeps.h"

@@ -1516,24 +1521,36 @@ data ReadUState a = ReadUState

1516toReadUState :: Array a -> ReadUState a

1521toReadUState :: Array a -> ReadUState a

1517toReadUState (Array contents start end _) = ReadUState contents end start

1522toReadUState (Array contents start end _) = ReadUState contents end start

1518

1523

1519−-- | Resumable unfold of an array.

1520−--

1521−{-# INLINE_NORMAL producer #-}

1524+{-# INLINE_NORMAL producerWith #-}

1525+producerWith ::

1522−producer :: forall m a. (MonadIO m, Storable a) => Producer m (Array a) a

1526+ forall m a. (Monad m, Storable a)

1527+ => (forall b. IO b -> m b) -> Producer m (Array a) a

1523−producer = Producer step (return . toReadUState) extract

1528+producerWith liftio = Producer step (return . toReadUState) extract

1524 where

1529 where

1525

1530

1526 {-# INLINE_LATE step #-}

1531 {-# INLINE_LATE step #-}

1527 step (ReadUState contents end cur)

1532 step (ReadUState contents end cur)

1528− | assert (cur <= end) (cur == end) = do

1533+ | assert (cur <= end) (cur == end) = liftio $ do

1534+ -- We bind it strictly here as we don't want this to be optimized

1535+ -- out in a purely lazy monad like Identity.

1529− liftIO $ touch contents

1536+ !_ <- touch contents

1530 return D.Stop

1537 return D.Stop

1531 step (ReadUState contents end cur) = do

1538 step (ReadUState contents end cur) = do

1539+ -- When we use a purely lazy Monad like Identity, we need to force a

1540+ -- few actions for correctness and execution order sanity. We want

1541+ -- the peek to occur right here and not lazily at some later point

1542+ -- because we want the peek to be ordered with respect to the touch.

1532− !x <- liftIO $ peek cur

1543+ !x <- liftio $ peek cur

1533 return $ D.Yield x (ReadUState contents end (PTR_NEXT(cur,a)))

1544 return $ D.Yield x (ReadUState contents end (PTR_NEXT(cur,a)))

1534

1545

1535 extract (ReadUState contents end cur) = return $ Array contents cur end end

1546 extract (ReadUState contents end cur) = return $ Array contents cur end end

1536

1547

1548+-- | Resumable unfold of an array.

1549+--

1550+{-# INLINE_NORMAL producer #-}

1551+producer :: forall m a. (MonadIO m, Storable a) => Producer m (Array a) a

1552+producer = producerWith liftIO

1553+

1537-- | Unfold an array into a stream.

1554-- | Unfold an array into a stream.

1538--

1555--

1539-- @since 0.7.0

1556-- @since 0.7.0

@@ -1541,11 +1558,11 @@ producer = Producer step (return . toReadUState) extract

1541read :: forall m a. (MonadIO m, Storable a) => Unfold m (Array a) a

1558read :: forall m a. (MonadIO m, Storable a) => Unfold m (Array a) a

1542read = Producer.simplify producer

1559read = Producer.simplify producer

1543

1560

1544−{-# INLINE_NORMAL readRev_ #-}

1561+{-# INLINE_NORMAL readRevWith #-}

1545−readRev_ ::

1562+readRevWith ::

1546 forall m a. (Monad m, Storable a)

1563 forall m a. (Monad m, Storable a)

1547 => (forall b. IO b -> m b) -> Unfold m (Array a) a

1564 => (forall b. IO b -> m b) -> Unfold m (Array a) a

1548−readRev_ liftio = Unfold step inject

1565+readRevWith liftio = Unfold step inject

1549 where

1566 where

1550

1567

1551 inject (Array contents start end _) =

1568 inject (Array contents start end _) =

@@ -1553,16 +1570,11 @@ readRev_ liftio = Unfold step inject

1553 in return $ ReadUState contents start p

1570 in return $ ReadUState contents start p

1554

1571

1555 {-# INLINE_LATE step #-}

1572 {-# INLINE_LATE step #-}

1556− step (ReadUState contents start p) | p < start = do

1573+ step (ReadUState contents start p) | p < start = liftio $ do

1557− -- Is the following replicated properly?

1558− -- let x = unsafeInlineIO $ touch contents

1559− -- in x `seq` return D.Stop

1560− liftio $ touch contents

1574+ !_ <- touch contents

1561 return D.Stop

1575 return D.Stop

1562− -- unit <- liftio $ touch contents

1563− -- unit `seq` return D.Stop

1564 step (ReadUState contents start p) = do

1576 step (ReadUState contents start p) = do

1565− x <- liftio $ peek p

1577+ !x <- liftio $ peek p

1566 return $ D.Yield x (ReadUState contents start (PTR_PREV(p,a)))

1578 return $ D.Yield x (ReadUState contents start (PTR_PREV(p,a)))

1567

1579

1568-- | Unfold an array into a stream in reverse order.

1580-- | Unfold an array into a stream in reverse order.

@@ -1570,7 +1582,7 @@ readRev_ liftio = Unfold step inject

1570-- /Pre-release/

1582-- /Pre-release/

1571{-# INLINE_NORMAL readRev #-}

1583{-# INLINE_NORMAL readRev #-}

1572readRev :: forall m a. (MonadIO m, Storable a) => Unfold m (Array a) a

1584readRev :: forall m a. (MonadIO m, Storable a) => Unfold m (Array a) a

1573−readRev = readRev_ liftIO

1585+readRev = readRevWith liftIO

1574

1586

1575-------------------------------------------------------------------------------

1587-------------------------------------------------------------------------------

1576-- to Lists and streams

1588-- to Lists and streams

@@ -1616,11 +1628,11 @@ toList Array{..} = liftIO $ go arrStart

1616 touch arrContents

1628 touch arrContents

1617 (:) x <$> go (PTR_NEXT(p,a))

1629 (:) x <$> go (PTR_NEXT(p,a))

1618

1630

1619−{-# INLINE_NORMAL toStreamD_ #-}

1631+{-# INLINE_NORMAL toStreamDWith #-}

1620−toStreamD_ ::

1632+toStreamDWith ::

1621 forall m a. (Monad m, Storable a)

1633 forall m a. (Monad m, Storable a)

1622 => (forall b. IO b -> m b) -> Array a -> D.Stream m a

1634 => (forall b. IO b -> m b) -> Array a -> D.Stream m a

1623−toStreamD_ liftio Array{..} = D.Stream step arrStart

1635+toStreamDWith liftio Array{..} = D.Stream step arrStart

1624

1636

1625 where

1637 where

1626

1638

@@ -1638,13 +1650,13 @@ toStreamD_ liftio Array{..} = D.Stream step arrStart

1638-- We can try this if the unfold has any performance issues.

1650-- We can try this if the unfold has any performance issues.

1639{-# INLINE_NORMAL toStreamD #-}

1651{-# INLINE_NORMAL toStreamD #-}

1640toStreamD :: forall m a. (MonadIO m, Storable a) => Array a -> D.Stream m a

1652toStreamD :: forall m a. (MonadIO m, Storable a) => Array a -> D.Stream m a

1641−toStreamD arr = toStreamD_ liftIO arr

1653+toStreamD = toStreamDWith liftIO

1642

1654

1643−{-# INLINE toStreamK_ #-}

1655+{-# INLINE toStreamKWith #-}

1644−toStreamK_ ::

1656+toStreamKWith ::

1645 forall m a. (Monad m, Storable a)

1657 forall m a. (Monad m, Storable a)

1646 => (forall b. IO b -> m b) -> Array a -> K.Stream m a

1658 => (forall b. IO b -> m b) -> Array a -> K.Stream m a

1647−toStreamK_ liftio Array{..} = go arrStart

1659+toStreamKWith liftio Array{..} = go arrStart

1648

1660

1649 where

1661 where

1650

1662

@@ -1658,13 +1670,13 @@ toStreamK_ liftio Array{..} = go arrStart

1658

1670

1659{-# INLINE toStreamK #-}

1671{-# INLINE toStreamK #-}

1660toStreamK :: forall m a. (MonadIO m, Storable a) => Array a -> K.Stream m a

1672toStreamK :: forall m a. (MonadIO m, Storable a) => Array a -> K.Stream m a

1661−toStreamK arr = toStreamK_ liftIO arr

1673+toStreamK = toStreamKWith liftIO

1662

1674

1663−{-# INLINE_NORMAL toStreamDRev_ #-}

1675+{-# INLINE_NORMAL toStreamDRevWith #-}

1664−toStreamDRev_ ::

1676+toStreamDRevWith ::

1665 forall m a. (Monad m, Storable a)

1677 forall m a. (Monad m, Storable a)

1666 => (forall b. IO b -> m b) -> Array a -> D.Stream m a

1678 => (forall b. IO b -> m b) -> Array a -> D.Stream m a

1667−toStreamDRev_ liftio Array{..} =

1679+toStreamDRevWith liftio Array{..} =

1668 let p = PTR_PREV(aEnd,a)

1680 let p = PTR_PREV(aEnd,a)

1669 in D.Stream step p

1681 in D.Stream step p

1670

1682

@@ -1684,13 +1696,13 @@ toStreamDRev_ liftio Array{..} =

1684-- We can try this if the unfold has any perf issues.

1696-- We can try this if the unfold has any perf issues.

1685{-# INLINE_NORMAL toStreamDRev #-}

1697{-# INLINE_NORMAL toStreamDRev #-}

1686toStreamDRev :: forall m a. (MonadIO m, Storable a) => Array a -> D.Stream m a

1698toStreamDRev :: forall m a. (MonadIO m, Storable a) => Array a -> D.Stream m a

1687−toStreamDRev arr = toStreamDRev_ liftIO arr

1699+toStreamDRev = toStreamDRevWith liftIO

1688

1700

1689−{-# INLINE toStreamKRev_ #-}

1701+{-# INLINE toStreamKRevWith #-}

1690−toStreamKRev_ ::

1702+toStreamKRevWith ::

1691 forall m a. (Monad m, Storable a)

1703 forall m a. (Monad m, Storable a)

1692 => (forall b. IO b -> m b) -> Array a -> K.Stream m a

1704 => (forall b. IO b -> m b) -> Array a -> K.Stream m a

1693−toStreamKRev_ liftio Array {..} =

1705+toStreamKRevWith liftio Array {..} =

1694 let p = PTR_PREV(aEnd,a)

1706 let p = PTR_PREV(aEnd,a)

1695 in go p

1707 in go p

1696

1708

@@ -1706,7 +1718,7 @@ toStreamKRev_ liftio Array {..} =

1706

1718

1707{-# INLINE toStreamKRev #-}

1719{-# INLINE toStreamKRev #-}

1708toStreamKRev :: forall m a. (MonadIO m, Storable a) => Array a -> K.Stream m a

1720toStreamKRev :: forall m a. (MonadIO m, Storable a) => Array a -> K.Stream m a

1709−toStreamKRev arr = toStreamKRev_ liftIO arr

1721+toStreamKRev = toStreamKRevWith liftIO

1710

1722

1711-------------------------------------------------------------------------------

1723-------------------------------------------------------------------------------

1712-- Folding

1724-- Folding

core/src/Streamly/Internal/Data/Array/Foreign/Type.hs+7−5

@@ -474,23 +474,25 @@ length arr = MA.length (unsafeThaw arr)

474-- @since 0.8.0

475{-# INLINE_NORMAL readRev #-}

476readRev :: forall m a. (Monad m, Storable a) => Unfold m (Array a) a

477−readRev = Unfold.lmap unsafeThaw $ MA.readRev_ (return . unsafeInlineIO)

477+readRev = Unfold.lmap unsafeThaw $ MA.readRevWith (return . unsafeInlineIO)

478

479{-# INLINE_NORMAL toStreamD #-}

480toStreamD :: forall m a. (Monad m, Storable a) => Array a -> D.Stream m a

481−toStreamD arr = MA.toStreamD_ (return . unsafeInlineIO) (unsafeThaw arr)

481+toStreamD arr = MA.toStreamDWith (return . unsafeInlineIO) (unsafeThaw arr)

482

483{-# INLINE toStreamK #-}

484toStreamK :: forall m a. (Monad m, Storable a) => Array a -> K.Stream m a

485−toStreamK arr = MA.toStreamK_ (return . unsafeInlineIO) (unsafeThaw arr)

485+toStreamK arr = MA.toStreamKWith (return . unsafeInlineIO) (unsafeThaw arr)

486

487{-# INLINE_NORMAL toStreamDRev #-}

488toStreamDRev :: forall m a. (Monad m, Storable a) => Array a -> D.Stream m a

489+toStreamDRev arr =

489−toStreamDRev arr = MA.toStreamDRev_ (return . unsafeInlineIO) (unsafeThaw arr)

490+ MA.toStreamDRevWith (return . unsafeInlineIO) (unsafeThaw arr)

490

491

491{-# INLINE toStreamKRev #-}

492{-# INLINE toStreamKRev #-}

492toStreamKRev :: forall m a. (Monad m, Storable a) => Array a -> K.Stream m a

493toStreamKRev :: forall m a. (Monad m, Storable a) => Array a -> K.Stream m a

494+toStreamKRev arr =

493−toStreamKRev arr = MA.toStreamKRev_ (return . unsafeInlineIO) (unsafeThaw arr)

495+ MA.toStreamKRevWith (return . unsafeInlineIO) (unsafeThaw arr)

494

496

495-- | Convert an 'Array' into a stream.

497-- | Convert an 'Array' into a stream.

496--

498--

core/src/Streamly/Internal/Data/Producer/Type.hs+4−4

@@ -23,7 +23,7 @@ module Streamly.Internal.Data.Producer.Type

23 -- * Mapping

24 , translate

25 , lmap

26− , lmap_

26+ , lmap2

28 -- * Nesting

29 , NestedLoop (..)

@@ -116,9 +116,9 @@ lmap f (Producer step inject extract) = Producer step (inject . f) extract

116-- | Map the producer input using the provided translation functions

117--

118-- /Pre-release/

119−{-# INLINE_NORMAL lmap_ #-}

119+{-# INLINE_NORMAL lmap2 #-}

120−lmap_ :: Functor m => (a -> c) -> (c -> a) -> Producer m c b -> Producer m a b

120+lmap2 :: Functor m => (a -> c) -> (c -> a) -> Producer m c b -> Producer m a b

121−lmap_ f g (Producer step inject extract) = Producer step (inject . f) (fmap g . extract)

121+lmap2 f g (Producer step inject extract) = Producer step (inject . f) (fmap g . extract)

122

123------------------------------------------------------------------------------

124-- Functor