git.iko.soy/streamly
on master help sign in
Fix writeN dropping one element at the end of Fold
f1a8faeHarendra Kumar
src/Streamly/Internal/Data/Array.hs+20−11
compare with master
unifiedsplit
@@ -75,20 +75,29 @@ bottomElement = undefined
75-- Construction - Folds
75-- Construction - Folds
76-------------------------------------------------------------------------------
76-------------------------------------------------------------------------------
77
77
78−-- writeN n = S.evertM (fromStreamDN n)
79{-# INLINE_NORMAL writeN #-}
78{-# INLINE_NORMAL writeN #-}
80writeN :: MonadIO m => Int -> Fold m a (Array a)
79writeN :: MonadIO m => Int -> Fold m a (Array a)
81−writeN limit = Fold step initial extract
80+writeN len = Fold step initial extract
81+
82− where
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+
83 initial = do
92 initial = do
84− marr <- liftIO $ newArray limit bottomElement
93+ marr <- liftIO $ newArray len bottomElement
85− return $ FL.Partial (Tuple' marr 0)
94+ next marr (-1)
95+
86− step st@(Tuple' marr i) x
96+ step (Tuple' marr i) x = do
87− | i == limit = fmap FL.Done $ extract st
88− | otherwise = do
89− liftIO $ writeArray marr i x
97+ liftIO $ writeArray marr i x
90− return $ FL.Partial $ Tuple' marr (i + 1)
98+ next marr i
99+
91− extract (Tuple' marr len) = liftIO $ freezeArray marr 0 len
100+ extract (Tuple' marr l) = liftIO $ freezeArray marr 0 l
92
101
93{-# INLINE_NORMAL write #-}
102{-# INLINE_NORMAL write #-}
94write :: MonadIO m => Fold m a (Array a)
103write :: 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
769{-# INLINE_NORMAL writeNAllocWith #-}
769{-# INLINE_NORMAL writeNAllocWith #-}
770writeNAllocWith :: forall m a. (MonadIO m, Storable a)
770writeNAllocWith :: forall m a. (MonadIO m, Storable a)
771 => (Int -> m (Array a)) -> Int -> Fold m a (Array a)
771 => (Int -> m (Array a)) -> Int -> Fold m a (Array a)
772−writeNAllocWith alloc n = Fold step initial extract
772+writeNAllocWith alloc n = FL.take n (writeNUnsafeWith alloc n)
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
784
773
785-- | @writeN n@ folds a maximum of @n@ elements from the input stream to an
774-- | @writeN n@ folds a maximum of @n@ elements from the input stream to an
786-- 'Array'.
775-- 'Array'.
787--
776--
777+-- >>> import qualified Streamly.Data.Fold as Fold
788−-- @writeN n = Fold.take n writeNUnsafe@
778+-- >>> writeN n = Fold.take n (writeNUnsafe n)
789--
779--
790-- @since 0.7.0
780-- @since 0.7.0
791{-# INLINE_NORMAL writeN #-}
781{-# INLINE_NORMAL writeN #-}
@@ -819,22 +809,21 @@ data ArrayUnsafe a = ArrayUnsafe
819 {-# UNPACK #-} !(ForeignPtr a) -- first address
809 {-# UNPACK #-} !(ForeignPtr a) -- first address
820 {-# UNPACK #-} !(Ptr a) -- first unused address
810 {-# UNPACK #-} !(Ptr a) -- first unused address
821
811
822−-- | Like 'writeN' but does not check the array bounds when writing. The fold
812+-- XXX We can carry bound as well in the state to make sure we do not lose the
823−-- driver must not call the step function more than 'n' times otherwise it will
813+-- remaining capacity. Need to check perf impact.
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.
827--
814--
828−-- @since 0.7.0
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+--
829−{-# INLINE_NORMAL writeNUnsafe #-}
818+{-# INLINE_NORMAL writeNUnsafeWith #-}
830−writeNUnsafe :: forall m a. (MonadIO m, Storable a)
819+writeNUnsafeWith :: forall m a. (MonadIO m, Storable a)
831− => Int -> Fold m a (Array a)
820+ => (Int -> m (Array a)) -> Int -> Fold m a (Array a)
832−writeNUnsafe n = Fold step initial extract
821+writeNUnsafeWith alloc n = Fold step initial extract
833
822
834 where
823 where
835
824
836 initial = do
825 initial = do
837− (Array start end _) <- liftIO $ newArray (max n 0)
826+ (Array start end _) <- alloc (max n 0)
838 return $ FL.Partial $ ArrayUnsafe start end
827 return $ FL.Partial $ ArrayUnsafe start end
839
828
840 step (ArrayUnsafe start end) x = do
829 step (ArrayUnsafe start end) x = do
@@ -845,6 +834,18 @@ writeNUnsafe n = Fold step initial extract
845
834
846 extract (ArrayUnsafe start end) = return $ Array start end end -- liftIO . shrinkToFit
835 extract (ArrayUnsafe start end) = return $ Array start end end -- liftIO . shrinkToFit
847
836
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+
848-- XXX Buffer to a list instead?
849-- XXX Buffer to a list instead?
849--
850--
850-- | Buffer a stream into a stream of arrays.
851-- | 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
163-- /Pre-release/
163-- /Pre-release/
164{-# INLINE_NORMAL writeN #-}
164{-# INLINE_NORMAL writeN #-}
165writeN :: (MonadIO m, Prim a) => Int -> Fold m a (Array a)
165writeN :: (MonadIO m, Prim a) => Int -> Fold m a (Array a)
166−writeN limit = Fold step initial extract
166+writeN len = Fold step initial extract
167
167
168 where
168 where
169
169
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+
170 initial = do
178 initial = do
171− marr <- newArray limit
179+ marr <- liftIO $ newArray len
172− return $ FL.Partial $ Tuple' marr 0
180+ next marr (-1)
173
181
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
175
185
176− step s@(Tuple' marr i) x
186+ extract (Tuple' marr l) = shrinkArray marr l >> return marr
177− | i == limit = FL.Done <$> extract s
178− | otherwise = do
179− unsafeWriteIndex marr i x
180− return $ FL.Partial $ Tuple' marr (i + 1)
181
187
182-- Use Tuple' instead?
188-- Use Tuple' instead?
183data ArrayUnsafe a = ArrayUnsafe
189data 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
105-- of elements use an 'Array' from either "Streamly.Data.Array" or "Streamly.Data.Array.Foreign".
105-- of elements use an 'Array' from either "Streamly.Data.Array" or "Streamly.Data.Array.Foreign".
106{-# INLINE_NORMAL writeN #-}
106{-# INLINE_NORMAL writeN #-}
107writeN :: MonadIO m => Int -> Fold m a (SmallArray a)
107writeN :: MonadIO m => Int -> Fold m a (SmallArray a)
108−writeN limit = Fold step initial extract
108+writeN len = Fold step initial extract
109
109
110 where
110 where
111
111
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+
112 initial = do
120 initial = do
113− marr <- liftIO $ newSmallArray limit bottomElement
121+ marr <- liftIO $ newSmallArray len bottomElement
122+ next marr (-1)
123+
114− return $ FL.Partial (Tuple' marr 0)
124+ step (Tuple' marr i) x = do
125+ liftIO $ writeSmallArray marr i x
126+ next marr i
115
127
116− step st@(Tuple' marr i) x
128+ extract (Tuple' marr l) = liftIO $ freezeSmallArray marr 0 l
117− | i == limit = FL.Done <$> extract st
118− | otherwise = do
119− liftIO $ writeSmallArray marr i x
120− return $ FL.Partial (Tuple' marr (i + 1))
121
129
122− extract (Tuple' marr len) = liftIO $ freezeSmallArray marr 0 len
123
130
124{-# INLINE_NORMAL fromStreamDN #-}
131{-# INLINE_NORMAL fromStreamDN #-}
125fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (SmallArray a)
132fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (SmallArray a)