@@ -9,50 +9,156 @@
99
1010module Streamly.Internal.Data.Ring
1111 ( Ring(..)
12+
13+ -- * Generation
1214 , createRing
13− , unsafeInsertRing
15+ , writeLastN
16+
17+ -- * Modification
18+ , seek
19+ , unsafeInsertRingWith
20+
21+ -- * Conversion
22+ , toMutArray
23+ , toStreamWith
1424 ) where
1525
16−import Control.Monad (when)
17−import Data.IORef (newIORef, readIORef, writeIORef, IORef)
26+#include "assert.hs"
27+
28+import Control.Monad.IO.Class (liftIO, MonadIO)
29+import Streamly.Internal.Data.Stream.StreamD.Type (Stream)
30+import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))
31+import Streamly.Internal.Data.Fold.Type (Fold(..))
1832import Streamly.Internal.Data.Array.Generic.Mut.Type
1933 ( MutArray(..)
2034 , new
2135 , uninit
2236 , putIndexUnsafe
37+ , putSliceUnsafe
2338 )
39+-- import qualified Streamly.Internal.Data.Stream.StreamD.Type as Stream
40+import qualified Streamly.Internal.Data.Fold.Type as Fold
2441
42+-- XXX Use MutableArray rather than keeping a MutArray here.
2543data Ring a = Ring
2644 { 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.
31− , ringHead :: IORef Int -- current index to be over-written
32− , ringMax :: !Int -- first index beyond allocated memory
45+ -- XXX We can keep the current fill amount, Or we can keep a count of total
46+ -- elements inserted and compute ring head as well using mod on that,
47+ -- assuming it won't overflow. But mod could be expensive.
48+ , ringHead :: !Int -- current index to be over-written
49+ , ringMax :: !Int -- first index beyond allocated memory
3350 }
3451
52+-------------------------------------------------------------------------------
53+-- Generation
54+-------------------------------------------------------------------------------
55+
3556-- XXX If we align the ringMax to nearest power of two then computation of the
3657-- index to write could be cheaper.
3758{-# INLINE createRing #-}
38−createRing :: Int -> IO (Ring a)
39−createRing count = do
59+createRing :: MonadIO m => Int -> m (Ring a)
60+createRing count = liftIO $ do
4061 arr <- new count
4162 arr1 <- uninit arr count
42− ref <- newIORef 0
4363 return (Ring
4464 { ringArr = arr1
45− , ringHead = ref
65+ , ringHead = 0
4666 , ringMax = count
4767 })
4868
69+
70+{-# INLINE writeLastN #-}
71+writeLastN :: MonadIO m => Int -> Fold m a (Ring a)
72+writeLastN n = Fold step initial extract
73+
74+ where
75+
76+ initial = do
77+ if n <= 0
78+ then Fold.Done <$> createRing 0
79+ else do
80+ rb <- createRing n
81+ return $ Fold.Partial $ Tuple' rb (0 :: Int)
82+
83+ step (Tuple' rb cnt) x = do
84+ rh1 <- liftIO $ unsafeInsertRingWith rb x
85+ return $ Fold.Partial $ Tuple' (rb {ringHead = rh1}) (cnt + 1)
86+
87+ extract (Tuple' rb@Ring{..} cnt) =
88+ return $
89+ if cnt < ringMax
90+ then Ring ringArr 0 ringHead
91+ else rb
92+
93+-------------------------------------------------------------------------------
94+-- Modification
95+-------------------------------------------------------------------------------
96+
4997-- XXX This is safe
50−{-# INLINE unsafeInsertRing #-}
51−unsafeInsertRing :: Ring a -> a -> IO ()
52−unsafeInsertRing Ring{..} x = do
53− when (ringMax > 0) $ do
98+-- Take the ring head and return the new ring head.
99+{-# INLINE unsafeInsertRingWith #-}
100+unsafeInsertRingWith :: Ring a -> a -> IO Int
101+unsafeInsertRingWith Ring{..} x = do
102+ assertM(ringMax >= 1)
103+ assertM(ringHead < ringMax)
104+ putIndexUnsafe ringHead ringArr x
105+ let rh1 = ringHead + 1
106+ next = if rh1 == ringMax then 0 else rh1
107+ return next
108+
109+-- | Move the ring head clockwise (+ve adj) or counter clockwise (-ve adj) by
110+-- the given amount.
111+{-# INLINE seek #-}
112+seek :: MonadIO m => Int -> Ring a -> m (Ring a)
113+seek adj rng@Ring{..}
114+ | ringMax > 0 = liftIO $ do
115+ -- XXX try avoiding mod when in bounds
116+ let idx1 = ringHead + adj
117+ next = mod idx1 ringMax
118+ return $ Ring ringArr next ringMax
119+ | otherwise = pure rng
120+
121+-------------------------------------------------------------------------------
122+-- Conversion
123+-------------------------------------------------------------------------------
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.
128+{-# INLINE toMutArray #-}
129+toMutArray :: MonadIO m => Int -> Int -> Ring a -> m (MutArray a)
130+toMutArray adj n Ring{..} = do
131+ let len = min ringMax n
132+ let idx = mod (ringHead + adj) ringMax
133+ end = idx + len
134+ if end <= ringMax
135+ then
136+ -- putSliceUnsafe ringArr idx arr1 0 len
137+ return $ ringArr { arrStart = idx, arrLen = len }
138+ else do
139+ -- XXX Just swap the elements in the existing ring and return the
140+ -- same array without reallocation.
141+ arr <- liftIO $ new len
142+ arr1 <- uninit arr len
143+ putSliceUnsafe ringArr idx arr1 0 (ringMax - idx)
144+ putSliceUnsafe ringArr 0 arr1 (ringMax - idx) (end - ringMax)
145+ return arr1
146+
147+-- This would be theoretically slower than toMutArray because of a branch
148+-- introduced for each element in the second half of the ring.
149+
150+-- | Seek by n and then read the entire ring. Use 'take' on the stream to
151+-- restrict the reads.
152+toStreamWith :: Int -> Ring a -> Stream m a
153+toStreamWith = undefined
154+{-
155+toStreamWith n Ring{..}
156+ | ringMax > 0 = concatEffect $ liftIO $ do
54157 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
158+ let idx1 = idx + adj
159+ next = mod idx1 ringMax
160+ s1 = undefined -- stream initial slice
161+ s2 = undefined -- stream next slice
162+ return (s1 `Stream.append` s2)
163+ | otherwise = Stream.nil
164+-}