{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Data.Vector.Algorithms.Intro
(
sort
, sortBy
, sortByBounds
, select
, selectBy
, selectByBounds
, partialSort
, partialSortBy
, partialSortByBounds
, 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, midPoint)
import qualified Data.Vector.Algorithms.Insertion as I
import qualified Data.Vector.Algorithms.Optimal as O
import qualified Data.Vector.Algorithms.Heap as H
sort :: (PrimMonad m, MVector v e, Ord e) => v (PrimState m) e -> m ()
sort :: 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 #-}
sortBy :: (PrimMonad m, MVector v e) => Comparison e -> v (PrimState m) e -> m ()
sortBy :: Comparison e -> v (PrimState m) e -> m ()
sortBy cmp :: Comparison e
cmp a :: 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 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 #-}
sortByBounds
:: (PrimMonad m, MVector v e)
=> Comparison e
-> v (PrimState m) e
-> Int
-> Int
-> m ()
sortByBounds :: Comparison e -> v (PrimState m) e -> Int -> Int -> m ()
sortByBounds cmp :: Comparison e
cmp a :: v (PrimState m) e
a l :: Int
l u :: Int
u
| Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< 2 = () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
| Int
len Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== 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
== 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
== 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 -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> Int -> Int -> m ()
introsort Comparison e
cmp v (PrimState m) e
a (Int -> Int
ilg Int
len) Int
l Int
u
where len :: Int
len = Int
u Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l
{-# INLINE sortByBounds #-}
introsort :: (PrimMonad m, MVector v e)
=> Comparison e -> v (PrimState m) e -> Int -> Int -> Int -> m ()
introsort :: Comparison e -> v (PrimState m) e -> Int -> Int -> Int -> m ()
introsort cmp :: Comparison e
cmp a :: v (PrimState m) e
a i :: Int
i l :: Int
l u :: Int
u = Int -> Int -> Int -> m ()
sort Int
i Int
l Int
u m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> 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 ()
I.sortByBounds Comparison e
cmp v (PrimState m) e
a Int
l Int
u
where
sort :: Int -> Int -> Int -> m ()
sort 0 l :: Int
l u :: 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 ()
H.sortByBounds Comparison e
cmp v (PrimState m) e
a Int
l Int
u
sort d :: Int
d l :: Int
l u :: Int
u
| Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
threshold = () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
| 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 ()
O.sort3ByIndex Comparison e
cmp v (PrimState m) e
a Int
c Int
l (Int
uInt -> Int -> Int
forall a. Num a => a -> a -> a
-1)
e
p <- 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
Int
mid <- Comparison e -> v (PrimState m) e -> e -> Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> e -> Int -> Int -> m Int
partitionBy Comparison e
cmp v (PrimState m) e
a e
p (Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
+1) Int
u
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
mid Int -> Int -> Int
forall a. Num a => a -> a -> a
- 1)
Int -> Int -> Int -> m ()
sort (Int
dInt -> Int -> Int
forall a. Num a => a -> a -> a
-1) Int
mid Int
u
Int -> Int -> Int -> m ()
sort (Int
dInt -> Int -> Int
forall a. Num a => a -> a -> a
-1) Int
l (Int
mid Int -> Int -> Int
forall a. Num a => a -> a -> a
- 1)
where
len :: Int
len = Int
u Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l
c :: Int
c = Int -> Int -> Int
midPoint Int
u Int
l
{-# INLINE introsort #-}
select
:: (PrimMonad m, MVector v e, Ord e)
=> v (PrimState m) e
-> Int
-> m ()
select :: 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 #-}
selectBy
:: (PrimMonad m, MVector v e)
=> Comparison e
-> v (PrimState m) e
-> Int
-> m ()
selectBy :: Comparison e -> v (PrimState m) e -> Int -> m ()
selectBy cmp :: Comparison e
cmp a :: v (PrimState m) e
a k :: 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 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 #-}
selectByBounds
:: (PrimMonad m, MVector v e)
=> Comparison e
-> v (PrimState m) e
-> Int
-> Int
-> Int
-> m ()
selectByBounds :: Comparison e -> v (PrimState m) e -> Int -> Int -> Int -> m ()
selectByBounds cmp :: Comparison e
cmp a :: v (PrimState m) e
a k :: Int
k l :: Int
l u :: Int
u
| Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
u = () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
| Bool
otherwise = Int -> Int -> Int -> Int -> m ()
go (Int -> Int
ilg Int
len) Int
l (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k) Int
u
where
len :: Int
len = Int
u Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l
go :: Int -> Int -> Int -> Int -> m ()
go 0 l :: Int
l m :: Int
m u :: 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 ()
H.selectByBounds Comparison e
cmp v (PrimState m) e
a (Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l) Int
l Int
u
go n :: Int
n l :: Int
l m :: Int
m u :: Int
u = 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 ()
O.sort3ByIndex Comparison e
cmp v (PrimState m) e
a Int
c Int
l (Int
uInt -> Int -> Int
forall a. Num a => a -> a -> a
-1)
e
p <- 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
Int
mid <- Comparison e -> v (PrimState m) e -> e -> Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> e -> Int -> Int -> m Int
partitionBy Comparison e
cmp v (PrimState m) e
a e
p (Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
+1) Int
u
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
mid Int -> Int -> Int
forall a. Num a => a -> a -> a
- 1)
if Int
m Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
mid
then Int -> Int -> Int -> Int -> m ()
go (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-1) Int
mid Int
m Int
u
else if Int
m Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
mid Int -> Int -> Int
forall a. Num a => a -> a -> a
- 1
then Int -> Int -> Int -> Int -> m ()
go (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-1) Int
l Int
m (Int
mid Int -> Int -> Int
forall a. Num a => a -> a -> a
- 1)
else () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
where c :: Int
c = Int -> Int -> Int
midPoint Int
u Int
l
{-# INLINE selectByBounds #-}
partialSort
:: (PrimMonad m, MVector v e, Ord e)
=> v (PrimState m) e
-> Int
-> m ()
partialSort :: 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 #-}
partialSortBy
:: (PrimMonad m, MVector v e)
=> Comparison e
-> v (PrimState m) e
-> Int
-> m ()
partialSortBy :: Comparison e -> v (PrimState m) e -> Int -> m ()
partialSortBy cmp :: Comparison e
cmp a :: v (PrimState m) e
a k :: 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 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 #-}
partialSortByBounds
:: (PrimMonad m, MVector v e)
=> Comparison e
-> v (PrimState m) e
-> Int
-> Int
-> Int
-> m ()
partialSortByBounds :: Comparison e -> v (PrimState m) e -> Int -> Int -> Int -> m ()
partialSortByBounds cmp :: Comparison e
cmp a :: v (PrimState m) e
a k :: Int
k l :: Int
l u :: Int
u
| Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
u = () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
| Bool
otherwise = Int -> Int -> Int -> Int -> m ()
go (Int -> Int
ilg Int
len) Int
l (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k) Int
u
where
isort :: Int -> Int -> Int -> m ()
isort = 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 ()
introsort Comparison e
cmp v (PrimState m) e
a
{-# INLINE [1] isort #-}
len :: Int
len = Int
u Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l
go :: Int -> Int -> Int -> Int -> m ()
go 0 l :: Int
l m :: Int
m n :: Int
n = 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 ()
H.partialSortByBounds Comparison e
cmp v (PrimState m) e
a (Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l) Int
l Int
u
go n :: Int
n l :: Int
l m :: Int
m u :: Int
u
| Int
l Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
m = () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
| 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 ()
O.sort3ByIndex Comparison e
cmp v (PrimState m) e
a Int
c Int
l (Int
uInt -> Int -> Int
forall a. Num a => a -> a -> a
-1)
e
p <- 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
Int
mid <- Comparison e -> v (PrimState m) e -> e -> Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> e -> Int -> Int -> m Int
partitionBy Comparison e
cmp v (PrimState m) e
a e
p (Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
+1) Int
u
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
mid Int -> Int -> Int
forall a. Num a => a -> a -> a
- 1)
case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
m Int
mid of
GT -> do Int -> Int -> Int -> m ()
isort (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-1) Int
l (Int
mid Int -> Int -> Int
forall a. Num a => a -> a -> a
- 1)
Int -> Int -> Int -> Int -> m ()
go (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-1) Int
mid Int
m Int
u
EQ -> Int -> Int -> Int -> m ()
isort (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-1) Int
l Int
m
LT -> Int -> Int -> Int -> Int -> m ()
go Int
n Int
l Int
m (Int
mid Int -> Int -> Int
forall a. Num a => a -> a -> a
- 1)
where c :: Int
c = Int -> Int -> Int
midPoint Int
u Int
l
{-# INLINE partialSortByBounds #-}
partitionBy :: forall m v e. (PrimMonad m, MVector v e)
=> Comparison e -> v (PrimState m) e -> e -> Int -> Int -> m Int
partitionBy :: Comparison e -> v (PrimState m) e -> e -> Int -> Int -> m Int
partitionBy cmp :: Comparison e
cmp a :: v (PrimState m) e
a = e -> Int -> Int -> m Int
partUp
where
partUp :: e -> Int -> Int -> m Int
partUp :: e -> Int -> Int -> m Int
partUp p :: e
p l :: Int
l u :: Int
u
| Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
u = do e
e <- 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
case Comparison e
cmp e
e e
p of
LT -> e -> Int -> Int -> m Int
partUp e
p (Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
+1) Int
u
_ -> e -> Int -> Int -> m Int
partDown e
p Int
l (Int
uInt -> Int -> Int
forall a. Num a => a -> a -> a
-1)
| Bool
otherwise = Int -> m Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
l
partDown :: e -> Int -> Int -> m Int
partDown :: e -> Int -> Int -> m Int
partDown p :: e
p l :: Int
l u :: Int
u
| Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
u = do e
e <- 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
p e
e of
LT -> e -> Int -> Int -> m Int
partDown e
p Int
l (Int
uInt -> Int -> Int
forall a. Num a => a -> a -> a
-1)
_ -> 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
u m () -> m Int -> m Int
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> e -> Int -> Int -> m Int
partUp e
p (Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
+1) Int
u
| Bool
otherwise = Int -> m Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
l
{-# INLINE partitionBy #-}
ilg :: Int -> Int
ilg :: Int -> Int
ilg m :: Int
m = 2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int -> Int -> Int
forall t t. (Num t, Num t, Bits t) => t -> t -> t
loop Int
m 0
where
loop :: t -> t -> t
loop 0 !t
k = t
k t -> t -> t
forall a. Num a => a -> a -> a
- 1
loop n :: t
n !t
k = t -> t -> t
loop (t
n t -> Int -> t
forall a. Bits a => a -> Int -> a
`shiftR` 1) (t
kt -> t -> t
forall a. Num a => a -> a -> a
+1)
threshold :: Int
threshold :: Int
threshold = 18
{-# INLINE threshold #-}