{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ExistentialQuantification #-}

-- |
-- Module      : Streamly.Memory.Ring
-- Copyright   : (c) 2019 Composewell Technologies
-- License     : BSD3
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC
--

module Streamly.Memory.Ring
    ( Ring(..)

    -- * Construction
    , new

    -- * Modification
    , unsafeInsert

    -- * Folds
    , unsafeFoldRing
    , unsafeFoldRingM
    , unsafeFoldRingFullM

    -- * Fast Byte Comparisons
    , unsafeEqArray
    , unsafeEqArrayN
    ) where

import Control.Exception (assert)
import Foreign.ForeignPtr (ForeignPtr, withForeignPtr, touchForeignPtr)
import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
import Foreign.Ptr (plusPtr, minusPtr, castPtr)
import Foreign.Storable (Storable(..))
import GHC.ForeignPtr (mallocPlainForeignPtrAlignedBytes)
import GHC.Ptr (Ptr(..))
import Prelude hiding (length, concat)

import Control.Monad.IO.Class (MonadIO(..))

import qualified Streamly.Internal.Memory.Array.Types as A

-- | A ring buffer is a mutable array of fixed size. Initially the array is
-- empty, with ringStart pointing at the start of allocated memory. We call the
-- next location to be written in the ring as ringHead. Initially ringHead ==
-- ringStart. When the first item is added, ringHead points to ringStart +
-- sizeof item. When the buffer becomes full ringHead would wrap around to
-- ringStart. When the buffer is full, ringHead always points at the oldest
-- item in the ring and the newest item added always overwrites the oldest
-- item.
--
-- When using it we should keep in mind that a ringBuffer is a mutable data
-- structure. We should not leak out references to it for immutable use.
--
data Ring a = Ring
    { Ring a -> ForeignPtr a
ringStart :: !(ForeignPtr a) -- first address
    , Ring a -> Ptr a
ringBound :: !(Ptr a)        -- first address beyond allocated memory
    }

-- | Create a new ringbuffer and return the ring buffer and the ringHead.
-- Returns the ring and the ringHead, the ringHead is same as ringStart.
{-# INLINE new #-}
new :: forall a. Storable a => Int -> IO (Ring a, Ptr a)
new :: Int -> IO (Ring a, Ptr a)
new Int
count = do
    let size :: Int
size = Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
* a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)
    ForeignPtr a
fptr <- Int -> Int -> IO (ForeignPtr a)
forall a. Int -> Int -> IO (ForeignPtr a)
mallocPlainForeignPtrAlignedBytes Int
size (a -> Int
forall a. Storable a => a -> Int
alignment (a
forall a. HasCallStack => a
undefined :: a))
    let p :: Ptr a
p = ForeignPtr a -> Ptr a
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr a
fptr
    (Ring a, Ptr a) -> IO (Ring a, Ptr a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Ring :: forall a. ForeignPtr a -> Ptr a -> Ring a
Ring
        { ringStart :: ForeignPtr a
ringStart = ForeignPtr a
fptr
        , ringBound :: Ptr a
ringBound = Ptr a
p Ptr a -> Int -> Ptr a
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
size
        }, Ptr a
p)

-- | Advance the ringHead by 1 item, wrap around if we hit the end of the
-- array.
{-# INLINE advance #-}
advance :: forall a. Storable a => Ring a -> Ptr a -> Ptr a
advance :: Ring a -> Ptr a -> Ptr a
advance Ring{Ptr a
ForeignPtr a
ringBound :: Ptr a
ringStart :: ForeignPtr a
ringBound :: forall a. Ring a -> Ptr a
ringStart :: forall a. Ring a -> ForeignPtr a
..} Ptr a
ringHead =
    let ptr :: Ptr b
ptr = Ptr a
ringHead Ptr a -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)
    in if Ptr a
forall b. Ptr b
ptr Ptr a -> Ptr a -> Bool
forall a. Ord a => a -> a -> Bool
<  Ptr a
ringBound
       then Ptr a
forall b. Ptr b
ptr
       else ForeignPtr a -> Ptr a
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr a
ringStart

-- | Insert an item at the head of the ring, when the ring is full this
-- replaces the oldest item in the ring with the new item. This is unsafe
-- beause ringHead supplied is not verified to be within the Ring. Also,
-- the ringStart foreignPtr must be guaranteed to be alive by the caller.
{-# INLINE unsafeInsert #-}
unsafeInsert :: Storable a => Ring a -> Ptr a -> a -> IO (Ptr a)
unsafeInsert :: Ring a -> Ptr a -> a -> IO (Ptr a)
unsafeInsert Ring a
rb Ptr a
ringHead a
newVal = do
    Ptr a -> a -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr a
ringHead a
newVal
    -- touchForeignPtr (ringStart rb)
    Ptr a -> IO (Ptr a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Ptr a -> IO (Ptr a)) -> Ptr a -> IO (Ptr a)
forall a b. (a -> b) -> a -> b
$ Ring a -> Ptr a -> Ptr a
forall a. Storable a => Ring a -> Ptr a -> Ptr a
advance Ring a
rb Ptr a
ringHead

-- XXX remove all usage of A.unsafeInlineIO
--
-- | Like 'unsafeEqArray' but compares only N bytes instead of entire length of
-- the ring buffer. This is unsafe because the ringHead Ptr is not checked to
-- be in range.
{-# INLINE unsafeEqArrayN #-}
unsafeEqArrayN :: Ring a -> Ptr a -> A.Array a -> Int -> Bool
unsafeEqArrayN :: Ring a -> Ptr a -> Array a -> Int -> Bool
unsafeEqArrayN Ring{Ptr a
ForeignPtr a
ringBound :: Ptr a
ringStart :: ForeignPtr a
ringBound :: forall a. Ring a -> Ptr a
ringStart :: forall a. Ring a -> ForeignPtr a
..} Ptr a
rh A.Array{Ptr a
ForeignPtr a
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
aStart :: forall a. Array a -> ForeignPtr a
aBound :: Ptr a
aEnd :: Ptr a
aStart :: ForeignPtr a
..} Int
n =
    let !res :: Bool
res = IO Bool -> Bool
forall a. IO a -> a
A.unsafeInlineIO (IO Bool -> Bool) -> IO Bool -> Bool
forall a b. (a -> b) -> a -> b
$ do
            let rs :: Ptr a
rs = ForeignPtr a -> Ptr a
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr a
ringStart
                as :: Ptr a
as = ForeignPtr a -> Ptr a
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr a
aStart
            Bool -> IO () -> IO ()
forall a. HasCallStack => Bool -> a -> a
assert (Ptr a
aBound Ptr a -> Ptr a -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
as Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Ptr a
ringBound Ptr a -> Ptr a -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
rs) (() -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
            let len :: Int
len = Ptr a
ringBound Ptr a -> Ptr a -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
rh
            Bool
r1 <- Ptr Word8 -> Ptr Word8 -> Int -> IO Bool
A.memcmp (Ptr a -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr a
rh) (Ptr a -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr a
as) (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
len Int
n)
            Bool
r2 <- if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
len
                then Ptr Word8 -> Ptr Word8 -> Int -> IO Bool
A.memcmp (Ptr a -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr a
rs) (Ptr Any -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr (Ptr a
as Ptr a -> Int -> Ptr Any
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
len))
                              (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Ptr a
rh Ptr a -> Ptr a -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
rs) (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
len))
                else Bool -> IO Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
            -- XXX enable these, check perf impact
            -- touchForeignPtr ringStart
            -- touchForeignPtr aStart
            Bool -> IO Bool
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool
r1 Bool -> Bool -> Bool
&& Bool
r2)
    in Bool
res

-- | Byte compare the entire length of ringBuffer with the given array,
-- starting at the supplied ringHead pointer.  Returns true if the Array and
-- the ringBuffer have identical contents.
--
-- This is unsafe because the ringHead Ptr is not checked to be in range. The
-- supplied array must be equal to or bigger than the ringBuffer, ARRAY BOUNDS
-- ARE NOT CHECKED.
{-# INLINE unsafeEqArray #-}
unsafeEqArray :: Ring a -> Ptr a -> A.Array a -> Bool
unsafeEqArray :: Ring a -> Ptr a -> Array a -> Bool
unsafeEqArray Ring{Ptr a
ForeignPtr a
ringBound :: Ptr a
ringStart :: ForeignPtr a
ringBound :: forall a. Ring a -> Ptr a
ringStart :: forall a. Ring a -> ForeignPtr a
..} Ptr a
rh A.Array{Ptr a
ForeignPtr a
aBound :: Ptr a
aEnd :: Ptr a
aStart :: ForeignPtr a
aBound :: forall a. Array a -> Ptr a
aEnd :: forall a. Array a -> Ptr a
aStart :: forall a. Array a -> ForeignPtr a
..} =
    let !res :: Bool
res = IO Bool -> Bool
forall a. IO a -> a
A.unsafeInlineIO (IO Bool -> Bool) -> IO Bool -> Bool
forall a b. (a -> b) -> a -> b
$ do
            let rs :: Ptr a
rs = ForeignPtr a -> Ptr a
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr a
ringStart
            let as :: Ptr a
as = ForeignPtr a -> Ptr a
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr a
aStart
            Bool -> IO () -> IO ()
forall a. HasCallStack => Bool -> a -> a
assert (Ptr a
aBound Ptr a -> Ptr a -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
as Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Ptr a
ringBound Ptr a -> Ptr a -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
rs)
                   (() -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
            let len :: Int
len = Ptr a
ringBound Ptr a -> Ptr a -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
rh
            Bool
r1 <- Ptr Word8 -> Ptr Word8 -> Int -> IO Bool
A.memcmp (Ptr a -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr a
rh) (Ptr a -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr a
as) Int
len
            Bool
r2 <- Ptr Word8 -> Ptr Word8 -> Int -> IO Bool
A.memcmp (Ptr a -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr a
rs) (Ptr Any -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr (Ptr a
as Ptr a -> Int -> Ptr Any
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
len))
                           (Ptr a
rh Ptr a -> Ptr a -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
rs)
            -- XXX enable these, check perf impact
            -- touchForeignPtr ringStart
            -- touchForeignPtr aStart
            Bool -> IO Bool
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool
r1 Bool -> Bool -> Bool
&& Bool
r2)
    in Bool
res

-- XXX use MonadIO
--
-- | Fold the buffer starting from ringStart up to the given 'Ptr' using a pure
-- step function. This is useful to fold the items in the ring when the ring is
-- not full. The supplied pointer is usually the end of the ring.
--
-- Unsafe because the supplied Ptr is not checked to be in range.
{-# INLINE unsafeFoldRing #-}
unsafeFoldRing :: forall a b. Storable a
    => Ptr a -> (b -> a -> b) -> b -> Ring a -> b
unsafeFoldRing :: Ptr a -> (b -> a -> b) -> b -> Ring a -> b
unsafeFoldRing Ptr a
ptr b -> a -> b
f b
z Ring{Ptr a
ForeignPtr a
ringBound :: Ptr a
ringStart :: ForeignPtr a
ringBound :: forall a. Ring a -> Ptr a
ringStart :: forall a. Ring a -> ForeignPtr a
..} =
    let !res :: b
res = IO b -> b
forall a. IO a -> a
A.unsafeInlineIO (IO b -> b) -> IO b -> b
forall a b. (a -> b) -> a -> b
$ ForeignPtr a -> (Ptr a -> IO b) -> IO b
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr a
ringStart ((Ptr a -> IO b) -> IO b) -> (Ptr a -> IO b) -> IO b
forall a b. (a -> b) -> a -> b
$ \Ptr a
p ->
                    b -> Ptr a -> Ptr a -> IO b
go b
z Ptr a
p Ptr a
ptr
    in b
res
    where
      go :: b -> Ptr a -> Ptr a -> IO b
go !b
acc !Ptr a
p !Ptr a
q
        | Ptr a
p Ptr a -> Ptr a -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr a
q = b -> IO b
forall (m :: * -> *) a. Monad m => a -> m a
return b
acc
        | Bool
otherwise = do
            a
x <- Ptr a -> IO a
forall a. Storable a => Ptr a -> IO a
peek Ptr a
p
            b -> Ptr a -> Ptr a -> IO b
go (b -> a -> b
f b
acc a
x) (Ptr a
p Ptr a -> Int -> Ptr a
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)) Ptr a
q

-- XXX Can we remove MonadIO here?
withForeignPtrM :: MonadIO m => ForeignPtr a -> (Ptr a -> m b) -> m b
withForeignPtrM :: ForeignPtr a -> (Ptr a -> m b) -> m b
withForeignPtrM ForeignPtr a
fp Ptr a -> m b
fn = do
    b
r <- Ptr a -> m b
fn (Ptr a -> m b) -> Ptr a -> m b
forall a b. (a -> b) -> a -> b
$ ForeignPtr a -> Ptr a
forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr a
fp
    IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ ForeignPtr a -> IO ()
forall a. ForeignPtr a -> IO ()
touchForeignPtr ForeignPtr a
fp
    b -> m b
forall (m :: * -> *) a. Monad m => a -> m a
return b
r

-- | Like unsafeFoldRing but with a monadic step function.
{-# INLINE unsafeFoldRingM #-}
unsafeFoldRingM :: forall m a b. (MonadIO m, Storable a)
    => Ptr a -> (b -> a -> m b) -> b -> Ring a -> m b
unsafeFoldRingM :: Ptr a -> (b -> a -> m b) -> b -> Ring a -> m b
unsafeFoldRingM Ptr a
ptr b -> a -> m b
f b
z Ring {Ptr a
ForeignPtr a
ringBound :: Ptr a
ringStart :: ForeignPtr a
ringBound :: forall a. Ring a -> Ptr a
ringStart :: forall a. Ring a -> ForeignPtr a
..} =
    ForeignPtr a -> (Ptr a -> m b) -> m b
forall (m :: * -> *) a b.
MonadIO m =>
ForeignPtr a -> (Ptr a -> m b) -> m b
withForeignPtrM ForeignPtr a
ringStart ((Ptr a -> m b) -> m b) -> (Ptr a -> m b) -> m b
forall a b. (a -> b) -> a -> b
$ \Ptr a
x -> b -> Ptr a -> Ptr a -> m b
go b
z Ptr a
x Ptr a
ptr
  where
    go :: b -> Ptr a -> Ptr a -> m b
go !b
acc !Ptr a
start !Ptr a
end
        | Ptr a
start Ptr a -> Ptr a -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr a
end = b -> m b
forall (m :: * -> *) a. Monad m => a -> m a
return b
acc
        | Bool
otherwise = do
            let !x :: a
x = IO a -> a
forall a. IO a -> a
A.unsafeInlineIO (IO a -> a) -> IO a -> a
forall a b. (a -> b) -> a -> b
$ Ptr a -> IO a
forall a. Storable a => Ptr a -> IO a
peek Ptr a
start
            b
acc' <- b -> a -> m b
f b
acc a
x
            b -> Ptr a -> Ptr a -> m b
go b
acc' (Ptr a
start Ptr a -> Int -> Ptr a
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)) Ptr a
end

-- | Fold the entire length of a ring buffer starting at the supplied ringHead
-- pointer.  Assuming the supplied ringHead pointer points to the oldest item,
-- this would fold the ring starting from the oldest item to the newest item in
-- the ring.
{-# INLINE unsafeFoldRingFullM #-}
unsafeFoldRingFullM :: forall m a b. (MonadIO m, Storable a)
    => Ptr a -> (b -> a -> m b) -> b -> Ring a -> m b
unsafeFoldRingFullM :: Ptr a -> (b -> a -> m b) -> b -> Ring a -> m b
unsafeFoldRingFullM Ptr a
rh b -> a -> m b
f b
z rb :: Ring a
rb@Ring {Ptr a
ForeignPtr a
ringBound :: Ptr a
ringStart :: ForeignPtr a
ringBound :: forall a. Ring a -> Ptr a
ringStart :: forall a. Ring a -> ForeignPtr a
..} =
    ForeignPtr a -> (Ptr a -> m b) -> m b
forall (m :: * -> *) a b.
MonadIO m =>
ForeignPtr a -> (Ptr a -> m b) -> m b
withForeignPtrM ForeignPtr a
ringStart ((Ptr a -> m b) -> m b) -> (Ptr a -> m b) -> m b
forall a b. (a -> b) -> a -> b
$ \Ptr a
_ -> b -> Ptr a -> m b
go b
z Ptr a
rh
  where
    go :: b -> Ptr a -> m b
go !b
acc !Ptr a
start = do
        let !x :: a
x = IO a -> a
forall a. IO a -> a
A.unsafeInlineIO (IO a -> a) -> IO a -> a
forall a b. (a -> b) -> a -> b
$ Ptr a -> IO a
forall a. Storable a => Ptr a -> IO a
peek Ptr a
start
        b
acc' <- b -> a -> m b
f b
acc a
x
        let ptr :: Ptr a
ptr = Ring a -> Ptr a -> Ptr a
forall a. Storable a => Ring a -> Ptr a -> Ptr a
advance Ring a
rb Ptr a
start
        if Ptr a
ptr Ptr a -> Ptr a -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr a
rh
            then b -> m b
forall (m :: * -> *) a. Monad m => a -> m a
return b
acc'
            else b -> Ptr a -> m b
go b
acc' Ptr a
ptr