aad6d62 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Fix a bug in writeLastN, and failing assert in putSliceUnsafe

aad6d62Harendra Kumar4 years ago

core/src/Streamly/Internal/Data/Array/Generic.hs+16−13

compare with master

unifiedsplit

@@ -256,20 +256,23 @@ writeLastN n

256 rb <- liftIO $ RB.createRing n

257 return $ FL.Partial $ Tuple' rb (0 :: Int)

258

259− step (Tuple' rb rh) x = do

259+ step (Tuple' rb cnt) x = do

260− liftIO $ RB.unsafeInsertRing rb rh x

260+ liftIO $ RB.unsafeInsertRing rb x

261− return $ FL.Partial $ Tuple' rb (rh + 1)

261+ return $ FL.Partial $ Tuple' rb (cnt + 1)

262−

262+

263− done (Tuple' rb rh) = do

263+ done (Tuple' rb cnt) = do

264+ let len = min cnt n

264− arr' <- liftIO $ MArray.new (min rh n)

265+ arr <- liftIO $ MArray.new len

266+ arr1 <- MArray.uninit arr len

267+ -- XXX Should use an API for retrieving the ring head

265− ref <- liftIO $ readIORef $ RB.ringHead rb

268+ rhead <- liftIO $ readIORef $ RB.ringHead rb

266− if rh < n

269+ if cnt <= n

267− then

270+ then do

268− MArray.putSliceUnsafe (RB.arr rb) 0 arr' 0 ref

271+ MArray.putSliceUnsafe (RB.ringArr rb) 0 arr1 0 cnt

269 else do

272 else do

270− MArray.putSliceUnsafe (RB.arr rb) ref arr' 0 (n - ref)

273+ MArray.putSliceUnsafe (RB.ringArr rb) rhead arr1 0 (n - rhead)

271− MArray.putSliceUnsafe (RB.arr rb) 0 arr' (n - ref) ref

274+ MArray.putSliceUnsafe (RB.ringArr rb) 0 arr1 (n - rhead) rhead

272− return $ unsafeFreeze arr'

275+ return $ unsafeFreeze arr1

273

276

274sliceUnsafe :: Int -> Int -> Array a -> Array a

277sliceUnsafe :: Int -> Int -> Array a -> Array a

275sliceUnsafe offset len (Array cont off1 _) = Array cont (off1 + offset) len

278sliceUnsafe 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

62 -- ** Reallocation

63 , realloc

64+ , uninit

65 -- ** Appending elements

66 -- ** Appending elements

66 , snocWith

67 , snocWith

@@ -229,7 +230,9 @@ bottomElement =

229

230

230-- XXX Would be nice if GHC can provide something like newUninitializedArray# so

231-- XXX Would be nice if GHC can provide something like newUninitializedArray# so

231-- that we do not have to write undefined or error in the whole array.

232-- that we do not have to write undefined or error in the whole array.

233+

232−-- | @new count@ allocates an empty array that can hold 'count' items.

234+-- | @new count@ allocates a zero length array that can be extended to hold

235+-- up to 'count' items without reallocating.

233--

236--

234-- /Pre-release/

237-- /Pre-release/

235{-# INLINE new #-}

238{-# INLINE new #-}

@@ -254,13 +257,14 @@ new n@(I# n#) =

254{-# INLINE putIndexUnsafe #-}

257{-# INLINE putIndexUnsafe #-}

255putIndexUnsafe :: forall m a. MonadIO m => Int -> MutArray a -> a -> m ()

258putIndexUnsafe :: forall m a. MonadIO m => Int -> MutArray a -> a -> m ()

256putIndexUnsafe i MutArray {..} x =

259putIndexUnsafe i MutArray {..} x =

260+ assert (i >= 0 && i < arrLen)

257− liftIO

261+ (liftIO

258 $ IO

262 $ IO

259 $ \s# ->

263 $ \s# ->

260 case i + arrStart of

264 case i + arrStart of

261 I# n# ->

265 I# n# ->

262 let s1# = writeArray# arrContents# n# x s#

266 let s1# = writeArray# arrContents# n# x s#

263− in (# s1#, () #)

267+ in (# s1#, () #))

264

268

265invalidIndex :: String -> Int -> a

269invalidIndex :: String -> Int -> a

266invalidIndex label i =

270invalidIndex label i =

@@ -410,6 +414,17 @@ snocWith sizer arr@MutArray {..} x = do

410snoc :: MonadIO m => MutArray a -> a -> m (MutArray a)

414snoc :: MonadIO m => MutArray a -> a -> m (MutArray a)

411snoc = snocWith (* 2)

415snoc = snocWith (* 2)

412

416

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+

413-------------------------------------------------------------------------------

428-------------------------------------------------------------------------------

414-- Random reads

429-- Random reads

415-------------------------------------------------------------------------------

430-------------------------------------------------------------------------------

@@ -586,7 +601,8 @@ reader = Producer.simplify producer

586-- | Put a sub range of a source array into a subrange of a destination array.

601-- | Put a sub range of a source array into a subrange of a destination array.

587-- This is not safe as it does not check the bounds.

602-- This is not safe as it does not check the bounds.

588{-# INLINE putSliceUnsafe #-}

603{-# INLINE putSliceUnsafe #-}

604+putSliceUnsafe :: MonadIO m =>

589−putSliceUnsafe :: MonadIO m => MutArray a -> Int -> MutArray a -> Int -> Int -> m ()

605+ MutArray a -> Int -> MutArray a -> Int -> Int -> m ()

590putSliceUnsafe src srcStart dst dstStart len = liftIO $ do

606putSliceUnsafe src srcStart dst dstStart len = liftIO $ do

591 assertM(len <= arrLen dst)

607 assertM(len <= arrLen dst)

592 assertM(len <= arrLen src)

608 assertM(len <= arrLen src)

core/src/Streamly/Internal/Data/Array/Mut/Type.hs+5−2

@@ -459,8 +459,11 @@ new =

459 (\s _ -> liftIO $ Unboxed.newUnpinnedBytes s)

460 (error "new: alignment is not used in unpinned arrays.")

461

462+-- XXX This should create a full length uninitialzed array so that the pointer

463+-- can be used.

464+

462−-- | Allocate an MutArray of the given size and run an IO action passing the array

465+-- | Allocate a pinned MutArray of the given size and run an IO action passing

463−-- start pointer.

466+-- the array start pointer.

464--

467--

465-- /Internal/

468-- /Internal/

466{-# INLINE withNewArrayUnsafe #-}

469{-# INLINE withNewArrayUnsafe #-}

core/src/Streamly/Internal/Data/Ring.hs+24−13

@@ -13,35 +13,46 @@ module Streamly.Internal.Data.Ring

13 , unsafeInsertRing

14 ) where

16+import Control.Monad (when)

16−import Data.IORef (modifyIORef', newIORef, readIORef, writeIORef, IORef)

17+import Data.IORef (newIORef, readIORef, writeIORef, IORef)

17import Streamly.Internal.Data.Array.Generic.Mut.Type

18import Streamly.Internal.Data.Array.Generic.Mut.Type

18 ( MutArray(..)

19 ( MutArray(..)

19 , new

20 , new

21+ , uninit

20 , putIndexUnsafe

22 , putIndexUnsafe

21 )

23 )

23data Ring a = Ring

25data 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.

25 , ringHead :: IORef Int -- current index to be over-written

31 , ringHead :: IORef Int -- current index to be over-written

26 , ringMax :: !Int -- first index beyond allocated memory

32 , ringMax :: !Int -- first index beyond allocated memory

27 }

33 }

35+-- XXX If we align the ringMax to nearest power of two then computation of the

36+-- index to write could be cheaper.

29{-# INLINE createRing #-}

37{-# INLINE createRing #-}

30createRing :: Int -> IO (Ring a)

38createRing :: Int -> IO (Ring a)

31createRing count = do

39createRing count = do

32− arr' <- new count

40+ arr <- new count

41+ arr1 <- uninit arr count

33− head' <- newIORef 0

42+ ref <- newIORef 0

34 return (Ring

43 return (Ring

35− { arr = arr'

44+ { ringArr = arr1

36− , ringHead = head'

45+ , ringHead = ref

37 , ringMax = count

46 , ringMax = count

38 })

47 })

49+-- XXX This is safe

40{-# INLINE unsafeInsertRing #-}

50{-# INLINE unsafeInsertRing #-}

41−unsafeInsertRing :: Ring a -> Int -> a -> IO ()

51+unsafeInsertRing :: Ring a -> a -> IO ()

42−unsafeInsertRing Ring{..} idx x = do

52+unsafeInsertRing Ring{..} x = do

43− putIndexUnsafe (mod idx ringMax) arr x

53+ when (ringMax > 0) $ do

44− ref <- readIORef ringHead

54+ idx <- readIORef ringHead

55+ putIndexUnsafe idx ringArr x

56+ let idx1 = idx + 1

45− if (ref+1) < ringMax

57+ next = if idx1 == ringMax then 0 else idx1

46− then modifyIORef' ringHead ( + 1)

47− else writeIORef ringHead 0

58+ writeIORef ringHead next