{-# 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 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 #-}
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 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 (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 -> 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 Comparison e
cmp v (PrimState m) e
a Int
i Int
l 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 Int
0 Int
l 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 Int
d Int
l 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
-Int
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
+Int
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
- Int
1)
Int -> Int -> Int -> m ()
sort (Int
dInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Int
mid Int
u
Int -> Int -> Int -> m ()
sort (Int
dInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Int
l (Int
mid Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
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 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 #-}
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 Comparison e
cmp v (PrimState m) e
a Int
k Int
l 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 Int
0 Int
l Int
m 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 Int
n Int
l Int
m 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
-Int
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
+Int
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
- Int
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
-Int
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
- Int
1
then Int -> Int -> Int -> Int -> m ()
go (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Int
l Int
m (Int
mid Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
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 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 #-}
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 Comparison e
cmp v (PrimState m) e
a Int
k Int
l 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 = let k' :: Int
k' = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Int
uInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
l) Int
k
in 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 Int
0 Int
l Int
m 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 Int
n Int
l Int
m 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
-Int
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
+Int
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
- Int
1)
case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
m Int
mid of
Ordering
GT -> do Int -> Int -> Int -> m ()
isort (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Int
l (Int
mid Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
Int -> Int -> Int -> Int -> m ()
go (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Int
mid Int
m Int
u
Ordering
EQ -> Int -> Int -> Int -> m ()
isort (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Int
l Int
m
Ordering
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
- Int
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 Comparison e
cmp 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 e
p Int
l 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
Ordering
LT -> e -> Int -> Int -> m Int
partUp e
p (Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
u
Ordering
_ -> e -> Int -> Int -> m Int
partDown e
p Int
l (Int
uInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
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 e
p Int
l 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
Ordering
LT -> e -> Int -> Int -> m Int
partDown e
p Int
l (Int
uInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
Ordering
_ -> 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
+Int
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 Int
m = Int
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 Int
0
where
loop :: t -> t -> t
loop t
0 !t
k = t
k t -> t -> t
forall a. Num a => a -> a -> a
- t
1
loop t
n !t
k = t -> t -> t
loop (t
n t -> Int -> t
forall a. Bits a => a -> Int -> a
`shiftR` Int
1) (t
kt -> t -> t
forall a. Num a => a -> a -> a
+t
1)
threshold :: Int
threshold :: Int
threshold = Int
18
{-# INLINE threshold #-}