{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ExistentialQuantification #-}
module Streamly.Memory.Ring
( Ring(..)
, new
, unsafeInsert
, unsafeFoldRing
, unsafeFoldRingM
, unsafeFoldRingFullM
, 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
data Ring a = Ring
{ forall a. Ring a -> ForeignPtr a
ringStart :: !(ForeignPtr a)
, forall a. Ring a -> Ptr a
ringBound :: !(Ptr a)
}
{-# INLINE new #-}
new :: forall a. Storable a => Int -> IO (Ring a, Ptr a)
new :: forall a. Storable a => Int -> IO (Ring a, Ptr a)
new Int
count = do
let size :: Int
size = Int
count forall a. Num a => a -> a -> a
* forall a. Storable a => a -> Int
sizeOf (forall a. HasCallStack => a
undefined :: a)
ForeignPtr a
fptr <- forall a. Int -> Int -> IO (ForeignPtr a)
mallocPlainForeignPtrAlignedBytes Int
size (forall a. Storable a => a -> Int
alignment (forall a. HasCallStack => a
undefined :: a))
let p :: Ptr a
p = forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr a
fptr
forall (m :: * -> *) a. Monad m => a -> m a
return (Ring
{ ringStart :: ForeignPtr a
ringStart = ForeignPtr a
fptr
, ringBound :: Ptr a
ringBound = Ptr a
p forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
size
}, Ptr a
p)
{-# INLINE advance #-}
advance :: forall a. Storable a => Ring a -> Ptr a -> Ptr a
advance :: forall a. Storable a => 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 forall a b. Ptr a -> Int -> Ptr b
`plusPtr` forall a. Storable a => a -> Int
sizeOf (forall a. HasCallStack => a
undefined :: a)
in if forall {b}. Ptr b
ptr forall a. Ord a => a -> a -> Bool
< Ptr a
ringBound
then forall {b}. Ptr b
ptr
else forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr a
ringStart
{-# INLINE unsafeInsert #-}
unsafeInsert :: Storable a => Ring a -> Ptr a -> a -> IO (Ptr a)
unsafeInsert :: forall a. Storable a => Ring a -> Ptr a -> a -> IO (Ptr a)
unsafeInsert Ring a
rb Ptr a
ringHead a
newVal = do
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr a
ringHead a
newVal
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. Storable a => Ring a -> Ptr a -> Ptr a
advance Ring a
rb Ptr a
ringHead
{-# INLINE unsafeEqArrayN #-}
unsafeEqArrayN :: Ring a -> Ptr a -> A.Array a -> Int -> Bool
unsafeEqArrayN :: forall a. 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 = forall a. IO a -> a
A.unsafeInlineIO forall a b. (a -> b) -> a -> b
$ do
let rs :: Ptr a
rs = forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr a
ringStart
as :: Ptr a
as = forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr a
aStart
forall a. HasCallStack => Bool -> a -> a
assert (Ptr a
aBound forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
as forall a. Ord a => a -> a -> Bool
>= Ptr a
ringBound forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
rs) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
let len :: Int
len = Ptr a
ringBound forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
rh
Bool
r1 <- Ptr Word8 -> Ptr Word8 -> Int -> IO Bool
A.memcmp (forall a b. Ptr a -> Ptr b
castPtr Ptr a
rh) (forall a b. Ptr a -> Ptr b
castPtr Ptr a
as) (forall a. Ord a => a -> a -> a
min Int
len Int
n)
Bool
r2 <- if Int
n forall a. Ord a => a -> a -> Bool
> Int
len
then Ptr Word8 -> Ptr Word8 -> Int -> IO Bool
A.memcmp (forall a b. Ptr a -> Ptr b
castPtr Ptr a
rs) (forall a b. Ptr a -> Ptr b
castPtr (Ptr a
as forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
len))
(forall a. Ord a => a -> a -> a
min (Ptr a
rh forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
rs) (Int
n forall a. Num a => a -> a -> a
- Int
len))
else forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool
r1 Bool -> Bool -> Bool
&& Bool
r2)
in Bool
res
{-# INLINE unsafeEqArray #-}
unsafeEqArray :: Ring a -> Ptr a -> A.Array a -> Bool
unsafeEqArray :: forall a. 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 = forall a. IO a -> a
A.unsafeInlineIO forall a b. (a -> b) -> a -> b
$ do
let rs :: Ptr a
rs = forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr a
ringStart
let as :: Ptr a
as = forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr a
aStart
forall a. HasCallStack => Bool -> a -> a
assert (Ptr a
aBound forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
as forall a. Ord a => a -> a -> Bool
>= Ptr a
ringBound forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
rs)
(forall (m :: * -> *) a. Monad m => a -> m a
return ())
let len :: Int
len = Ptr a
ringBound forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
rh
Bool
r1 <- Ptr Word8 -> Ptr Word8 -> Int -> IO Bool
A.memcmp (forall a b. Ptr a -> Ptr b
castPtr Ptr a
rh) (forall a b. Ptr a -> Ptr b
castPtr Ptr a
as) Int
len
Bool
r2 <- Ptr Word8 -> Ptr Word8 -> Int -> IO Bool
A.memcmp (forall a b. Ptr a -> Ptr b
castPtr Ptr a
rs) (forall a b. Ptr a -> Ptr b
castPtr (Ptr a
as forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
len))
(Ptr a
rh forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr a
rs)
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool
r1 Bool -> Bool -> Bool
&& Bool
r2)
in Bool
res
{-# INLINE unsafeFoldRing #-}
unsafeFoldRing :: forall a b. Storable a
=> Ptr a -> (b -> a -> b) -> b -> Ring a -> b
unsafeFoldRing :: forall a b.
Storable a =>
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 = forall a. IO a -> a
A.unsafeInlineIO forall a b. (a -> b) -> a -> b
$ forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr a
ringStart 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 forall a. Eq a => a -> a -> Bool
== Ptr a
q = forall (m :: * -> *) a. Monad m => a -> m a
return b
acc
| Bool
otherwise = do
a
x <- 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 forall a b. Ptr a -> Int -> Ptr b
`plusPtr` forall a. Storable a => a -> Int
sizeOf (forall a. HasCallStack => a
undefined :: a)) Ptr a
q
withForeignPtrM :: MonadIO m => ForeignPtr a -> (Ptr a -> m b) -> m b
withForeignPtrM :: forall (m :: * -> *) a b.
MonadIO m =>
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 forall a b. (a -> b) -> a -> b
$ forall a. ForeignPtr a -> Ptr a
unsafeForeignPtrToPtr ForeignPtr a
fp
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. ForeignPtr a -> IO ()
touchForeignPtr ForeignPtr a
fp
forall (m :: * -> *) a. Monad m => a -> m a
return b
r
{-# INLINE unsafeFoldRingM #-}
unsafeFoldRingM :: forall m a b. (MonadIO m, Storable a)
=> Ptr a -> (b -> a -> m b) -> b -> Ring a -> m b
unsafeFoldRingM :: forall (m :: * -> *) a b.
(MonadIO m, Storable a) =>
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
..} =
forall (m :: * -> *) a b.
MonadIO m =>
ForeignPtr a -> (Ptr a -> m b) -> m b
withForeignPtrM ForeignPtr a
ringStart 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 forall a. Eq a => a -> a -> Bool
== Ptr a
end = forall (m :: * -> *) a. Monad m => a -> m a
return b
acc
| Bool
otherwise = do
let !x :: a
x = forall a. IO a -> a
A.unsafeInlineIO forall a b. (a -> b) -> a -> b
$ 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 forall a b. Ptr a -> Int -> Ptr b
`plusPtr` forall a. Storable a => a -> Int
sizeOf (forall a. HasCallStack => a
undefined :: a)) Ptr a
end
{-# INLINE unsafeFoldRingFullM #-}
unsafeFoldRingFullM :: forall m a b. (MonadIO m, Storable a)
=> Ptr a -> (b -> a -> m b) -> b -> Ring a -> m b
unsafeFoldRingFullM :: forall (m :: * -> *) a b.
(MonadIO m, Storable a) =>
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
..} =
forall (m :: * -> *) a b.
MonadIO m =>
ForeignPtr a -> (Ptr a -> m b) -> m b
withForeignPtrM ForeignPtr a
ringStart 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 = forall a. IO a -> a
A.unsafeInlineIO forall a b. (a -> b) -> a -> b
$ 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 = forall a. Storable a => Ring a -> Ptr a -> Ptr a
advance Ring a
rb Ptr a
start
if Ptr a
ptr forall a. Eq a => a -> a -> Bool
== Ptr a
rh
then forall (m :: * -> *) a. Monad m => a -> m a
return b
acc'
else b -> Ptr a -> m b
go b
acc' Ptr a
ptr