git.iko.soy/streamly
on master help sign in
Fix writeN dropping one element at the end of Fold
commitf1a8fae977b51a501bcae9237cfd20a67eb71b73 authorHarendra Kumar <harendra@composewell.com> dateWed 6 Oct 2021, 11:17 (5 years ago) parent8139397
4 files changed+76−53
compare with master
unifiedsplit
src/Streamly/Internal/Data/Array.hs+20−11 src/Streamly/Internal/Data/Array/Foreign/Mut/Type.hs+25−24 src/Streamly/Internal/Data/Array/Prim/MutTypesInclude.hs+15−9 src/Streamly/Internal/Data/SmallArray.hs+16−9
src/Streamly/Internal/Data/Array.hs+20−11
@@ -75,20 +75,29 @@ bottomElement = undefined
7575-- Construction - Folds
7676-------------------------------------------------------------------------------
7777
78−-- writeN n = S.evertM (fromStreamDN n)
7978{-# INLINE_NORMAL writeN #-}
8079writeN :: MonadIO m => Int -> Fold m a (Array a)
81−writeN limit = Fold step initial extract
82− where
80+writeN len = Fold step initial extract
81+
82+ where
83+
84+ {-# INLINE next #-}
85+ next marr i = do
86+ let i1 = i + 1
87+ st = Tuple' marr i1
88+ if len > i1
89+ then return $ FL.Partial st
90+ else fmap FL.Done $ extract st
91+
8392 initial = do
84− marr <- liftIO $ newArray limit bottomElement
85− return $ FL.Partial (Tuple' marr 0)
86− step st@(Tuple' marr i) x
87− | i == limit = fmap FL.Done $ extract st
88− | otherwise = do
89− liftIO $ writeArray marr i x
90− return $ FL.Partial $ Tuple' marr (i + 1)
91− extract (Tuple' marr len) = liftIO $ freezeArray marr 0 len
93+ marr <- liftIO $ newArray len bottomElement
94+ next marr (-1)
95+
96+ step (Tuple' marr i) x = do
97+ liftIO $ writeArray marr i x
98+ next marr i
99+
100+ extract (Tuple' marr l) = liftIO $ freezeArray marr 0 l
92101
93102{-# INLINE_NORMAL write #-}
94103write :: MonadIO m => Fold m a (Array a)
src/Streamly/Internal/Data/Array/Foreign/Mut/Type.hs+25−24
@@ -769,23 +769,13 @@ foldr f z arr = runIdentity $ D.foldr f z $ toStreamD arr
769769{-# INLINE_NORMAL writeNAllocWith #-}
770770writeNAllocWith :: forall m a. (MonadIO m, Storable a)
771771 => (Int -> m (Array a)) -> Int -> Fold m a (Array a)
772−writeNAllocWith alloc n = Fold step initial extract
773−
774− where
775−
776− initial = FL.Partial <$> alloc (max n 0)
777− step arr@(Array _ end bound) _ | end == bound = return $ FL.Done arr
778− step (Array start end bound) x = do
779− liftIO $ poke end x
780− return $ FL.Partial $ Array start (end `plusPtr` sizeOf (undefined :: a)) bound
781− -- XXX note that shirkToFit does not maintain alignment, in case we are
782− -- using aligned allocation.
783− extract = return -- liftIO . shrinkToFit
772+writeNAllocWith alloc n = FL.take n (writeNUnsafeWith alloc n)
784773
785774-- | @writeN n@ folds a maximum of @n@ elements from the input stream to an
786775-- 'Array'.
787776--
788−-- @writeN n = Fold.take n writeNUnsafe@
777+-- >>> import qualified Streamly.Data.Fold as Fold
778+-- >>> writeN n = Fold.take n (writeNUnsafe n)
789779--
790780-- @since 0.7.0
791781{-# INLINE_NORMAL writeN #-}
@@ -819,22 +809,21 @@ data ArrayUnsafe a = ArrayUnsafe
819809 {-# UNPACK #-} !(ForeignPtr a) -- first address
820810 {-# UNPACK #-} !(Ptr a) -- first unused address
821811
822−-- | Like 'writeN' but does not check the array bounds when writing. The fold
823−-- driver must not call the step function more than 'n' times otherwise it will
824−-- corrupt the memory and crash. This function exists mainly because any
825−-- conditional in the step function blocks fusion causing 10x performance
826−-- slowdown.
812+-- XXX We can carry bound as well in the state to make sure we do not lose the
813+-- remaining capacity. Need to check perf impact.
827814--
828−-- @since 0.7.0
829−{-# INLINE_NORMAL writeNUnsafe #-}
830−writeNUnsafe :: forall m a. (MonadIO m, Storable a)
831− => Int -> Fold m a (Array a)
832−writeNUnsafe n = Fold step initial extract
815+-- | This API is only useful when we do not want to use the remaining space in
816+-- the array to append anything else later as that space is lost.
817+--
818+{-# INLINE_NORMAL writeNUnsafeWith #-}
819+writeNUnsafeWith :: forall m a. (MonadIO m, Storable a)
820+ => (Int -> m (Array a)) -> Int -> Fold m a (Array a)
821+writeNUnsafeWith alloc n = Fold step initial extract
833822
834823 where
835824
836825 initial = do
837− (Array start end _) <- liftIO $ newArray (max n 0)
826+ (Array start end _) <- alloc (max n 0)
838827 return $ FL.Partial $ ArrayUnsafe start end
839828
840829 step (ArrayUnsafe start end) x = do
@@ -845,6 +834,18 @@ writeNUnsafe n = Fold step initial extract
845834
846835 extract (ArrayUnsafe start end) = return $ Array start end end -- liftIO . shrinkToFit
847836
837+-- | Like 'writeN' but does not check the array bounds when writing. The fold
838+-- driver must not call the step function more than 'n' times otherwise it will
839+-- corrupt the memory and crash. This function exists mainly because any
840+-- conditional in the step function blocks fusion causing 10x performance
841+-- slowdown.
842+--
843+-- @since 0.7.0
844+{-# INLINE_NORMAL writeNUnsafe #-}
845+writeNUnsafe :: forall m a. (MonadIO m, Storable a)
846+ => Int -> Fold m a (Array a)
847+writeNUnsafe = writeNUnsafeWith newArray
848+
848849-- XXX Buffer to a list instead?
849850--
850851-- | Buffer a stream into a stream of arrays.
src/Streamly/Internal/Data/Array/Prim/MutTypesInclude.hs+15−9
@@ -163,21 +163,27 @@ write = FL.rmapM extract $ FL.foldlM' step initial
163163-- /Pre-release/
164164{-# INLINE_NORMAL writeN #-}
165165writeN :: (MonadIO m, Prim a) => Int -> Fold m a (Array a)
166−writeN limit = Fold step initial extract
166+writeN len = Fold step initial extract
167167
168168 where
169169
170+ {-# INLINE next #-}
171+ next marr i = do
172+ let i1 = i + 1
173+ st = Tuple' marr i1
174+ if len > i1
175+ then return $ FL.Partial st
176+ else fmap FL.Done $ extract st
177+
170178 initial = do
171− marr <- newArray limit
172− return $ FL.Partial $ Tuple' marr 0
179+ marr <- liftIO $ newArray len
180+ next marr (-1)
173181
174− extract (Tuple' marr len) = shrinkArray marr len >> return marr
182+ step (Tuple' marr i) x = do
183+ liftIO $ unsafeWriteIndex marr i x
184+ next marr i
175185
176− step s@(Tuple' marr i) x
177− | i == limit = FL.Done <$> extract s
178− | otherwise = do
179− unsafeWriteIndex marr i x
180− return $ FL.Partial $ Tuple' marr (i + 1)
186+ extract (Tuple' marr l) = shrinkArray marr l >> return marr
181187
182188-- Use Tuple' instead?
183189data ArrayUnsafe a = ArrayUnsafe
src/Streamly/Internal/Data/SmallArray.hs+16−9
@@ -105,21 +105,28 @@ foldr f z arr = runIdentity $ D.foldr f z $ toStreamD arr
105105-- of elements use an 'Array' from either "Streamly.Data.Array" or "Streamly.Data.Array.Foreign".
106106{-# INLINE_NORMAL writeN #-}
107107writeN :: MonadIO m => Int -> Fold m a (SmallArray a)
108−writeN limit = Fold step initial extract
108+writeN len = Fold step initial extract
109109
110110 where
111111
112+ {-# INLINE next #-}
113+ next marr i = do
114+ let i1 = i + 1
115+ st = Tuple' marr i1
116+ if len > i1
117+ then return $ FL.Partial st
118+ else fmap FL.Done $ extract st
119+
112120 initial = do
113− marr <- liftIO $ newSmallArray limit bottomElement
114− return $ FL.Partial (Tuple' marr 0)
121+ marr <- liftIO $ newSmallArray len bottomElement
122+ next marr (-1)
123+
124+ step (Tuple' marr i) x = do
125+ liftIO $ writeSmallArray marr i x
126+ next marr i
115127
116− step st@(Tuple' marr i) x
117− | i == limit = FL.Done <$> extract st
118− | otherwise = do
119− liftIO $ writeSmallArray marr i x
120− return $ FL.Partial (Tuple' marr (i + 1))
128+ extract (Tuple' marr l) = liftIO $ freezeSmallArray marr 0 l
121129
122− extract (Tuple' marr len) = liftIO $ freezeSmallArray marr 0 len
123130
124131{-# INLINE_NORMAL fromStreamDN #-}
125132fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (SmallArray a)