{-# LANGUAGE RecordWildCards #-}

module Data.Heap.Binary where

import Control.Monad
import Control.Monad.Primitive
import Data.Bits
import Data.Coerce
import Data.Function
import Data.Functor.Identity
import Data.Kind
import Data.Ord
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as UM

--
import My.Prelude (rev)

data BinaryHeap (f :: Type -> Type) s a = BinaryHeap
  { forall (f :: * -> *) s a. BinaryHeap f s a -> a -> f a
priorityBH :: a -> f a
  , forall (f :: * -> *) s a. BinaryHeap f s a -> MVector s Int
intVarsBH :: !(UM.MVector s Int)
  , forall (f :: * -> *) s a. BinaryHeap f s a -> MVector s a
internalVecBH :: !(UM.MVector s a)
  }

_sizeBH :: Int
_sizeBH :: Int
_sizeBH = Int
0
{-# INLINE _sizeBH #-}

type MinBinaryHeap s a = BinaryHeap Identity s a
type MaxBinaryHeap s a = BinaryHeap Down s a

newBinaryHeap :: (U.Unbox a, PrimMonad m) => (a -> f a) -> Int -> m (BinaryHeap f (PrimState m) a)
newBinaryHeap :: forall a (m :: * -> *) (f :: * -> *).
(Unbox a, PrimMonad m) =>
(a -> f a) -> Int -> m (BinaryHeap f (PrimState m) a)
newBinaryHeap a -> f a
prio Int
n = (a -> f a)
-> MVector (PrimState m) Int
-> MVector (PrimState m) a
-> BinaryHeap f (PrimState m) a
forall (f :: * -> *) s a.
(a -> f a) -> MVector s Int -> MVector s a -> BinaryHeap f s a
BinaryHeap a -> f a
prio (MVector (PrimState m) Int
 -> MVector (PrimState m) a -> BinaryHeap f (PrimState m) a)
-> m (MVector (PrimState m) Int)
-> m (MVector (PrimState m) a -> BinaryHeap f (PrimState m) a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> m (MVector (PrimState m) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
UM.replicate Int
1 Int
0 m (MVector (PrimState m) a -> BinaryHeap f (PrimState m) a)
-> m (MVector (PrimState m) a) -> m (BinaryHeap f (PrimState m) a)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> m (MVector (PrimState m) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
UM.unsafeNew Int
n

newMinBinaryHeap :: (U.Unbox a, PrimMonad m) => Int -> m (MinBinaryHeap (PrimState m) a)
newMinBinaryHeap :: forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Int -> m (MinBinaryHeap (PrimState m) a)
newMinBinaryHeap = (a -> Identity a) -> Int -> m (BinaryHeap Identity (PrimState m) a)
forall a (m :: * -> *) (f :: * -> *).
(Unbox a, PrimMonad m) =>
(a -> f a) -> Int -> m (BinaryHeap f (PrimState m) a)
newBinaryHeap a -> Identity a
forall a. a -> Identity a
Identity

newMaxBinaryHeap :: (U.Unbox a, PrimMonad m) => Int -> m (MaxBinaryHeap (PrimState m) a)
newMaxBinaryHeap :: forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Int -> m (MaxBinaryHeap (PrimState m) a)
newMaxBinaryHeap = (a -> Down a) -> Int -> m (BinaryHeap Down (PrimState m) a)
forall a (m :: * -> *) (f :: * -> *).
(Unbox a, PrimMonad m) =>
(a -> f a) -> Int -> m (BinaryHeap f (PrimState m) a)
newBinaryHeap a -> Down a
forall a. a -> Down a
Down

getBinaryHeapSize :: (PrimMonad m) => BinaryHeap f (PrimState m) a -> m Int
getBinaryHeapSize :: forall (m :: * -> *) (f :: * -> *) a.
PrimMonad m =>
BinaryHeap f (PrimState m) a -> m Int
getBinaryHeapSize BinaryHeap{MVector (PrimState m) a
MVector (PrimState m) Int
a -> f a
priorityBH :: forall (f :: * -> *) s a. BinaryHeap f s a -> a -> f a
intVarsBH :: forall (f :: * -> *) s a. BinaryHeap f s a -> MVector s Int
internalVecBH :: forall (f :: * -> *) s a. BinaryHeap f s a -> MVector s a
priorityBH :: a -> f a
intVarsBH :: MVector (PrimState m) Int
internalVecBH :: MVector (PrimState m) a
..} = MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
UM.unsafeRead MVector (PrimState m) Int
intVarsBH Int
_sizeBH
{-# INLINE getBinaryHeapSize #-}

siftUpBy ::
  (U.Unbox a, PrimMonad m) =>
  (a -> a -> Ordering) ->
  Int ->
  UM.MVector (PrimState m) a ->
  m ()
siftUpBy :: forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
(a -> a -> Ordering) -> Int -> MVector (PrimState m) a -> m ()
siftUpBy a -> a -> Ordering
cmp Int
k MVector (PrimState m) a
vec = do
  a
x <- MVector (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
UM.unsafeRead MVector (PrimState m) a
vec Int
k
  (((Int -> m ()) -> Int -> m ()) -> Int -> m ())
-> Int -> ((Int -> m ()) -> Int -> m ()) -> m ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Int -> m ()) -> Int -> m ()) -> Int -> m ()
forall a. (a -> a) -> a
fix Int
k (((Int -> m ()) -> Int -> m ()) -> m ())
-> ((Int -> m ()) -> Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int -> m ()
loop !Int
i ->
    if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
      then do
        let parent :: Int
parent = (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
1
        a
p <- MVector (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
UM.unsafeRead MVector (PrimState m) a
vec Int
parent
        case a -> a -> Ordering
cmp a
p a
x of
          Ordering
GT -> MVector (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
UM.unsafeWrite MVector (PrimState m) a
vec Int
i a
p m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> m ()
loop Int
parent
          Ordering
_ -> MVector (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
UM.unsafeWrite MVector (PrimState m) a
vec Int
i a
x
      else MVector (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
UM.unsafeWrite MVector (PrimState m) a
vec Int
0 a
x
{-# INLINE siftUpBy #-}

siftDownBy ::
  (U.Unbox a, PrimMonad m) =>
  (a -> a -> Ordering) ->
  Int ->
  UM.MVector (PrimState m) a ->
  m ()
siftDownBy :: forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
(a -> a -> Ordering) -> Int -> MVector (PrimState m) a -> m ()
siftDownBy a -> a -> Ordering
cmp Int
k MVector (PrimState m) a
vec = do
  a
x <- MVector (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
UM.unsafeRead MVector (PrimState m) a
vec Int
k
  let !n :: Int
n = MVector (PrimState m) a -> Int
forall a s. Unbox a => MVector s a -> Int
UM.length MVector (PrimState m) a
vec
  (((Int -> m ()) -> Int -> m ()) -> Int -> m ())
-> Int -> ((Int -> m ()) -> Int -> m ()) -> m ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Int -> m ()) -> Int -> m ()) -> Int -> m ()
forall a. (a -> a) -> a
fix Int
k (((Int -> m ()) -> Int -> m ()) -> m ())
-> ((Int -> m ()) -> Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int -> m ()
loop !Int
i -> do
    let l :: Int
l = Int -> Int -> Int
forall a. Bits a => a -> Int -> a
unsafeShiftL Int
i Int
1 Int -> Int -> Int
forall a. Bits a => a -> a -> a
.|. Int
1
    let r :: Int
r = Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
    if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
l
      then MVector (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
UM.unsafeWrite MVector (PrimState m) a
vec Int
i a
x
      else do
        a
vl <- MVector (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
UM.unsafeRead MVector (PrimState m) a
vec Int
l
        if Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n
          then do
            a
vr <- MVector (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
UM.unsafeRead MVector (PrimState m) a
vec Int
r
            case a -> a -> Ordering
cmp a
vr a
vl of
              Ordering
LT -> case a -> a -> Ordering
cmp a
x a
vr of
                Ordering
GT -> MVector (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
UM.unsafeWrite MVector (PrimState m) a
vec Int
i a
vr m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> m ()
loop Int
r
                Ordering
_ -> MVector (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
UM.unsafeWrite MVector (PrimState m) a
vec Int
i a
x
              Ordering
_ -> case a -> a -> Ordering
cmp a
x a
vl of
                Ordering
GT -> MVector (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
UM.unsafeWrite MVector (PrimState m) a
vec Int
i a
vl m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> m ()
loop Int
l
                Ordering
_ -> MVector (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
UM.unsafeWrite MVector (PrimState m) a
vec Int
i a
x
          else case a -> a -> Ordering
cmp a
x a
vl of
            Ordering
GT -> MVector (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
UM.unsafeWrite MVector (PrimState m) a
vec Int
i a
vl m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> m ()
loop Int
l
            Ordering
_ -> MVector (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
UM.unsafeWrite MVector (PrimState m) a
vec Int
i a
x
{-# INLINE siftDownBy #-}

heapifyBy ::
  (U.Unbox a, PrimMonad m) =>
  (a -> a -> Ordering) ->
  UM.MVector (PrimState m) a ->
  m ()
heapifyBy :: forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
(a -> a -> Ordering) -> MVector (PrimState m) a -> m ()
heapifyBy a -> a -> Ordering
cmp MVector (PrimState m) a
vec = do
  Int -> (Int -> m ()) -> m ()
forall (m :: * -> *). Monad m => Int -> (Int -> m ()) -> m ()
rev (MVector (PrimState m) a -> Int
forall a s. Unbox a => MVector s a -> Int
UM.length MVector (PrimState m) a
vec Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
2) ((Int -> m ()) -> m ()) -> (Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int
i -> do
    (a -> a -> Ordering) -> Int -> MVector (PrimState m) a -> m ()
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
(a -> a -> Ordering) -> Int -> MVector (PrimState m) a -> m ()
siftDownBy a -> a -> Ordering
cmp Int
i MVector (PrimState m) a
vec
{-# INLINE heapifyBy #-}

class OrdVia f a where
  compareVia :: (a -> f a) -> a -> a -> Ordering

instance (Ord a) => OrdVia Identity a where
  compareVia :: (a -> Identity a) -> a -> a -> Ordering
compareVia a -> Identity a
_ = (Identity a -> Identity a -> Ordering) -> a -> a -> Ordering
forall a b. Coercible a b => a -> b
coerce (Identity a -> Identity a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare :: Identity a -> Identity a -> Ordering)
  {-# INLINE compareVia #-}

instance (Ord a) => OrdVia Down a where
  compareVia :: (a -> Down a) -> a -> a -> Ordering
compareVia a -> Down a
_ = (Down a -> Down a -> Ordering) -> a -> a -> Ordering
forall a b. Coercible a b => a -> b
coerce (Down a -> Down a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare :: Down a -> Down a -> Ordering)
  {-# INLINE compareVia #-}

buildBinaryHeapVia ::
  (OrdVia f a, U.Unbox a, PrimMonad m) =>
  (a -> f a) ->
  U.Vector a ->
  m (BinaryHeap f (PrimState m) a)
buildBinaryHeapVia :: forall (f :: * -> *) a (m :: * -> *).
(OrdVia f a, Unbox a, PrimMonad m) =>
(a -> f a) -> Vector a -> m (BinaryHeap f (PrimState m) a)
buildBinaryHeapVia a -> f a
priorityBH Vector a
vec = do
  MVector (PrimState m) Int
intVarsBH <- Int -> Int -> m (MVector (PrimState m) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
UM.replicate Int
1 (Int -> m (MVector (PrimState m) Int))
-> Int -> m (MVector (PrimState m) Int)
forall a b. (a -> b) -> a -> b
$ Vector a -> Int
forall a. Unbox a => Vector a -> Int
U.length Vector a
vec
  MVector (PrimState m) a
internalVecBH <- Vector a -> m (MVector (PrimState m) a)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
U.thaw Vector a
vec
  (a -> a -> Ordering) -> MVector (PrimState m) a -> m ()
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
(a -> a -> Ordering) -> MVector (PrimState m) a -> m ()
heapifyBy ((a -> f a) -> a -> a -> Ordering
forall (f :: * -> *) a.
OrdVia f a =>
(a -> f a) -> a -> a -> Ordering
compareVia a -> f a
priorityBH) MVector (PrimState m) a
internalVecBH
  BinaryHeap f (PrimState m) a -> m (BinaryHeap f (PrimState m) a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (BinaryHeap f (PrimState m) a -> m (BinaryHeap f (PrimState m) a))
-> BinaryHeap f (PrimState m) a -> m (BinaryHeap f (PrimState m) a)
forall a b. (a -> b) -> a -> b
$! BinaryHeap{MVector (PrimState m) a
MVector (PrimState m) Int
a -> f a
priorityBH :: a -> f a
intVarsBH :: MVector (PrimState m) Int
internalVecBH :: MVector (PrimState m) a
priorityBH :: a -> f a
intVarsBH :: MVector (PrimState m) Int
internalVecBH :: MVector (PrimState m) a
..}
{-# INLINE buildBinaryHeapVia #-}

buildMinBinaryHeap ::
  (Ord a, U.Unbox a, PrimMonad m) =>
  U.Vector a ->
  m (BinaryHeap Identity (PrimState m) a)
buildMinBinaryHeap :: forall a (m :: * -> *).
(Ord a, Unbox a, PrimMonad m) =>
Vector a -> m (BinaryHeap Identity (PrimState m) a)
buildMinBinaryHeap = (a -> Identity a)
-> Vector a -> m (BinaryHeap Identity (PrimState m) a)
forall (f :: * -> *) a (m :: * -> *).
(OrdVia f a, Unbox a, PrimMonad m) =>
(a -> f a) -> Vector a -> m (BinaryHeap f (PrimState m) a)
buildBinaryHeapVia a -> Identity a
forall a. a -> Identity a
Identity
{-# INLINE buildMinBinaryHeap #-}

buildMaxBinaryHeap ::
  (Ord a, U.Unbox a, PrimMonad m) =>
  U.Vector a ->
  m (BinaryHeap Down (PrimState m) a)
buildMaxBinaryHeap :: forall a (m :: * -> *).
(Ord a, Unbox a, PrimMonad m) =>
Vector a -> m (BinaryHeap Down (PrimState m) a)
buildMaxBinaryHeap = (a -> Down a) -> Vector a -> m (BinaryHeap Down (PrimState m) a)
forall (f :: * -> *) a (m :: * -> *).
(OrdVia f a, Unbox a, PrimMonad m) =>
(a -> f a) -> Vector a -> m (BinaryHeap f (PrimState m) a)
buildBinaryHeapVia a -> Down a
forall a. a -> Down a
Down
{-# INLINE buildMaxBinaryHeap #-}

unsafeViewBH ::
  (U.Unbox a, PrimMonad m) =>
  BinaryHeap f (PrimState m) a ->
  m a
unsafeViewBH :: forall a (m :: * -> *) (f :: * -> *).
(Unbox a, PrimMonad m) =>
BinaryHeap f (PrimState m) a -> m a
unsafeViewBH BinaryHeap{MVector (PrimState m) a
MVector (PrimState m) Int
a -> f a
priorityBH :: forall (f :: * -> *) s a. BinaryHeap f s a -> a -> f a
intVarsBH :: forall (f :: * -> *) s a. BinaryHeap f s a -> MVector s Int
internalVecBH :: forall (f :: * -> *) s a. BinaryHeap f s a -> MVector s a
priorityBH :: a -> f a
intVarsBH :: MVector (PrimState m) Int
internalVecBH :: MVector (PrimState m) a
..} = MVector (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
UM.unsafeRead MVector (PrimState m) a
internalVecBH Int
0
{-# INLINE unsafeViewBH #-}

viewBH ::
  (U.Unbox a, PrimMonad m) =>
  BinaryHeap f (PrimState m) a ->
  m (Maybe a)
viewBH :: forall a (m :: * -> *) (f :: * -> *).
(Unbox a, PrimMonad m) =>
BinaryHeap f (PrimState m) a -> m (Maybe a)
viewBH BinaryHeap f (PrimState m) a
bh = do
  Int
size <- BinaryHeap f (PrimState m) a -> m Int
forall (m :: * -> *) (f :: * -> *) a.
PrimMonad m =>
BinaryHeap f (PrimState m) a -> m Int
getBinaryHeapSize BinaryHeap f (PrimState m) a
bh
  if Int
size Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
    then a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> m a -> m (Maybe a)
forall (m :: * -> *) a b. Monad m => (a -> b) -> m a -> m b
<$!> BinaryHeap f (PrimState m) a -> m a
forall a (m :: * -> *) (f :: * -> *).
(Unbox a, PrimMonad m) =>
BinaryHeap f (PrimState m) a -> m a
unsafeViewBH BinaryHeap f (PrimState m) a
bh
    else 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
{-# INLINE viewBH #-}

insertBH ::
  (OrdVia f a, U.Unbox a, PrimMonad m) =>
  a ->
  BinaryHeap f (PrimState m) a ->
  m ()
insertBH :: forall (f :: * -> *) a (m :: * -> *).
(OrdVia f a, Unbox a, PrimMonad m) =>
a -> BinaryHeap f (PrimState m) a -> m ()
insertBH a
x BinaryHeap{MVector (PrimState m) a
MVector (PrimState m) Int
a -> f a
priorityBH :: forall (f :: * -> *) s a. BinaryHeap f s a -> a -> f a
intVarsBH :: forall (f :: * -> *) s a. BinaryHeap f s a -> MVector s Int
internalVecBH :: forall (f :: * -> *) s a. BinaryHeap f s a -> MVector s a
priorityBH :: a -> f a
intVarsBH :: MVector (PrimState m) Int
internalVecBH :: MVector (PrimState m) a
..} = do
  Int
size <- MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
UM.unsafeRead MVector (PrimState m) Int
intVarsBH Int
_sizeBH
  MVector (PrimState m) Int -> Int -> Int -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
UM.unsafeWrite MVector (PrimState m) Int
intVarsBH Int
_sizeBH (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  MVector (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
UM.unsafeWrite MVector (PrimState m) a
internalVecBH Int
size a
x
  (a -> a -> Ordering) -> Int -> MVector (PrimState m) a -> m ()
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
(a -> a -> Ordering) -> Int -> MVector (PrimState m) a -> m ()
siftUpBy ((a -> f a) -> a -> a -> Ordering
forall (f :: * -> *) a.
OrdVia f a =>
(a -> f a) -> a -> a -> Ordering
compareVia a -> f a
priorityBH) Int
size MVector (PrimState m) a
internalVecBH
{-# INLINE insertBH #-}

unsafeDeleteBH ::
  (OrdVia f a, U.Unbox a, PrimMonad m) =>
  BinaryHeap f (PrimState m) a ->
  m ()
unsafeDeleteBH :: forall (f :: * -> *) a (m :: * -> *).
(OrdVia f a, Unbox a, PrimMonad m) =>
BinaryHeap f (PrimState m) a -> m ()
unsafeDeleteBH BinaryHeap{MVector (PrimState m) a
MVector (PrimState m) Int
a -> f a
priorityBH :: forall (f :: * -> *) s a. BinaryHeap f s a -> a -> f a
intVarsBH :: forall (f :: * -> *) s a. BinaryHeap f s a -> MVector s Int
internalVecBH :: forall (f :: * -> *) s a. BinaryHeap f s a -> MVector s a
priorityBH :: a -> f a
intVarsBH :: MVector (PrimState m) Int
internalVecBH :: MVector (PrimState m) a
..} = do
  Int
size' <- Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
1 (Int -> Int) -> m Int -> m Int
forall (m :: * -> *) a b. Monad m => (a -> b) -> m a -> m b
<$!> MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
UM.unsafeRead MVector (PrimState m) Int
intVarsBH Int
_sizeBH
  MVector (PrimState m) Int -> Int -> Int -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
UM.unsafeWrite MVector (PrimState m) Int
intVarsBH Int
_sizeBH Int
size'
  MVector (PrimState m) a -> Int -> Int -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> Int -> m ()
UM.unsafeSwap MVector (PrimState m) a
internalVecBH Int
0 Int
size'
  (a -> a -> Ordering) -> Int -> MVector (PrimState m) a -> m ()
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
(a -> a -> Ordering) -> Int -> MVector (PrimState m) a -> m ()
siftDownBy ((a -> f a) -> a -> a -> Ordering
forall (f :: * -> *) a.
OrdVia f a =>
(a -> f a) -> a -> a -> Ordering
compareVia a -> f a
priorityBH) Int
0 (Int -> MVector (PrimState m) a -> MVector (PrimState m) a
forall a s. Unbox a => Int -> MVector s a -> MVector s a
UM.unsafeTake Int
size' MVector (PrimState m) a
internalVecBH)
{-# INLINE unsafeDeleteBH #-}

modifyTopBH ::
  (OrdVia f a, U.Unbox a, PrimMonad m) =>
  (a -> a) ->
  BinaryHeap f (PrimState m) a ->
  m ()
modifyTopBH :: forall (f :: * -> *) a (m :: * -> *).
(OrdVia f a, Unbox a, PrimMonad m) =>
(a -> a) -> BinaryHeap f (PrimState m) a -> m ()
modifyTopBH a -> a
f BinaryHeap{MVector (PrimState m) a
MVector (PrimState m) Int
a -> f a
priorityBH :: forall (f :: * -> *) s a. BinaryHeap f s a -> a -> f a
intVarsBH :: forall (f :: * -> *) s a. BinaryHeap f s a -> MVector s Int
internalVecBH :: forall (f :: * -> *) s a. BinaryHeap f s a -> MVector s a
priorityBH :: a -> f a
intVarsBH :: MVector (PrimState m) Int
internalVecBH :: MVector (PrimState m) a
..} = do
  MVector (PrimState m) a -> (a -> a) -> Int -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> (a -> a) -> Int -> m ()
UM.unsafeModify MVector (PrimState m) a
internalVecBH a -> a
f Int
0
  Int
size <- MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
UM.unsafeRead MVector (PrimState m) Int
intVarsBH Int
_sizeBH
  (a -> a -> Ordering) -> Int -> MVector (PrimState m) a -> m ()
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
(a -> a -> Ordering) -> Int -> MVector (PrimState m) a -> m ()
siftDownBy ((a -> f a) -> a -> a -> Ordering
forall (f :: * -> *) a.
OrdVia f a =>
(a -> f a) -> a -> a -> Ordering
compareVia a -> f a
priorityBH) Int
0 (Int -> MVector (PrimState m) a -> MVector (PrimState m) a
forall a s. Unbox a => Int -> MVector s a -> MVector s a
UM.unsafeTake Int
size MVector (PrimState m) a
internalVecBH)
{-# INLINE modifyTopBH #-}

deleteFindTopBH ::
  (OrdVia f a, U.Unbox a, PrimMonad m) =>
  BinaryHeap f (PrimState m) a ->
  m (Maybe a)
deleteFindTopBH :: forall (f :: * -> *) a (m :: * -> *).
(OrdVia f a, Unbox a, PrimMonad m) =>
BinaryHeap f (PrimState m) a -> m (Maybe a)
deleteFindTopBH BinaryHeap f (PrimState m) a
bh = do
  Int
size <- BinaryHeap f (PrimState m) a -> m Int
forall (m :: * -> *) (f :: * -> *) a.
PrimMonad m =>
BinaryHeap f (PrimState m) a -> m Int
getBinaryHeapSize BinaryHeap f (PrimState m) a
bh
  if Int
size Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
    then do
      !a
top <- BinaryHeap f (PrimState m) a -> m a
forall a (m :: * -> *) (f :: * -> *).
(Unbox a, PrimMonad m) =>
BinaryHeap f (PrimState m) a -> m a
unsafeViewBH BinaryHeap f (PrimState m) a
bh m a -> m () -> m a
forall a b. m a -> m b -> m a
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<* BinaryHeap f (PrimState m) a -> m ()
forall (f :: * -> *) a (m :: * -> *).
(OrdVia f a, Unbox a, PrimMonad m) =>
BinaryHeap f (PrimState m) a -> m ()
unsafeDeleteBH BinaryHeap f (PrimState m) a
bh
      Maybe a -> m (Maybe a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe a -> m (Maybe a)) -> Maybe a -> m (Maybe a)
forall a b. (a -> b) -> a -> b
$ a -> Maybe a
forall a. a -> Maybe a
Just a
top
    else 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
{-# INLINE deleteFindTopBH #-}

clearBH :: (PrimMonad m) => BinaryHeap f (PrimState m) a -> m ()
clearBH :: forall (m :: * -> *) (f :: * -> *) a.
PrimMonad m =>
BinaryHeap f (PrimState m) a -> m ()
clearBH BinaryHeap{MVector (PrimState m) a
MVector (PrimState m) Int
a -> f a
priorityBH :: forall (f :: * -> *) s a. BinaryHeap f s a -> a -> f a
intVarsBH :: forall (f :: * -> *) s a. BinaryHeap f s a -> MVector s Int
internalVecBH :: forall (f :: * -> *) s a. BinaryHeap f s a -> MVector s a
priorityBH :: a -> f a
intVarsBH :: MVector (PrimState m) Int
internalVecBH :: MVector (PrimState m) a
..} = MVector (PrimState m) Int -> Int -> Int -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
UM.unsafeWrite MVector (PrimState m) Int
intVarsBH Int
0 Int
0

freezeInternalVecBH ::
  (U.Unbox a, PrimMonad m) =>
  BinaryHeap f (PrimState m) a ->
  m (U.Vector a)
freezeInternalVecBH :: forall a (m :: * -> *) (f :: * -> *).
(Unbox a, PrimMonad m) =>
BinaryHeap f (PrimState m) a -> m (Vector a)
freezeInternalVecBH BinaryHeap{MVector (PrimState m) a
MVector (PrimState m) Int
a -> f a
priorityBH :: forall (f :: * -> *) s a. BinaryHeap f s a -> a -> f a
intVarsBH :: forall (f :: * -> *) s a. BinaryHeap f s a -> MVector s Int
internalVecBH :: forall (f :: * -> *) s a. BinaryHeap f s a -> MVector s a
priorityBH :: a -> f a
intVarsBH :: MVector (PrimState m) Int
internalVecBH :: MVector (PrimState m) a
..} = do
  Int
size <- MVector (PrimState m) Int -> Int -> m Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
UM.unsafeRead MVector (PrimState m) Int
intVarsBH Int
_sizeBH
  MVector (PrimState m) a -> m (Vector a)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
U.unsafeFreeze (Int -> MVector (PrimState m) a -> MVector (PrimState m) a
forall a s. Unbox a => Int -> MVector s a -> MVector s a
UM.unsafeTake Int
size MVector (PrimState m) a
internalVecBH)