module DataSketches.Quantiles.RelativeErrorQuantile.Internal.Auxiliary
( ReqAuxiliary(..)
, MReqAuxiliary (..)
, mkAuxiliary
, getQuantile
, mergeSortIn
) where
import GHC.TypeLits
import Control.Monad (when)
import Control.Monad.Primitive
import Data.Bits (shiftL)
import Data.Word
import Data.Primitive.MutVar
import Data.Vector.Algorithms.Search
import qualified Data.Vector as Vector
import qualified Data.Vector.Unboxed.Mutable as MUVector
import DataSketches.Quantiles.RelativeErrorQuantile.Types
import DataSketches.Quantiles.RelativeErrorQuantile.Internal.Compactor (ReqCompactor)
import qualified DataSketches.Quantiles.RelativeErrorQuantile.Internal.Compactor as Compactor
import DataSketches.Quantiles.RelativeErrorQuantile.Internal.DoubleBuffer (DoubleBuffer)
import qualified DataSketches.Quantiles.RelativeErrorQuantile.Internal.DoubleBuffer as DoubleBuffer
import qualified Data.Vector.Unboxed as U
import Control.Monad.ST
import DataSketches.Quantiles.RelativeErrorQuantile.Internal.InequalitySearch (find)
import qualified DataSketches.Quantiles.RelativeErrorQuantile.Internal.InequalitySearch as IS
import Debug.Trace
import qualified Data.Vector.Generic.Mutable as MG
data ReqAuxiliary = ReqAuxiliary
{ ReqAuxiliary -> Vector (Double, Word64)
raWeightedItems :: {-# UNPACK #-} !(U.Vector (Double, Word64))
, ReqAuxiliary -> RankAccuracy
raHighRankAccuracy :: !RankAccuracy
, ReqAuxiliary -> Word64
raSize :: {-# UNPACK #-} !Word64
}
deriving (Int -> ReqAuxiliary -> ShowS
[ReqAuxiliary] -> ShowS
ReqAuxiliary -> String
(Int -> ReqAuxiliary -> ShowS)
-> (ReqAuxiliary -> String)
-> ([ReqAuxiliary] -> ShowS)
-> Show ReqAuxiliary
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ReqAuxiliary] -> ShowS
$cshowList :: [ReqAuxiliary] -> ShowS
show :: ReqAuxiliary -> String
$cshow :: ReqAuxiliary -> String
showsPrec :: Int -> ReqAuxiliary -> ShowS
$cshowsPrec :: Int -> ReqAuxiliary -> ShowS
Show, ReqAuxiliary -> ReqAuxiliary -> Bool
(ReqAuxiliary -> ReqAuxiliary -> Bool)
-> (ReqAuxiliary -> ReqAuxiliary -> Bool) -> Eq ReqAuxiliary
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: ReqAuxiliary -> ReqAuxiliary -> Bool
$c/= :: ReqAuxiliary -> ReqAuxiliary -> Bool
== :: ReqAuxiliary -> ReqAuxiliary -> Bool
$c== :: ReqAuxiliary -> ReqAuxiliary -> Bool
Eq)
data MReqAuxiliary s = MReqAuxiliary
{ MReqAuxiliary s -> MutVar s (MVector s (Double, Word64))
mraWeightedItems :: {-# UNPACK #-} !(MutVar s (MUVector.MVector s (Double, Word64)))
, MReqAuxiliary s -> RankAccuracy
mraHighRankAccuracy :: !RankAccuracy
, MReqAuxiliary s -> Word64
mraSize :: {-# UNPACK #-} !Word64
}
mkAuxiliary :: (PrimMonad m, s ~ PrimState m) => RankAccuracy -> Word64 -> Int -> Vector.Vector (ReqCompactor s) -> m ReqAuxiliary
mkAuxiliary :: RankAccuracy
-> Word64 -> Int -> Vector (ReqCompactor s) -> m ReqAuxiliary
mkAuxiliary RankAccuracy
rankAccuracy Word64
totalN Int
retainedItems Vector (ReqCompactor s)
compactors = do
MutVar s (MVector s (Double, Word64))
items <- MVector s (Double, Word64)
-> m (MutVar s (MVector s (Double, Word64)))
forall (m :: * -> *) a.
PrimMonad m =>
a -> m (MutVar (PrimState m) a)
newMutVar (MVector s (Double, Word64)
-> m (MutVar s (MVector s (Double, Word64))))
-> m (MVector s (Double, Word64))
-> m (MutVar s (MVector s (Double, Word64)))
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Int
-> (Double, Word64) -> m (MVector (PrimState m) (Double, Word64))
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
MUVector.replicate Int
retainedItems (Double
0, Word64
0)
let this :: MReqAuxiliary s
this = MReqAuxiliary :: forall s.
MutVar s (MVector s (Double, Word64))
-> RankAccuracy -> Word64 -> MReqAuxiliary s
MReqAuxiliary
{ mraWeightedItems :: MutVar s (MVector s (Double, Word64))
mraWeightedItems = MutVar s (MVector s (Double, Word64))
items
, mraHighRankAccuracy :: RankAccuracy
mraHighRankAccuracy = RankAccuracy
rankAccuracy
, mraSize :: Word64
mraSize = Word64
totalN
}
(Int -> ReqCompactor s -> m Int)
-> Int -> Vector (ReqCompactor s) -> m ()
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> m a) -> a -> Vector b -> m ()
Vector.foldM_ (MReqAuxiliary (PrimState m)
-> Int -> ReqCompactor (PrimState m) -> m Int
forall (m :: * -> *).
PrimMonad m =>
MReqAuxiliary (PrimState m)
-> Int -> ReqCompactor (PrimState m) -> m Int
mergeBuffers MReqAuxiliary s
MReqAuxiliary (PrimState m)
this) Int
0 Vector (ReqCompactor s)
compactors
MReqAuxiliary (PrimState m) -> m ()
forall (m :: * -> *).
PrimMonad m =>
MReqAuxiliary (PrimState m) -> m ()
createCumulativeWeights MReqAuxiliary s
MReqAuxiliary (PrimState m)
this
MReqAuxiliary (PrimState m) -> m ()
forall (m :: * -> *).
PrimMonad m =>
MReqAuxiliary (PrimState m) -> m ()
dedup MReqAuxiliary s
MReqAuxiliary (PrimState m)
this
Vector (Double, Word64)
items' <- MVector s (Double, Word64) -> m (Vector (Double, Word64))
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze (MVector s (Double, Word64) -> m (Vector (Double, Word64)))
-> m (MVector s (Double, Word64)) -> m (Vector (Double, Word64))
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< MutVar (PrimState m) (MVector s (Double, Word64))
-> m (MVector s (Double, Word64))
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar MutVar s (MVector s (Double, Word64))
MutVar (PrimState m) (MVector s (Double, Word64))
items
ReqAuxiliary -> m ReqAuxiliary
forall (f :: * -> *) a. Applicative f => a -> f a
pure ReqAuxiliary :: Vector (Double, Word64) -> RankAccuracy -> Word64 -> ReqAuxiliary
ReqAuxiliary
{ raWeightedItems :: Vector (Double, Word64)
raWeightedItems = Vector (Double, Word64)
items'
, raHighRankAccuracy :: RankAccuracy
raHighRankAccuracy = RankAccuracy
rankAccuracy
, raSize :: Word64
raSize = Word64
totalN
}
where
mergeBuffers :: MReqAuxiliary (PrimState m)
-> Int -> ReqCompactor (PrimState m) -> m Int
mergeBuffers MReqAuxiliary (PrimState m)
this Int
auxCount ReqCompactor (PrimState m)
compactor = do
DoubleBuffer (PrimState m)
buff <- ReqCompactor (PrimState m) -> m (DoubleBuffer (PrimState m))
forall (m :: * -> *).
PrimMonad m =>
ReqCompactor (PrimState m) -> m (DoubleBuffer (PrimState m))
Compactor.getBuffer ReqCompactor (PrimState m)
compactor
Int
buffSize <- DoubleBuffer (PrimState m) -> m Int
forall (m :: * -> *).
PrimMonad m =>
DoubleBuffer (PrimState m) -> m Int
DoubleBuffer.getCount DoubleBuffer (PrimState m)
buff
let lgWeight :: Word8
lgWeight = ReqCompactor (PrimState m) -> Word8
forall s. ReqCompactor s -> Word8
Compactor.getLgWeight ReqCompactor (PrimState m)
compactor
weight :: Word64
weight = Word64
1 Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftL` Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
lgWeight
MReqAuxiliary (PrimState m)
-> DoubleBuffer (PrimState m) -> Word64 -> Int -> m ()
forall (m :: * -> *).
PrimMonad m =>
MReqAuxiliary (PrimState m)
-> DoubleBuffer (PrimState m) -> Word64 -> Int -> m ()
mergeSortIn MReqAuxiliary (PrimState m)
this DoubleBuffer (PrimState m)
buff Word64
weight Int
auxCount
Int -> m Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> m Int) -> Int -> m Int
forall a b. (a -> b) -> a -> b
$ Int
auxCount Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
buffSize
getWeightedItems :: PrimMonad m => MReqAuxiliary (PrimState m) -> m (MUVector.MVector (PrimState m) (Double, Word64))
getWeightedItems :: MReqAuxiliary (PrimState m)
-> m (MVector (PrimState m) (Double, Word64))
getWeightedItems = MutVar (PrimState m) (MVector (PrimState m) (Double, Word64))
-> m (MVector (PrimState m) (Double, Word64))
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar (MutVar (PrimState m) (MVector (PrimState m) (Double, Word64))
-> m (MVector (PrimState m) (Double, Word64)))
-> (MReqAuxiliary (PrimState m)
-> MutVar (PrimState m) (MVector (PrimState m) (Double, Word64)))
-> MReqAuxiliary (PrimState m)
-> m (MVector (PrimState m) (Double, Word64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MReqAuxiliary (PrimState m)
-> MutVar (PrimState m) (MVector (PrimState m) (Double, Word64))
forall s. MReqAuxiliary s -> MutVar s (MVector s (Double, Word64))
mraWeightedItems
getItems :: PrimMonad m => MReqAuxiliary (PrimState m) -> m (MUVector.MVector (PrimState m) Double)
getItems :: MReqAuxiliary (PrimState m) -> m (MVector (PrimState m) Double)
getItems = (MVector (PrimState m) (Double, Word64)
-> MVector (PrimState m) Double)
-> m (MVector (PrimState m) (Double, Word64))
-> m (MVector (PrimState m) Double)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((MVector (PrimState m) Double, MVector (PrimState m) Word64)
-> MVector (PrimState m) Double
forall a b. (a, b) -> a
fst ((MVector (PrimState m) Double, MVector (PrimState m) Word64)
-> MVector (PrimState m) Double)
-> (MVector (PrimState m) (Double, Word64)
-> (MVector (PrimState m) Double, MVector (PrimState m) Word64))
-> MVector (PrimState m) (Double, Word64)
-> MVector (PrimState m) Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MVector (PrimState m) (Double, Word64)
-> (MVector (PrimState m) Double, MVector (PrimState m) Word64)
forall a b s.
(Unbox a, Unbox b) =>
MVector s (a, b) -> (MVector s a, MVector s b)
MUVector.unzip) (m (MVector (PrimState m) (Double, Word64))
-> m (MVector (PrimState m) Double))
-> (MReqAuxiliary (PrimState m)
-> m (MVector (PrimState m) (Double, Word64)))
-> MReqAuxiliary (PrimState m)
-> m (MVector (PrimState m) Double)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MReqAuxiliary (PrimState m)
-> m (MVector (PrimState m) (Double, Word64))
forall (m :: * -> *).
PrimMonad m =>
MReqAuxiliary (PrimState m)
-> m (MVector (PrimState m) (Double, Word64))
getWeightedItems
getWeights :: PrimMonad m => MReqAuxiliary (PrimState m) -> m (MUVector.MVector (PrimState m) Word64)
getWeights :: MReqAuxiliary (PrimState m) -> m (MVector (PrimState m) Word64)
getWeights = (MVector (PrimState m) (Double, Word64)
-> MVector (PrimState m) Word64)
-> m (MVector (PrimState m) (Double, Word64))
-> m (MVector (PrimState m) Word64)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((MVector (PrimState m) Double, MVector (PrimState m) Word64)
-> MVector (PrimState m) Word64
forall a b. (a, b) -> b
snd ((MVector (PrimState m) Double, MVector (PrimState m) Word64)
-> MVector (PrimState m) Word64)
-> (MVector (PrimState m) (Double, Word64)
-> (MVector (PrimState m) Double, MVector (PrimState m) Word64))
-> MVector (PrimState m) (Double, Word64)
-> MVector (PrimState m) Word64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MVector (PrimState m) (Double, Word64)
-> (MVector (PrimState m) Double, MVector (PrimState m) Word64)
forall a b s.
(Unbox a, Unbox b) =>
MVector s (a, b) -> (MVector s a, MVector s b)
MUVector.unzip) (m (MVector (PrimState m) (Double, Word64))
-> m (MVector (PrimState m) Word64))
-> (MReqAuxiliary (PrimState m)
-> m (MVector (PrimState m) (Double, Word64)))
-> MReqAuxiliary (PrimState m)
-> m (MVector (PrimState m) Word64)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MReqAuxiliary (PrimState m)
-> m (MVector (PrimState m) (Double, Word64))
forall (m :: * -> *).
PrimMonad m =>
MReqAuxiliary (PrimState m)
-> m (MVector (PrimState m) (Double, Word64))
getWeightedItems
getQuantile :: ReqAuxiliary -> Double -> Criterion -> Double
getQuantile :: ReqAuxiliary -> Double -> Criterion -> Double
getQuantile ReqAuxiliary
this Double
normalRank Criterion
ltEq = (Double, Word64) -> Double
forall a b. (a, b) -> a
fst (Vector (Double, Word64)
weightedItems Vector (Double, Word64) -> Int -> (Double, Word64)
forall a. Unbox a => Vector a -> Int -> a
U.! Int
ix)
where
ix :: Int
ix = if Int
searchResult Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Vector (Double, Word64) -> Int
forall a. Unbox a => Vector a -> Int
U.length Vector (Double, Word64)
weightedItems
then Int
searchResult Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
else Int
searchResult
searchResult :: Int
searchResult = (forall s. ST s Int) -> Int
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s Int) -> Int) -> (forall s. ST s Int) -> Int
forall a b. (a -> b) -> a -> b
$ do
MVector s Word64
v <- Vector Word64 -> ST s (MVector (PrimState (ST s)) Word64)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
U.unsafeThaw (Vector Word64 -> ST s (MVector (PrimState (ST s)) Word64))
-> Vector Word64 -> ST s (MVector (PrimState (ST s)) Word64)
forall a b. (a -> b) -> a -> b
$ (Vector Double, Vector Word64) -> Vector Word64
forall a b. (a, b) -> b
snd ((Vector Double, Vector Word64) -> Vector Word64)
-> (Vector Double, Vector Word64) -> Vector Word64
forall a b. (a -> b) -> a -> b
$ Vector (Double, Word64) -> (Vector Double, Vector Word64)
forall a b.
(Unbox a, Unbox b) =>
Vector (a, b) -> (Vector a, Vector b)
U.unzip Vector (Double, Word64)
weightedItems
let search :: MVector s Word64 -> Int -> Int -> Word64 -> ST s Int
search = case Criterion
ltEq of
Criterion
(:<) -> (:>)
-> MVector (PrimState (ST s)) Word64
-> Int
-> Int
-> Word64
-> ST s Int
forall s (m :: * -> *) (v :: * -> * -> *) a.
(InequalitySearch s, PrimMonad m, MVector v a, Ord a) =>
s -> v (PrimState m) a -> Int -> Int -> a -> m Int
find (:>)
(IS.:>)
Criterion
(:<=) -> (:>=)
-> MVector (PrimState (ST s)) Word64
-> Int
-> Int
-> Word64
-> ST s Int
forall s (m :: * -> *) (v :: * -> * -> *) a.
(InequalitySearch s, PrimMonad m, MVector v a, Ord a) =>
s -> v (PrimState m) a -> Int -> Int -> a -> m Int
find (:>=)
(IS.:>=)
MVector s Word64 -> Int -> Int -> Word64 -> ST s Int
search MVector s Word64
v Int
0 (Int
weightsSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Word64
rank
weightedItems :: Vector (Double, Word64)
weightedItems = ReqAuxiliary -> Vector (Double, Word64)
raWeightedItems ReqAuxiliary
this
weightsSize :: Int
weightsSize = Vector (Double, Word64) -> Int
forall a. Unbox a => Vector a -> Int
U.length Vector (Double, Word64)
weightedItems
rank :: Word64
rank = Double -> Word64
forall a b. (RealFrac a, Integral b) => a -> b
floor (Double
normalRank Double -> Double -> Double
forall a. Num a => a -> a -> a
* Word64 -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ReqAuxiliary -> Word64
raSize ReqAuxiliary
this))
createCumulativeWeights :: PrimMonad m => MReqAuxiliary (PrimState m) -> m ()
createCumulativeWeights :: MReqAuxiliary (PrimState m) -> m ()
createCumulativeWeights MReqAuxiliary (PrimState m)
this = do
MVector (PrimState m) Word64
weights <- MReqAuxiliary (PrimState m) -> m (MVector (PrimState m) Word64)
forall (m :: * -> *).
PrimMonad m =>
MReqAuxiliary (PrimState m) -> m (MVector (PrimState m) Word64)
getWeights MReqAuxiliary (PrimState m)
this
let size :: Int
size = MVector (PrimState m) Word64 -> Int
forall a s. Unbox a => MVector s a -> Int
MUVector.length MVector (PrimState m) Word64
weights
let accumulateM :: Int -> Word64 -> m ()
accumulateM Int
i Word64
weight = do
Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
Word64
prevWeight <- MVector (PrimState m) Word64 -> Int -> m Word64
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MUVector.read MVector (PrimState m) Word64
weights (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
MVector (PrimState m) Word64 -> Int -> Word64 -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MUVector.unsafeWrite MVector (PrimState m) Word64
weights Int
i (Word64
weight Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
prevWeight)
MVector (PrimState m) Word64 -> (Int -> m ()) -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a b.
(Monad m, MVector v a) =>
v (PrimState m) a -> (Int -> m b) -> m ()
forI_ MVector (PrimState m) Word64
weights (\Int
i -> MVector (PrimState m) Word64 -> Int -> m Word64
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MUVector.read MVector (PrimState m) Word64
weights Int
i m Word64 -> (Word64 -> m ()) -> m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Word64
x -> Int -> Word64 -> m ()
accumulateM Int
i Word64
x)
Word64
lastWeight <- MVector (PrimState m) Word64 -> Int -> m Word64
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MUVector.read MVector (PrimState m) Word64
weights (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Word64
lastWeight Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
/= MReqAuxiliary (PrimState m) -> Word64
forall s. MReqAuxiliary s -> Word64
mraSize MReqAuxiliary (PrimState m)
this) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
String -> m ()
forall a. HasCallStack => String -> a
error String
"invariant violated: lastWeight does not equal raSize"
where
forI_ :: (Monad m, MG.MVector v a) => v (PrimState m) a -> (Int -> m b) -> m ()
{-# INLINE forI_ #-}
forI_ :: v (PrimState m) a -> (Int -> m b) -> m ()
forI_ v (PrimState m) a
v Int -> m b
f = Int -> m ()
loop Int
0
where
loop :: Int -> m ()
loop Int
i
| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
| Bool
otherwise = Int -> m b
f Int
i m b -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> m ()
loop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
n :: Int
n = v (PrimState m) a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
MG.length v (PrimState m) a
v
dedup :: PrimMonad m => MReqAuxiliary (PrimState m) -> m ()
dedup :: MReqAuxiliary (PrimState m) -> m ()
dedup MReqAuxiliary (PrimState m)
this = do
MVector (PrimState m) (Double, Word64)
weightedItems <- MReqAuxiliary (PrimState m)
-> m (MVector (PrimState m) (Double, Word64))
forall (m :: * -> *).
PrimMonad m =>
MReqAuxiliary (PrimState m)
-> m (MVector (PrimState m) (Double, Word64))
getWeightedItems MReqAuxiliary (PrimState m)
this
let size :: Int
size = MVector (PrimState m) (Double, Word64) -> Int
forall a s. Unbox a => MVector s a -> Int
MUVector.length MVector (PrimState m) (Double, Word64)
weightedItems
MVector (PrimState m) (Double, Word64)
weightedItemsB <- Int
-> (Double, Word64) -> m (MVector (PrimState m) (Double, Word64))
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
MUVector.replicate Int
size (Double
0, Word64
0)
Int
bi <- MVector (PrimState m) (Double, Word64)
-> Int
-> MVector (PrimState m) (Double, Word64)
-> Int
-> Int
-> m Int
forall (f :: * -> *) a b.
(PrimMonad f, Unbox a, Unbox b, Eq a) =>
MVector (PrimState f) (a, b)
-> Int -> MVector (PrimState f) (a, b) -> Int -> Int -> f Int
doDedup MVector (PrimState m) (Double, Word64)
weightedItems Int
size MVector (PrimState m) (Double, Word64)
weightedItemsB Int
0 Int
0
MutVar (PrimState m) (MVector (PrimState m) (Double, Word64))
-> MVector (PrimState m) (Double, Word64) -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> a -> m ()
writeMutVar (MReqAuxiliary (PrimState m)
-> MutVar (PrimState m) (MVector (PrimState m) (Double, Word64))
forall s. MReqAuxiliary s -> MutVar s (MVector s (Double, Word64))
mraWeightedItems MReqAuxiliary (PrimState m)
this) (MVector (PrimState m) (Double, Word64) -> m ())
-> MVector (PrimState m) (Double, Word64) -> m ()
forall a b. (a -> b) -> a -> b
$ Int
-> Int
-> MVector (PrimState m) (Double, Word64)
-> MVector (PrimState m) (Double, Word64)
forall a s. Unbox a => Int -> Int -> MVector s a -> MVector s a
MUVector.slice Int
0 Int
bi MVector (PrimState m) (Double, Word64)
weightedItemsB
where
doDedup :: MVector (PrimState f) (a, b)
-> Int -> MVector (PrimState f) (a, b) -> Int -> Int -> f Int
doDedup MVector (PrimState f) (a, b)
weightedItems Int
itemsSize MVector (PrimState f) (a, b)
weightedItemsB = Int -> Int -> f Int
go
where
go :: Int -> Int -> f Int
go !Int
i !Int
bi
| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
itemsSize = Int -> f Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
bi
| Bool
otherwise = do
let j :: Int
j = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
hidup :: Int
hidup = Int
j
countDups :: Int -> Int -> f (Int, Int)
countDups !Int
j !Int
hidup = if Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
itemsSize
then do
(a
itemI, b
_) <- MVector (PrimState f) (a, b) -> Int -> f (a, b)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MUVector.read MVector (PrimState f) (a, b)
weightedItems Int
i
(a
itemJ, b
_) <- MVector (PrimState f) (a, b) -> Int -> f (a, b)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MUVector.read MVector (PrimState f) (a, b)
weightedItems Int
j
if a
itemI a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
itemJ
then Int -> Int -> f (Int, Int)
countDups (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
j
else (Int, Int) -> f (Int, Int)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
j, Int
hidup)
else (Int, Int) -> f (Int, Int)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
j, Int
hidup)
(Int
j', Int
hidup') <- Int -> Int -> f (Int, Int)
countDups Int
j Int
hidup
if Int
j' Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1
then do
(a
item, b
weight) <- MVector (PrimState f) (a, b) -> Int -> f (a, b)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MUVector.read MVector (PrimState f) (a, b)
weightedItems Int
i
MVector (PrimState f) (a, b) -> Int -> (a, b) -> f ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MUVector.unsafeWrite MVector (PrimState f) (a, b)
weightedItemsB Int
bi (a
item, b
weight)
Int -> Int -> f Int
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
bi Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
else do
(a
item, b
weight) <- MVector (PrimState f) (a, b) -> Int -> f (a, b)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MUVector.read MVector (PrimState f) (a, b)
weightedItems Int
hidup'
MVector (PrimState f) (a, b) -> Int -> (a, b) -> f ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MUVector.unsafeWrite MVector (PrimState f) (a, b)
weightedItemsB Int
bi (a
item, b
weight)
Int -> Int -> f Int
go Int
j' (Int
bi Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
mergeSortIn :: PrimMonad m => MReqAuxiliary (PrimState m) -> DoubleBuffer (PrimState m) -> Word64 -> Int -> m ()
mergeSortIn :: MReqAuxiliary (PrimState m)
-> DoubleBuffer (PrimState m) -> Word64 -> Int -> m ()
mergeSortIn MReqAuxiliary (PrimState m)
this DoubleBuffer (PrimState m)
bufIn Word64
defaultWeight Int
auxCount = do
DoubleBuffer (PrimState m) -> m ()
forall (m :: * -> *).
PrimMonad m =>
DoubleBuffer (PrimState m) -> m ()
DoubleBuffer.sort DoubleBuffer (PrimState m)
bufIn
MVector (PrimState m) (Double, Word64)
weightedItems <- MReqAuxiliary (PrimState m)
-> m (MVector (PrimState m) (Double, Word64))
forall (m :: * -> *).
PrimMonad m =>
MReqAuxiliary (PrimState m)
-> m (MVector (PrimState m) (Double, Word64))
getWeightedItems MReqAuxiliary (PrimState m)
this
MVector (PrimState m) Double
otherItems <- DoubleBuffer (PrimState m) -> m (MVector (PrimState m) Double)
forall (m :: * -> *) s.
(PrimMonad m, PrimState m ~ s) =>
DoubleBuffer s -> m (MVector s Double)
DoubleBuffer.getVector DoubleBuffer (PrimState m)
bufIn
Int
otherBuffSize <- DoubleBuffer (PrimState m) -> m Int
forall (m :: * -> *).
PrimMonad m =>
DoubleBuffer (PrimState m) -> m Int
DoubleBuffer.getCount DoubleBuffer (PrimState m)
bufIn
Int
otherBuffCapacity <- DoubleBuffer (PrimState m) -> m Int
forall (m :: * -> *).
PrimMonad m =>
DoubleBuffer (PrimState m) -> m Int
DoubleBuffer.getCapacity DoubleBuffer (PrimState m)
bufIn
let totalSize :: Int
totalSize = Int
otherBuffSize Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
auxCount Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
height :: Int
height = case MReqAuxiliary (PrimState m) -> RankAccuracy
forall s. MReqAuxiliary s -> RankAccuracy
mraHighRankAccuracy MReqAuxiliary (PrimState m)
this of
RankAccuracy
HighRanksAreAccurate -> Int
otherBuffCapacity Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
RankAccuracy
LowRanksAreAccurate -> Int
otherBuffSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
Int
-> MVector (PrimState m) (Double, Word64)
-> MVector (PrimState m) Double
-> Int
-> Int
-> Int
-> m ()
merge Int
totalSize MVector (PrimState m) (Double, Word64)
weightedItems MVector (PrimState m) Double
otherItems (Int
auxCount Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Int
otherBuffSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int
height
where
merge :: Int
-> MVector (PrimState m) (Double, Word64)
-> MVector (PrimState m) Double
-> Int
-> Int
-> Int
-> m ()
merge Int
totalSize MVector (PrimState m) (Double, Word64)
weightedItems MVector (PrimState m) Double
otherItems = Int -> Int -> Int -> Int -> m ()
go Int
totalSize
where
go :: Int -> Int -> Int -> Int -> m ()
go !Int
k !Int
i !Int
j !Int
h
| Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = () -> m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 = do
(Double
item, Word64
weight) <- MVector (PrimState m) (Double, Word64) -> Int -> m (Double, Word64)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MUVector.read MVector (PrimState m) (Double, Word64)
weightedItems Int
i
Double
otherItem <- MVector (PrimState m) Double -> Int -> m Double
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MUVector.read MVector (PrimState m) Double
otherItems Int
h
if Double
item Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
otherItem
then do
MVector (PrimState m) (Double, Word64)
-> Int -> (Double, Word64) -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MUVector.unsafeWrite MVector (PrimState m) (Double, Word64)
weightedItems Int
k (Double
item, Word64
weight)
Int -> Int -> Int -> m ()
continue (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int
j Int
h
else do
MVector (PrimState m) (Double, Word64)
-> Int -> (Double, Word64) -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MUVector.unsafeWrite MVector (PrimState m) (Double, Word64)
weightedItems Int
k (Double
otherItem, Word64
defaultWeight)
Int -> Int -> Int -> m ()
continue Int
i (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Int
h Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 = do
MVector (PrimState m) (Double, Word64) -> Int -> m (Double, Word64)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MUVector.read MVector (PrimState m) (Double, Word64)
weightedItems Int
i m (Double, Word64) -> ((Double, Word64) -> m ()) -> m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= MVector (PrimState m) (Double, Word64)
-> Int -> (Double, Word64) -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MUVector.write MVector (PrimState m) (Double, Word64)
weightedItems Int
k
Int -> Int -> Int -> m ()
continue (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int
j Int
h
| Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 = do
Double
otherItem <- MVector (PrimState m) Double -> Int -> m Double
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MUVector.read MVector (PrimState m) Double
otherItems Int
h
MVector (PrimState m) (Double, Word64)
-> Int -> (Double, Word64) -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MUVector.unsafeWrite MVector (PrimState m) (Double, Word64)
weightedItems Int
k (Double
otherItem, Word64
defaultWeight)
Int -> Int -> Int -> m ()
continue Int
i (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Int
h Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
| Bool
otherwise = () -> m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
where
continue :: Int -> Int -> Int -> m ()
continue = Int -> Int -> Int -> Int -> m ()
go (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)