git.iko.soy/streamly
on master help sign in
Implement a polymorphic ring buffer
commit95b8d904430e8ac8b03734b4f22717e907bcfab9 authorRanjeet Kumar Ranjan <ranjeetr@composewell.com> dateWed 22 Sep 2021, 21:00 (5 years ago) parent57ab52c
4 files changed+84−2
compare with master
unifiedsplit
benchmark/Streamly/Benchmark/Data/Array.hs+9−2 src/Streamly/Internal/Data/Array.hs+30 src/Streamly/Internal/Data/Ring.hs+44 streamly.cabal+1
benchmark/Streamly/Benchmark/Data/Array.hs+9−2
@@ -13,7 +13,12 @@ import System.Random (randomRIO)
1313
1414import qualified Streamly.Benchmark.Data.ArrayOps as Ops
1515import qualified Streamly.Data.Array.Foreign as A
16+
17+#ifdef DATA_ARRAY
18+import qualified Streamly.Internal.Data.Array as IA
19+#else
1620import qualified Streamly.Internal.Data.Array.Foreign as IA
21+#endif
1722import qualified Streamly.Prelude as S
1823
1924import Gauge
@@ -117,8 +122,10 @@ o_1_space_elimination value =
117122 , benchIOSink value "foldl'" Ops.pureFoldl'
118123 , benchIOSink value "read" Ops.unfoldReadDrain
119124 , benchIOSink value "toStreamRev" Ops.toStreamRevDrain
120− , benchFold "writeLastN.1" (S.fold (IA.writeLastN 1)) (P.sourceUnfoldrM value)
121− , benchFold "writeLastN.10" (S.fold (IA.writeLastN 10)) (P.sourceUnfoldrM value)
125+ , benchFold "writeLastN.1"
126+ (S.fold (IA.writeLastN 1)) (P.sourceUnfoldrM value)
127+ , benchFold "writeLastN.10"
128+ (S.fold (IA.writeLastN 10)) (P.sourceUnfoldrM value)
122129#if !defined(DATA_ARRAY_PRIM) && !defined(DATA_ARRAY_PRIM_PINNED)
123130#ifdef DEVBUILD
124131 , benchPureSink value "foldable/foldl'" Ops.foldableFoldl'
src/Streamly/Internal/Data/Array.hs+30
@@ -18,6 +18,7 @@ module Streamly.Internal.Data.Array
1818 -- * Construction
1919 , writeN
2020 , write
21+ , writeLastN
2122
2223 , fromStreamDN
2324 , fromStreamD
@@ -52,6 +53,7 @@ import Control.DeepSeq (NFData(..))
5253import Control.Monad (when)
5354import Control.Monad.IO.Class (liftIO, MonadIO)
5455import Data.Functor.Identity (runIdentity)
56+import Data.IORef
5557import GHC.Base (Int(..))
5658import GHC.IO (unsafePerformIO)
5759
@@ -62,6 +64,7 @@ import Streamly.Internal.Data.Unfold.Type (Unfold(..))
6264
6365import qualified GHC.Exts as Exts
6466import qualified Streamly.Internal.Data.Fold.Type as FL
67+import qualified Streamly.Internal.Data.Ring as RB
6568import qualified Streamly.Internal.Data.Stream.StreamD as D
6669
6770import Data.Primitive.Array hiding (fromList, fromListN)
@@ -248,3 +251,30 @@ streamFold f arr = f (toStream arr)
248251{-# INLINE getIndexUnsafe #-}
249252getIndexUnsafe :: Array a -> Int -> a
250253getIndexUnsafe = indexArray
254+
255+{-# INLINE writeLastN #-}
256+writeLastN :: MonadIO m => Int -> Fold m a (Array a)
257+writeLastN n
258+ | n <= 0 = fmap (const mempty) FL.drain
259+ | otherwise = Fold step initial done
260+
261+ where
262+
263+ initial = do
264+ rb <- liftIO $ RB.createRing n
265+ return $ FL.Partial $ Tuple' rb (0 :: Int)
266+
267+ step (Tuple' rb rh) x = do
268+ liftIO $ RB.unsafeInsertRing rb rh x
269+ return $ FL.Partial $ Tuple' rb (rh + 1)
270+
271+ done (Tuple' rb rh) = do
272+ arr' <- liftIO $ newArray (min rh n) (undefined :: a)
273+ ref <- liftIO $ readIORef $ RB.ringHead rb
274+ if rh < n
275+ then
276+ liftIO $ copyMutableArray arr' 0 (RB.arr rb) 0 ref
277+ else do
278+ liftIO $ copyMutableArray arr' 0 (RB.arr rb) ref (n - ref)
279+ liftIO $ copyMutableArray arr' (n - ref) (RB.arr rb) 0 ref
280+ liftIO $ unsafeFreezeArray arr'
src/Streamly/Internal/Data/Ring.hsnew file+44
@@ -0,0 +1,44 @@
1+-- |
2+-- Module : Streamly.Internal.Data.Ring
3+-- Copyright : (c) 2021 Composewell Technologies
4+-- License : BSD-3-Clause
5+-- Maintainer : streamly@composewell.com
6+-- Stability : experimental
7+-- Portability : GHC
8+--
9+
10+module Streamly.Internal.Data.Ring
11+ ( Ring(..)
12+ , createRing
13+ , unsafeInsertRing
14+ ) where
15+
16+import Control.Monad.Primitive (PrimMonad(PrimState))
17+import Data.IORef (modifyIORef', newIORef, readIORef, writeIORef, IORef)
18+import Data.Primitive.Array (newArray, writeArray, MutableArray)
19+
20+data Ring a = Ring
21+ { arr :: MutableArray (PrimState IO) a
22+ , ringHead :: IORef Int -- current index to be over-written
23+ , ringMax :: !Int -- first index beyond allocated memory
24+ }
25+
26+{-# INLINE createRing #-}
27+createRing :: Int -> IO (Ring a)
28+createRing count = do
29+ arr' <- newArray count (undefined :: a)
30+ head' <- newIORef 0
31+ return (Ring
32+ { arr = arr'
33+ , ringHead = head'
34+ , ringMax = count
35+ })
36+
37+{-# INLINE unsafeInsertRing #-}
38+unsafeInsertRing :: Ring a -> Int -> a -> IO ()
39+unsafeInsertRing Ring{..} idx x = do
40+ writeArray arr (mod idx ringMax) x
41+ ref <- readIORef ringHead
42+ if (ref+1) < ringMax
43+ then modifyIORef' ringHead ( + 1)
44+ else writeIORef ringHead 0
streamly.cabal+1
@@ -477,6 +477,7 @@ library
477477
478478 -- Memory storage
479479 , Streamly.Internal.Ring.Foreign
480+ , Streamly.Internal.Data.Ring
480481
481482 -- IOVec (depends on arrays/streams)
482483 , Streamly.Internal.System.IOVec