--------------------------------------------------------------------------------
-- |
-- Module      :  HGeometry.Foldable.Sort
-- Copyright   :  (C) Frank Staals
-- License     :  see the LICENSE file
-- Maintainer  :  Frank Staals
--
-- Sorting foldable collections by using fast sorting functions from
-- vector-algorithms.
--
--------------------------------------------------------------------------------
module HGeometry.Foldable.Sort
  ( sortBy
  , sort
  , sortOn
  , sortOnCheap
  ) where

import           Control.Monad.ST
import           Data.Ord (comparing)
import qualified Data.Vector.Algorithms.Intro as Intro
import qualified Data.Vector.Generic as Vector
import qualified Data.Vector as BoxedVector
import qualified VectorBuilder.Builder as Builder
import qualified VectorBuilder.MVector as MVectorBuilder
import qualified VectorBuilder.Vector as VectorBuilder

--------------------------------------------------------------------------------

-- | Sort the given collection using intro sort.
--
-- \(O(n\log n)\)
sort :: forall vector f a. ( Foldable f
                           , Vector.Vector vector a
                           , Ord a
                           )
     => f a -> vector a
sort :: forall (vector :: * -> *) (f :: * -> *) a.
(Foldable f, Vector vector a, Ord a) =>
f a -> vector a
sort = (a -> a -> Ordering) -> f a -> vector a
forall (vector :: * -> *) (f :: * -> *) a.
(Foldable f, Vector vector a) =>
(a -> a -> Ordering) -> f a -> vector a
sortBy a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare
{-# INLINE sort #-}


-- | Sort the collection using the given "expensive" function; i.e. the
-- function f is called only once and cached.
--
-- \(O(Tn + n\log n)\), where \(T\) is the time required to evaluate the function
sortOn      :: forall f a b.
               (Foldable f, Ord b)
            => (a -> b) -> f a -> BoxedVector.Vector a
sortOn :: forall (f :: * -> *) a b.
(Foldable f, Ord b) =>
(a -> b) -> f a -> Vector a
sortOn a -> b
f = (Helper b a -> a) -> Vector (Helper b a) -> Vector a
forall a b. (a -> b) -> Vector a -> Vector b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(Helper b
_ a
x) -> a
x)
         (Vector (Helper b a) -> Vector a)
-> (f a -> Vector (Helper b a)) -> f a -> Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (vector :: * -> *) (f :: * -> *) a b.
(Foldable f, Vector vector a, Ord b) =>
(a -> b) -> f a -> vector a
sortOnCheap @_ @BoxedVector.Vector (\(Helper b
b a
_) -> b
b)
         (Vector (Helper b a) -> Vector (Helper b a))
-> (f a -> Vector (Helper b a)) -> f a -> Vector (Helper b a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Helper b a) -> Vector a -> Vector (Helper b a)
forall a b. (a -> b) -> Vector a -> Vector b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap  (\a
x -> b -> a -> Helper b a
forall b a. b -> a -> Helper b a
Helper (a -> b
f a
x) a
x)
         (Vector a -> Vector (Helper b a))
-> (f a -> Vector a) -> f a -> Vector (Helper b a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Builder a -> Vector a
forall (vector :: * -> *) element.
Vector vector element =>
Builder element -> vector element
VectorBuilder.build (Builder a -> Vector a) -> (f a -> Builder a) -> f a -> Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f a -> Builder a
forall (foldable :: * -> *) element.
Foldable foldable =>
foldable element -> Builder element
Builder.foldable
{-# INLINABLE sortOn #-}

data Helper b a = Helper !b a

-- | Sort the collection using the given "cheap" function; i.e. the
-- function f is called at every comparison
--
-- \(O(Tn\log n)\), where \(T\) is the time required to evaluate the function
sortOnCheap   :: forall vector f a b. ( Foldable f
                                 , Vector.Vector vector a
                                 , Ord b
                                 )
              => (a -> b) -> f a -> vector a
sortOnCheap :: forall (vector :: * -> *) (f :: * -> *) a b.
(Foldable f, Vector vector a, Ord b) =>
(a -> b) -> f a -> vector a
sortOnCheap a -> b
f = (a -> a -> Ordering) -> f a -> vector a
forall (vector :: * -> *) (f :: * -> *) a.
(Foldable f, Vector vector a) =>
(a -> a -> Ordering) -> f a -> vector a
sortBy ((a -> b) -> a -> a -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing a -> b
f)
{-# INLINE sortOnCheap #-}

-- | Sort a collection using the given comparator using intro-sort (essentially quicksort).
--
-- \(O(T n\log n)\), where \(T\) is the time to do a comparison.
sortBy        :: forall vector f a. ( Foldable f
                                    , Vector.Vector vector a)
              => (a -> a -> Ordering) -> f a -> vector a
sortBy :: forall (vector :: * -> *) (f :: * -> *) a.
(Foldable f, Vector vector a) =>
(a -> a -> Ordering) -> f a -> vector a
sortBy a -> a -> Ordering
cmp f a
xs = (forall s. ST s (vector a)) -> vector a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (vector a)) -> vector a)
-> (forall s. ST s (vector a)) -> vector a
forall a b. (a -> b) -> a -> b
$ do v <- Builder a -> ST s (Mutable vector s a)
forall (vector :: * -> * -> *) element s.
MVector vector element =>
Builder element -> ST s (vector s element)
MVectorBuilder.build (Builder a -> ST s (Mutable vector s a))
-> Builder a -> ST s (Mutable vector s a)
forall a b. (a -> b) -> a -> b
$ f a -> Builder a
forall (foldable :: * -> *) element.
Foldable foldable =>
foldable element -> Builder element
Builder.foldable f a
xs
                           Intro.sortBy cmp v
                           Vector.unsafeFreeze v
{-# INLINABLE sortBy #-}