-----------------------------------------------------------------------------
-- |
-- Module      :  Instrument.Sampler
-- Copyright   :  Soostone Inc
-- License     :  BSD3
--
-- Maintainer  :  Ozgun Ataman
-- Stability   :  experimental
--
-- A container that can capture actual numeric values, efficiently and
-- concurrently buffering them in memory. We can impose a cap on how
-- many values are captured at a time.
----------------------------------------------------------------------------

module Instrument.Sampler
    ( Sampler (..)
    , new
    , sample
    , get
    , reset
    ) where

-------------------------------------------------------------------------------
import           Control.Concurrent.MVar
import           Control.Exception       (mask_, onException)
import           Control.Monad
import qualified Data.Vector.Mutable     as V
-------------------------------------------------------------------------------


-- |'BoundedChan' is an abstract data type representing a bounded channel.
data Buffer a = B {
       Buffer a -> Int
_size     :: Int
     , Buffer a -> IOVector a
_contents :: V.IOVector a
     , Buffer a -> MVar Int
_writePos :: MVar Int
     }

-- Versions of modifyMVar and withMVar that do not 'restore' the previous mask state when running
-- 'io', with added modification strictness.  The lack of 'restore' may make these perform better
-- than the normal version.  Moving strictness here makes using them more pleasant.
{-# INLINE modifyMVar_mask #-}
modifyMVar_mask :: MVar a -> (a -> IO (a,b)) -> IO b
modifyMVar_mask :: MVar a -> (a -> IO (a, b)) -> IO b
modifyMVar_mask MVar a
m a -> IO (a, b)
io =
  IO b -> IO b
forall a. IO a -> IO a
mask_ (IO b -> IO b) -> IO b -> IO b
forall a b. (a -> b) -> a -> b
$ do
    a
a <- MVar a -> IO a
forall a. MVar a -> IO a
takeMVar MVar a
m
    (a
a',b
b) <- a -> IO (a, b)
io a
a IO (a, b) -> IO () -> IO (a, b)
forall a b. IO a -> IO b -> IO a
`onException` MVar a -> a -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar a
m a
a
    MVar a -> a -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar a
m (a -> IO ()) -> a -> IO ()
forall a b. (a -> b) -> a -> b
$! a
a'
    b -> IO b
forall (m :: * -> *) a. Monad m => a -> m a
return b
b

{-# INLINE modifyMVar_mask_ #-}
modifyMVar_mask_ :: MVar a -> (a -> IO a) -> IO ()
modifyMVar_mask_ :: MVar a -> (a -> IO a) -> IO ()
modifyMVar_mask_ MVar a
m a -> IO a
io =
  IO () -> IO ()
forall a. IO a -> IO a
mask_ (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
    a
a <- MVar a -> IO a
forall a. MVar a -> IO a
takeMVar MVar a
m
    a
a' <- a -> IO a
io a
a IO a -> IO () -> IO a
forall a b. IO a -> IO b -> IO a
`onException` MVar a -> a -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar a
m a
a
    MVar a -> a -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar a
m (a -> IO ()) -> a -> IO ()
forall a b. (a -> b) -> a -> b
$! a
a'


-------------------------------------------------------------------------------
newBuffer :: Int -> IO (Buffer a)
newBuffer :: Int -> IO (Buffer a)
newBuffer Int
lim = do
  MVar Int
pos  <- Int -> IO (MVar Int)
forall a. a -> IO (MVar a)
newMVar Int
0
  IOVector a
entries <- Int -> IO (MVector (PrimState IO) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> m (MVector (PrimState m) a)
V.new Int
lim
  Buffer a -> IO (Buffer a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> IOVector a -> MVar Int -> Buffer a
forall a. Int -> IOVector a -> MVar Int -> Buffer a
B Int
lim IOVector a
entries MVar Int
pos)


-- | Write an element to the channel. If the channel is full, nothing
-- will happen and the function will return immediately. We don't want
-- to disturb production code.
writeBuffer :: Buffer a -> a -> IO ()
writeBuffer :: Buffer a -> a -> IO ()
writeBuffer (B Int
size IOVector a
contents MVar Int
wposMV) a
x = MVar Int -> (Int -> IO Int) -> IO ()
forall a. MVar a -> (a -> IO a) -> IO ()
modifyMVar_mask_ MVar Int
wposMV ((Int -> IO Int) -> IO ()) -> (Int -> IO Int) -> IO ()
forall a b. (a -> b) -> a -> b
$
  \Int
wpos ->
    case Int
wpos Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
size of
      Bool
True -> Int -> IO Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
wpos       -- buffer full, don't do anything
      Bool
False -> do
        MVector (PrimState IO) a -> Int -> a -> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
V.write IOVector a
MVector (PrimState IO) a
contents Int
wpos a
x
        Int -> IO Int
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> Int
forall a. Enum a => a -> a
succ Int
wpos)


-------------------------------------------------------------------------------
getBuffer :: Buffer a -> IO [a]
getBuffer :: Buffer a -> IO [a]
getBuffer (B Int
_size IOVector a
contents MVar Int
pos) = do
  Int
wpos <- MVar Int -> (Int -> IO (Int, Int)) -> IO Int
forall a b. MVar a -> (a -> IO (a, b)) -> IO b
modifyMVar_mask MVar Int
pos ((Int -> IO (Int, Int)) -> IO Int)
-> (Int -> IO (Int, Int)) -> IO Int
forall a b. (a -> b) -> a -> b
$ \ Int
wpos -> (Int, Int) -> IO (Int, Int)
forall (m :: * -> *) a. Monad m => a -> m a
return (Int
wpos, Int
wpos)
  [Int] -> (Int -> IO a) -> IO [a]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Int
0.. (Int
wpos Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)] ((Int -> IO a) -> IO [a]) -> (Int -> IO a) -> IO [a]
forall a b. (a -> b) -> a -> b
$ \ Int
i -> (MVector (PrimState IO) a -> Int -> IO a
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m a
V.read IOVector a
MVector (PrimState IO) a
contents Int
i)



-------------------------------------------------------------------------------
resetBuffer :: Buffer a -> IO ()
resetBuffer :: Buffer a -> IO ()
resetBuffer (B Int
_size IOVector a
_els MVar Int
pos) = MVar Int -> (Int -> IO Int) -> IO ()
forall a. MVar a -> (a -> IO a) -> IO ()
modifyMVar_mask_ MVar Int
pos (IO Int -> Int -> IO Int
forall a b. a -> b -> a
const (Int -> IO Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
0))




-- | An in-memory, bounded buffer for measurement samples.
newtype Sampler = S { Sampler -> Buffer Double
unS :: Buffer Double }


-- | Create a new, empty sampler
new :: Int -> IO Sampler
new :: Int -> IO Sampler
new Int
i = Buffer Double -> Sampler
S (Buffer Double -> Sampler) -> IO (Buffer Double) -> IO Sampler
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` Int -> IO (Buffer Double)
forall a. Int -> IO (Buffer a)
newBuffer Int
i


-------------------------------------------------------------------------------
sample :: Double -> Sampler -> IO ()
sample :: Double -> Sampler -> IO ()
sample Double
v Sampler
s = Buffer Double -> Double -> IO ()
forall a. Buffer a -> a -> IO ()
writeBuffer (Sampler -> Buffer Double
unS Sampler
s) Double
v


-------------------------------------------------------------------------------
get :: Sampler -> IO [Double]
get :: Sampler -> IO [Double]
get (S Buffer Double
buffer) = Buffer Double -> IO [Double]
forall a. Buffer a -> IO [a]
getBuffer Buffer Double
buffer


-------------------------------------------------------------------------------
reset :: Sampler -> IO ()
reset :: Sampler -> IO ()
reset (S Buffer Double
buf) = Buffer Double -> IO ()
forall a. Buffer a -> IO ()
resetBuffer Buffer Double
buf