git.iko.soy/streamly
on cedc81c help sign in
streamly/core/src/Streamly/Internal/Data/Ring/Generic.hs
185 lines 6.3 kB history raw
1-- |
2-- Module : Streamly.Internal.Data.Ring.Generic
3-- Copyright : (c) 2021 Composewell Technologies
4-- License : BSD-3-Clause
5-- Maintainer : streamly@composewell.com
6-- Stability : experimental
7-- Portability : GHC
8--
10module Streamly.Internal.Data.Ring.Generic
11 ( Ring(..)
13 -- * Generation
14 , createRing
15 , writeLastN
17 -- * Modification
18 , seek
19 , unsafeInsertRingWith
21 -- * Conversion
22 , toMutArray
23 , copyToMutArray
24 , toStreamWith
25 ) where
27#include "assert.hs"
29import Control.Monad.IO.Class (liftIO, MonadIO)
30import Streamly.Internal.Data.Stream.Type (Stream)
31import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))
32import Streamly.Internal.Data.Fold.Type (Fold(..))
33import Streamly.Internal.Data.MutArray.Generic (MutArray(..))
35-- import qualified Streamly.Internal.Data.Stream.Type as Stream
36import qualified Streamly.Internal.Data.Fold.Type as Fold
37import qualified Streamly.Internal.Data.MutArray.Generic as MutArray
39-- XXX Use MutableArray rather than keeping a MutArray here.
40data Ring a = Ring
41 { ringArr :: MutArray a
42 -- XXX We can keep the current fill amount, Or we can keep a count of total
43 -- elements inserted and compute ring head as well using mod on that,
44 -- assuming it won't overflow. But mod could be expensive.
45 , ringHead :: !Int -- current index to be over-written
46 , ringMax :: !Int -- first index beyond allocated memory
47 }
49-------------------------------------------------------------------------------
50-- Generation
51-------------------------------------------------------------------------------
53-- XXX If we align the ringMax to nearest power of two then computation of the
54-- index to write could be cheaper.
55{-# INLINE createRing #-}
56createRing :: MonadIO m => Int -> m (Ring a)
57createRing count = liftIO $ do
58 arr <- MutArray.new count
59 arr1 <- MutArray.uninit arr count
60 return (Ring
61 { ringArr = arr1
62 , ringHead = 0
63 , ringMax = count
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
74 where
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)
83 step (Tuple' rb cnt) x = do
84 rh1 <- liftIO $ unsafeInsertRingWith rb x
85 return $ Fold.Partial $ Tuple' (rb {ringHead = rh1}) (cnt + 1)
87 extract (Tuple' rb@Ring{..} cnt) =
88 return $
89 if cnt < ringMax
90 then Ring ringArr 0 ringHead
91 else rb
93-------------------------------------------------------------------------------
94-- Modification
95-------------------------------------------------------------------------------
97-- XXX This is safe
98-- Take the ring head and return the new ring head.
99{-# INLINE unsafeInsertRingWith #-}
100unsafeInsertRingWith :: Ring a -> a -> IO Int
101unsafeInsertRingWith Ring{..} x = do
102 assertM(ringMax >= 1)
103 assertM(ringHead < ringMax)
104 MutArray.putIndexUnsafe ringHead ringArr x
105 let rh1 = ringHead + 1
106 next = if rh1 == ringMax then 0 else rh1
107 return next
109-- | Move the ring head clockwise (+ve adj) or counter clockwise (-ve adj) by
110-- the given amount.
111{-# INLINE seek #-}
112seek :: MonadIO m => Int -> Ring a -> m (Ring a)
113seek 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
121-------------------------------------------------------------------------------
122-- Conversion
123-------------------------------------------------------------------------------
125-- | @toMutArray rignHeadAdjustment lengthToRead ring@.
126-- Convert the ring into a boxed mutable array. Note that the returned MutArray
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--
130{-# INLINE toMutArray #-}
131toMutArray :: MonadIO m => Int -> Int -> Ring a -> m (MutArray a)
132toMutArray 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
137 let len = min ringMax n
138 let idx = mod (ringHead + adj) ringMax
139 end = idx + len
140 if end <= ringMax
141 then
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
152-- | Copy out the mutable ring to a mutable Array.
153{-# INLINE copyToMutArray #-}
154copyToMutArray :: MonadIO m => Int -> Int -> Ring a -> m (MutArray a)
155copyToMutArray adj n Ring{..} = do
156 if ringMax <= 0
157 then MutArray.nil
158 else do
159 let len = min ringMax n
160 let idx = mod (ringHead + adj) ringMax
161 end = idx + len
162 arr <- MutArray.new len
163 arr1 <- MutArray.uninit arr len
164 MutArray.putSliceUnsafe ringArr idx arr1 0 (ringMax - idx)
165 MutArray.putSliceUnsafe ringArr 0 arr1 (ringMax - idx) (end - ringMax)
166 return arr1
168-- This would be theoretically slower than toMutArray because of a branch
169-- introduced for each element in the second half of the ring.
171-- | Seek by n and then read the entire ring. Use 'take' on the stream to
172-- restrict the reads.
173toStreamWith :: Int -> Ring a -> Stream m a
174toStreamWith = undefined
175{-
176toStreamWith n Ring{..}
177 | ringMax > 0 = concatEffect $ liftIO $ do
178 idx <- readIORef ringHead
179 let idx1 = idx + adj
180 next = mod idx1 ringMax
181 s1 = undefined -- stream initial slice
182 s2 = undefined -- stream next slice
183 return (s1 `Stream.append` s2)
184 | otherwise = Stream.nil
185-}