git.iko.soy/streamly
on 3e53339 help sign in
streamly/src/Streamly/Memory/Ring.hs
250 lines 9.3 kB history raw
1-- |
2-- Module : Streamly.Memory.Ring
3-- Copyright : (c) 2019 Composewell Technologies
4-- License : BSD3
5-- Maintainer : streamly@composewell.com
6-- Stability : experimental
7-- Portability : GHC
8--
10module Streamly.Memory.Ring
11 ( Ring(..)
13 -- * Construction
14 , new
15 , advance
16 , moveBy
17 , startOf
19 -- * Modification
20 , unsafeInsert
22 -- * Folds
23 , unsafeFoldRing
24 , unsafeFoldRingM
25 , unsafeFoldRingFullM
26 , unsafeFoldRingNM
28 -- * Fast Byte Comparisons
29 , unsafeEqArray
30 , unsafeEqArrayN
31 ) where
33import Control.Exception (assert)
34import Foreign.ForeignPtr (ForeignPtr, withForeignPtr, touchForeignPtr)
35import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
36import Foreign.Ptr (plusPtr, minusPtr, castPtr)
37import Foreign.Storable (Storable(..))
38import GHC.ForeignPtr (mallocPlainForeignPtrAlignedBytes)
39import GHC.Ptr (Ptr(..))
40import Prelude hiding (length, concat)
42import Control.Monad.IO.Class (MonadIO(..))
44import qualified Streamly.Internal.Data.Array.Storable.Foreign.Types as A
46-- | A ring buffer is a mutable array of fixed size. Initially the array is
47-- empty, with ringStart pointing at the start of allocated memory. We call the
48-- next location to be written in the ring as ringHead. Initially ringHead ==
49-- ringStart. When the first item is added, ringHead points to ringStart +
50-- sizeof item. When the buffer becomes full ringHead would wrap around to
51-- ringStart. When the buffer is full, ringHead always points at the oldest
52-- item in the ring and the newest item added always overwrites the oldest
53-- item.
54--
55-- When using it we should keep in mind that a ringBuffer is a mutable data
56-- structure. We should not leak out references to it for immutable use.
57--
58data Ring a = Ring
59 { ringStart :: {-# UNPACK #-} !(ForeignPtr a) -- first address
60 , ringBound :: {-# UNPACK #-} !(Ptr a) -- first address beyond allocated memory
61 }
63-- | Get the first address of the ring as a pointer.
64startOf :: Ring a -> Ptr a
65startOf = unsafeForeignPtrToPtr . ringStart
67-- | Create a new ringbuffer and return the ring buffer and the ringHead.
68-- Returns the ring and the ringHead, the ringHead is same as ringStart.
69{-# INLINE new #-}
70new :: forall a. Storable a => Int -> IO (Ring a, Ptr a)
71new count = do
72 let size = count * sizeOf (undefined :: a)
73 fptr <- mallocPlainForeignPtrAlignedBytes size (alignment (undefined :: a))
74 let p = unsafeForeignPtrToPtr fptr
75 return (Ring
76 { ringStart = fptr
77 , ringBound = p `plusPtr` size
78 }, p)
80-- | Advance the ringHead by 1 item, wrap around if we hit the end of the
81-- array.
82{-# INLINE advance #-}
83advance :: forall a. Storable a => Ring a -> Ptr a -> Ptr a
84advance Ring{..} ringHead =
85 let ptr = ringHead `plusPtr` sizeOf (undefined :: a)
86 in if ptr < ringBound
87 then ptr
88 else unsafeForeignPtrToPtr ringStart
90-- | Move the ringHead by n items. The direction depends on the sign on whether
91-- n is positive or negative. Wrap around if we hit the beginning or end of the
92-- array.
93{-# INLINE moveBy #-}
94moveBy :: forall a. Storable a => Int -> Ring a -> Ptr a -> Ptr a
95moveBy by Ring {..} ringHead = ringStartPtr `plusPtr` advanceFromHead
97 where
99 elemSize = sizeOf (undefined :: a)
100 ringStartPtr = unsafeForeignPtrToPtr ringStart
101 lenInBytes = ringBound `minusPtr` ringStartPtr
102 offInBytes = ringHead `minusPtr` ringStartPtr
103 len = assert (lenInBytes `mod` elemSize == 0) $ lenInBytes `div` elemSize
104 off = assert (offInBytes `mod` elemSize == 0) $ offInBytes `div` elemSize
105 advanceFromHead = (off + by `mod` len) * elemSize
107-- | Insert an item at the head of the ring, when the ring is full this
108-- replaces the oldest item in the ring with the new item. This is unsafe
109-- beause ringHead supplied is not verified to be within the Ring. Also,
110-- the ringStart foreignPtr must be guaranteed to be alive by the caller.
111{-# INLINE unsafeInsert #-}
112unsafeInsert :: Storable a => Ring a -> Ptr a -> a -> IO (Ptr a)
113unsafeInsert rb ringHead newVal = do
114 poke ringHead newVal
115 -- touchForeignPtr (ringStart rb)
116 return $ advance rb ringHead
118-- XXX remove all usage of A.unsafeInlineIO
119--
120-- | Like 'unsafeEqArray' but compares only N bytes instead of entire length of
121-- the ring buffer. This is unsafe because the ringHead Ptr is not checked to
122-- be in range.
123{-# INLINE unsafeEqArrayN #-}
124unsafeEqArrayN :: Ring a -> Ptr a -> A.Array a -> Int -> Bool
125unsafeEqArrayN Ring{..} rh A.Array{..} n =
126 let !res = A.unsafeInlineIO $ do
127 let rs = unsafeForeignPtrToPtr ringStart
128 as = unsafeForeignPtrToPtr aStart
129 assert (aEnd `minusPtr` as >= ringBound `minusPtr` rs) (return ())
130 let len = ringBound `minusPtr` rh
131 r1 <- A.memcmp (castPtr rh) (castPtr as) (min len n)
132 r2 <- if n > len
133 then A.memcmp (castPtr rs) (castPtr (as `plusPtr` len))
134 (min (rh `minusPtr` rs) (n - len))
135 else return True
136 -- XXX enable these, check perf impact
137 -- touchForeignPtr ringStart
138 -- touchForeignPtr aStart
139 return (r1 && r2)
140 in res
142-- | Byte compare the entire length of ringBuffer with the given array,
143-- starting at the supplied ringHead pointer. Returns true if the Array and
144-- the ringBuffer have identical contents.
145--
146-- This is unsafe because the ringHead Ptr is not checked to be in range. The
147-- supplied array must be equal to or bigger than the ringBuffer, ARRAY BOUNDS
148-- ARE NOT CHECKED.
149{-# INLINE unsafeEqArray #-}
150unsafeEqArray :: Ring a -> Ptr a -> A.Array a -> Bool
151unsafeEqArray Ring{..} rh A.Array{..} =
152 let !res = A.unsafeInlineIO $ do
153 let rs = unsafeForeignPtrToPtr ringStart
154 let as = unsafeForeignPtrToPtr aStart
155 assert (aEnd `minusPtr` as >= ringBound `minusPtr` rs)
156 (return ())
157 let len = ringBound `minusPtr` rh
158 r1 <- A.memcmp (castPtr rh) (castPtr as) len
159 r2 <- A.memcmp (castPtr rs) (castPtr (as `plusPtr` len))
160 (rh `minusPtr` rs)
161 -- XXX enable these, check perf impact
162 -- touchForeignPtr ringStart
163 -- touchForeignPtr aStart
164 return (r1 && r2)
165 in res
167-- XXX use MonadIO
168--
169-- | Fold the buffer starting from ringStart up to the given 'Ptr' using a pure
170-- step function. This is useful to fold the items in the ring when the ring is
171-- not full. The supplied pointer is usually the end of the ring.
172--
173-- Unsafe because the supplied Ptr is not checked to be in range.
174{-# INLINE unsafeFoldRing #-}
175unsafeFoldRing :: forall a b. Storable a
176 => Ptr a -> (b -> a -> b) -> b -> Ring a -> b
177unsafeFoldRing ptr f z Ring{..} =
178 let !res = A.unsafeInlineIO $ withForeignPtr ringStart $ \p ->
179 go z p ptr
180 in res
181 where
182 go !acc !p !q
183 | p == q = return acc
184 | otherwise = do
185 x <- peek p
186 go (f acc x) (p `plusPtr` sizeOf (undefined :: a)) q
188-- XXX Can we remove MonadIO here?
189withForeignPtrM :: MonadIO m => ForeignPtr a -> (Ptr a -> m b) -> m b
190withForeignPtrM fp fn = do
191 r <- fn $ unsafeForeignPtrToPtr fp
192 liftIO $ touchForeignPtr fp
193 return r
195-- | Like unsafeFoldRing but with a monadic step function.
196{-# INLINE unsafeFoldRingM #-}
197unsafeFoldRingM :: forall m a b. (MonadIO m, Storable a)
198 => Ptr a -> (b -> a -> m b) -> b -> Ring a -> m b
199unsafeFoldRingM ptr f z Ring {..} =
200 withForeignPtrM ringStart $ \x -> go z x ptr
201 where
202 go !acc !start !end
203 | start == end = return acc
204 | otherwise = do
205 let !x = A.unsafeInlineIO $ peek start
206 acc' <- f acc x
207 go acc' (start `plusPtr` sizeOf (undefined :: a)) end
209-- | Fold the entire length of a ring buffer starting at the supplied ringHead
210-- pointer. Assuming the supplied ringHead pointer points to the oldest item,
211-- this would fold the ring starting from the oldest item to the newest item in
212-- the ring.
213--
214-- Note, this will crash on ring of 0 size.
215--
216{-# INLINE unsafeFoldRingFullM #-}
217unsafeFoldRingFullM :: forall m a b. (MonadIO m, Storable a)
218 => Ptr a -> (b -> a -> m b) -> b -> Ring a -> m b
219unsafeFoldRingFullM rh f z rb@Ring {..} =
220 withForeignPtrM ringStart $ \_ -> go z rh
221 where
222 go !acc !start = do
223 let !x = A.unsafeInlineIO $ peek start
224 acc' <- f acc x
225 let ptr = advance rb start
226 if ptr == rh
227 then return acc'
228 else go acc' ptr
230-- | Fold @Int@ items in the ring starting at @Ptr a@. Won't fold more
231-- than the length of the ring.
232--
233-- Note, this will crash on ring of 0 size.
234--
235{-# INLINE unsafeFoldRingNM #-}
236unsafeFoldRingNM :: forall m a b. (MonadIO m, Storable a)
237 => Int -> Ptr a -> (b -> a -> m b) -> b -> Ring a -> m b
238unsafeFoldRingNM count rh f z rb@Ring {..} =
239 withForeignPtrM ringStart $ \_ -> go count z rh
241 where
243 go 0 acc _ = return acc
244 go !n !acc !start = do
245 let !x = A.unsafeInlineIO $ peek start
246 acc' <- f acc x
247 let ptr = advance rb start
248 if ptr == rh || n == 0
249 then return acc'
250 else go (n - 1) acc' ptr