module Data.Metrics.Reservoir.Uniform (
UniformReservoir,
reservoir,
unsafeReservoir,
clear,
unsafeClear,
size,
snapshot,
update,
unsafeUpdate
) where
import Control.Lens
import Control.Lens.TH
import Control.Monad.ST
import Data.Metrics.Internal
import Data.Time.Clock
import qualified Data.Metrics.Reservoir as R
import qualified Data.Metrics.Snapshot as S
import Data.Primitive.MutVar
import System.Random.MWC
import qualified Data.Vector.Unboxed as I
import qualified Data.Vector.Unboxed.Mutable as V
data UniformReservoir = UniformReservoir
{ uniformReservoirCount :: !Int
, uniformReservoirInnerReservoir :: !(I.Vector Double)
, uniformReservoirSeed :: !Seed
}
makeFields ''UniformReservoir
reservoir :: Seed
-> Int
-> R.Reservoir
reservoir g r = R.Reservoir
{ R.reservoirClear = clear
, R.reservoirSize = size
, R.reservoirSnapshot = snapshot
, R.reservoirUpdate = update
, R.reservoirState = UniformReservoir 0 (I.replicate r 0) g
}
unsafeReservoir :: Seed -> Int -> R.Reservoir
unsafeReservoir g r = R.Reservoir
{ R.reservoirClear = unsafeClear
, R.reservoirSize = size
, R.reservoirSnapshot = snapshot
, R.reservoirUpdate = unsafeUpdate
, R.reservoirState = UniformReservoir 0 (I.replicate r 0) g
}
clear :: NominalDiffTime -> UniformReservoir -> UniformReservoir
clear = go
where
go _ c = c & count .~ 0 & innerReservoir %~ newRes
newRes v = runST $ do
v' <- I.thaw v
V.set v' 0
I.unsafeFreeze v'
unsafeClear :: NominalDiffTime -> UniformReservoir -> UniformReservoir
unsafeClear = go
where
go _ c = c & count .~ 0 & innerReservoir %~ newRes
newRes v = runST $ do
v' <- I.unsafeThaw v
V.set v' 0
I.unsafeFreeze v'
size :: UniformReservoir -> Int
size = go
where
go c = min (c ^. count) (I.length $ c ^. innerReservoir)
snapshot :: UniformReservoir -> S.Snapshot
snapshot = go
where
go c = runST $ do
v' <- I.unsafeThaw $ c ^. innerReservoir
S.takeSnapshot $ V.slice 0 (size c) v'
update :: Double -> NominalDiffTime -> UniformReservoir -> UniformReservoir
update = go
where
go x _ c = c & count .~ newCount & innerReservoir .~ newRes & seed .~ newSeed
where
newCount = c ^. count . to succ
(newSeed, newRes) = runST $ do
v' <- I.thaw $ c ^. innerReservoir
g <- restore $ c ^. seed
if newCount <= V.length v'
then V.unsafeWrite v' (c ^. count) x
else do
i <- uniformR (0, newCount) g
if i < V.length v'
then V.unsafeWrite v' i x
else return ()
v'' <- I.unsafeFreeze v'
s <- save g
return (s, v'')
unsafeUpdate :: Double -> NominalDiffTime -> UniformReservoir -> UniformReservoir
unsafeUpdate = go
where
go x _ c = c & count .~ newCount & innerReservoir .~ newRes & seed .~ newSeed
where
newCount = c ^. count . to succ
(newSeed, newRes) = runST $ do
v' <- I.unsafeThaw $ c ^. innerReservoir
g <- restore (uniformReservoirSeed c)
if newCount <= V.length v'
then V.unsafeWrite v' (c ^. count) x
else do
i <- uniformR (0, newCount) g
if i < V.length v'
then V.unsafeWrite v' i x
else return ()
v'' <- I.unsafeFreeze v'
s <- save g
return (s, v'')