{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE CPP                 #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_GHC -fexpose-all-unfoldings #-}

module KMerge.LoserTree (
    MutableLoserTree,
    newLoserTree,
    replace,
    remove,
) where

import           Control.Monad.Primitive (PrimMonad (PrimState), RealWorld)
import qualified Control.Monad.ST as Lazy
import qualified Control.Monad.ST as Strict
import           Data.Bits (unsafeShiftR)
import           Data.List.NonEmpty (NonEmpty (..))
import           Data.Primitive (MutablePrimArray, SmallMutableArray,
                     newPrimArray, newSmallArray, readPrimArray, readSmallArray,
                     setPrimArray, writePrimArray, writeSmallArray)
import           Data.Primitive.PrimVar (PrimVar, newPrimVar, readPrimVar,
                     writePrimVar)
import           Unsafe.Coerce (unsafeCoerce)

-- | Mutable Loser Tree.
data MutableLoserTree s a = MLT
    !(PrimVar s Int)                 -- ^ element count, i.e. size.
    !(PrimVar s Int)                 -- ^ index of the hole (i.e. winner's initial index)
    !(MutablePrimArray s Int)        -- ^ indices, we store the index of first match. -1 if there is no match.
    !(SmallMutableArray s a)         -- ^ values

placeholder :: a
placeholder :: forall a. a
placeholder = () -> a
forall a b. a -> b
unsafeCoerce ()

-- | Create new 'MutableLoserTree'.
--
-- The second half of a pair is the winner value (only losers are stored in the tree).
--
newLoserTree :: forall a m. (PrimMonad m, Ord a) => NonEmpty a -> m (MutableLoserTree (PrimState m) a, a)
newLoserTree :: forall a (m :: * -> *).
(PrimMonad m, Ord a) =>
NonEmpty a -> m (MutableLoserTree (PrimState m) a, a)
newLoserTree (a
x0 :| [a]
xs0) = do
    -- allocate array, we need one less than there are elements.
    -- one of the elements will be the winner.
    MutablePrimArray (PrimState m) Int
ids <- Int -> m (MutablePrimArray (PrimState m) Int)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray  Int
len
    SmallMutableArray (PrimState m) a
arr <- Int -> a -> m (SmallMutableArray (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
len a
forall a. a
placeholder
    case [a]
xs0 of
      [] -> do
        PrimVar (PrimState m) Int
sizeRef <- Int -> m (PrimVar (PrimState m) Int)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
a -> m (PrimVar (PrimState m) a)
newPrimVar Int
0
        PrimVar (PrimState m) Int
holeRef <- Int -> m (PrimVar (PrimState m) Int)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
a -> m (PrimVar (PrimState m) a)
newPrimVar Int
0
        (MutableLoserTree (PrimState m) a, a)
-> m (MutableLoserTree (PrimState m) a, a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ((MutableLoserTree (PrimState m) a, a)
 -> m (MutableLoserTree (PrimState m) a, a))
-> (MutableLoserTree (PrimState m) a, a)
-> m (MutableLoserTree (PrimState m) a, a)
forall a b. (a -> b) -> a -> b
$! (PrimVar (PrimState m) Int
-> PrimVar (PrimState m) Int
-> MutablePrimArray (PrimState m) Int
-> SmallMutableArray (PrimState m) a
-> MutableLoserTree (PrimState m) a
forall s a.
PrimVar s Int
-> PrimVar s Int
-> MutablePrimArray s Int
-> SmallMutableArray s a
-> MutableLoserTree s a
MLT PrimVar (PrimState m) Int
sizeRef PrimVar (PrimState m) Int
holeRef MutablePrimArray (PrimState m) Int
ids SmallMutableArray (PrimState m) a
arr, a
x0)
      [a]
_ -> do
        MutablePrimArray (PrimState m) Int -> Int -> Int -> Int -> m ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> Int -> a -> m ()
setPrimArray MutablePrimArray (PrimState m) Int
ids Int
0 Int
len (-Int
1)
        MutablePrimArray (PrimState m) Int
-> SmallMutableArray (PrimState m) a
-> Int
-> NonEmpty a
-> m (MutableLoserTree (PrimState m) a, a)
loop MutablePrimArray (PrimState m) Int
ids SmallMutableArray (PrimState m) a
arr Int
len (NonEmpty a -> m (MutableLoserTree (PrimState m) a, a))
-> NonEmpty a -> m (MutableLoserTree (PrimState m) a, a)
forall a b. (a -> b) -> a -> b
$ a
x0 a -> [a] -> NonEmpty a
forall a. a -> [a] -> NonEmpty a
:| [a]
xs0
  where
    !len :: Int
len = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs0

    loop :: MutablePrimArray (PrimState m) Int -> SmallMutableArray (PrimState m) a -> Int -> NonEmpty a -> m (MutableLoserTree (PrimState m) a, a)
    loop :: MutablePrimArray (PrimState m) Int
-> SmallMutableArray (PrimState m) a
-> Int
-> NonEmpty a
-> m (MutableLoserTree (PrimState m) a, a)
loop  MutablePrimArray (PrimState m) Int
ids  SmallMutableArray (PrimState m) a
arr  Int
idx (a
x :| [a]
xs) = do
        MutablePrimArray (PrimState m) Int
-> SmallMutableArray (PrimState m) a
-> Int
-> Int
-> a
-> Int
-> [a]
-> m (MutableLoserTree (PrimState m) a, a)
sift MutablePrimArray (PrimState m) Int
ids SmallMutableArray (PrimState m) a
arr (Int -> Int
parentOf Int
idx) (Int -> Int
parentOf Int
idx) a
x Int
idx [a]
xs

    sift :: MutablePrimArray (PrimState m) Int -> SmallMutableArray (PrimState m) a -> Int -> Int -> a -> Int -> [a] -> m (MutableLoserTree (PrimState m) a, a)
    sift :: MutablePrimArray (PrimState m) Int
-> SmallMutableArray (PrimState m) a
-> Int
-> Int
-> a
-> Int
-> [a]
-> m (MutableLoserTree (PrimState m) a, a)
sift !MutablePrimArray (PrimState m) Int
ids !SmallMutableArray (PrimState m) a
arr !Int
idxX !Int
j !a
x !Int
idx0 [a]
xs = do
        !Int
idxY <- MutablePrimArray (PrimState m) Int -> Int -> m Int
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> m a
readPrimArray MutablePrimArray (PrimState m) Int
ids Int
j
        a
y     <- SmallMutableArray (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> m a
readSmallArray SmallMutableArray (PrimState m) a
arr Int
j
        -- NOTE: The length of xs is equal to number of uninitialised entries
        -- from this we can deduce that an entry at j is uninitialised implies
        -- that xs cannot be empty.
        -- We check this invariant here and throw an exception
        -- with a descriptive error message if it is violated.
        if Int
idxY Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0
        then case [a]
xs of
          [] -> [Char] -> m (MutableLoserTree (PrimState m) a, a)
forall a. HasCallStack => [Char] -> a
error ([Char] -> m (MutableLoserTree (PrimState m) a, a))
-> [Char] -> m (MutableLoserTree (PrimState m) a, a)
forall a b. (a -> b) -> a -> b
$ [[Char]] -> [Char]
unlines
            [ [Char]
"Error in KMerge.LoserTree.newLoserTree"
            , [[Char]] -> [Char]
unwords [ [Char]
"Invariant violated at entry # j =", Int -> [Char]
forall a. Show a => a -> [Char]
show Int
j, [Char]
"with xs = [] and idxY =", Int -> [Char]
forall a. Show a => a -> [Char]
show Int
idxY ]
            ]
          a
e:[a]
es -> do
            MutablePrimArray (PrimState m) Int -> Int -> Int -> m ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray  MutablePrimArray (PrimState m) Int
ids Int
j Int
idxX
            SmallMutableArray (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray (PrimState m) a
arr Int
j a
x
            MutablePrimArray (PrimState m) Int
-> SmallMutableArray (PrimState m) a
-> Int
-> NonEmpty a
-> m (MutableLoserTree (PrimState m) a, a)
loop MutablePrimArray (PrimState m) Int
ids SmallMutableArray (PrimState m) a
arr (Int
idx0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (NonEmpty a -> m (MutableLoserTree (PrimState m) a, a))
-> NonEmpty a -> m (MutableLoserTree (PrimState m) a, a)
forall a b. (a -> b) -> a -> b
$ a
e a -> [a] -> NonEmpty a
forall a. a -> [a] -> NonEmpty a
:| [a]
es
        else
            if Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0
            then do
                    if a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
y
                    then do
                        PrimVar (PrimState m) Int
sizeRef <- Int -> m (PrimVar (PrimState m) Int)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
a -> m (PrimVar (PrimState m) a)
newPrimVar Int
len
                        PrimVar (PrimState m) Int
holeRef <- Int -> m (PrimVar (PrimState m) Int)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
a -> m (PrimVar (PrimState m) a)
newPrimVar Int
idxX
                        (MutableLoserTree (PrimState m) a, a)
-> m (MutableLoserTree (PrimState m) a, a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (PrimVar (PrimState m) Int
-> PrimVar (PrimState m) Int
-> MutablePrimArray (PrimState m) Int
-> SmallMutableArray (PrimState m) a
-> MutableLoserTree (PrimState m) a
forall s a.
PrimVar s Int
-> PrimVar s Int
-> MutablePrimArray s Int
-> SmallMutableArray s a
-> MutableLoserTree s a
MLT PrimVar (PrimState m) Int
sizeRef PrimVar (PrimState m) Int
holeRef MutablePrimArray (PrimState m) Int
ids SmallMutableArray (PrimState m) a
arr, a
x)
                    else do
                        MutablePrimArray (PrimState m) Int -> Int -> Int -> m ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray  MutablePrimArray (PrimState m) Int
ids Int
j Int
idxX
                        SmallMutableArray (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray (PrimState m) a
arr Int
j a
x
                        PrimVar (PrimState m) Int
sizeRef <- Int -> m (PrimVar (PrimState m) Int)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
a -> m (PrimVar (PrimState m) a)
newPrimVar Int
len
                        PrimVar (PrimState m) Int
holeRef <- Int -> m (PrimVar (PrimState m) Int)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
a -> m (PrimVar (PrimState m) a)
newPrimVar Int
idxY
                        (MutableLoserTree (PrimState m) a, a)
-> m (MutableLoserTree (PrimState m) a, a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (PrimVar (PrimState m) Int
-> PrimVar (PrimState m) Int
-> MutablePrimArray (PrimState m) Int
-> SmallMutableArray (PrimState m) a
-> MutableLoserTree (PrimState m) a
forall s a.
PrimVar s Int
-> PrimVar s Int
-> MutablePrimArray s Int
-> SmallMutableArray s a
-> MutableLoserTree s a
MLT PrimVar (PrimState m) Int
sizeRef PrimVar (PrimState m) Int
holeRef MutablePrimArray (PrimState m) Int
ids SmallMutableArray (PrimState m) a
arr, a
y)
            else do
                    if a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
y
                    then do
                        MutablePrimArray (PrimState m) Int
-> SmallMutableArray (PrimState m) a
-> Int
-> Int
-> a
-> Int
-> [a]
-> m (MutableLoserTree (PrimState m) a, a)
sift MutablePrimArray (PrimState m) Int
ids SmallMutableArray (PrimState m) a
arr Int
idxX (Int -> Int
parentOf Int
j) a
x Int
idx0 [a]
xs
                    else do
                        MutablePrimArray (PrimState m) Int -> Int -> Int -> m ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray  MutablePrimArray (PrimState m) Int
ids Int
j Int
idxX
                        SmallMutableArray (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray (PrimState m) a
arr Int
j a
x
                        MutablePrimArray (PrimState m) Int
-> SmallMutableArray (PrimState m) a
-> Int
-> Int
-> a
-> Int
-> [a]
-> m (MutableLoserTree (PrimState m) a, a)
sift MutablePrimArray (PrimState m) Int
ids SmallMutableArray (PrimState m) a
arr Int
idxY (Int -> Int
parentOf Int
j) a
y Int
idx0 [a]
xs

{-# SPECIALIZE newLoserTree :: forall a.   Ord a => NonEmpty a -> IO          (MutableLoserTree RealWorld a, a) #-}
{-# SPECIALIZE newLoserTree :: forall a s. Ord a => NonEmpty a -> Strict.ST s (MutableLoserTree s         a, a) #-}
{-# SPECIALIZE newLoserTree :: forall a s. Ord a => NonEmpty a -> Lazy.ST   s (MutableLoserTree s         a, a) #-}

{-------------------------------------------------------------------------------
  Updates
-------------------------------------------------------------------------------}

{-# SPECIALIZE replace :: forall a.   Ord a => MutableLoserTree RealWorld a -> a -> IO          a #-}
{-# SPECIALIZE replace :: forall a s. Ord a => MutableLoserTree s         a -> a -> Strict.ST s a #-}
{-# SPECIALIZE replace :: forall a s. Ord a => MutableLoserTree s         a -> a -> Lazy.ST s   a #-}

{-# SPECIALIZE remove :: forall a.   Ord a => MutableLoserTree RealWorld a -> IO          (Maybe a) #-}
{-# SPECIALIZE remove :: forall a s. Ord a => MutableLoserTree s         a -> Strict.ST s (Maybe a) #-}
{-# SPECIALIZE remove :: forall a s. Ord a => MutableLoserTree s         a -> Lazy.ST s   (Maybe a) #-}

-- | Don't fill the winner "hole". Return a next winner of (smaller) tournament.
remove :: forall a m. (PrimMonad m, Ord a) => MutableLoserTree (PrimState m) a -> m (Maybe a)
remove :: forall a (m :: * -> *).
(PrimMonad m, Ord a) =>
MutableLoserTree (PrimState m) a -> m (Maybe a)
remove (MLT PrimVar (PrimState m) Int
sizeRef PrimVar (PrimState m) Int
holeRef MutablePrimArray (PrimState m) Int
ids SmallMutableArray (PrimState m) a
arr) = do
    Int
size <- PrimVar (PrimState m) Int -> m Int
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimVar (PrimState m) a -> m a
readPrimVar PrimVar (PrimState m) Int
sizeRef
    if Int
size Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0
    then Maybe a -> m (Maybe a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
forall a. Maybe a
Nothing
    else do
        PrimVar (PrimState m) Int -> Int -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimVar (PrimState m) a -> a -> m ()
writePrimVar PrimVar (PrimState m) Int
sizeRef (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
        Int
hole <- PrimVar (PrimState m) Int -> m Int
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimVar (PrimState m) a -> m a
readPrimVar PrimVar (PrimState m) Int
holeRef
        Int -> m (Maybe a)
siftEmpty Int
hole
  where
    siftEmpty :: Int -> m (Maybe a)
    siftEmpty :: Int -> m (Maybe a)
siftEmpty !Int
j = do
        !Int
idxY <- MutablePrimArray (PrimState m) Int -> Int -> m Int
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> m a
readPrimArray MutablePrimArray (PrimState m) Int
ids Int
j
        a
y     <- SmallMutableArray (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> m a
readSmallArray SmallMutableArray (PrimState m) a
arr Int
j
        if Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0
        then if Int
idxY Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0
            then Maybe a -> m (Maybe a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
forall a. Maybe a
Nothing
            else do
                MutablePrimArray (PrimState m) Int -> Int -> Int -> m ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray  MutablePrimArray (PrimState m) Int
ids Int
j (-Int
1)
                SmallMutableArray (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray (PrimState m) a
arr Int
j a
forall a. a
placeholder
                PrimVar (PrimState m) Int -> Int -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimVar (PrimState m) a -> a -> m ()
writePrimVar PrimVar (PrimState m) Int
holeRef Int
idxY
                Maybe a -> m (Maybe a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> Maybe a
forall a. a -> Maybe a
Just a
y)
        else if Int
idxY Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0
            then
                Int -> m (Maybe a)
siftEmpty (Int -> Int
parentOf Int
j)
            else do
                MutablePrimArray (PrimState m) Int -> Int -> Int -> m ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray  MutablePrimArray (PrimState m) Int
ids Int
j (-Int
1)
                SmallMutableArray (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray (PrimState m) a
arr Int
j a
forall a. a
placeholder
                a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> m a -> m (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MutablePrimArray (PrimState m) Int
-> SmallMutableArray (PrimState m) a
-> PrimVar (PrimState m) Int
-> Int
-> Int
-> a
-> m a
forall a (m :: * -> *).
(PrimMonad m, Ord a) =>
MutablePrimArray (PrimState m) Int
-> SmallMutableArray (PrimState m) a
-> PrimVar (PrimState m) Int
-> Int
-> Int
-> a
-> m a
siftUp MutablePrimArray (PrimState m) Int
ids SmallMutableArray (PrimState m) a
arr PrimVar (PrimState m) Int
holeRef (Int -> Int
parentOf Int
j) Int
idxY a
y

-- | Fill the winner "hole" with a new element. Return a new tournament winner.
replace :: forall a m. (PrimMonad m, Ord a) => MutableLoserTree (PrimState m) a -> a -> m a
replace :: forall a (m :: * -> *).
(PrimMonad m, Ord a) =>
MutableLoserTree (PrimState m) a -> a -> m a
replace (MLT PrimVar (PrimState m) Int
sizeRef PrimVar (PrimState m) Int
holeRef MutablePrimArray (PrimState m) Int
ids SmallMutableArray (PrimState m) a
arr) a
val = do
    Int
size <- PrimVar (PrimState m) Int -> m Int
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimVar (PrimState m) a -> m a
readPrimVar PrimVar (PrimState m) Int
sizeRef
    if Int
size Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0
    then a -> m a
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
val
    else do
        Int
hole <- PrimVar (PrimState m) Int -> m Int
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimVar (PrimState m) a -> m a
readPrimVar PrimVar (PrimState m) Int
holeRef
        MutablePrimArray (PrimState m) Int
-> SmallMutableArray (PrimState m) a
-> PrimVar (PrimState m) Int
-> Int
-> Int
-> a
-> m a
forall a (m :: * -> *).
(PrimMonad m, Ord a) =>
MutablePrimArray (PrimState m) Int
-> SmallMutableArray (PrimState m) a
-> PrimVar (PrimState m) Int
-> Int
-> Int
-> a
-> m a
siftUp MutablePrimArray (PrimState m) Int
ids SmallMutableArray (PrimState m) a
arr PrimVar (PrimState m) Int
holeRef Int
hole Int
hole a
val

{-# SPECIALIZE siftUp :: forall a.   Ord a => MutablePrimArray RealWorld Int -> SmallMutableArray RealWorld a -> PrimVar RealWorld Int -> Int -> Int -> a -> IO          a #-}
{-# SPECIALIZE siftUp :: forall a s. Ord a => MutablePrimArray s Int         -> SmallMutableArray s         a -> PrimVar s         Int -> Int -> Int -> a -> Strict.ST s a #-}
{-# SPECIALIZE siftUp :: forall a s. Ord a => MutablePrimArray s Int         -> SmallMutableArray s         a -> PrimVar s         Int -> Int -> Int -> a -> Lazy.ST s   a #-}

siftUp :: forall a m. (PrimMonad m, Ord a) => MutablePrimArray (PrimState m) Int -> SmallMutableArray (PrimState m) a -> PrimVar (PrimState m) Int -> Int -> Int -> a -> m a
siftUp :: forall a (m :: * -> *).
(PrimMonad m, Ord a) =>
MutablePrimArray (PrimState m) Int
-> SmallMutableArray (PrimState m) a
-> PrimVar (PrimState m) Int
-> Int
-> Int
-> a
-> m a
siftUp MutablePrimArray (PrimState m) Int
ids SmallMutableArray (PrimState m) a
arr PrimVar (PrimState m) Int
holeRef = Int -> Int -> a -> m a
sift
  where
    sift :: Int -> Int -> a -> m a
    sift :: Int -> Int -> a -> m a
sift !Int
j !Int
idxX !a
x = do
        !Int
idxY <- MutablePrimArray (PrimState m) Int -> Int -> m Int
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> m a
readPrimArray MutablePrimArray (PrimState m) Int
ids Int
j
        a
y     <- SmallMutableArray (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> m a
readSmallArray SmallMutableArray (PrimState m) a
arr Int
j
        if Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0
        then if Int
idxY Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0
            then do
                PrimVar (PrimState m) Int -> Int -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimVar (PrimState m) a -> a -> m ()
writePrimVar PrimVar (PrimState m) Int
holeRef Int
idxX
                a -> m a
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
x
            else do
                if a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
y
                then do
                    PrimVar (PrimState m) Int -> Int -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimVar (PrimState m) a -> a -> m ()
writePrimVar PrimVar (PrimState m) Int
holeRef Int
idxX
                    a -> m a
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
x
                else do
                    MutablePrimArray (PrimState m) Int -> Int -> Int -> m ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray  MutablePrimArray (PrimState m) Int
ids Int
j Int
idxX
                    SmallMutableArray (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray (PrimState m) a
arr Int
j a
x
                    PrimVar (PrimState m) Int -> Int -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimVar (PrimState m) a -> a -> m ()
writePrimVar PrimVar (PrimState m) Int
holeRef Int
idxY
                    a -> m a
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
y
        else if Int
idxY Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0
            then Int -> Int -> a -> m a
sift (Int -> Int
parentOf Int
j) Int
idxX a
x
            else do
                if a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
y
                then do
                    Int -> Int -> a -> m a
sift (Int -> Int
parentOf Int
j) Int
idxX a
x
                else do
                    MutablePrimArray (PrimState m) Int -> Int -> Int -> m ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray  MutablePrimArray (PrimState m) Int
ids Int
j Int
idxX
                    SmallMutableArray (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray (PrimState m) a
arr Int
j a
x
                    Int -> Int -> a -> m a
sift (Int -> Int
parentOf Int
j) Int
idxY a
y

{-------------------------------------------------------------------------------
  Helpers
-------------------------------------------------------------------------------}

halfOf :: Int -> Int
halfOf :: Int -> Int
halfOf Int
i = Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftR Int
i Int
1
{-# INLINE halfOf #-}

parentOf :: Int -> Int
parentOf :: Int -> Int
parentOf Int
i = Int -> Int
halfOf (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
{-# INLINE parentOf #-}