{-# LANGUAGE MultiParamTypeClasses #-}
module Basement.Alg.Mutable
    ( inplaceSortBy
    ) where

import           GHC.Types
import           GHC.Prim
import           Basement.Compat.Base
import           Basement.Numerical.Additive
import           Basement.Numerical.Multiplicative
import           Basement.Types.OffsetSize
import           Basement.PrimType
import           Basement.Monad
import           Basement.Alg.Class

inplaceSortBy :: (PrimMonad prim, RandomAccess container prim ty) 
              => (ty -> ty -> Ordering)
              -- ^ Function defining the ordering relationship
              -> (Offset ty) -- ^ Offset to first element to sort
              -> (CountOf ty) -- ^ Number of elements to sort
              -> container -- ^ Data to be sorted
              -> prim ()
inplaceSortBy :: forall (prim :: * -> *) container ty.
(PrimMonad prim, RandomAccess container prim ty) =>
(ty -> ty -> Ordering)
-> Offset ty -> CountOf ty -> container -> prim ()
inplaceSortBy ty -> ty -> Ordering
ford Offset ty
start CountOf ty
len container
mvec
    = Offset ty -> Offset ty -> prim ()
qsort Offset ty
start (Offset ty
start forall ty. Offset ty -> CountOf ty -> Offset ty
`offsetPlusE` CountOf ty
len forall a. Offset a -> Offset a -> Offset a
`offsetSub` Offset ty
1)
    where
        qsort :: Offset ty -> Offset ty -> prim ()
qsort Offset ty
lo Offset ty
hi
            | Offset ty
lo forall a. Ord a => a -> a -> Bool
>= Offset ty
hi  = forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            | Bool
otherwise = do
                Offset ty
p <- Offset ty -> Offset ty -> prim (Offset ty)
partition Offset ty
lo Offset ty
hi
                Offset ty -> Offset ty -> prim ()
qsort Offset ty
lo (forall a. Enum a => a -> a
pred Offset ty
p)
                Offset ty -> Offset ty -> prim ()
qsort (Offset ty
pforall a. Additive a => a -> a -> a
+Offset ty
1) Offset ty
hi
        pivotStrategy :: Offset ty -> Offset ty -> prim ty
pivotStrategy (Offset Int
low) hi :: Offset ty
hi@(Offset Int
high) = do
            let mid :: Offset ty
mid = forall ty. Int -> Offset ty
Offset forall a b. (a -> b) -> a -> b
$ (Int
low forall a. Additive a => a -> a -> a
+ Int
high) forall a. IDivisible a => a -> a -> a
`div` Int
2
            ty
pivot <- forall container (prim :: * -> *) ty.
RandomAccess container prim ty =>
container -> Offset ty -> prim ty
read container
mvec Offset ty
mid
            forall container (prim :: * -> *) ty.
RandomAccess container prim ty =>
container -> Offset ty -> prim ty
read container
mvec Offset ty
hi forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall container (prim :: * -> *) ty.
RandomAccess container prim ty =>
container -> Offset ty -> ty -> prim ()
write container
mvec Offset ty
mid
            forall container (prim :: * -> *) ty.
RandomAccess container prim ty =>
container -> Offset ty -> ty -> prim ()
write container
mvec Offset ty
hi ty
pivot -- move pivot @ pivotpos := hi
            forall (f :: * -> *) a. Applicative f => a -> f a
pure ty
pivot
        partition :: Offset ty -> Offset ty -> prim (Offset ty)
partition Offset ty
lo Offset ty
hi = do
            ty
pivot <- Offset ty -> Offset ty -> prim ty
pivotStrategy Offset ty
lo Offset ty
hi
            -- RETURN: index of pivot with [<pivot | pivot | >=pivot]
            -- INVARIANT: i & j are valid array indices; pivotpos==hi
            let go :: Offset ty -> Offset ty -> prim (Offset ty)
go Offset ty
i Offset ty
j = do
                    -- INVARIANT: k <= pivotpos
                    let fw :: Offset ty -> prim (Offset ty, ty)
fw Offset ty
k = do ty
ak <- forall container (prim :: * -> *) ty.
RandomAccess container prim ty =>
container -> Offset ty -> prim ty
read container
mvec Offset ty
k
                                  if ty -> ty -> Ordering
ford ty
ak ty
pivot forall a. Eq a => a -> a -> Bool
== Ordering
LT 
                                    then Offset ty -> prim (Offset ty, ty)
fw (Offset ty
kforall a. Additive a => a -> a -> a
+Offset ty
1)
                                    else forall (f :: * -> *) a. Applicative f => a -> f a
pure (Offset ty
k, ty
ak)
                    (Offset ty
i, ty
ai) <- Offset ty -> prim (Offset ty, ty)
fw Offset ty
i -- POST: ai >= pivot
                    -- INVARIANT: k >= i
                    let bw :: Offset ty -> prim (Offset ty, ty)
bw Offset ty
k | Offset ty
kforall a. Eq a => a -> a -> Bool
==Offset ty
i = forall (f :: * -> *) a. Applicative f => a -> f a
pure (Offset ty
i, ty
ai)
                             | Bool
otherwise = do ty
ak <- forall container (prim :: * -> *) ty.
RandomAccess container prim ty =>
container -> Offset ty -> prim ty
read container
mvec Offset ty
k
                                              if ty -> ty -> Ordering
ford ty
ak ty
pivot forall a. Eq a => a -> a -> Bool
/= Ordering
LT
                                                then Offset ty -> prim (Offset ty, ty)
bw (forall a. Enum a => a -> a
pred Offset ty
k)
                                                else forall (f :: * -> *) a. Applicative f => a -> f a
pure (Offset ty
k, ty
ak)
                    (Offset ty
j, ty
aj) <- Offset ty -> prim (Offset ty, ty)
bw Offset ty
j -- POST: i==j OR (aj<pivot AND j<pivotpos)
                    -- POST: ai>=pivot AND (i==j OR aj<pivot AND (j<pivotpos))
                    if Offset ty
i forall a. Ord a => a -> a -> Bool
< Offset ty
j
                        then do -- (ai>=p AND aj<p) AND (i<j<pivotpos)
                            -- swap two non-pivot elements and proceed
                            forall container (prim :: * -> *) ty.
RandomAccess container prim ty =>
container -> Offset ty -> ty -> prim ()
write container
mvec Offset ty
i ty
aj
                            forall container (prim :: * -> *) ty.
RandomAccess container prim ty =>
container -> Offset ty -> ty -> prim ()
write container
mvec Offset ty
j ty
ai
                            -- POST: (ai < pivot <= aj)
                            Offset ty -> Offset ty -> prim (Offset ty)
go (Offset ty
iforall a. Additive a => a -> a -> a
+Offset ty
1) (forall a. Enum a => a -> a
pred Offset ty
j)
                        else do -- ai >= pivot 
                            -- complete partitioning by swapping pivot to the center
                            forall container (prim :: * -> *) ty.
RandomAccess container prim ty =>
container -> Offset ty -> ty -> prim ()
write container
mvec Offset ty
hi ty
ai 
                            forall container (prim :: * -> *) ty.
RandomAccess container prim ty =>
container -> Offset ty -> ty -> prim ()
write container
mvec Offset ty
i ty
pivot
                            forall (f :: * -> *) a. Applicative f => a -> f a
pure Offset ty
i
            Offset ty -> Offset ty -> prim (Offset ty)
go Offset ty
lo Offset ty
hi
{-# INLINE inplaceSortBy #-}