-- |
-- Module      : Streamly.Internal.Data.Ring
-- Copyright   : (c) 2021 Composewell Technologies
-- License     : BSD-3-Clause
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC
--

module Streamly.Internal.Data.Ring
    ( Ring(..)

    -- * Generation
    , createRing
    , writeLastN

    -- * Modification
    , seek
    , unsafeInsertRingWith

    -- * Conversion
    , toMutArray
    , toStreamWith
    ) where

#include "assert.hs"

import Control.Monad.IO.Class (liftIO, MonadIO)
import Streamly.Internal.Data.Stream.StreamD.Type (Stream)
import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))
import Streamly.Internal.Data.Fold.Type (Fold(..))
import Streamly.Internal.Data.Array.Generic.Mut.Type
    ( MutArray(..)
    , new
    , uninit
    , putIndexUnsafe
    , putSliceUnsafe
    )
-- import qualified Streamly.Internal.Data.Stream.StreamD.Type as Stream
import qualified Streamly.Internal.Data.Fold.Type as Fold

-- XXX Use MutableArray rather than keeping a MutArray here.
data Ring a = Ring
    { forall a. Ring a -> MutArray a
ringArr :: MutArray a
    -- XXX We can keep the current fill amount, Or we can keep a count of total
    -- elements inserted and compute ring head as well using mod on that,
    -- assuming it won't overflow. But mod could be expensive.
    , forall a. Ring a -> Int
ringHead :: !Int -- current index to be over-written
    , forall a. Ring a -> Int
ringMax :: !Int  -- first index beyond allocated memory
    }

-------------------------------------------------------------------------------
-- Generation
-------------------------------------------------------------------------------

-- XXX If we align the ringMax to nearest power of two then computation of the
-- index to write could be cheaper.
{-# INLINE createRing #-}
createRing :: MonadIO m => Int -> m (Ring a)
createRing :: forall (m :: * -> *) a. MonadIO m => Int -> m (Ring a)
createRing Int
count = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
    MutArray a
arr <- forall (m :: * -> *) a. MonadIO m => Int -> m (MutArray a)
new Int
count
    MutArray a
arr1 <- forall (m :: * -> *) a.
MonadIO m =>
MutArray a -> Int -> m (MutArray a)
uninit MutArray a
arr Int
count
    forall (m :: * -> *) a. Monad m => a -> m a
return (Ring
        { ringArr :: MutArray a
ringArr = MutArray a
arr1
        , ringHead :: Int
ringHead = Int
0
        , ringMax :: Int
ringMax = Int
count
        })


{-# INLINE writeLastN #-}
writeLastN :: MonadIO m => Int -> Fold m a (Ring a)
writeLastN :: forall (m :: * -> *) a. MonadIO m => Int -> Fold m a (Ring a)
writeLastN Int
n = forall (m :: * -> *) a b s.
(s -> a -> m (Step s b))
-> m (Step s b) -> (s -> m b) -> Fold m a b
Fold forall {m :: * -> *} {b} {a} {b}.
(MonadIO m, Num b) =>
Tuple' (Ring a) b -> a -> m (Step (Tuple' (Ring a) b) b)
step forall {a} {a}. m (Step (Tuple' (Ring a) Int) (Ring a))
initial forall {m :: * -> *} {a}.
Monad m =>
Tuple' (Ring a) Int -> m (Ring a)
extract

    where

    initial :: m (Step (Tuple' (Ring a) Int) (Ring a))
initial = do
        if Int
n forall a. Ord a => a -> a -> Bool
<= Int
0
        then forall s b. b -> Step s b
Fold.Done forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a. MonadIO m => Int -> m (Ring a)
createRing Int
0
        else do
            Ring a
rb <- forall (m :: * -> *) a. MonadIO m => Int -> m (Ring a)
createRing Int
n
            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s b. s -> Step s b
Fold.Partial forall a b. (a -> b) -> a -> b
$ forall a b. a -> b -> Tuple' a b
Tuple' Ring a
rb (Int
0 :: Int)

    step :: Tuple' (Ring a) b -> a -> m (Step (Tuple' (Ring a) b) b)
step (Tuple' Ring a
rb b
cnt) a
x = do
        Int
rh1 <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Ring a -> a -> IO Int
unsafeInsertRingWith Ring a
rb a
x
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s b. s -> Step s b
Fold.Partial forall a b. (a -> b) -> a -> b
$ forall a b. a -> b -> Tuple' a b
Tuple' (Ring a
rb {ringHead :: Int
ringHead = Int
rh1}) (b
cnt forall a. Num a => a -> a -> a
+ b
1)

    extract :: Tuple' (Ring a) Int -> m (Ring a)
extract (Tuple' rb :: Ring a
rb@Ring{Int
MutArray a
ringMax :: Int
ringHead :: Int
ringArr :: MutArray a
ringMax :: forall a. Ring a -> Int
ringHead :: forall a. Ring a -> Int
ringArr :: forall a. Ring a -> MutArray a
..} Int
cnt) =
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$
            if Int
cnt forall a. Ord a => a -> a -> Bool
< Int
ringMax
            then forall a. MutArray a -> Int -> Int -> Ring a
Ring MutArray a
ringArr Int
0 Int
ringHead
            else Ring a
rb

-------------------------------------------------------------------------------
-- Modification
-------------------------------------------------------------------------------

-- XXX This is safe
-- Take the ring head and return the new ring head.
{-# INLINE unsafeInsertRingWith #-}
unsafeInsertRingWith :: Ring a -> a -> IO Int
unsafeInsertRingWith :: forall a. Ring a -> a -> IO Int
unsafeInsertRingWith Ring{Int
MutArray a
ringMax :: Int
ringHead :: Int
ringArr :: MutArray a
ringMax :: forall a. Ring a -> Int
ringHead :: forall a. Ring a -> Int
ringArr :: forall a. Ring a -> MutArray a
..} a
x = do
    assertM(Int
ringMax forall a. Ord a => a -> a -> Bool
>= Int
1)
    assertM(Int
ringHead forall a. Ord a => a -> a -> Bool
< Int
ringMax)
    forall (m :: * -> *) a. MonadIO m => Int -> MutArray a -> a -> m ()
putIndexUnsafe Int
ringHead MutArray a
ringArr a
x
    let rh1 :: Int
rh1 = Int
ringHead forall a. Num a => a -> a -> a
+ Int
1
        next :: Int
next = if Int
rh1 forall a. Eq a => a -> a -> Bool
== Int
ringMax then Int
0 else Int
rh1
    forall (m :: * -> *) a. Monad m => a -> m a
return Int
next

-- | Move the ring head clockwise (+ve adj) or counter clockwise (-ve adj) by
-- the given amount.
{-# INLINE seek #-}
seek :: MonadIO m => Int -> Ring a -> m (Ring a)
seek :: forall (m :: * -> *) a. MonadIO m => Int -> Ring a -> m (Ring a)
seek Int
adj rng :: Ring a
rng@Ring{Int
MutArray a
ringMax :: Int
ringHead :: Int
ringArr :: MutArray a
ringMax :: forall a. Ring a -> Int
ringHead :: forall a. Ring a -> Int
ringArr :: forall a. Ring a -> MutArray a
..}
    | Int
ringMax forall a. Ord a => a -> a -> Bool
> Int
0 = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
        -- XXX try avoiding mod when in bounds
        let idx1 :: Int
idx1 = Int
ringHead forall a. Num a => a -> a -> a
+ Int
adj
            next :: Int
next = forall a. Integral a => a -> a -> a
mod Int
idx1 Int
ringMax
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. MutArray a -> Int -> Int -> Ring a
Ring MutArray a
ringArr Int
next Int
ringMax
    | Bool
otherwise = forall (f :: * -> *) a. Applicative f => a -> f a
pure Ring a
rng

-------------------------------------------------------------------------------
-- Conversion
-------------------------------------------------------------------------------

-- | @toMutArray rignHeadAdjustment lengthToRead ring@.
-- Convert the ring into a boxed mutable array. Note that the returned MutArray
-- may share the same underlying memory as the Ring.
{-# INLINE toMutArray #-}
toMutArray :: MonadIO m => Int -> Int -> Ring a -> m (MutArray a)
toMutArray :: forall (m :: * -> *) a.
MonadIO m =>
Int -> Int -> Ring a -> m (MutArray a)
toMutArray Int
adj Int
n Ring{Int
MutArray a
ringMax :: Int
ringHead :: Int
ringArr :: MutArray a
ringMax :: forall a. Ring a -> Int
ringHead :: forall a. Ring a -> Int
ringArr :: forall a. Ring a -> MutArray a
..} = do
    let len :: Int
len = forall a. Ord a => a -> a -> a
min Int
ringMax Int
n
    let idx :: Int
idx = forall a. Integral a => a -> a -> a
mod (Int
ringHead forall a. Num a => a -> a -> a
+ Int
adj) Int
ringMax
        end :: Int
end = Int
idx forall a. Num a => a -> a -> a
+ Int
len
    if Int
end forall a. Ord a => a -> a -> Bool
<= Int
ringMax
    then
        -- putSliceUnsafe ringArr idx arr1 0 len
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ MutArray a
ringArr { arrStart :: Int
arrStart = Int
idx, arrLen :: Int
arrLen = Int
len }
    else do
        -- XXX Just swap the elements in the existing ring and return the
        -- same array without reallocation.
        MutArray a
arr <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. MonadIO m => Int -> m (MutArray a)
new Int
len
        MutArray a
arr1 <- forall (m :: * -> *) a.
MonadIO m =>
MutArray a -> Int -> m (MutArray a)
uninit MutArray a
arr Int
len
        forall (m :: * -> *) a.
MonadIO m =>
MutArray a -> Int -> MutArray a -> Int -> Int -> m ()
putSliceUnsafe MutArray a
ringArr Int
idx MutArray a
arr1 Int
0 (Int
ringMax forall a. Num a => a -> a -> a
- Int
idx)
        forall (m :: * -> *) a.
MonadIO m =>
MutArray a -> Int -> MutArray a -> Int -> Int -> m ()
putSliceUnsafe MutArray a
ringArr Int
0 MutArray a
arr1 (Int
ringMax forall a. Num a => a -> a -> a
- Int
idx) (Int
end forall a. Num a => a -> a -> a
- Int
ringMax)
        forall (m :: * -> *) a. Monad m => a -> m a
return MutArray a
arr1

-- This would be theoretically slower than toMutArray because of a branch
-- introduced for each element in the second half of the ring.

-- | Seek by n and then read the entire ring. Use 'take' on the stream to
-- restrict the reads.
toStreamWith :: Int -> Ring a -> Stream m a
toStreamWith :: forall a (m :: * -> *). Int -> Ring a -> Stream m a
toStreamWith = forall a. (?callStack::CallStack) => a
undefined
{-
toStreamWith n Ring{..}
    | ringMax > 0 = concatEffect $ liftIO $ do
        idx <- readIORef ringHead
        let idx1 = idx + adj
            next = mod idx1 ringMax
            s1 = undefined  -- stream initial slice
            s2 = undefined  -- stream next slice
        return (s1 `Stream.append` s2)
    | otherwise = Stream.nil
-}