{-# language BangPatterns #-}
{-# language TupleSections #-}
{-# language TypeFamilies #-}

module RingBuffers.Internal
  ( RingBuffer(..)
  , RingState(..)
  , withRing
  , new
  , clear
  , capacity
  , filledLength
  , latest
  , unsafeLatest
  , advance
  , extend
  , append
  , foldMap
  , toList
  ) where

import qualified Data.Primitive.Contiguous as Contiguous

data RingBuffer arr a = RingBuffer
  { RingBuffer arr a -> Mutable arr RealWorld a
_ringBufferBuffer :: !(Mutable arr RealWorld a)
  , RingBuffer arr a -> MVar RingState
_ringBufferState :: {-# UNPACK #-} !(MVar RingState)
  }

data RingState = RingState
  { RingState -> Bool
_ringStateFull :: !Bool -- ^ Is the ring buffer full?
  , RingState -> Int
_ringStateHead :: !Int  -- ^ next entry to be written
  }

ringState0 :: RingState
ringState0 :: RingState
ringState0 = Bool -> Int -> RingState
RingState Bool
False Int
0
{-# inline ringState0 #-}

withRing :: (Contiguous arr, Element arr a)
  => RingBuffer arr a
  -> (Mutable arr RealWorld a -> RingState -> IO (RingState, r))
  -> IO r
withRing :: RingBuffer arr a
-> (Mutable arr RealWorld a -> RingState -> IO (RingState, r))
-> IO r
withRing (RingBuffer Mutable arr RealWorld a
ba MVar RingState
bs) Mutable arr RealWorld a -> RingState -> IO (RingState, r)
f = MVar RingState -> (RingState -> IO (RingState, r)) -> IO r
forall a b. MVar a -> (a -> IO (a, b)) -> IO b
modifyMVar MVar RingState
bs (Mutable arr RealWorld a -> RingState -> IO (RingState, r)
f Mutable arr RealWorld a
ba)
{-# inline withRing #-}

new :: (Contiguous arr, Element arr a)
  => Int
  -> IO (RingBuffer arr a)
new :: Int -> IO (RingBuffer arr a)
new !Int
sz = do
  Mutable arr RealWorld a
ba <- Int -> IO (Mutable arr (PrimState IO) a)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Int -> m (Mutable arr (PrimState m) b)
Contiguous.new Int
sz
  MVar RingState
s0 <- RingState -> IO (MVar RingState)
forall a. a -> IO (MVar a)
newMVar RingState
ringState0
  RingBuffer arr a -> IO (RingBuffer arr a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Mutable arr RealWorld a -> MVar RingState -> RingBuffer arr a
forall (arr :: * -> *) a.
Mutable arr RealWorld a -> MVar RingState -> RingBuffer arr a
RingBuffer Mutable arr RealWorld a
ba MVar RingState
s0)
{-# inlineable new #-}

clear :: (Contiguous arr, Element arr a)
  => RingBuffer arr a
  -> IO ()
clear :: RingBuffer arr a -> IO ()
clear RingBuffer arr a
rb = RingBuffer arr a
-> (Mutable arr RealWorld a -> RingState -> IO (RingState, ()))
-> IO ()
forall (arr :: * -> *) a r.
(Contiguous arr, Element arr a) =>
RingBuffer arr a
-> (Mutable arr RealWorld a -> RingState -> IO (RingState, r))
-> IO r
withRing RingBuffer arr a
rb ((Mutable arr RealWorld a -> RingState -> IO (RingState, ()))
 -> IO ())
-> (Mutable arr RealWorld a -> RingState -> IO (RingState, ()))
-> IO ()
forall a b. (a -> b) -> a -> b
$ \Mutable arr RealWorld a
_ RingState
_ -> (RingState, ()) -> IO (RingState, ())
forall (f :: * -> *) a. Applicative f => a -> f a
pure (RingState
ringState0,())
{-# inline clear #-}

capacity :: (Contiguous arr, Element arr a)
  => RingBuffer arr a
  -> IO Int
capacity :: RingBuffer arr a -> IO Int
capacity (RingBuffer Mutable arr RealWorld a
buf MVar RingState
_) = Mutable arr (PrimState IO) a -> IO Int
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m Int
Contiguous.sizeMutable Mutable arr RealWorld a
Mutable arr (PrimState IO) a
buf
{-# inline capacity #-}

filledLength :: (Contiguous arr, Element arr a)
  => RingBuffer arr a
  -> IO Int
filledLength :: RingBuffer arr a -> IO Int
filledLength RingBuffer arr a
rb = do
  RingBuffer arr a
-> (Mutable arr RealWorld a -> RingState -> IO (RingState, Int))
-> IO Int
forall (arr :: * -> *) a r.
(Contiguous arr, Element arr a) =>
RingBuffer arr a
-> (Mutable arr RealWorld a -> RingState -> IO (RingState, r))
-> IO r
withRing RingBuffer arr a
rb ((Mutable arr RealWorld a -> RingState -> IO (RingState, Int))
 -> IO Int)
-> (Mutable arr RealWorld a -> RingState -> IO (RingState, Int))
-> IO Int
forall a b. (a -> b) -> a -> b
$ \Mutable arr RealWorld a
_ rs :: RingState
rs@(RingState Bool
full Int
pos) ->
    if Bool
full
      then do
        Int
cap <- RingBuffer arr a -> IO Int
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
RingBuffer arr a -> IO Int
capacity RingBuffer arr a
rb
        (RingState, Int) -> IO (RingState, Int)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (RingState
rs,Int
cap)
      else (RingState, Int) -> IO (RingState, Int)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (RingState
rs,Int
pos)
{-# inline filledLength #-}

latest :: (Contiguous arr, Element arr a)
  => RingBuffer arr a
  -> Int
  -> IO (Maybe a)
latest :: RingBuffer arr a -> Int -> IO (Maybe a)
latest RingBuffer arr a
rb Int
n = do
  RingBuffer arr a
-> (Mutable arr RealWorld a
    -> RingState -> IO (RingState, Maybe a))
-> IO (Maybe a)
forall (arr :: * -> *) a r.
(Contiguous arr, Element arr a) =>
RingBuffer arr a
-> (Mutable arr RealWorld a -> RingState -> IO (RingState, r))
-> IO r
withRing RingBuffer arr a
rb ((Mutable arr RealWorld a -> RingState -> IO (RingState, Maybe a))
 -> IO (Maybe a))
-> (Mutable arr RealWorld a
    -> RingState -> IO (RingState, Maybe a))
-> IO (Maybe a)
forall a b. (a -> b) -> a -> b
$ \Mutable arr RealWorld a
ra rs :: RingState
rs@(RingState Bool
full Int
pos) -> do
    Int
cap <- RingBuffer arr a -> IO Int
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
RingBuffer arr a -> IO Int
capacity RingBuffer arr a
rb
    let len :: Int
len = if Bool
full then Int
cap else Int
pos
    if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
len
      then do
        (RingState, Maybe a) -> IO (RingState, Maybe a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (RingState
rs, Maybe a
forall a. Maybe a
Nothing)
      else do
        let ix :: Int
ix = (Int
pos Int -> Int -> Int
forall a. Ring a => a -> a -> a
- Int
n Int -> Int -> Int
forall a. Ring a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
cap
        a
a <- Mutable arr (PrimState IO) a -> Int -> IO a
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> m b
Contiguous.read Mutable arr RealWorld a
Mutable arr (PrimState IO) a
ra Int
ix
        (RingState, Maybe a) -> IO (RingState, Maybe a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((RingState, Maybe a) -> IO (RingState, Maybe a))
-> (RingState, Maybe a) -> IO (RingState, Maybe a)
forall a b. (a -> b) -> a -> b
$ (RingState
rs, a -> Maybe a
forall a. a -> Maybe a
Just a
a)
{-# inline latest #-}

unsafeLatest :: (Contiguous arr, Element arr a)
  => RingBuffer arr a
  -> Int
  -> IO a
unsafeLatest :: RingBuffer arr a -> Int -> IO a
unsafeLatest RingBuffer arr a
rb Int
n = do
  Int
cap <- RingBuffer arr a -> IO Int
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
RingBuffer arr a -> IO Int
capacity RingBuffer arr a
rb
  RingBuffer arr a
-> (Mutable arr RealWorld a -> RingState -> IO (RingState, a))
-> IO a
forall (arr :: * -> *) a r.
(Contiguous arr, Element arr a) =>
RingBuffer arr a
-> (Mutable arr RealWorld a -> RingState -> IO (RingState, r))
-> IO r
withRing RingBuffer arr a
rb ((Mutable arr RealWorld a -> RingState -> IO (RingState, a))
 -> IO a)
-> (Mutable arr RealWorld a -> RingState -> IO (RingState, a))
-> IO a
forall a b. (a -> b) -> a -> b
$ \Mutable arr RealWorld a
ba bs :: RingState
bs@(RingState Bool
_ Int
hd) -> do
    let idx :: Int
idx = (Int
hd Int -> Int -> Int
forall a. Ring a => a -> a -> a
- Int
n Int -> Int -> Int
forall a. Ring a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
cap
    (RingState
bs,) (a -> (RingState, a)) -> IO a -> IO (RingState, a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mutable arr (PrimState IO) a -> Int -> IO a
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> m b
Contiguous.read Mutable arr RealWorld a
Mutable arr (PrimState IO) a
ba Int
idx

advance :: (Contiguous arr, Element arr a)
  => Int
  -> (Mutable arr RealWorld a -> RingState -> IO (RingState, ()))
advance :: Int -> Mutable arr RealWorld a -> RingState -> IO (RingState, ())
advance Int
n = \Mutable arr RealWorld a
ba (RingState Bool
full Int
pos) -> do
  Int
cap <- Mutable arr (PrimState IO) a -> IO Int
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m Int
Contiguous.sizeMutable Mutable arr RealWorld a
Mutable arr (PrimState IO) a
ba
  let (Int
a,Int
pos') = (Int
pos Int -> Int -> Int
forall a. Semiring a => a -> a -> a
+ Int
n) Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
`divMod` Int
cap
  (RingState, ()) -> IO (RingState, ())
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> Int -> RingState
RingState (Bool
full Bool -> Bool -> Bool
|| Int
a Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) Int
pos', ())
{-# inline advance #-}

append :: (Contiguous arr, Element arr a)
  => a
  -> RingBuffer arr a
  -> IO ()
append :: a -> RingBuffer arr a -> IO ()
append a
x RingBuffer arr a
rb = RingBuffer arr a
-> (Mutable arr RealWorld a -> RingState -> IO (RingState, ()))
-> IO ()
forall (arr :: * -> *) a r.
(Contiguous arr, Element arr a) =>
RingBuffer arr a
-> (Mutable arr RealWorld a -> RingState -> IO (RingState, r))
-> IO r
withRing RingBuffer arr a
rb ((Mutable arr RealWorld a -> RingState -> IO (RingState, ()))
 -> IO ())
-> (Mutable arr RealWorld a -> RingState -> IO (RingState, ()))
-> IO ()
forall a b. (a -> b) -> a -> b
$ \Mutable arr RealWorld a
ba RingState
bs -> do
  Mutable arr (PrimState IO) a -> Int -> a -> IO ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
Contiguous.write Mutable arr RealWorld a
Mutable arr (PrimState IO) a
ba (RingState -> Int
_ringStateHead RingState
bs) a
x
  Int -> Mutable arr RealWorld a -> RingState -> IO (RingState, ())
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
Int -> Mutable arr RealWorld a -> RingState -> IO (RingState, ())
advance Int
1 Mutable arr RealWorld a
ba RingState
bs
{-# inline append #-}

extend :: (Contiguous arr, Element arr a)
  => arr a
  -> RingBuffer arr a
  -> IO ()
extend :: arr a -> RingBuffer arr a -> IO ()
extend arr a
xs RingBuffer arr a
rb = RingBuffer arr a
-> (Mutable arr RealWorld a -> RingState -> IO (RingState, ()))
-> IO ()
forall (arr :: * -> *) a r.
(Contiguous arr, Element arr a) =>
RingBuffer arr a
-> (Mutable arr RealWorld a -> RingState -> IO (RingState, r))
-> IO r
withRing RingBuffer arr a
rb ((Mutable arr RealWorld a -> RingState -> IO (RingState, ()))
 -> IO ())
-> (Mutable arr RealWorld a -> RingState -> IO (RingState, ()))
-> IO ()
forall a b. (a -> b) -> a -> b
$ \Mutable arr RealWorld a
ba RingState
bs -> do
  Int
cap <- RingBuffer arr a -> IO Int
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
RingBuffer arr a -> IO Int
capacity RingBuffer arr a
rb
  let extensionLength :: Int
extensionLength = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (arr a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
Contiguous.size arr a
xs) Int
cap
  let currentHead :: Int
currentHead = RingState -> Int
_ringStateHead RingState
bs
  let go :: Int -> IO ()
go !Int
ix = Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
extensionLength) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
        a
atIx <- arr a -> Int -> IO a
forall (arr :: * -> *) b (m :: * -> *).
(Contiguous arr, Element arr b, Monad m) =>
arr b -> Int -> m b
Contiguous.indexM arr a
xs Int
ix
        Mutable arr (PrimState IO) a -> Int -> a -> IO ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
Contiguous.write Mutable arr RealWorld a
Mutable arr (PrimState IO) a
ba (Int
currentHead Int -> Int -> Int
forall a. Semiring a => a -> a -> a
+ Int
ix) a
atIx
        Int -> IO ()
go (Int
ix Int -> Int -> Int
forall a. Semiring a => a -> a -> a
+ Int
1)
  Int -> IO ()
go Int
0
  Int -> Mutable arr RealWorld a -> RingState -> IO (RingState, ())
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
Int -> Mutable arr RealWorld a -> RingState -> IO (RingState, ())
advance Int
extensionLength Mutable arr RealWorld a
ba RingState
bs
{-# inlineable extend #-}

foldMap :: (Contiguous arr, Element arr a, Monoid b)
  => RingBuffer arr a
  -> (a -> IO b)
  -> IO b
foldMap :: RingBuffer arr a -> (a -> IO b) -> IO b
foldMap RingBuffer arr a
rb a -> IO b
action = do
  RingBuffer arr a
-> (Mutable arr RealWorld a -> RingState -> IO (RingState, b))
-> IO b
forall (arr :: * -> *) a r.
(Contiguous arr, Element arr a) =>
RingBuffer arr a
-> (Mutable arr RealWorld a -> RingState -> IO (RingState, r))
-> IO r
withRing RingBuffer arr a
rb ((Mutable arr RealWorld a -> RingState -> IO (RingState, b))
 -> IO b)
-> (Mutable arr RealWorld a -> RingState -> IO (RingState, b))
-> IO b
forall a b. (a -> b) -> a -> b
$ \Mutable arr RealWorld a
ba bs :: RingState
bs@(RingState Bool
full Int
pos) -> do
    Int
n <- do
      Int
cap <- RingBuffer arr a -> IO Int
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
RingBuffer arr a -> IO Int
capacity RingBuffer arr a
rb
      Int -> IO Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> IO Int) -> Int -> IO Int
forall a b. (a -> b) -> a -> b
$ if Bool
full then Int
cap else Int
pos
    let go :: Int -> b -> IO (RingState, b)
go !Int
ix !b
acc = if Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n
          then do
            a
v <- Mutable arr (PrimState IO) a -> Int -> IO a
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> m b
Contiguous.read Mutable arr RealWorld a
Mutable arr (PrimState IO) a
ba Int
ix
            b
m <- a -> IO b
action a
v
            Int -> b -> IO (RingState, b)
go (Int
ix Int -> Int -> Int
forall a. Semiring a => a -> a -> a
+ Int
1) (b
acc b -> b -> b
forall a. Semigroup a => a -> a -> a
<> b
m)
          else
            (RingState, b) -> IO (RingState, b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (RingState
bs, b
acc)
    Int -> b -> IO (RingState, b)
go Int
0 b
forall a. Monoid a => a
mempty
{-# inline foldMap #-}

toList :: (Contiguous arr, Element arr a)
  => RingBuffer arr a
  -> IO [a]
toList :: RingBuffer arr a -> IO [a]
toList RingBuffer arr a
rb = do
  Int
len <- RingBuffer arr a -> IO Int
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
RingBuffer arr a -> IO Int
filledLength RingBuffer arr a
rb
  (Int -> IO a) -> [Int] -> IO [a]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (RingBuffer arr a -> Int -> IO a
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
RingBuffer arr a -> Int -> IO a
unsafeLatest RingBuffer arr a
rb) [Int
0..Int
lenInt -> Int -> Int
forall a. Ring a => a -> a -> a
-Int
1]