module Data.Histogram.ST (
MHistogram
, newMHistogram
, fillOne
, fillOneW
, fillMonoid
, unsafeFreezeHist
, freezeHist
) where
import Control.Monad.Primitive
import Data.Monoid
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as M
import Data.Histogram.Generic
data MHistogram s v bin a =
MHistogram
!Int
!bin
!(v s a)
newMHistogram :: (PrimMonad m, Bin bin, M.MVector v a) => a -> bin -> m (MHistogram (PrimState m) v bin a)
newMHistogram zero bin = do
let n = nBins bin
a <- M.replicate (n + 2) zero
return $ MHistogram n bin a
fill :: (PrimMonad m, M.MVector v a, Bin bin) => MHistogram (PrimState m) v bin a -> BinValue bin -> (a -> a) -> m ()
fill (MHistogram n bin arr) !x f
| i < 0 = M.unsafeWrite arr n . f =<< M.unsafeRead arr n
| i >= n = M.unsafeWrite arr (n+1) . f =<< M.unsafeRead arr (n+1)
| otherwise = M.unsafeWrite arr i . f =<< M.unsafeRead arr i
where
i = toIndex bin x
fillOne :: (PrimMonad m, Num a, M.MVector v a, Bin bin) => MHistogram (PrimState m) v bin a -> BinValue bin -> m ()
fillOne h !x = fill h x (+1)
fillOneW :: (PrimMonad m, Num a, M.MVector v a, Bin bin) => MHistogram (PrimState m) v bin a -> (BinValue bin, a) -> m ()
fillOneW h (!x,!w) = fill h x (+w)
fillMonoid :: (PrimMonad m, Monoid a, M.MVector v a, Bin bin) => MHistogram (PrimState m) v bin a -> (BinValue bin, a) -> m ()
fillMonoid h (!x,!m) = fill h x (`mappend` m)
unsafeFreezeHist :: (PrimMonad m, G.Vector v a, Bin bin)
=> MHistogram (PrimState m) (G.Mutable v) bin a
-> m (Histogram v bin a)
unsafeFreezeHist (MHistogram n bin arr) = do
u <- M.unsafeRead arr n
o <- M.unsafeRead arr (n+1)
a <- G.unsafeFreeze $ M.slice 0 n arr
return $ histogramUO bin (Just (u,o)) a
freezeHist :: (PrimMonad m, G.Vector v a, Bin bin)
=> MHistogram (PrimState m) (G.Mutable v) bin a
-> m (Histogram v bin a)
freezeHist (MHistogram n bin arr) = do
u <- M.unsafeRead arr n
o <- M.unsafeRead arr (n+1)
a <- G.freeze $ M.slice 0 n arr
return $ histogramUO bin (Just (u,o)) a