{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}

-- ---------------------------------------------------------------------------
-- |
-- Module      : Data.Vector.Algorithms.Heap
-- Copyright   : (c) 2008-2015 Dan Doel
-- Maintainer  : Dan Doel <dan.doel@gmail.com>
-- Stability   : Experimental
-- Portability : Non-portable (type operators)
--
-- This module implements operations for working with a quaternary heap stored
-- in an unboxed array. Most heapsorts are defined in terms of a binary heap,
-- in which each internal node has at most two children. By contrast, a
-- quaternary heap has internal nodes with up to four children. This reduces
-- the number of comparisons in a heapsort slightly, and improves locality
-- (again, slightly) by flattening out the heap.

module Data.Vector.Algorithms.Heap
       ( -- * Sorting
         sort
       , sortUniq
       , sortBy
       , sortUniqBy
       , sortByBounds
         -- * Selection
       , select
       , selectBy
       , selectByBounds
         -- * Partial sorts
       , partialSort
       , partialSortBy
       , partialSortByBounds
         -- * Heap operations
       , heapify
       , pop
       , popTo
       , sortHeap
       , heapInsert
       , Comparison
       ) where

import Prelude hiding (read, length)

import Control.Monad
import Control.Monad.Primitive

import Data.Bits

import Data.Vector.Generic.Mutable

import Data.Vector.Algorithms.Common (Comparison, uniqueMutableBy)

import qualified Data.Vector.Algorithms.Optimal as O

-- | Sorts an entire array using the default ordering.
sort :: (PrimMonad m, MVector v e, Ord e) => v (PrimState m) e -> m ()
sort :: forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e, Ord e) =>
v (PrimState m) e -> m ()
sort = Comparison e -> v (PrimState m) e -> m ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> m ()
sortBy Comparison e
forall a. Ord a => a -> a -> Ordering
compare
{-# INLINABLE sort #-}

-- | A variant on `sort` that returns a vector of unique elements.
sortUniq :: (PrimMonad m, MVector v e, Ord e) => v (PrimState m) e -> m (v (PrimState m) e)
sortUniq :: forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e, Ord e) =>
v (PrimState m) e -> m (v (PrimState m) e)
sortUniq = Comparison e -> v (PrimState m) e -> m (v (PrimState m) e)
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> m (v (PrimState m) e)
sortUniqBy Comparison e
forall a. Ord a => a -> a -> Ordering
compare
{-# INLINABLE sortUniq #-}

-- | Sorts an entire array using a custom ordering.
sortBy :: (PrimMonad m, MVector v e) => Comparison e -> v (PrimState m) e -> m ()
sortBy :: forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> m ()
sortBy Comparison e
cmp v (PrimState m) e
a = Comparison e -> v (PrimState m) e -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> Int -> m ()
sortByBounds Comparison e
cmp v (PrimState m) e
a Int
0 (v (PrimState m) e -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
length v (PrimState m) e
a)
{-# INLINE sortBy #-}

-- | A variant on `sortBy` which returns a vector of unique elements.
sortUniqBy :: (PrimMonad m, MVector v e)
  => Comparison e -> v (PrimState m) e -> m (v (PrimState m) e)
sortUniqBy :: forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> m (v (PrimState m) e)
sortUniqBy Comparison e
cmp v (PrimState m) e
a = do
  Comparison e -> v (PrimState m) e -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> Int -> m ()
sortByBounds Comparison e
cmp v (PrimState m) e
a Int
0 (v (PrimState m) e -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
length v (PrimState m) e
a)
  Comparison e -> v (PrimState m) e -> m (v (PrimState m) e)
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> m (v (PrimState m) e)
uniqueMutableBy Comparison e
cmp v (PrimState m) e
a
{-# INLINE sortUniqBy #-}

-- | Sorts a portion of an array [l,u) using a custom ordering
sortByBounds
  :: (PrimMonad m, MVector v e)
  => Comparison e
  -> v (PrimState m) e
  -> Int -- ^ lower index, l
  -> Int -- ^ upper index, u
  -> m ()
sortByBounds :: forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> Int -> m ()
sortByBounds Comparison e
cmp v (PrimState m) e
a Int
l Int
u
  | Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
2   = () -> m ()
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
  | Int
len Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2  = Comparison e -> v (PrimState m) e -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> m ()
O.sort2ByOffset Comparison e
cmp v (PrimState m) e
a Int
l
  | Int
len Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
3  = Comparison e -> v (PrimState m) e -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> m ()
O.sort3ByOffset Comparison e
cmp v (PrimState m) e
a Int
l
  | Int
len Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
4  = Comparison e -> v (PrimState m) e -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> m ()
O.sort4ByOffset Comparison e
cmp v (PrimState m) e
a Int
l
  | Bool
otherwise = Comparison e -> v (PrimState m) e -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> Int -> m ()
heapify Comparison e
cmp v (PrimState m) e
a Int
l Int
u m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Comparison e -> v (PrimState m) e -> Int -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> Int -> Int -> m ()
sortHeap Comparison e
cmp v (PrimState m) e
a Int
l (Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
4) Int
u m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Comparison e -> v (PrimState m) e -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> m ()
O.sort4ByOffset Comparison e
cmp v (PrimState m) e
a Int
l
 where len :: Int
len = Int
u Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l
{-# INLINE sortByBounds #-}

-- | Moves the lowest k elements to the front of the array.
-- The elements will be in no particular order.
select
  :: (PrimMonad m, MVector v e, Ord e)
  => v (PrimState m) e
  -> Int -- ^ number of elements to select, k
  -> m ()
select :: forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e, Ord e) =>
v (PrimState m) e -> Int -> m ()
select = Comparison e -> v (PrimState m) e -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> m ()
selectBy Comparison e
forall a. Ord a => a -> a -> Ordering
compare
{-# INLINE select #-}

-- | Moves the lowest (as defined by the comparison) k elements
-- to the front of the array. The elements will be in no particular
-- order.
selectBy
  :: (PrimMonad m, MVector v e)
  => Comparison e
  -> v (PrimState m) e
  -> Int -- ^ number of elements to select, k
  -> m ()
selectBy :: forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> m ()
selectBy Comparison e
cmp v (PrimState m) e
a Int
k = Comparison e -> v (PrimState m) e -> Int -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> Int -> Int -> m ()
selectByBounds Comparison e
cmp v (PrimState m) e
a Int
k Int
0 (v (PrimState m) e -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
length v (PrimState m) e
a)
{-# INLINE selectBy #-}

-- | Moves the 'lowest' k elements in the portion [l,u) of the
-- array into the positions [l,k+l). The elements will be in
-- no particular order.
selectByBounds
  :: (PrimMonad m, MVector v e)
  => Comparison e
  -> v (PrimState m) e
  -> Int -- ^ number of elements to select, k
  -> Int -- ^ lower index, l
  -> Int -- ^ upper index, u
  -> m ()
selectByBounds :: forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> Int -> Int -> m ()
selectByBounds Comparison e
cmp v (PrimState m) e
a Int
k Int
l Int
u
  | Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
u = Comparison e -> v (PrimState m) e -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> Int -> m ()
heapify Comparison e
cmp v (PrimState m) e
a Int
l (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k) 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 -> Int -> Int -> m ()
go Int
l (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k) (Int
u Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
  | Bool
otherwise  = () -> m ()
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
 where
 go :: Int -> Int -> Int -> m ()
go Int
l Int
m Int
u
   | Int
u Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
m      = () -> m ()
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
   | Bool
otherwise  = do e
el <- v (PrimState m) e -> Int -> m e
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
unsafeRead v (PrimState m) e
a Int
l
                     e
eu <- v (PrimState m) e -> Int -> m e
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
unsafeRead v (PrimState m) e
a Int
u
                     case Comparison e
cmp e
eu e
el of
                       Ordering
LT -> Comparison e -> v (PrimState m) e -> Int -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> Int -> Int -> m ()
popTo Comparison e
cmp v (PrimState m) e
a Int
l Int
m Int
u
                       Ordering
_  -> () -> m ()
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
                     Int -> Int -> Int -> m ()
go Int
l Int
m (Int
u Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
{-# INLINE selectByBounds #-}

-- | Moves the lowest k elements to the front of the array, sorted.
--
-- The remaining values of the array will be in no particular order.
partialSort
  :: (PrimMonad m, MVector v e, Ord e)
  => v (PrimState m) e
  -> Int -- ^ number of elements to sort, k
  -> m ()
partialSort :: forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e, Ord e) =>
v (PrimState m) e -> Int -> m ()
partialSort = Comparison e -> v (PrimState m) e -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> m ()
partialSortBy Comparison e
forall a. Ord a => a -> a -> Ordering
compare
{-# INLINE partialSort #-}

-- | Moves the lowest k elements (as defined by the comparison) to
-- the front of the array, sorted.
--
-- The remaining values of the array will be in no particular order.
partialSortBy
  :: (PrimMonad m, MVector v e)
  => Comparison e
  -> v (PrimState m) e
  -> Int -- ^ number of elements to sort, k
  -> m ()
partialSortBy :: forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> m ()
partialSortBy Comparison e
cmp v (PrimState m) e
a Int
k = Comparison e -> v (PrimState m) e -> Int -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> Int -> Int -> m ()
partialSortByBounds Comparison e
cmp v (PrimState m) e
a Int
k Int
0 (v (PrimState m) e -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
length v (PrimState m) e
a)
{-# INLINE partialSortBy #-}

-- | Moves the lowest k elements in the portion [l,u) of the array
-- into positions [l,k+l), sorted.
--
-- The remaining values in [l,u) will be in no particular order. Values outside
-- the range [l,u) will be unaffected.
partialSortByBounds
  :: (PrimMonad m, MVector v e)
  => Comparison e
  -> v (PrimState m) e
  -> Int -- ^ number of elements to sort, k
  -> Int -- ^ lower index, l
  -> Int -- ^ upper index, u
  -> m ()
partialSortByBounds :: forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> Int -> Int -> m ()
partialSortByBounds Comparison e
cmp v (PrimState m) e
a Int
k Int
l Int
u
  -- this potentially does more work than absolutely required,
  -- but using a heap to find the least 2 of 4 elements
  -- seems unlikely to be better than just sorting all of them
  -- with an optimal sort, and the latter is obviously index
  -- correct.
  | Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<  Int
2   = () -> m ()
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
  | Int
len Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2   = Comparison e -> v (PrimState m) e -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> m ()
O.sort2ByOffset Comparison e
cmp v (PrimState m) e
a Int
l
  | Int
len Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
3   = Comparison e -> v (PrimState m) e -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> m ()
O.sort3ByOffset Comparison e
cmp v (PrimState m) e
a Int
l
  | Int
len Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
4   = Comparison e -> v (PrimState m) e -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> m ()
O.sort4ByOffset Comparison e
cmp v (PrimState m) e
a Int
l
  | Int
u Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k = Comparison e -> v (PrimState m) e -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> Int -> m ()
sortByBounds Comparison e
cmp v (PrimState m) e
a Int
l Int
u
  | Bool
otherwise  = do Comparison e -> v (PrimState m) e -> Int -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> Int -> Int -> m ()
selectByBounds Comparison e
cmp v (PrimState m) e
a Int
k Int
l Int
u
                    Comparison e -> v (PrimState m) e -> Int -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> Int -> Int -> m ()
sortHeap Comparison e
cmp v (PrimState m) e
a Int
l (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
4) (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k)
                    Comparison e -> v (PrimState m) e -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> m ()
O.sort4ByOffset Comparison e
cmp v (PrimState m) e
a Int
l
 where
 len :: Int
len = Int
u Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l
{-# INLINE partialSortByBounds #-}

-- | Constructs a heap in a portion of an array [l, u), using the values therein.
--
-- Note: 'heapify' is more efficient than constructing a heap by repeated
-- insertion. Repeated insertion has complexity O(n*log n) while 'heapify' is able
-- to construct a heap in O(n), where n is the number of elements in the heap.
heapify
  :: (PrimMonad m, MVector v e)
  => Comparison e
  -> v (PrimState m) e
  -> Int -- ^ lower index, l
  -> Int -- ^ upper index, u
  -> m ()
heapify :: forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> Int -> m ()
heapify Comparison e
cmp v (PrimState m) e
a Int
l Int
u = Int -> m ()
loop (Int -> m ()) -> Int -> m ()
forall a b. (a -> b) -> a -> b
$ (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
2
  where
 len :: Int
len = Int
u Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l
 loop :: Int -> m ()
loop Int
k
   | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0     = () -> m ()
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
   | Bool
otherwise = v (PrimState m) e -> Int -> m e
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
unsafeRead v (PrimState m) e
a (Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
k) m e -> (e -> m ()) -> m ()
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \e
e ->
                   Comparison e -> v (PrimState m) e -> e -> Int -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> e -> Int -> Int -> Int -> m ()
siftByOffset Comparison e
cmp v (PrimState m) e
a e
e Int
l Int
k Int
len 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
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
{-# INLINE heapify #-}

-- | Given a heap stored in a portion of an array [l,u), swaps the
-- top of the heap with the element at u and rebuilds the heap.
pop
  :: (PrimMonad m, MVector v e)
  => Comparison e
  -> v (PrimState m) e
  -> Int -- ^ lower heap index, l
  -> Int -- ^ upper heap index, u
  -> m ()
pop :: forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> Int -> m ()
pop Comparison e
cmp v (PrimState m) e
a Int
l Int
u = Comparison e -> v (PrimState m) e -> Int -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> Int -> Int -> m ()
popTo Comparison e
cmp v (PrimState m) e
a Int
l Int
u Int
u
{-# INLINE pop #-}

-- | Given a heap stored in a portion of an array [l,u) swaps the top
-- of the heap with the element at position t, and rebuilds the heap.
popTo
  :: (PrimMonad m, MVector v e)
  => Comparison e
  -> v (PrimState m) e
  -> Int -- ^ lower heap index, l
  -> Int -- ^ upper heap index, u
  -> Int -- ^ index to pop to, t
  -> m ()
popTo :: forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> Int -> Int -> m ()
popTo Comparison e
cmp v (PrimState m) e
a Int
l Int
u Int
t = do e
al <- v (PrimState m) e -> Int -> m e
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
unsafeRead v (PrimState m) e
a Int
l
                       e
at <- v (PrimState m) e -> Int -> m e
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
unsafeRead v (PrimState m) e
a Int
t
                       v (PrimState m) e -> Int -> e -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
unsafeWrite v (PrimState m) e
a Int
t e
al
                       Comparison e -> v (PrimState m) e -> e -> Int -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> e -> Int -> Int -> Int -> m ()
siftByOffset Comparison e
cmp v (PrimState m) e
a e
at Int
l Int
0 (Int
u Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l)
{-# INLINE popTo #-}

-- | Given a heap stored in a portion of an array [l,u), sorts the
-- highest values into [m,u). The elements in [l,m) are not in any
-- particular order.
sortHeap
  :: (PrimMonad m, MVector v e)
  => Comparison e
  -> v (PrimState m) e
  -> Int -- ^ lower heap index, l
  -> Int -- ^ lower bound of final sorted portion, m
  -> Int -- ^ upper heap index, u
  -> m ()
sortHeap :: forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> Int -> Int -> m ()
sortHeap Comparison e
cmp v (PrimState m) e
a Int
l Int
m Int
u = Int -> m ()
loop (Int
uInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> v (PrimState m) e -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> Int -> m ()
unsafeSwap v (PrimState m) e
a Int
l Int
m
 where
 loop :: Int -> m ()
loop Int
k
   | Int
m Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
k     = Comparison e -> v (PrimState m) e -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> Int -> m ()
pop Comparison e
cmp v (PrimState m) e
a Int
l Int
k 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
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
   | Bool
otherwise = () -> m ()
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
{-# INLINE sortHeap #-}

-- | Given a heap stored in a portion of an array [l,u) and an element e,
-- inserts the element into the heap, resulting in a heap in [l,u].
--
-- Note: it is best to only use this operation when incremental construction of
-- a heap is required. 'heapify' is capable of building a heap in O(n) time,
-- while repeated insertion takes O(n*log n) time.
heapInsert
  :: (PrimMonad m, MVector v e)
  => Comparison e
  -> v (PrimState m) e
  -> Int -- ^ lower heap index, l
  -> Int -- ^ upper heap index, u
  -> e -- ^ element to be inserted, e
  -> m ()
heapInsert :: forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> Int -> e -> m ()
heapInsert Comparison e
cmp v (PrimState m) e
v Int
l Int
u e
e = Int -> m ()
sift (Int
u Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l)
 where
 sift :: Int -> m ()
sift Int
k
   | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0    = v (PrimState m) e -> Int -> e -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
unsafeWrite v (PrimState m) e
v Int
l e
e
   | Bool
otherwise = let pi :: Int
pi = Int -> Int -> Int
forall a. Bits a => a -> Int -> a
shiftR (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Int
2
                  in v (PrimState m) e -> Int -> m e
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
unsafeRead v (PrimState m) e
v (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
pi) m e -> (e -> m ()) -> m ()
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \e
p -> case Comparison e
cmp e
p e
e of
                       Ordering
LT -> v (PrimState m) e -> Int -> e -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
unsafeWrite v (PrimState m) e
v (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k) e
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 ()
sift Int
pi
                       Ordering
_  -> v (PrimState m) e -> Int -> e -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
unsafeWrite v (PrimState m) e
v (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k) e
e
{-# INLINE heapInsert #-}

-- Rebuilds a heap with a hole in it from start downwards. Afterward,
-- the heap property should apply for [start + off, len + off). val
-- is the new value to be put in the hole.
siftByOffset :: (PrimMonad m, MVector v e)
             => Comparison e -> v (PrimState m) e -> e -> Int -> Int -> Int -> m ()
siftByOffset :: forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> e -> Int -> Int -> Int -> m ()
siftByOffset Comparison e
cmp v (PrimState m) e
a e
val Int
off Int
start Int
len = e -> Int -> Int -> m ()
sift e
val Int
start Int
len
 where
 sift :: e -> Int -> Int -> m ()
sift e
val Int
root Int
len
   | Int
child Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
len = do (Int
child', e
ac) <- Comparison e
-> v (PrimState m) e -> Int -> Int -> Int -> m (Int, e)
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e
-> v (PrimState m) e -> Int -> Int -> Int -> m (Int, e)
maximumChild Comparison e
cmp v (PrimState m) e
a Int
off Int
child Int
len
                      case Comparison e
cmp e
val e
ac of
                        Ordering
LT -> v (PrimState m) e -> Int -> e -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
unsafeWrite v (PrimState m) e
a (Int
root Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
off) e
ac m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> e -> Int -> Int -> m ()
sift e
val Int
child' Int
len
                        Ordering
_  -> v (PrimState m) e -> Int -> e -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
unsafeWrite v (PrimState m) e
a (Int
root Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
off) e
val
   | Bool
otherwise = v (PrimState m) e -> Int -> e -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
unsafeWrite v (PrimState m) e
a (Int
root Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
off) e
val
  where child :: Int
child = Int
root Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
{-# INLINE siftByOffset #-}

-- Finds the maximum child of a heap node, given the indx of the first child.
maximumChild :: (PrimMonad m, MVector v e)
             => Comparison e -> v (PrimState m) e -> Int -> Int -> Int -> m (Int,  e)
maximumChild :: forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e
-> v (PrimState m) e -> Int -> Int -> Int -> m (Int, e)
maximumChild Comparison e
cmp v (PrimState m) e
a Int
off Int
child1 Int
len
  | Int
child4 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
len = do e
ac1 <- v (PrimState m) e -> Int -> m e
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
unsafeRead v (PrimState m) e
a (Int
child1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
off)
                      e
ac2 <- v (PrimState m) e -> Int -> m e
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
unsafeRead v (PrimState m) e
a (Int
child2 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
off)
                      e
ac3 <- v (PrimState m) e -> Int -> m e
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
unsafeRead v (PrimState m) e
a (Int
child3 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
off)
                      e
ac4 <- v (PrimState m) e -> Int -> m e
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
unsafeRead v (PrimState m) e
a (Int
child4 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
off)
                      (Int, e) -> m (Int, e)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ((Int, e) -> m (Int, e)) -> (Int, e) -> m (Int, e)
forall a b. (a -> b) -> a -> b
$ case Comparison e
cmp e
ac1 e
ac2 of
                                 Ordering
LT -> case Comparison e
cmp e
ac2 e
ac3 of
                                         Ordering
LT -> case Comparison e
cmp e
ac3 e
ac4 of
                                                 Ordering
LT -> (Int
child4, e
ac4)
                                                 Ordering
_  -> (Int
child3, e
ac3)
                                         Ordering
_  -> case Comparison e
cmp e
ac2 e
ac4 of
                                                 Ordering
LT -> (Int
child4, e
ac4)
                                                 Ordering
_  -> (Int
child2, e
ac2)
                                 Ordering
_  -> case Comparison e
cmp e
ac1 e
ac3 of
                                         Ordering
LT -> case Comparison e
cmp e
ac3 e
ac4 of
                                                 Ordering
LT -> (Int
child4, e
ac4)
                                                 Ordering
_  -> (Int
child3, e
ac3)
                                         Ordering
_  -> case Comparison e
cmp e
ac1 e
ac4 of
                                                 Ordering
LT -> (Int
child4, e
ac4)
                                                 Ordering
_  -> (Int
child1, e
ac1)
  | Int
child3 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
len = do e
ac1 <- v (PrimState m) e -> Int -> m e
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
unsafeRead v (PrimState m) e
a (Int
child1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
off)
                      e
ac2 <- v (PrimState m) e -> Int -> m e
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
unsafeRead v (PrimState m) e
a (Int
child2 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
off)
                      e
ac3 <- v (PrimState m) e -> Int -> m e
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
unsafeRead v (PrimState m) e
a (Int
child3 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
off)
                      (Int, e) -> m (Int, e)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ((Int, e) -> m (Int, e)) -> (Int, e) -> m (Int, e)
forall a b. (a -> b) -> a -> b
$ case Comparison e
cmp e
ac1 e
ac2 of
                                 Ordering
LT -> case Comparison e
cmp e
ac2 e
ac3 of
                                         Ordering
LT -> (Int
child3, e
ac3)
                                         Ordering
_  -> (Int
child2, e
ac2)
                                 Ordering
_  -> case Comparison e
cmp e
ac1 e
ac3 of
                                         Ordering
LT -> (Int
child3, e
ac3)
                                         Ordering
_  -> (Int
child1, e
ac1)
  | Int
child2 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
len = do e
ac1 <- v (PrimState m) e -> Int -> m e
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
unsafeRead v (PrimState m) e
a (Int
child1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
off)
                      e
ac2 <- v (PrimState m) e -> Int -> m e
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
unsafeRead v (PrimState m) e
a (Int
child2 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
off)
                      (Int, e) -> m (Int, e)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ((Int, e) -> m (Int, e)) -> (Int, e) -> m (Int, e)
forall a b. (a -> b) -> a -> b
$ case Comparison e
cmp e
ac1 e
ac2 of
                                 Ordering
LT -> (Int
child2, e
ac2)
                                 Ordering
_  -> (Int
child1, e
ac1)
  | Bool
otherwise    = do e
ac1 <- v (PrimState m) e -> Int -> m e
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
unsafeRead v (PrimState m) e
a (Int
child1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
off) ; (Int, e) -> m (Int, e)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int
child1, e
ac1)
 where
 child2 :: Int
child2 = Int
child1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
 child3 :: Int
child3 = Int
child1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2
 child4 :: Int
child4 = Int
child1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
3
{-# INLINE maximumChild #-}