    {-# 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 ford start len mvec
        = qsort start (start `offsetPlusE` len `offsetSub` 1)
        where
            qsort lo hi
                | lo >= hi  = pure ()
                | otherwise = do
                    p <- partition lo hi
                    qsort lo (pred p)
                    qsort (p+1) hi
            pivotStrategy (Offset low) hi@(Offset high) = do
                let mid = Offset $ (low + high) `div` 2
                pivot <- read mvec mid
                read mvec hi >>= write mvec mid
                write mvec hi pivot -- move pivot @ pivotpos := hi
                pure pivot
            partition lo hi = do
                pivot <- pivotStrategy lo hi
                -- RETURN: index of pivot with [<pivot | pivot | >=pivot]
                -- INVARIANT: i & j are valid array indices; pivotpos==hi
                let go i j = do
                        -- INVARIANT: k <= pivotpos
                        let fw k = do ak <- read mvec k
                                      if ford ak pivot == LT
                                        then fw (k+1)
                                        else pure (k, ak)
                        (i, ai) <- fw i -- POST: ai >= pivot
                        -- INVARIANT: k >= i
                        let bw k | k==i = pure (i, ai)
                                 | otherwise = do ak <- read mvec k
                                                  if ford ak pivot /= LT
                                                    then bw (pred k)
                                                    else pure (k, ak)
                        (j, aj) <- bw j -- POST: i==j OR (aj<pivot AND j<pivotpos)
                        -- POST: ai>=pivot AND (i==j OR aj<pivot AND (j<pivotpos))
                        if i < j
                            then do -- (ai>=p AND aj<p) AND (i<j<pivotpos)
                                -- swap two non-pivot elements and proceed
                                write mvec i aj
                                write mvec j ai
                                -- POST: (ai < pivot <= aj)
                                go (i+1) (pred j)
                            else do -- ai >= pivot
                                -- complete partitioning by swapping pivot to the center
                                write mvec hi ai
                                write mvec i pivot
                                pure i
                go lo hi
    {-# INLINE inplaceSortBy #-}