liquidhaskell-0.8.10.7: benchmarks/vector-algorithms-0.5.4.2/Data/Vector/Algorithms/Intro.hs
{-# LANGUAGE PartialTypeSignatures, TypeOperators, BangPatterns, ScopedTypeVariables #-}
-- ---------------------------------------------------------------------------
-- |
-- Module : Data.Vector.Algorithms.Intro
-- Copyright : (c) 2008-2011 Dan Doel
-- Maintainer : Dan Doel <dan.doel@gmail.com>
-- Stability : Experimental
-- Portability : Non-portable (type operators, bang patterns)
--
-- This module implements various algorithms based on the introsort algorithm,
-- originally described by David R. Musser in the paper /Introspective Sorting
-- and Selection Algorithms/. It is also in widespread practical use, as the
-- standard unstable sort used in the C++ Standard Template Library.
--
-- Introsort is at its core a quicksort. The version implemented here has the
-- following optimizations that make it perform better in practice:
--
-- * Small segments of the array are left unsorted until a final insertion
-- sort pass. This is faster than recursing all the way down to
-- one-element arrays.
--
-- * The pivot for segment [l,u) is chosen as the median of the elements at
-- l, u-1 and (u+l)/2. This yields good behavior on mostly sorted (or
-- reverse-sorted) arrays.
--
-- * The algorithm tracks its recursion depth, and if it decides it is
-- taking too long (depth greater than 2 * lg n), it switches to a heap
-- sort to maintain O(n lg n) worst case behavior. (This is what makes the
-- algorithm introsort).
module Data.Vector.Algorithms.Intro
( -- * Sorting
sort
, sortBy
, sortByBounds
-- * Selecting
, select
, selectBy
, selectByBounds
-- * Partial sorting
, 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, shiftRI)
import qualified Data.Vector.Algorithms.Insertion as I
import qualified Data.Vector.Algorithms.Optimal as O
import qualified Data.Vector.Algorithms.Heap as H
-- | Sorts an entire array using the default ordering.
{-@ sort :: (PrimMonad m, MVector v e, Ord e) => {v: (v (PrimState m) e) | 0 < (vsize v)} -> m () @-}
sort :: (PrimMonad m, MVector v e, Ord e) => v (PrimState m) e -> m ()
sort = sortBy compare
{-# INLINABLE sort #-}
-- | Sorts an entire array using a custom ordering.
{-@ sortBy :: (PrimMonad m, MVector v e) => Comparison e -> {v: (v (PrimState m) e) | 0 < (vsize v)} -> m () @-}
sortBy :: (PrimMonad m, MVector v e) => Comparison e -> v (PrimState m) e -> m ()
sortBy cmp a = sortByBounds cmp a 0 (length a)
{-# INLINE sortBy #-}
-- | Sorts a portion of an array [l,u) using a custom ordering
{-@ sortByBounds :: (PrimMonad m, MVector v e)
=> Comparison e -> vec:(v (PrimState m) e)
-> l:(OkIdx vec) -> u:{v:Nat | (InRngL v l (vsize vec))}
-> m ()
@-}
sortByBounds :: (PrimMonad m, MVector v e)
=> Comparison e -> v (PrimState m) e -> Int -> Int -> m ()
sortByBounds cmp a l u
| len < 2 = return ()
| len == 2 = O.sort2ByOffset cmp a l
| len == 3 = O.sort3ByOffset cmp a l
| len == 4 = O.sort4ByOffset cmp a l
| otherwise = introsort cmp a (ilg len) l u
where len = u - l
{-# INLINE sortByBounds #-}
-- Internal version of the introsort loop which allows partial
-- sort functions to call with a specified bound on iterations.
{-@ introsort :: (PrimMonad m, MVector v e)
=> Comparison e -> vec:(v (PrimState m) e)
-> Nat -> l:(OkIdx vec) -> u:{v:Nat | (InRng v l (vsize vec))}
-> m ()
@-}
introsort :: (PrimMonad m, MVector v e)
=> Comparison e -> v (PrimState m) e -> Int -> Int -> Int -> m ()
introsort cmp a i l u = sort i l u >> I.sortByBounds cmp a l u
where
sort :: Int -> Int -> Int -> _
sort 0 l u = H.sortByBounds cmp a l u
{- LIQUID WITNESS -}
sort (d :: Int) l u
| len < threshold = return ()
| otherwise = do O.sort3ByIndex cmp a c l (u-1) -- sort the median into the lowest position
p <- unsafeRead a l
mid <- partitionBy cmp a p (l+1) u
unsafeSwap a l (mid - 1)
sort (d-1) mid u
sort (d-1) l (mid - 1)
where
len = u - l
c = (u + l) `shiftRI` 1 -- `div` 2
{-# INLINE introsort #-}
-- | Moves the least k elements to the front of the array in
-- no particular order.
{-@ select :: (PrimMonad m, MVector v e, Ord e) => (NeVec v m e) -> Pos -> m () @-}
select :: (PrimMonad m, MVector v e, Ord e) => v (PrimState m) e -> Int -> m ()
select = selectBy compare
{-# INLINE select #-}
-- | Moves the least k elements (as defined by the comparison) to
-- the front of the array in no particular order.
{-@ selectBy :: (PrimMonad m, MVector v e) => (Comparison e) -> (NeVec v m e) -> Pos -> m () @-}
selectBy :: (PrimMonad m, MVector v e)
=> Comparison e -> v (PrimState m) e -> Int -> m ()
selectBy cmp a k = selectByBounds cmp a k 0 (length a)
{-# INLINE selectBy #-}
-- | Moves the least k elements in the interval [l,u) to the positions
-- [l,k+l) in no particular order.
{-@ selectByBounds :: (PrimMonad m, MVector v e)
=> Comparison e -> vec:(NeVec v m e)
-> Pos -> l:(OkIdx vec) -> u:{v:Nat | (InRngL v l (vsize vec))}
-> m ()
@-}
selectByBounds :: (PrimMonad m, MVector v e)
=> Comparison e -> v (PrimState m) e -> Int -> Int -> Int -> m ()
selectByBounds cmp a k l u
| m >= u = return () -- LIQUID: changed, possible bugfix! (or did I add a bug?)
| otherwise = go (ilg len) l u
where
len = u - l
m = l + k
{- LIQUID WITNESS -}
go (0 :: Int) l u = H.selectByBounds cmp a k l u
go n l u = do O.sort3ByIndex cmp a c l (u-1)
p <- unsafeRead a l
mid <- partitionBy cmp a p (l+1) u
unsafeSwap a l (mid - 1)
if m > mid
then go (n-1) mid u
else if m < mid - 1
then go (n-1) l (mid - 1)
else return ()
where c = (u + l) `shiftRI` 1 -- `div` 2
{-# INLINE selectByBounds #-}
-- | Moves the least k elements to the front of the array, sorted.
{-@ partialSort :: (PrimMonad m, MVector v e, Ord e) => vec:(NeVec v m e) -> Pos -> m () @-}
partialSort :: (PrimMonad m, MVector v e, Ord e) => v (PrimState m) e -> Int -> m ()
partialSort = partialSortBy compare
{-# INLINE partialSort #-}
-- | Moves the least k elements (as defined by the comparison) to
-- the front of the array, sorted.
{-@ partialSortBy :: (PrimMonad m, MVector v e) => (Comparison e) -> vec:(NeVec v m e) -> Pos -> m () @-}
partialSortBy :: (PrimMonad m, MVector v e)
=> Comparison e -> v (PrimState m) e -> Int -> m ()
partialSortBy cmp a k = partialSortByBounds cmp a k 0 (length a)
{-# INLINE partialSortBy #-}
-- | Moves the least k elements in the interval [l,u) to the positions
-- [l,k+l), sorted.
{-@ partialSortByBounds :: (PrimMonad m, MVector v e)
=> Comparison e -> vec:(NeVec v m e)
-> k:Pos
-> l:(OkIdx vec)
-> u:{v:Nat | (InRngL v l (vsize vec))}
-> m ()
@-}
partialSortByBounds :: (PrimMonad m, MVector v e)
=> Comparison e -> v (PrimState m) e -> Int -> Int -> Int -> m ()
partialSortByBounds cmp a k l u
| m0 >= u = return () -- LIQUID: changed, possible bugfix! (or did I add a bug?)
| otherwise = go (ilg len) l u m0
where
isort = introsort cmp a
{-# INLINE [1] isort #-}
len = u - l
m0 = l + k
{-@ decrease go 1 3 @-}
go 0 l n _ = H.partialSortByBounds cmp a k l u
go n l u (m :: Int)
| l == m = return ()
| otherwise = do O.sort3ByIndex cmp a c l (u-1)
p <- unsafeRead a l
mid <- partitionBy cmp a p (l+1) u
unsafeSwap a l (mid - 1)
case compare m mid of
GT -> do isort (n-1) l (mid - 1)
go (n-1) mid u m
EQ -> isort (n-1) l m
LT -> go n l (mid - 1) m
where c = (u + l) `shiftRI` 1 -- `div` 2
{-# INLINE partialSortByBounds #-}
{-@ partitionBy :: forall m v e. (PrimMonad m, MVector v e)
=> Comparison e -> vec:(v (PrimState m) e) -> e
-> l:(OkIdx vec) -> u:{v:Nat | (InRng v l (vsize vec))}
-> m {v:Int | (InRng v l u)}
@-}
partitionBy :: forall m v e. (PrimMonad m, MVector v e)
=> Comparison e -> v (PrimState m) e -> e -> Int -> Int -> m Int
partitionBy cmp a p l u = partUp p l u (u-l)
where
-- 6.10 panics without the signatures for partUp and partDown, 6.12 and later
-- versions don't need them
{-@ decrease partUp 4 @-}
{-@ decrease partDown 4 @-}
partUp :: e -> Int -> Int -> Int -> m Int
partUp p l u _
| l < u = do e <- unsafeRead a l
case cmp e p of
LT -> partUp p (l+1) u (u-l-1)
_ -> partDown p l (u-1) (u-l-1)
| otherwise = return l
partDown :: e -> Int -> Int -> Int -> m Int
partDown p l u _
| l < u = do e <- unsafeRead a u
case cmp p e of
LT -> partDown p l (u-1) (u-l-1)
_ -> unsafeSwap a l u >> partUp p (l+1) u (u-l-1)
| otherwise = return l
{-# INLINE partitionBy #-}
-- computes the number of recursive calls after which heapsort should
-- be invoked given the lower and upper indices of the array to be sorted
{-@ ilg :: Pos -> Nat @-}
ilg :: Int -> Int
ilg m = 2 * loop m 0
where
{-@ loop :: n:Nat -> {v:Nat | ((n = 0) => (v > 0))} -> Nat @-}
loop :: Int -> Int -> Int
loop 0 !k = k - 1
loop n !k = loop (n `shiftRI` 1) (k+1)
-- the size of array at which the introsort algorithm switches to insertion sort
{-@ threshold :: {v:Int | v = 18} @-}
threshold :: Int
threshold = 18
{-# INLINE threshold #-}