git.iko.soy/streamly
on efe0be3 help sign in
streamly/core/src/Streamly/Internal/Data/Ring.hs
164 lines 5.3 kB history raw
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--
10module Streamly.Internal.Data.Ring
11 ( Ring(..)
13 -- * Generation
14 , createRing
15 , writeLastN
17 -- * Modification
18 , seek
19 , unsafeInsertRingWith
21 -- * Conversion
22 , toMutArray
23 , toStreamWith
24 ) where
26#include "assert.hs"
28import Control.Monad.IO.Class (liftIO, MonadIO)
29import Streamly.Internal.Data.Stream.StreamD.Type (Stream)
30import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))
31import Streamly.Internal.Data.Fold.Type (Fold(..))
32import Streamly.Internal.Data.Array.Generic.Mut.Type
33 ( MutArray(..)
34 , new
35 , uninit
36 , putIndexUnsafe
37 , putSliceUnsafe
38 )
39-- import qualified Streamly.Internal.Data.Stream.StreamD.Type as Stream
40import qualified Streamly.Internal.Data.Fold.Type as Fold
42-- XXX Use MutableArray rather than keeping a MutArray here.
43data Ring a = Ring
44 { ringArr :: MutArray a
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
50 }
52-------------------------------------------------------------------------------
53-- Generation
54-------------------------------------------------------------------------------
56-- XXX If we align the ringMax to nearest power of two then computation of the
57-- index to write could be cheaper.
58{-# INLINE createRing #-}
59createRing :: MonadIO m => Int -> m (Ring a)
60createRing count = liftIO $ do
61 arr <- new count
62 arr1 <- uninit arr count
63 return (Ring
64 { ringArr = arr1
65 , ringHead = 0
66 , ringMax = count
67 })
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 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-- may share the same underlying memory as the Ring.
128{-# INLINE toMutArray #-}
129toMutArray :: MonadIO m => Int -> Int -> Ring a -> m (MutArray a)
130toMutArray 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
147-- This would be theoretically slower than toMutArray because of a branch
148-- introduced for each element in the second half of the ring.
150-- | Seek by n and then read the entire ring. Use 'take' on the stream to
151-- restrict the reads.
152toStreamWith :: Int -> Ring a -> Stream m a
153toStreamWith = undefined
154{-
155toStreamWith n Ring{..}
156 | ringMax > 0 = concatEffect $ liftIO $ do
157 idx <- readIORef ringHead
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-}