src/Data/Sequence/Internal/Sorting.hs

    {-# LANGUAGE BangPatterns #-}

    {-# OPTIONS_HADDOCK not-home #-}

    -- |
    --
    -- = WARNING
    --
    -- This module is considered __internal__.
    --
    -- The Package Versioning Policy __does not apply__.
    --
    -- The contents of this module may change __in any way whatsoever__
    -- and __without any warning__ between minor versions of this package.
    --
    -- Authors importing this module are expected to track development
    -- closely.
    --
    -- = Description
    --
    -- This module provides the various sorting implementations for
    -- "Data.Sequence". Further notes are available in the file sorting.md
    -- (in this directory).

    module Data.Sequence.Internal.Sorting
      (
       -- * Sort Functions
       sort
      ,sortBy
      ,sortOn
      ,unstableSort
      ,unstableSortBy
      ,unstableSortOn
      ,
       -- * Heaps
       -- $heaps
       Queue(..)
      ,QList(..)
      ,IndexedQueue(..)
      ,IQList(..)
      ,TaggedQueue(..)
      ,TQList(..)
      ,IndexedTaggedQueue(..)
      ,ITQList(..)
      ,
       -- * Merges
       -- $merges
       mergeQ
      ,mergeIQ
      ,mergeTQ
      ,mergeITQ
      ,
       -- * popMin
       -- $popMin
       popMinQ
      ,popMinIQ
      ,popMinTQ
      ,popMinITQ
      ,
       -- * Building
       -- $building
       buildQ
      ,buildIQ
      ,buildTQ
      ,buildITQ
      ,
       -- * Special folds
       -- $folds
       foldToMaybeTree
      ,foldToMaybeWithIndexTree)
      where

    import Data.Sequence.Internal
           (Elem(..), Seq(..), Node(..), Digit(..), Sized(..), FingerTree(..),
            replicateA, foldDigit, foldNode, foldWithIndexDigit,
            foldWithIndexNode)
    import Utils.Containers.Internal.State (State(..), execState)
    -- | \( O(n \log n) \).  'sort' sorts the specified 'Seq' by the natural
    -- ordering of its elements.  The sort is stable.  If stability is not
    -- required, 'unstableSort' can be slightly faster.
    --
    -- @since 0.3.0
    sort :: Ord a => Seq a -> Seq a
    sort = sortBy compare

    -- | \( O(n \log n) \).  'sortBy' sorts the specified 'Seq' according to the
    -- specified comparator.  The sort is stable.  If stability is not required,
    -- 'unstableSortBy' can be slightly faster.
    --
    -- @since 0.3.0
    sortBy :: (a -> a -> Ordering) -> Seq a -> Seq a
    sortBy cmp (Seq xs) =
        maybe
            (Seq EmptyT)
            (execState (replicateA (size xs) (State (popMinIQ cmp))))
            (buildIQ cmp (\s (Elem x) -> IQ s x IQNil) 0 xs)

    -- | \( O(n \log n) \). 'sortOn' sorts the specified 'Seq' by comparing
    -- the results of a key function applied to each element. @'sortOn' f@ is
    -- equivalent to @'sortBy' ('compare' ``Data.Function.on`` f)@, but has the
    -- performance advantage of only evaluating @f@ once for each element in the
    -- input list. This is called the decorate-sort-undecorate paradigm, or
    -- Schwartzian transform.
    --
    -- An example of using 'sortOn' might be to sort a 'Seq' of strings
    -- according to their length:
    --
    -- > sortOn length (fromList ["alligator", "monkey", "zebra"]) == fromList ["zebra", "monkey", "alligator"]
    --
    -- If, instead, 'sortBy' had been used, 'length' would be evaluated on
    -- every comparison, giving \( O(n \log n) \) evaluations, rather than
    -- \( O(n) \).
    --
    -- If @f@ is very cheap (for example a record selector, or 'fst'),
    -- @'sortBy' ('compare' ``Data.Function.on`` f)@ will be faster than
    -- @'sortOn' f@.
    --
    -- @since 0.5.11
    sortOn :: Ord b => (a -> b) -> Seq a -> Seq a
    sortOn f (Seq xs) =
        maybe
           (Seq EmptyT)
           (execState (replicateA (size xs) (State (popMinITQ compare))))
           (buildITQ compare (\s (Elem x) -> ITQ s (f x) x ITQNil) 0 xs)

    -- | \( O(n \log n) \).  'unstableSort' sorts the specified 'Seq' by
    -- the natural ordering of its elements, but the sort is not stable.
    -- This algorithm is frequently faster and uses less memory than 'sort'.

    -- Notes on the implementation and choice of heap are available in
    -- the file sorting.md (in this directory).
    --
    -- @since 0.3.0
    unstableSort :: Ord a => Seq a -> Seq a
    unstableSort = unstableSortBy compare

    -- | \( O(n \log n) \).  A generalization of 'unstableSort', 'unstableSortBy'
    -- takes an arbitrary comparator and sorts the specified sequence.
    -- The sort is not stable.  This algorithm is frequently faster and
    -- uses less memory than 'sortBy'.
    --
    -- @since 0.3.0
    unstableSortBy :: (a -> a -> Ordering) -> Seq a -> Seq a
    unstableSortBy cmp (Seq xs) =
        maybe
            (Seq EmptyT)
            (execState (replicateA (size xs) (State (popMinQ cmp))))
            (buildQ cmp (\(Elem x) -> Q x Nil) xs)

    -- | \( O(n \log n) \). 'unstableSortOn' sorts the specified 'Seq' by
    -- comparing the results of a key function applied to each element.
    -- @'unstableSortOn' f@ is equivalent to @'unstableSortBy' ('compare' ``Data.Function.on`` f)@,
    -- but has the performance advantage of only evaluating @f@ once for each
    -- element in the input list. This is called the
    -- decorate-sort-undecorate paradigm, or Schwartzian transform.
    --
    -- An example of using 'unstableSortOn' might be to sort a 'Seq' of strings
    -- according to their length:
    --
    -- > unstableSortOn length (fromList ["alligator", "monkey", "zebra"]) == fromList ["zebra", "monkey", "alligator"]
    --
    -- If, instead, 'unstableSortBy' had been used, 'length' would be evaluated on
    -- every comparison, giving \( O(n \log n) \) evaluations, rather than
    -- \( O(n) \).
    --
    -- If @f@ is very cheap (for example a record selector, or 'fst'),
    -- @'unstableSortBy' ('compare' ``Data.Function.on`` f)@ will be faster than
    -- @'unstableSortOn' f@.
    --
    -- @since 0.5.11
    unstableSortOn :: Ord b => (a -> b) -> Seq a -> Seq a
    unstableSortOn f (Seq xs) =
        maybe
           (Seq EmptyT)
           (execState (replicateA (size xs) (State (popMinTQ compare))))
           (buildTQ compare (\(Elem x) -> TQ (f x) x TQNil) xs)

    ------------------------------------------------------------------------
    -- $heaps
    --
    -- The following are definitions for various specialized pairing heaps.
    --
    -- All of the heaps are defined to be non-empty, which speeds up the
    -- merge functions.
    ------------------------------------------------------------------------

    -- | A simple pairing heap.
    data Queue e = Q !e (QList e)
    data QList e
        = Nil
        | QCons {-# UNPACK #-} !(Queue e)
                (QList e)

    -- | A pairing heap tagged with the original position of elements,
    -- to allow for stable sorting.
    data IndexedQueue e =
        IQ {-# UNPACK #-} !Int !e (IQList e)
    data IQList e
        = IQNil
        | IQCons {-# UNPACK #-} !(IndexedQueue e)
                 (IQList e)

    -- | A pairing heap tagged with some key for sorting elements, for use
    -- in 'unstableSortOn'.
    data TaggedQueue a b =
        TQ !a b (TQList a b)
    data TQList a b
        = TQNil
        | TQCons {-# UNPACK #-} !(TaggedQueue a b)
                 (TQList a b)

    -- | A pairing heap tagged with both a key and the original position
    -- of its elements, for use in 'sortOn'.
    data IndexedTaggedQueue e a =
        ITQ {-# UNPACK #-} !Int !e a (ITQList e a)
    data ITQList e a
        = ITQNil
        | ITQCons {-# UNPACK #-} !(IndexedTaggedQueue e a)
                  (ITQList e a)

    infixr 8 `ITQCons`, `TQCons`, `QCons`, `IQCons`

    ------------------------------------------------------------------------
    -- $merges
    --
    -- The following are definitions for "merge" for each of the heaps
    -- above. Each takes a comparison function which is used to order the
    -- elements.
    ------------------------------------------------------------------------

    -- | 'mergeQ' merges two 'Queue's.
    mergeQ :: (a -> a -> Ordering) -> Queue a -> Queue a -> Queue a
    mergeQ cmp q1@(Q x1 ts1) q2@(Q x2 ts2)
      | cmp x1 x2 == GT = Q x2 (q1 `QCons` ts2)
      | otherwise       = Q x1 (q2 `QCons` ts1)

    -- | 'mergeTQ' merges two 'TaggedQueue's, based on the tag value.
    mergeTQ :: (a -> a -> Ordering)
            -> TaggedQueue a b
            -> TaggedQueue a b
            -> TaggedQueue a b
    mergeTQ cmp q1@(TQ x1 y1 ts1) q2@(TQ x2 y2 ts2)
      | cmp x1 x2 == GT = TQ x2 y2 (q1 `TQCons` ts2)
      | otherwise       = TQ x1 y1 (q2 `TQCons` ts1)

    -- | 'mergeIQ' merges two 'IndexedQueue's, taking into account the
    -- original position of the elements.
    mergeIQ :: (a -> a -> Ordering)
            -> IndexedQueue a
            -> IndexedQueue a
            -> IndexedQueue a
    mergeIQ cmp q1@(IQ i1 x1 ts1) q2@(IQ i2 x2 ts2) =
        case cmp x1 x2 of
            LT -> IQ i1 x1 (q2 `IQCons` ts1)
            EQ | i1 <= i2 -> IQ i1 x1 (q2 `IQCons` ts1)
            _ -> IQ i2 x2 (q1 `IQCons` ts2)

    -- | 'mergeITQ' merges two 'IndexedTaggedQueue's, based on the tag
    -- value, taking into account the original position of the elements.
    mergeITQ
        :: (a -> a -> Ordering)
        -> IndexedTaggedQueue a b
        -> IndexedTaggedQueue a b
        -> IndexedTaggedQueue a b
    mergeITQ cmp q1@(ITQ i1 x1 y1 ts1) q2@(ITQ i2 x2 y2 ts2) =
        case cmp x1 x2 of
            LT -> ITQ i1 x1 y1 (q2 `ITQCons` ts1)
            EQ | i1 <= i2 -> ITQ i1 x1 y1 (q2 `ITQCons` ts1)
            _ -> ITQ i2 x2 y2 (q1 `ITQCons` ts2)

    ------------------------------------------------------------------------
    -- $popMin
    --
    -- The following are definitions for @popMin@, a function which
    -- constructs a stateful action which pops the smallest element from the
    -- queue, where "smallest" is according to the supplied comparison
    -- function.
    --
    -- All of the functions fail on an empty queue.
    --
    -- Each of these functions is structured something like this:
    --
    -- @popMinQ cmp (Q x ts) = (mergeQs ts, x)@
    --
    -- The reason the call to @mergeQs@ is lazy is that it will be bottom
    -- for the last element in the queue, preventing us from evaluating the
    -- fully sorted sequence.
    ------------------------------------------------------------------------

    -- | Pop the smallest element from the queue, using the supplied
    -- comparator.
    popMinQ :: (e -> e -> Ordering) -> Queue e -> (Queue e, e)
    popMinQ cmp (Q x xs) = (mergeQs xs, x)
      where
        mergeQs (t `QCons` Nil) = t
        mergeQs (t1 `QCons` t2 `QCons` Nil) = t1 <+> t2
        mergeQs (t1 `QCons` t2 `QCons` ts) = (t1 <+> t2) <+> mergeQs ts
        mergeQs Nil = error "popMinQ: tried to pop from empty queue"
        (<+>) = mergeQ cmp

    -- | Pop the smallest element from the queue, using the supplied
    -- comparator, deferring to the item's original position when the
    -- comparator returns 'EQ'.
    popMinIQ :: (e -> e -> Ordering) -> IndexedQueue e -> (IndexedQueue e, e)
    popMinIQ cmp (IQ _ x xs) = (mergeQs xs, x)
      where
        mergeQs (t `IQCons` IQNil) = t
        mergeQs (t1 `IQCons` t2 `IQCons` IQNil) = t1 <+> t2
        mergeQs (t1 `IQCons` t2 `IQCons` ts) = (t1 <+> t2) <+> mergeQs ts
        mergeQs IQNil = error "popMinQ: tried to pop from empty queue"
        (<+>) = mergeIQ cmp

    -- | Pop the smallest element from the queue, using the supplied
    -- comparator on the tag.
    popMinTQ :: (a -> a -> Ordering) -> TaggedQueue a b -> (TaggedQueue a b, b)
    popMinTQ cmp (TQ _ x xs) = (mergeQs xs, x)
      where
        mergeQs (t `TQCons` TQNil) = t
        mergeQs (t1 `TQCons` t2 `TQCons` TQNil) = t1 <+> t2
        mergeQs (t1 `TQCons` t2 `TQCons` ts) = (t1 <+> t2) <+> mergeQs ts
        mergeQs TQNil = error "popMinQ: tried to pop from empty queue"
        (<+>) = mergeTQ cmp

    -- | Pop the smallest element from the queue, using the supplied
    -- comparator on the tag, deferring to the item's original position
    -- when the comparator returns 'EQ'.
    popMinITQ :: (e -> e -> Ordering)
              -> IndexedTaggedQueue e b
              -> (IndexedTaggedQueue e b, b)
    popMinITQ cmp (ITQ _ _ x xs) = (mergeQs xs, x)
      where
        mergeQs (t `ITQCons` ITQNil) = t
        mergeQs (t1 `ITQCons` t2 `ITQCons` ITQNil) = t1 <+> t2
        mergeQs (t1 `ITQCons` t2 `ITQCons` ts) = (t1 <+> t2) <+> mergeQs ts
        mergeQs ITQNil = error "popMinQ: tried to pop from empty queue"
        (<+>) = mergeITQ cmp

    ------------------------------------------------------------------------
    -- $building
    --
    -- The following are definitions for functions to build queues, given a
    -- comparison function.
    ------------------------------------------------------------------------

    buildQ :: (b -> b -> Ordering) -> (a -> Queue b) -> FingerTree a -> Maybe (Queue b)
    buildQ cmp = foldToMaybeTree (mergeQ cmp)

    buildIQ
        :: (b -> b -> Ordering)
        -> (Int -> Elem y -> IndexedQueue b)
        -> Int
        -> FingerTree (Elem y)
        -> Maybe (IndexedQueue b)
    buildIQ cmp = foldToMaybeWithIndexTree (mergeIQ cmp)

    buildTQ
        :: (b -> b -> Ordering)
        -> (a -> TaggedQueue b c)
        -> FingerTree a
        -> Maybe (TaggedQueue b c)
    buildTQ cmp = foldToMaybeTree (mergeTQ cmp)

    buildITQ
        :: (b -> b -> Ordering)
        -> (Int -> Elem y -> IndexedTaggedQueue b c)
        -> Int
        -> FingerTree (Elem y)
        -> Maybe (IndexedTaggedQueue b c)
    buildITQ cmp = foldToMaybeWithIndexTree (mergeITQ cmp)

    ------------------------------------------------------------------------
    -- $folds
    --
    -- A big part of what makes the heaps fast is that they're non empty,
    -- so the merge function can avoid an extra case match. To take
    -- advantage of this, though, we need specialized versions of 'foldMap'
    -- and 'Data.Sequence.foldMapWithIndex', which can alternate between
    -- calling the faster semigroup-like merge when folding over non empty
    -- structures (like 'Node' and 'Digit'), and the
    -- 'Data.Semirgroup.Option'-like mappend, when folding over structures
    -- which can be empty (like 'FingerTree').
    ------------------------------------------------------------------------

    -- | A 'foldMap'-like function, specialized to the
    -- 'Data.Semigroup.Option' monoid, which takes advantage of the
    -- internal structure of 'Seq' to avoid wrapping in 'Maybe' at certain
    -- points.
    foldToMaybeTree :: (b -> b -> b) -> (a -> b) -> FingerTree a -> Maybe b
    foldToMaybeTree _ _ EmptyT = Nothing
    foldToMaybeTree _ f (Single xs) = Just (f xs)
    foldToMaybeTree (<+>) f (Deep _ pr m sf) =
        Just (maybe (pr' <+> sf') ((pr' <+> sf') <+>) m')
      where
        pr' = foldDigit (<+>) f pr
        sf' = foldDigit (<+>) f sf
        m' = foldToMaybeTree (<+>) (foldNode (<+>) f) m

    -- | A 'Data.Sequence.foldMapWithIndex'-like function, specialized to the
    -- 'Data.Semigroup.Option' monoid, which takes advantage of the
    -- internal structure of 'Seq' to avoid wrapping in 'Maybe' at certain
    -- points.
    foldToMaybeWithIndexTree :: (b -> b -> b)
                             -> (Int -> Elem y -> b)
                             -> Int
                             -> FingerTree (Elem y)
                             -> Maybe b
    foldToMaybeWithIndexTree = foldToMaybeWithIndexTree'
      where
        {-# SPECIALISE foldToMaybeWithIndexTree' :: (b -> b -> b) -> (Int -> Elem y -> b) -> Int -> FingerTree (Elem y) -> Maybe b #-}
        {-# SPECIALISE foldToMaybeWithIndexTree' :: (b -> b -> b) -> (Int -> Node y -> b) -> Int -> FingerTree (Node y) -> Maybe b #-}
        foldToMaybeWithIndexTree'
            :: Sized a
            => (b -> b -> b) -> (Int -> a -> b) -> Int -> FingerTree a -> Maybe b
        foldToMaybeWithIndexTree' _ _ !_s EmptyT = Nothing
        foldToMaybeWithIndexTree' _ f s (Single xs) = Just (f s xs)
        foldToMaybeWithIndexTree' (<+>) f s (Deep _ pr m sf) =
            Just (maybe (pr' <+> sf') ((pr' <+> sf') <+>) m')
          where
            pr' = digit (<+>) f s pr
            sf' = digit (<+>) f sPsprm sf
            m' = foldToMaybeWithIndexTree' (<+>) (node (<+>) f) sPspr m
            !sPspr = s + size pr
            !sPsprm = sPspr + size m
        {-# SPECIALISE digit :: (b -> b -> b) -> (Int -> Elem y -> b) -> Int -> Digit (Elem y) -> b #-}
        {-# SPECIALISE digit :: (b -> b -> b) -> (Int -> Node y -> b) -> Int -> Digit (Node y) -> b #-}
        digit
            :: Sized a
            => (b -> b -> b) -> (Int -> a -> b) -> Int -> Digit a -> b
        digit = foldWithIndexDigit
        {-# SPECIALISE node :: (b -> b -> b) -> (Int -> Elem y -> b) -> Int -> Node (Elem y) -> b #-}
        {-# SPECIALISE node :: (b -> b -> b) -> (Int -> Node y -> b) -> Int -> Node (Node y) -> b #-}
        node
            :: Sized a
            => (b -> b -> b) -> (Int -> a -> b) -> Int -> Node a -> b
        node = foldWithIndexNode
    {-# INLINE foldToMaybeWithIndexTree #-}