git.iko.soy/streamly
on master help sign in
Fix a bug in writeLastN, and failing assert in putSliceUnsafe
commitaad6d62c3f3895e0ea6a2b76f4975aabcc3d90a1 authorHarendra Kumar <harendra@composewell.com> dateFri 10 Mar 2023, 02:42 (4 years ago) parent8fe249d
4 files changed+65−32
compare with master
unifiedsplit
core/src/Streamly/Internal/Data/Array/Generic.hs+16−13 core/src/Streamly/Internal/Data/Array/Generic/Mut/Type.hs+20−4 core/src/Streamly/Internal/Data/Array/Mut/Type.hs+5−2 core/src/Streamly/Internal/Data/Ring.hs+24−13
core/src/Streamly/Internal/Data/Array/Generic.hs+16−13
@@ -256,20 +256,23 @@ writeLastN n
256256 rb <- liftIO $ RB.createRing n
257257 return $ FL.Partial $ Tuple' rb (0 :: Int)
258258
259− step (Tuple' rb rh) x = do
260− liftIO $ RB.unsafeInsertRing rb rh x
261− return $ FL.Partial $ Tuple' rb (rh + 1)
262−
263− done (Tuple' rb rh) = do
264− arr' <- liftIO $ MArray.new (min rh n)
265− ref <- liftIO $ readIORef $ RB.ringHead rb
266− if rh < n
267− then
268− MArray.putSliceUnsafe (RB.arr rb) 0 arr' 0 ref
259+ step (Tuple' rb cnt) x = do
260+ liftIO $ RB.unsafeInsertRing rb x
261+ return $ FL.Partial $ Tuple' rb (cnt + 1)
262+
263+ done (Tuple' rb cnt) = do
264+ let len = min cnt n
265+ arr <- liftIO $ MArray.new len
266+ arr1 <- MArray.uninit arr len
267+ -- XXX Should use an API for retrieving the ring head
268+ rhead <- liftIO $ readIORef $ RB.ringHead rb
269+ if cnt <= n
270+ then do
271+ MArray.putSliceUnsafe (RB.ringArr rb) 0 arr1 0 cnt
269272 else do
270− MArray.putSliceUnsafe (RB.arr rb) ref arr' 0 (n - ref)
271− MArray.putSliceUnsafe (RB.arr rb) 0 arr' (n - ref) ref
272− return $ unsafeFreeze arr'
273+ MArray.putSliceUnsafe (RB.ringArr rb) rhead arr1 0 (n - rhead)
274+ MArray.putSliceUnsafe (RB.ringArr rb) 0 arr1 (n - rhead) rhead
275+ return $ unsafeFreeze arr1
273276
274277sliceUnsafe :: Int -> Int -> Array a -> Array a
275278sliceUnsafe offset len (Array cont off1 _) = Array cont (off1 + offset) len
core/src/Streamly/Internal/Data/Array/Generic/Mut/Type.hs+20−4
@@ -61,6 +61,7 @@ module Streamly.Internal.Data.Array.Generic.Mut.Type
6161
6262 -- ** Reallocation
6363 , realloc
64+ , uninit
6465
6566 -- ** Appending elements
6667 , snocWith
@@ -229,7 +230,9 @@ bottomElement =
229230
230231-- XXX Would be nice if GHC can provide something like newUninitializedArray# so
231232-- that we do not have to write undefined or error in the whole array.
232−-- | @new count@ allocates an empty array that can hold 'count' items.
233+
234+-- | @new count@ allocates a zero length array that can be extended to hold
235+-- up to 'count' items without reallocating.
233236--
234237-- /Pre-release/
235238{-# INLINE new #-}
@@ -254,13 +257,14 @@ new n@(I# n#) =
254257{-# INLINE putIndexUnsafe #-}
255258putIndexUnsafe :: forall m a. MonadIO m => Int -> MutArray a -> a -> m ()
256259putIndexUnsafe i MutArray {..} x =
257− liftIO
260+ assert (i >= 0 && i < arrLen)
261+ (liftIO
258262 $ IO
259263 $ \s# ->
260264 case i + arrStart of
261265 I# n# ->
262266 let s1# = writeArray# arrContents# n# x s#
263− in (# s1#, () #)
267+ in (# s1#, () #))
264268
265269invalidIndex :: String -> Int -> a
266270invalidIndex label i =
@@ -410,6 +414,17 @@ snocWith sizer arr@MutArray {..} x = do
410414snoc :: MonadIO m => MutArray a -> a -> m (MutArray a)
411415snoc = snocWith (* 2)
412416
417+-- | Make the uninitialized memory in the array available for use extending it
418+-- by the supplied length beyond the current length of the array. The array may
419+-- be reallocated.
420+--
421+{-# INLINE uninit #-}
422+uninit :: MonadIO m => MutArray a -> Int -> m (MutArray a)
423+uninit arr@MutArray{..} len =
424+ if arrStart + arrLen + len <= arrTrueLen
425+ then return $ arr {arrLen = arrLen + len}
426+ else realloc (arrLen + len) arr
427+
413428-------------------------------------------------------------------------------
414429-- Random reads
415430-------------------------------------------------------------------------------
@@ -586,7 +601,8 @@ reader = Producer.simplify producer
586601-- | Put a sub range of a source array into a subrange of a destination array.
587602-- This is not safe as it does not check the bounds.
588603{-# INLINE putSliceUnsafe #-}
589−putSliceUnsafe :: MonadIO m => MutArray a -> Int -> MutArray a -> Int -> Int -> m ()
604+putSliceUnsafe :: MonadIO m =>
605+ MutArray a -> Int -> MutArray a -> Int -> Int -> m ()
590606putSliceUnsafe src srcStart dst dstStart len = liftIO $ do
591607 assertM(len <= arrLen dst)
592608 assertM(len <= arrLen src)
core/src/Streamly/Internal/Data/Array/Mut/Type.hs+5−2
@@ -459,8 +459,11 @@ new =
459459 (\s _ -> liftIO $ Unboxed.newUnpinnedBytes s)
460460 (error "new: alignment is not used in unpinned arrays.")
461461
462−-- | Allocate an MutArray of the given size and run an IO action passing the array
463−-- start pointer.
462+-- XXX This should create a full length uninitialzed array so that the pointer
463+-- can be used.
464+
465+-- | Allocate a pinned MutArray of the given size and run an IO action passing
466+-- the array start pointer.
464467--
465468-- /Internal/
466469{-# INLINE withNewArrayUnsafe #-}
core/src/Streamly/Internal/Data/Ring.hs+24−13
@@ -13,35 +13,46 @@ module Streamly.Internal.Data.Ring
1313 , unsafeInsertRing
1414 ) where
1515
16−import Data.IORef (modifyIORef', newIORef, readIORef, writeIORef, IORef)
16+import Control.Monad (when)
17+import Data.IORef (newIORef, readIORef, writeIORef, IORef)
1718import Streamly.Internal.Data.Array.Generic.Mut.Type
1819 ( MutArray(..)
1920 , new
21+ , uninit
2022 , putIndexUnsafe
2123 )
2224
2325data Ring a = Ring
24− { arr :: MutArray a
26+ { ringArr :: MutArray a
27+ -- XXX This could be an unboxed IO ref
28+ -- If we want to rely on an external Int count using mod ringMax as index
29+ -- into the ring then we do not need to maintain this, assuming the counter
30+ -- would not overflow.
2531 , ringHead :: IORef Int -- current index to be over-written
2632 , ringMax :: !Int -- first index beyond allocated memory
2733 }
2834
35+-- XXX If we align the ringMax to nearest power of two then computation of the
36+-- index to write could be cheaper.
2937{-# INLINE createRing #-}
3038createRing :: Int -> IO (Ring a)
3139createRing count = do
32− arr' <- new count
33− head' <- newIORef 0
40+ arr <- new count
41+ arr1 <- uninit arr count
42+ ref <- newIORef 0
3443 return (Ring
35− { arr = arr'
36− , ringHead = head'
44+ { ringArr = arr1
45+ , ringHead = ref
3746 , ringMax = count
3847 })
3948
49+-- XXX This is safe
4050{-# INLINE unsafeInsertRing #-}
41−unsafeInsertRing :: Ring a -> Int -> a -> IO ()
42−unsafeInsertRing Ring{..} idx x = do
43− putIndexUnsafe (mod idx ringMax) arr x
44− ref <- readIORef ringHead
45− if (ref+1) < ringMax
46− then modifyIORef' ringHead ( + 1)
47− else writeIORef ringHead 0
51+unsafeInsertRing :: Ring a -> a -> IO ()
52+unsafeInsertRing Ring{..} x = do
53+ when (ringMax > 0) $ do
54+ idx <- readIORef ringHead
55+ putIndexUnsafe idx ringArr x
56+ let idx1 = idx + 1
57+ next = if idx1 == ringMax then 0 else idx1
58+ writeIORef ringHead next