313d178 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Fix Data.Array.Generic.writeLastN bug (#2567)

313d178Ranjeet Ranjan3 years ago

core/src/Streamly/Internal/Data/Array/Generic.hs+12−1

compare with master

unifiedsplit

@@ -242,6 +242,17 @@ getIndex i arr@Array {..} =

242 then Just $ getIndexUnsafe i arr

243 else Nothing

244

245+-- >>> import qualified Streamly.Data.Stream as Stream

246+-- >>> import qualified Streamly.Data.Fold as Fold

247+-- >>> import qualified Streamly.Internal.Data.Array.Generic as Array

248+-- >>> import Data.Function ((&))

249+-- >>> :{

250+-- Stream.fromList [1,2,3,4,5::Int]

251+-- & Stream.scan (Array.writeLastN 2)

252+-- & Stream.fold Fold.toList

253+-- :}

254+-- [fromList [],fromList [1],fromList [1,2],fromList [2,3],fromList [3,4],fromList [4,5]]

255+--

245{-# INLINE writeLastN #-}

256{-# INLINE writeLastN #-}

246writeLastN :: MonadIO m => Int -> Fold m a (Array a)

257writeLastN :: MonadIO m => Int -> Fold m a (Array a)

247writeLastN n = FL.rmapM f (RB.writeLastN n)

258writeLastN n = FL.rmapM f (RB.writeLastN n)

@@ -249,7 +260,7 @@ writeLastN n = FL.rmapM f (RB.writeLastN n)

249 where

260 where

250

261

251 f rb = do

262 f rb = do

252− arr <- RB.toMutArray 0 n rb

263+ arr <- RB.copyToMutArray 0 n rb

253 return $ unsafeFreeze arr

264 return $ unsafeFreeze arr

254

265

255{-# INLINE getSliceUnsafe #-}

266{-# INLINE getSliceUnsafe #-}

core/src/Streamly/Internal/Data/Ring/Generic.hs+46−25

@@ -20,6 +20,7 @@ module Streamly.Internal.Data.Ring.Generic

21 -- * Conversion

22 , toMutArray

23+ , copyToMutArray

23 , toStreamWith

24 , toStreamWith

24 ) where

25 ) where

@@ -29,15 +30,11 @@ import Control.Monad.IO.Class (liftIO, MonadIO)

29import Streamly.Internal.Data.Stream.Type (Stream)

30import Streamly.Internal.Data.Stream.Type (Stream)

30import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))

31import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))

31import Streamly.Internal.Data.Fold.Type (Fold(..))

32import Streamly.Internal.Data.Fold.Type (Fold(..))

32−import Streamly.Internal.Data.MutArray.Generic

33+import Streamly.Internal.Data.MutArray.Generic (MutArray(..))

33− ( MutArray(..)

34+

34− , new

35− , uninit

36− , putIndexUnsafe

37− , putSliceUnsafe

38− )

39-- import qualified Streamly.Internal.Data.Stream.Type as Stream

35-- import qualified Streamly.Internal.Data.Stream.Type as Stream

40import qualified Streamly.Internal.Data.Fold.Type as Fold

36import qualified Streamly.Internal.Data.Fold.Type as Fold

37+import qualified Streamly.Internal.Data.MutArray.Generic as MutArray

42-- XXX Use MutableArray rather than keeping a MutArray here.

39-- XXX Use MutableArray rather than keeping a MutArray here.

43data Ring a = Ring

40data Ring a = Ring

@@ -58,8 +55,8 @@ data Ring a = Ring

58{-# INLINE createRing #-}

55{-# INLINE createRing #-}

59createRing :: MonadIO m => Int -> m (Ring a)

56createRing :: MonadIO m => Int -> m (Ring a)

60createRing count = liftIO $ do

57createRing count = liftIO $ do

61− arr <- new count

58+ arr <- MutArray.new count

62− arr1 <- uninit arr count

59+ arr1 <- MutArray.uninit arr count

63 return (Ring

60 return (Ring

64 { ringArr = arr1

61 { ringArr = arr1

65 , ringHead = 0

62 , ringHead = 0

@@ -67,6 +64,9 @@ createRing count = liftIO $ do

67 })

64 })

67+-- | Note that it is not safe to return a reference to the mutable Ring using a

68+-- scan as the Ring is continuously getting mutated. You could however copy out

69+-- the Ring.

70{-# INLINE writeLastN #-}

71writeLastN :: MonadIO m => Int -> Fold m a (Ring a)

72writeLastN n = Fold step initial extract

@@ -101,7 +101,7 @@ unsafeInsertRingWith :: Ring a -> a -> IO Int

101unsafeInsertRingWith Ring{..} x = do

102 assertM(ringMax >= 1)

103 assertM(ringHead < ringMax)

104− putIndexUnsafe ringHead ringArr x

104+ MutArray.putIndexUnsafe ringHead ringArr x

105 let rh1 = ringHead + 1

106 next = if rh1 == ringMax then 0 else rh1

107 return next

@@ -124,24 +124,45 @@ seek adj rng@Ring{..}

124

125-- | @toMutArray rignHeadAdjustment lengthToRead ring@.

126-- Convert the ring into a boxed mutable array. Note that the returned MutArray

127−-- may share the same underlying memory as the Ring.

127+-- shares the same underlying memory as the Ring, the user of this API needs to

128+-- ensure that the ring is not mutated during and after the conversion.

129+--

128{-# INLINE toMutArray #-}

130{-# INLINE toMutArray #-}

129toMutArray :: MonadIO m => Int -> Int -> Ring a -> m (MutArray a)

131toMutArray :: MonadIO m => Int -> Int -> Ring a -> m (MutArray a)

130−toMutArray adj n Ring{..} = do

132+toMutArray adj n Ring{..} =

133+ -- XXX for empty Ring it will raise an Exception: divide by zero

134+ if ringMax <= 0

135+ then MutArray.nil

136+ else do

131− let len = min ringMax n

137+ let len = min ringMax n

132− let idx = mod (ringHead + adj) ringMax

138+ let idx = mod (ringHead + adj) ringMax

133− end = idx + len

139+ end = idx + len

134− if end <= ringMax

140+ if end <= ringMax

135− then

141+ then

136− -- putSliceUnsafe ringArr idx arr1 0 len

137− return $ ringArr { arrStart = idx, arrLen = len }

142+ return $ ringArr { arrStart = idx, arrLen = len }

143+ else do

144+ -- XXX Just swap the elements in the existing ring and return the

145+ -- same array without reallocation.

146+ arr <- liftIO $ MutArray.new len

147+ arr1 <- MutArray.uninit arr len

148+ MutArray.putSliceUnsafe ringArr idx arr1 0 (ringMax - idx)

149+ MutArray.putSliceUnsafe ringArr 0 arr1 (ringMax - idx) (end - ringMax)

150+ return arr1

151+

152+-- | Copy out the mutable ring to a mutable Array.

153+{-# INLINE copyToMutArray #-}

154+copyToMutArray :: MonadIO m => Int -> Int -> Ring a -> m (MutArray a)

155+copyToMutArray adj n Ring{..} = do

156+ if ringMax <= 0

157+ then MutArray.nil

138 else do

158 else do

139− -- XXX Just swap the elements in the existing ring and return the

159+ let len = min ringMax n

140− -- same array without reallocation.

160+ let idx = mod (ringHead + adj) ringMax

161+ end = idx + len

141− arr <- liftIO $ new len

162+ arr <- MutArray.new len

142− arr1 <- uninit arr len

163+ arr1 <- MutArray.uninit arr len

143− putSliceUnsafe ringArr idx arr1 0 (ringMax - idx)

164+ MutArray.putSliceUnsafe ringArr idx arr1 0 (ringMax - idx)

144− putSliceUnsafe ringArr 0 arr1 (ringMax - idx) (end - ringMax)

165+ MutArray.putSliceUnsafe ringArr 0 arr1 (ringMax - idx) (end - ringMax)

145 return arr1

166 return arr1

146

167

147-- This would be theoretically slower than toMutArray because of a branch

168-- This would be theoretically slower than toMutArray because of a branch