module DataSketches.Quantiles.RelativeErrorQuantile.Internal.Auxiliary
  ( ReqAuxiliary(..)
  , MReqAuxiliary (..)
  , mkAuxiliary
  , getQuantile
  -- | Really extra private, just needed for tests
  , 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

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
  MVector (PrimState m) Word64 -> (Int -> Word64 -> m ()) -> m ()
forall (m :: * -> *) a b.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> (Int -> a -> m b) -> m ()
MUVector.iforM_ MVector (PrimState m) Word64
weights ((Int -> Word64 -> m ()) -> m ())
-> (Int -> Word64 -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \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)
  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"

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 -- no dups
              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)