{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE PatternSynonyms     #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections       #-}
{-# LANGUAGE ViewPatterns        #-}

-- |
-- Module      : Data.Set.NonEmpty
-- Copyright   : (c) Justin Le 2018
-- License     : BSD3
--
-- Maintainer  : justin@jle.im
-- Stability   : experimental
-- Portability : non-portable
--
-- = Non-Empty Finite Sets
--
-- The @'NESet' e@ type represents a non-empty set of elements of type @e@.
-- Most operations require that @e@ be an instance of the 'Ord' class.
-- A 'NESet' is strict in its elements.
--
-- See documentation for 'NESet' for information on how to convert and
-- manipulate such non-empty set.
--
-- This module essentially re-imports the API of "Data.Set" and its 'Set'
-- type, along with semantics and asymptotics.  In most situations,
-- asymptotics are different only by a constant factor.  In some
-- situations, asmyptotics are even better (constant-time instead of
-- log-time).  All typeclass constraints are identical to their "Data.Set"
-- counterparts.
--
-- Because 'NESet' is implemented using 'Set', all of the caveats of using
-- 'Set' apply (such as the limitation of the maximum size of sets).
--
-- All functions take non-empty sets as inputs.  In situations where their
-- results can be guarunteed to also be non-empty, they also return
-- non-empty sets.  In situations where their results could potentially be
-- empty, 'Set' is returned instead.
--
-- Some functions ('partition', 'spanAntitone', 'split') have modified
-- return types to account for possible configurations of non-emptiness.
--
-- This module is intended to be imported qualified, to avoid name clashes
-- with "Prelude" and "Data.Set" functions:
--
-- > import qualified Data.Set.NonEmpty as NES
module Data.Set.NonEmpty (
  -- * Non-Empty Set Type
    NESet
  -- ** Conversions between empty and non-empty sets
  , pattern IsNonEmpty
  , pattern IsEmpty
  , nonEmptySet
  , toSet
  , withNonEmpty
  , insertSet
  , insertSetMin
  , insertSetMax
  , unsafeFromSet

  -- * Construction
  , singleton
  , fromList
  , fromAscList
  , fromDescList
  , fromDistinctAscList
  , fromDistinctDescList
  , powerSet

  -- * Insertion
  , insert

  -- * Deletion
  , delete

  -- * Query
  , member
  , notMember
  , lookupLT
  , lookupGT
  , lookupLE
  , lookupGE
  , size
  , isSubsetOf
  , isProperSubsetOf
  , disjoint

  -- * Combine
  , union
  , unions
  , difference
  , (\\)
  , intersection
  , cartesianProduct
  , disjointUnion

  -- * Filter
  , filter
  , takeWhileAntitone
  , dropWhileAntitone
  , spanAntitone
  , partition
  , split
  , splitMember
  , splitRoot

  -- * Indexed
  , lookupIndex
  , findIndex
  , elemAt
  , deleteAt
  , take
  , drop
  , splitAt

  -- * Map
  , map
  , mapMonotonic

  -- * Folds
  , foldr
  , foldl
  , foldr1
  , foldl1
  -- ** Strict folds
  , foldr'
  , foldl'
  , foldr1'
  , foldl1'

  -- * Min\/Max
  , findMin
  , findMax
  , deleteMin
  , deleteMax
  , deleteFindMin
  , deleteFindMax

  -- * Conversion

  -- ** List
  , elems
  , toList
  , toAscList
  , toDescList

  -- * Debugging
  , valid
  ) where

import           Control.Applicative
import           Data.Bifunctor
import           Data.List.NonEmpty         (NonEmpty(..))
import           Data.Maybe
import           Data.Set                   (Set)
import           Data.Set.NonEmpty.Internal
import           Data.These
import           Prelude hiding             (foldr, foldl, filter, map, take, drop, splitAt)
import qualified Data.List.NonEmpty         as NE
import qualified Data.Semigroup.Foldable    as F1
import qualified Data.Set                   as S

-- | /O(1)/ match, /O(log n)/ usage of contents. The 'IsNonEmpty' and
-- 'IsEmpty' patterns allow you to treat a 'Set' as if it were either
-- a @'IsNonEmpty' n@ (where @n@ is a 'NESet') or an 'IsEmpty'.
--
-- For example, you can pattern match on a 'Set':
--
-- @
-- myFunc :: 'Set' X -> Y
-- myFunc ('IsNonEmpty' n) =  -- here, the user provided a non-empty set, and @n@ is the 'NESet'
-- myFunc 'IsEmpty'        =  -- here, the user provided an empty set
-- @
--
-- Matching on @'IsNonEmpty' n@ means that the original 'Set' was /not/
-- empty, and you have a verified-non-empty 'NESet' @n@ to use.
--
-- Note that patching on this pattern is /O(1)/.  However, using the
-- contents requires a /O(log n)/ cost that is deferred until after the
-- pattern is matched on (and is not incurred at all if the contents are
-- never used).
--
-- A case statement handling both 'IsNonEmpty' and 'IsEmpty' provides
-- complete coverage.
--
-- This is a bidirectional pattern, so you can use 'IsNonEmpty' to convert
-- a 'NESet' back into a 'Set', obscuring its non-emptiness (see 'toSet').
pattern IsNonEmpty :: NESet a -> Set a
pattern IsNonEmpty n <- (nonEmptySet->Just n)
  where
    IsNonEmpty n = toSet n

-- | /O(1)/. The 'IsNonEmpty' and 'IsEmpty' patterns allow you to treat
-- a 'Set' as if it were either a @'IsNonEmpty' n@ (where @n@ is
-- a 'NESet') or an 'IsEmpty'.
--
-- Matching on 'IsEmpty' means that the original 'Set' was empty.
--
-- A case statement handling both 'IsNonEmpty' and 'IsEmpty' provides
-- complete coverage.
--
-- This is a bidirectional pattern, so you can use 'IsEmpty' as an
-- expression, and it will be interpreted as 'Data.Set.empty'.
--
-- See 'IsNonEmpty' for more information.
pattern IsEmpty :: Set a
pattern IsEmpty <- (S.null->True)
  where
    IsEmpty = S.empty

{-# COMPLETE IsNonEmpty, IsEmpty #-}

-- | /O(log n)/. Unsafe version of 'nonEmptySet'.  Coerces a 'Set' into an
-- 'NESet', but is undefined (throws a runtime exception when evaluation is
-- attempted) for an empty 'Set'.
unsafeFromSet
    :: Set a
    -> NESet a
unsafeFromSet = withNonEmpty e id
  where
    e = errorWithoutStackTrace "NESet.unsafeFromSet: empty set"
{-# INLINE unsafeFromSet #-}

-- | /O(log n)/. Convert a 'Set' into an 'NESet' by adding a value.
-- Because of this, we know that the set must have at least one
-- element, and so therefore cannot be empty.
--
-- See 'insertSetMin' for a version that is constant-time if the new value is
-- /strictly smaller than/ all values in the original set
--
-- > insertSet 4 (Data.Set.fromList [5, 3]) == fromList (3 :| [4, 5])
-- > insertSet 4 Data.Set.empty == singleton 4 "c"
insertSet :: Ord a => a -> Set a -> NESet a
insertSet x = withNonEmpty (singleton x) (insert x)
{-# INLINE insertSet #-}

-- | /O(1)/ Convert a 'Set' into an 'NESet' by adding a value where the
-- value is /strictly less than/ all values in the input set  The values in
-- the original map must all be /strictly greater than/ the new value.
-- /The precondition is not checked./
--
-- > insertSetMin 2 (Data.Set.fromList [5, 3]) == fromList (2 :| [3, 5])
-- > valid (insertSetMin 2 (Data.Set.fromList [5, 3])) == True
-- > valid (insertSetMin 7 (Data.Set.fromList [5, 3])) == False
-- > valid (insertSetMin 3 (Data.Set.fromList [5, 3])) == False
insertSetMin :: a -> Set a -> NESet a
insertSetMin = NESet
{-# INLINE insertSetMin #-}

-- | /O(log n)/ Convert a 'Set' into an 'NESet' by adding a value where the
-- value is /strictly less than/ all values in the input set  The values in
-- the original map must all be /strictly greater than/ the new value.
-- /The precondition is not checked./
--
-- While this has the same asymptotics as 'insertSet', it saves a constant
-- factor for key comparison (so may be helpful if comparison is expensive)
-- and also does not require an 'Ord' instance for the key type.
--
-- > insertSetMin 7 (Data.Set.fromList [5, 3]) == fromList (3 :| [5, 7])
-- > valid (insertSetMin 7 (Data.Set.fromList [5, 3])) == True
-- > valid (insertSetMin 2 (Data.Set.fromList [5, 3])) == False
-- > valid (insertSetMin 5 (Data.Set.fromList [5, 3])) == False
insertSetMax :: a -> Set a -> NESet a
insertSetMax x = withNonEmpty (singleton x) go
  where
    go (NESet x0 s0) = NESet x0 . insertMaxSet x $ s0
{-# INLINE insertSetMax #-}

-- | /O(n)/. Build a set from an ascending list in linear time.  /The
-- precondition (input list is ascending) is not checked./
fromAscList :: Eq a => NonEmpty a -> NESet a
fromAscList = fromDistinctAscList . combineEq
{-# INLINE fromAscList #-}

-- | /O(n)/. Build a set from an ascending list of distinct elements in linear time.
-- /The precondition (input list is strictly ascending) is not checked./
fromDistinctAscList :: NonEmpty a -> NESet a
fromDistinctAscList (x :| xs) = insertSetMin x
                              . S.fromDistinctAscList
                              $ xs
{-# INLINE fromDistinctAscList #-}

-- | /O(n)/. Build a set from a descending list in linear time.
-- /The precondition (input list is descending) is not checked./
fromDescList :: Eq a => NonEmpty a -> NESet a
fromDescList = fromDistinctDescList . combineEq
{-# INLINE fromDescList #-}

-- | /O(n)/. Build a set from a descending list of distinct elements in linear time.
-- /The precondition (input list is strictly descending) is not checked./
fromDistinctDescList :: NonEmpty a -> NESet a
fromDistinctDescList (x :| xs) = insertSetMax x
                               . S.fromDistinctDescList
                               $ xs
{-# INLINE fromDistinctDescList #-}

-- | Calculate the power set of a non-empty: the set of all its (non-empty)
-- subsets.
--
-- @
-- t ``member`` powerSet s == t ``isSubsetOf`` s
-- @
--
-- Example:
--
-- @
-- powerSet (fromList (1 :| [2,3])) =
--   fromList (singleton 1 :| [ singleton 2
--                            , singleton 3
--                            , fromList (1 :| [2])
--                            , fromList (1 :| [3])
--                            , fromList (2 :| [3])
--                            , fromList (1 :| [2,3])
--                            ]
--            )
-- @
--
-- We know that the result is non-empty because the result will always at
-- least contain the original set.
powerSet
    :: forall a. ()
    => NESet a
    -> NESet (NESet a)
powerSet (NESet x s0) = case nonEmptySet p1 of
    -- s0 was empty originally
    Nothing -> singleton (singleton x)
    -- s1 was not empty originally
    Just p2 -> mapMonotonic (insertSetMin x) p0
       `merge` p2
  where
    -- powerset should never be empty
    p0 :: NESet (Set a)
    p0@(NESet _ p0s) = forSure $ powerSetSet s0
    p1 :: Set (NESet a)
    p1 = S.mapMonotonic forSure p0s  -- only minimal element is empty, so the rest aren't
    forSure = withNonEmpty (errorWithoutStackTrace "NESet.powerSet: internal error")
                        id
{-# INLINABLE powerSet #-}

-- | /O(log n)/. Insert an element in a set.
-- If the set already contains an element equal to the given value,
-- it is replaced with the new value.
insert :: Ord a => a -> NESet a -> NESet a
insert x n@(NESet x0 s) = case compare x x0 of
    LT -> NESet x  $ toSet n
    EQ -> NESet x  s
    GT -> NESet x0 $ S.insert x s
{-# INLINE insert #-}

-- | /O(log n)/. Delete an element from a set.
delete :: Ord a => a -> NESet a -> Set a
delete x n@(NESet x0 s) = case compare x x0 of
    LT -> toSet n
    EQ -> s
    GT -> insertMinSet x0 . S.delete x $ s
{-# INLINE delete #-}

-- | /O(log n)/. Is the element in the set?
member :: Ord a => a -> NESet a -> Bool
member x (NESet x0 s) = case compare x x0 of
    LT -> False
    EQ -> True
    GT -> S.member x s
{-# INLINE member #-}

-- | /O(log n)/. Is the element not in the set?
notMember :: Ord a => a -> NESet a -> Bool
notMember x (NESet x0 s) = case compare x x0 of
    LT -> True
    EQ -> False
    GT -> S.notMember x s
{-# INLINE notMember #-}

-- | /O(log n)/. Find largest element smaller than the given one.
--
-- > lookupLT 3 (fromList (3 :| [5])) == Nothing
-- > lookupLT 5 (fromList (3 :| [5])) == Just 3
lookupLT :: Ord a => a -> NESet a -> Maybe a
lookupLT x (NESet x0 s) = case compare x x0 of
    LT -> Nothing
    EQ -> Nothing
    GT -> S.lookupLT x s <|> Just x0
{-# INLINE lookupLT #-}

-- | /O(log n)/. Find smallest element greater than the given one.
--
-- > lookupLT 4 (fromList (3 :| [5])) == Just 5
-- > lookupLT 5 (fromList (3 :| [5])) == Nothing
lookupGT :: Ord a => a -> NESet a -> Maybe a
lookupGT x (NESet x0 s) = case compare x x0 of
    LT -> Just x0
    EQ -> S.lookupMin s
    GT -> S.lookupGT x s
{-# INLINE lookupGT #-}

-- | /O(log n)/. Find largest element smaller or equal to the given one.
--
-- > lookupLT 2 (fromList (3 :| [5])) == Nothing
-- > lookupLT 4 (fromList (3 :| [5])) == Just 3
-- > lookupLT 5 (fromList (3 :| [5])) == Just 5
lookupLE :: Ord a => a -> NESet a -> Maybe a
lookupLE x (NESet x0 s) = case compare x x0 of
    LT -> Nothing
    EQ -> Just x0
    GT -> S.lookupLE x s <|> Just x0
{-# INLINE lookupLE #-}

-- | /O(log n)/. Find smallest element greater or equal to the given one.
--
-- > lookupLT 3 (fromList (3 :| [5])) == Just 3
-- > lookupLT 4 (fromList (3 :| [5])) == Just 5
-- > lookupLT 6 (fromList (3 :| [5])) == Nothing
lookupGE :: Ord a => a -> NESet a -> Maybe a
lookupGE x (NESet x0 s) = case compare x x0 of
    LT -> Just x0
    EQ -> Just x0
    GT -> S.lookupGE x s
{-# INLINE lookupGE #-}

-- | /O(n+m)/. Is this a subset?
-- @(s1 \`isSubsetOf\` s2)@ tells whether @s1@ is a subset of @s2@.
isSubsetOf
    :: Ord a
    => NESet a
    -> NESet a
    -> Bool
isSubsetOf (NESet x s0) (toSet->s1) = x `S.member` s1
                                   && s0 `S.isSubsetOf` s1
{-# INLINE isSubsetOf #-}

-- | /O(n+m)/. Is this a proper subset? (ie. a subset but not equal).
isProperSubsetOf
    :: Ord a
    => NESet a
    -> NESet a
    -> Bool
isProperSubsetOf s0 s1 = S.size (nesSet s0) < S.size (nesSet s1)
                      && s0 `isSubsetOf` s1
{-# INLINE isProperSubsetOf #-}

-- | /O(n+m)/. Check whether two sets are disjoint (i.e. their intersection
--   is empty).
--
-- > disjoint (fromList (2:|[4,6]))   (fromList (1:|[3]))     == True
-- > disjoint (fromList (2:|[4,6,8])) (fromList (2:|[3,5,7])) == False
-- > disjoint (fromList (1:|[2]))     (fromList (1:|[2,3,4])) == False
disjoint
    :: Ord a
    => NESet a
    -> NESet a
    -> Bool
disjoint n1@(NESet x1 s1) n2@(NESet x2 s2) = case compare x1 x2 of
    -- x1 is not in n2
    LT -> s1 `disjointSet` toSet n2
    -- k1 and k2 are a part of the result
    EQ -> False
    -- k2 is not in n1
    GT -> toSet n1 `disjointSet` s2
{-# INLINE disjoint #-}

-- | /O(m*log(n\/m + 1)), m <= n/. Difference of two sets.
--
-- Returns a potentially empty set ('Set') because the first set might be
-- a subset of the second set, and therefore have all of its elements
-- removed.
difference
    :: Ord a
    => NESet a
    -> NESet a
    -> Set a
difference n1@(NESet x1 s1) n2@(NESet x2 s2) = case compare x1 x2 of
    -- x1 is not in n2, so cannot be deleted
    LT -> insertMinSet x1 $ s1 `S.difference` toSet n2
    -- x2 deletes x1, and only x1
    EQ -> s1 `S.difference` s2
    -- x2 is not in n1, so cannot delete anything, so we can just difference n1 // s2.
    GT -> toSet n1 `S.difference` s2
{-# INLINE difference #-}

-- | Same as 'difference'.
(\\)
    :: Ord a
    => NESet a
    -> NESet a
    -> Set a
(\\) = difference
{-# INLINE (\\) #-}

-- | /O(m*log(n\/m + 1)), m <= n/. The intersection of two sets.
--
-- Returns a potentially empty set ('Set'), because the two sets might have
-- an empty intersection.
--
-- Elements of the result come from the first set, so for example
--
-- > import qualified Data.Set.NonEmpty as NES
-- > data AB = A | B deriving Show
-- > instance Ord AB where compare _ _ = EQ
-- > instance Eq AB where _ == _ = True
-- > main = print (NES.singleton A `NES.intersection` NES.singleton B,
-- >               NES.singleton B `NES.intersection` NES.singleton A)
--
-- prints @(fromList (A:|[]),fromList (B:|[]))@.
intersection
    :: Ord a
    => NESet a
    -> NESet a
    -> Set a
intersection n1@(NESet x1 s1) n2@(NESet x2 s2) = case compare x1 x2 of
    -- x1 is not in n2
    LT -> s1 `S.intersection` toSet n2
    -- x1 and x2 are a part of the result
    EQ -> insertMinSet x1 $ s1 `S.intersection` s2
    -- x2 is not in n1
    GT -> toSet n1 `S.intersection` s2
{-# INLINE intersection #-}

-- | Calculate the Cartesian product of two sets.
--
-- @
-- cartesianProduct xs ys = fromList $ liftA2 (,) (toList xs) (toList ys)
-- @
--
-- Example:
--
-- @
-- cartesianProduct (fromList (1:|[2])) (fromList (\'a\':|[\'b\'])) =
--   fromList ((1,\'a\') :| [(1,\'b\'), (2,\'a\'), (2,\'b\')])
-- @
cartesianProduct
    :: NESet a
    -> NESet b
    -> NESet (a, b)
cartesianProduct n1 n2 = getMergeNESet
                       . F1.foldMap1 (\x -> MergeNESet $ mapMonotonic (x,) n2)
                       $ n1
{-# INLINE cartesianProduct #-}

-- | Calculate the disjoint union of two sets.
--
-- @ disjointUnion xs ys = map Left xs ``union`` map Right ys @
--
-- Example:
--
-- @
-- disjointUnion (fromList (1:|[2])) (fromList ("hi":|["bye"])) =
--   fromList (Left 1 :| [Left 2, Right "hi", Right "bye"])
-- @
disjointUnion
    :: NESet a
    -> NESet b
    -> NESet (Either a b)
disjointUnion (NESet x1 s1) n2 = NESet (Left x1)
                                       (s1 `disjointUnionSet` toSet n2)
{-# INLINE disjointUnion #-}

-- | /O(n)/. Filter all elements that satisfy the predicate.
--
-- Returns a potentially empty set ('Set') because the predicate might
-- filter out all items in the original non-empty set.
filter
    :: (a -> Bool)
    -> NESet a
    -> Set a
filter f (NESet x s1)
    | f x       = insertMinSet x . S.filter f $ s1
    | otherwise = S.filter f s1
{-# INLINE filter #-}

-- | /O(log n)/. Take while a predicate on the elements holds.  The user is
-- responsible for ensuring that for all elements @j@ and @k@ in the set,
-- @j \< k ==\> p j \>= p k@. See note at 'spanAntitone'.
--
-- Returns a potentially empty set ('Set') because the predicate might fail
-- on the first input.
--
-- @
-- takeWhileAntitone p = Data.Set.fromDistinctAscList . Data.List.NonEmpty.takeWhile p . 'toList'
-- takeWhileAntitone p = 'filter' p
-- @
takeWhileAntitone
    :: (a -> Bool)
    -> NESet a
    -> Set a
takeWhileAntitone f (NESet x s)
    | f x       = insertMinSet x . S.takeWhileAntitone f $ s
    | otherwise = S.empty
{-# INLINE takeWhileAntitone #-}

-- | /O(log n)/. Drop while a predicate on the elements holds.  The user is
-- responsible for ensuring that for all elements @j@ and @k@ in the set,
-- @j \< k ==\> p j \>= p k@. See note at 'spanAntitone'.
--
-- Returns a potentially empty set ('Set') because the predicate might be
-- true for all items.
--
-- @
-- dropWhileAntitone p = Data.Set.fromDistinctAscList . Data.List.NonEmpty.dropWhile p . 'toList'
-- dropWhileAntitone p = 'filter' (not . p)
-- @
dropWhileAntitone
    :: (a -> Bool)
    -> NESet a
    -> Set a
dropWhileAntitone f n@(NESet x s)
    | f x       = S.dropWhileAntitone f s
    | otherwise = toSet n
{-# INLINE dropWhileAntitone #-}

-- | /O(log n)/. Divide a set at the point where a predicate on the
-- elements stops holding.  The user is responsible for ensuring that for
-- all elements @j@ and @k@ in the set, @j \< k ==\> p j \>= p k@.
--
-- Returns a 'These' with potentially two non-empty sets:
--
-- *   @'This' n1@ means that the predicate never failed for any item,
--     returning the original set
-- *   @'That' n2@ means that the predicate failed for the first item,
--     returning the original set
-- *   @'These' n1 n2@ gives @n1@ (the set up to the point where the
--     predicate stops holding) and @n2@ (the set starting from
--     the point where the predicate stops holding)
--
-- @
-- spanAntitone p xs = partition p xs
-- @
--
-- Note: if @p@ is not actually antitone, then @spanAntitone@ will split the set
-- at some /unspecified/ point where the predicate switches from holding to not
-- holding (where the predicate is seen to hold before the first element and to fail
-- after the last element).
spanAntitone
    :: (a -> Bool)
    -> NESet a
    -> These (NESet a) (NESet a)
spanAntitone f n@(NESet x s0)
    | f x       = case (nonEmptySet s1, nonEmptySet s2) of
        (Nothing, Nothing) -> This  n
        (Just _ , Nothing) -> This  n
        (Nothing, Just n2) -> These (singleton x)       n2
        (Just _ , Just n2) -> These (insertSetMin x s1) n2
    | otherwise = That n
  where
    (s1, s2) = S.spanAntitone f s0
{-# INLINABLE spanAntitone #-}

-- | /O(n)/. Partition the map according to a predicate.
--
-- Returns a 'These' with potentially two non-empty sets:
--
-- *   @'This' n1@ means that the predicate was true for all items.
-- *   @'That' n2@ means that the predicate was false for all items.
-- *   @'These' n1 n2@ gives @n1@ (all of the items that were true for the
--     predicate) and @n2@ (all of the items that were false for the
--     predicate).
--
-- See also 'split'.
--
-- > partition (> 3) (fromList (5 :| [3])) == These (singleton 5) (singleton 3)
-- > partition (< 7) (fromList (5 :| [3])) == This  (fromList (3 :| [5]))
-- > partition (> 7) (fromList (5 :| [3])) == That  (fromList (3 :| [5]))
partition
    :: (a -> Bool)
    -> NESet a
    -> These (NESet a) (NESet a)
partition f n@(NESet x s0) = case (nonEmptySet s1, nonEmptySet s2) of
    (Nothing, Nothing)
      | f x       -> This  n
      | otherwise -> That                      n
    (Just n1, Nothing)
      | f x       -> This  n
      | otherwise -> These n1                  (singleton x)
    (Nothing, Just n2)
      | f x       -> These (singleton x)       n2
      | otherwise -> That                      n
    (Just n1, Just n2)
      | f x       -> These (insertSetMin x s1) n2
      | otherwise -> These n1                  (insertSetMin x s2)
  where
    (s1, s2) = S.partition f s0
{-# INLINABLE partition #-}

-- | /O(log n)/. The expression (@'split' x set@) is potentially a 'These'
-- containing up to two 'NESet's based on splitting the set into sets
-- containing items before and after the value @x@.  It will never return
-- a set that contains @x@ itself.
--
-- *   'Nothing' means that @x@ was the only value in the the original set,
--     and so there are no items before or after it.
-- *   @'Just' ('This' n1)@ means @x@ was larger than or equal to all items
--     in the set, and @n1@ is the entire original set (minus @x@, if it
--     was present)
-- *   @'Just' ('That' n2)@ means @x@ was smaller than or equal to all
--     items in the set, and @n2@ is the entire original set (minus @x@, if
--     it was present)
-- *   @'Just' ('These' n1 n2)@ gives @n1@ (the set of all values from the
--     original set less than @x@) and @n2@ (the set of all values from the
--     original set greater than @x@).
--
-- > split 2 (fromList (5 :| [3])) == Just (That  (fromList (3 :| [5]))      )
-- > split 3 (fromList (5 :| [3])) == Just (That  (singleton 5)              )
-- > split 4 (fromList (5 :| [3])) == Just (These (singleton 3) (singleton 5))
-- > split 5 (fromList (5 :| [3])) == Just (This  (singleton 3)              )
-- > split 6 (fromList (5 :| [3])) == Just (This  (fromList (3 :| [5]))      )
-- > split 5 (singleton 5)         == Nothing
split
    :: Ord a
    => a
    -> NESet a
    -> Maybe (These (NESet a) (NESet a))
split x n@(NESet x0 s0) = case compare x x0 of
    LT -> Just $ That n
    EQ -> That <$> nonEmptySet s0
    GT -> case (nonEmptySet s1, nonEmptySet s2) of
      (Nothing, Nothing) -> Just $ This  (singleton x0)
      (Just _ , Nothing) -> Just $ This  (insertSetMin x0 s1)
      (Nothing, Just n2) -> Just $ These (singleton x0)       n2
      (Just _ , Just n2) -> Just $ These (insertSetMin x0 s1) n2
  where
    (s1, s2) = S.split x s0
{-# INLINABLE split #-}

-- | /O(log n)/. The expression (@'splitMember' x set@) splits a set just
-- like 'split' but also returns @'member' x set@ (whether or not @x@ was
-- in @set@)
--
-- > splitMember 2 (fromList (5 :| [3])) == (False, Just (That  (fromList (3 :| [5)]))))
-- > splitMember 3 (fromList (5 :| [3])) == (True , Just (That  (singleton 5)))
-- > splitMember 4 (fromList (5 :| [3])) == (False, Just (These (singleton 3) (singleton 5)))
-- > splitMember 5 (fromList (5 :| [3])) == (True , Just (This  (singleton 3))
-- > splitMember 6 (fromList (5 :| [3])) == (False, Just (This  (fromList (3 :| [5])))
-- > splitMember 5 (singleton 5)         == (True , Nothing)
splitMember
    :: Ord a
    => a
    -> NESet a
    -> (Bool, Maybe (These (NESet a) (NESet a)))
splitMember x n@(NESet x0 s0) = case compare x x0 of
    LT -> (False, Just $ That n)
    EQ -> (True , That <$> nonEmptySet s0)
    GT -> (mem  ,) $ case (nonEmptySet s1, nonEmptySet s2) of
      (Nothing, Nothing) -> Just $ This  (singleton x0)
      (Just _ , Nothing) -> Just $ This  (insertSetMin x0 s1)
      (Nothing, Just n2) -> Just $ These (singleton x0)       n2
      (Just _ , Just n2) -> Just $ These (insertSetMin x0 s1) n2
  where
    (s1, mem, s2) = S.splitMember x s0
{-# INLINABLE splitMember #-}

-- | /O(1)/.  Decompose a set into pieces based on the structure of the underlying
-- tree.  This function is useful for consuming a set in parallel.
--
-- No guarantee is made as to the sizes of the pieces; an internal, but
-- deterministic process determines this.  However, it is guaranteed that
-- the pieces returned will be in ascending order (all elements in the
-- first subset less than all elements in the second, and so on).
--
--  Note that the current implementation does not return more than four
--  subsets, but you should not depend on this behaviour because it can
--  change in the future without notice.
splitRoot
    :: NESet a
    -> NonEmpty (NESet a)
splitRoot (NESet x s) = singleton x
                     :| mapMaybe nonEmptySet (S.splitRoot s)
{-# INLINE splitRoot #-}

-- | /O(log n)/. Lookup the /index/ of an element, which is its zero-based
-- index in the sorted sequence of elements. The index is a number from /0/
-- up to, but not including, the 'size' of the set.
--
-- > isJust   (lookupIndex 2 (fromList (5:|[3]))) == False
-- > fromJust (lookupIndex 3 (fromList (5:|[3]))) == 0
-- > fromJust (lookupIndex 5 (fromList (5:|[3]))) == 1
-- > isJust   (lookupIndex 6 (fromList (5:|[3]))) == False
lookupIndex
    :: Ord a
    => a
    -> NESet a
    -> Maybe Int
lookupIndex x (NESet x0 s) = case compare x x0 of
    LT -> Nothing
    EQ -> Just 0
    GT -> (+ 1) <$> S.lookupIndex x s
{-# INLINE lookupIndex #-}

-- | /O(log n)/. Return the /index/ of an element, which is its zero-based
-- index in the sorted sequence of elements. The index is a number from /0/
-- up to, but not including, the 'size' of the set. Calls 'error' when the
-- element is not a 'member' of the set.
--
-- > findIndex 2 (fromList (5:|[3]))    Error: element is not in the set
-- > findIndex 3 (fromList (5:|[3])) == 0
-- > findIndex 5 (fromList (5:|[3])) == 1
-- > findIndex 6 (fromList (5:|[3]))    Error: element is not in the set
findIndex
    :: Ord a
    => a
    -> NESet a
    -> Int
findIndex k = fromMaybe e . lookupIndex k
  where
    e = error "NESet.findIndex: element is not in the set"
{-# INLINE findIndex #-}

-- | /O(log n)/. Retrieve an element by its /index/, i.e. by its zero-based
-- index in the sorted sequence of elements. If the /index/ is out of range
-- (less than zero, greater or equal to 'size' of the set), 'error' is
-- called.
--
-- > elemAt 0 (fromList (5:|[3])) == 3
-- > elemAt 1 (fromList (5:|[3])) == 5
-- > elemAt 2 (fromList (5:|[3]))    Error: index out of range
elemAt
    :: Int
    -> NESet a
    -> a
elemAt 0 (NESet x _) = x
elemAt i (NESet _ s) = S.elemAt (i - 1) s
{-# INLINE elemAt #-}

-- | /O(log n)/. Delete the element at /index/, i.e. by its zero-based
-- index in the sorted sequence of elements. If the /index/ is out of range
-- (less than zero, greater or equal to 'size' of the set), 'error' is
-- called.
--
-- Returns a potentially empty set ('Set'), because this could potentailly
-- delete the final element in a singleton set.
--
-- > deleteAt 0    (fromList (5:|[3])) == singleton 5
-- > deleteAt 1    (fromList (5:|[3])) == singleton 3
-- > deleteAt 2    (fromList (5:|[3]))    Error: index out of range
-- > deleteAt (-1) (fromList (5:|[3]))    Error: index out of range
deleteAt
    :: Int
    -> NESet a
    -> Set a
deleteAt 0 (NESet _ s) = s
deleteAt i (NESet x s) = insertMinSet x . S.deleteAt (i - 1) $ s
{-# INLINABLE deleteAt #-}

-- | Take a given number of elements in order, beginning
-- with the smallest ones.
--
-- Returns a potentailly empty set ('Set'), which can only happen when
-- calling @take 0@.
--
-- @
-- take n = Data.Set.fromDistinctAscList . Data.List.NonEmpty.take n . 'toAscList'
-- @
take
    :: Int
    -> NESet a
    -> Set a
take 0 (NESet _ _) = S.empty
take i (NESet x s) = insertMinSet x . S.take (i - 1) $ s
{-# INLINABLE take #-}

-- | Drop a given number of elements in order, beginning
-- with the smallest ones.
--
-- Returns a potentailly empty set ('Set'), in the case that 'drop' is
-- called with a number equal to or greater the number of items in the set,
-- and we drop every item.
--
-- @
-- drop n = Data.Set.fromDistinctAscList . Data.List.NonEmpty.drop n . 'toAscList'
-- @
drop
    :: Int
    -> NESet a
    -> Set a
drop 0 n           = toSet n
drop n (NESet _ s) = S.drop (n - 1) s
{-# INLINABLE drop #-}

-- | /O(log n)/. Split a set at a particular index @i@.
--
-- *   @'This' n1@ means that there are less than @i@ items in the set, and
--     @n1@ is the original set.
-- *   @'That' n2@ means @i@ was 0; we dropped 0 items, so @n2@ is the
--     original set.
-- *   @'These' n1 n2@ gives @n1@ (taking @i@ items from the original set)
--     and @n2@ (dropping @i@ items from the original set))
splitAt
    :: Int
    -> NESet a
    -> These (NESet a) (NESet a)
splitAt 0 n              = That n
splitAt i n@(NESet x s0) = case (nonEmptySet s1, nonEmptySet s2) of
    (Nothing, Nothing) -> This  (singleton x)
    (Just _ , Nothing) -> This  n
    (Nothing, Just n2) -> These (singleton x)       n2
    (Just _ , Just n2) -> These (insertSetMin x s1) n2
  where
    (s1, s2) = S.splitAt (i - 1) s0
{-# INLINABLE splitAt #-}

-- | /O(n*log n)/.
-- @'map' f s@ is the set obtained by applying @f@ to each element of @s@.
--
-- It's worth noting that the size of the result may be smaller if,
-- for some @(x,y)@, @x \/= y && f x == f y@
map :: Ord b
    => (a -> b)
    -> NESet a
    -> NESet b
map f (NESet x0 s) = fromList
                   . (f x0 :|)
                   . S.foldr (\x xs -> f x : xs) []
                   $ s
{-# INLINE map #-}

-- | /O(n)/.
-- @'mapMonotonic' f s == 'map' f s@, but works only when @f@ is strictly
-- increasing.  /The precondition is not checked./ Semi-formally, we have:
--
-- > and [x < y ==> f x < f y | x <- ls, y <- ls]
-- >                     ==> mapMonotonic f s == map f s
-- >     where ls = Data.Foldable.toList s
mapMonotonic
    :: (a -> b)
    -> NESet a
    -> NESet b
mapMonotonic f (NESet x s) = NESet (f x) (S.mapMonotonic f s)
{-# INLINE mapMonotonic #-}

-- | /O(n)/. A strict version of 'foldr1'. Each application of the operator
-- is evaluated before using the result in the next application. This
-- function is strict in the starting value.
foldr1' :: (a -> a -> a) -> NESet a -> a
foldr1' f (NESet x s) = case S.maxView s of
    Nothing      -> x
    Just (y, s') -> let !z = S.foldr' f y s' in x `f` z
{-# INLINE foldr1' #-}

-- | /O(n)/. A strict version of 'foldl1'. Each application of the operator
-- is evaluated before using the result in the next application. This
-- function is strict in the starting value.
foldl1' :: (a -> a -> a) -> NESet a -> a
foldl1' f (NESet x s) = S.foldl' f x s
{-# INLINE foldl1' #-}

-- | /O(1)/. The minimal element of a set.  Note that this is total, making
-- 'Data.Set.lookupMin' obsolete.  It is constant-time, so has better
-- asymptotics than @Data.Set.lookupMin@ and @Data.Map.findMin@ as well.
--
-- > findMin (fromList (5 :| [3])) == 3
findMin :: NESet a -> a
findMin (NESet x _) = x
{-# INLINE findMin #-}

-- | /O(log n)/. The maximal key of a set  Note that this is total,
-- making 'Data.Set.lookupMin' obsolete.
--
-- > findMax (fromList (5 :| [3])) == 5
findMax :: NESet a -> a
findMax (NESet x s) = fromMaybe x . S.lookupMax $ s
{-# INLINE findMax #-}

-- | /O(1)/. Delete the minimal element.  Returns a potentially empty set
-- ('Set'), because we might delete the final item in a singleton set.  It
-- is constant-time, so has better asymptotics than @Data.Set.deleteMin@.
--
-- > deleteMin (fromList (5 :| [3, 7])) == Data.Set.fromList [5, 7]
-- > deleteMin (singleton 5) == Data.Set.empty
deleteMin :: NESet a -> Set a
deleteMin (NESet _ s) = s
{-# INLINE deleteMin #-}

-- | /O(log n)/. Delete the maximal element.  Returns a potentially empty
-- set ('Set'), because we might delete the final item in a singleton set.
--
-- > deleteMax (fromList (5 :| [3, 7])) == Data.Set.fromList [3, 5]
-- > deleteMax (singleton 5) == Data.Set.empty
deleteMax :: NESet a -> Set a
deleteMax (NESet x s) = insertMinSet x . S.deleteMax $ s
{-# INLINE deleteMax #-}

-- | /O(1)/. Delete and find the minimal element.  It is constant-time, so
-- has better asymptotics that @Data.Set.minView@ for 'Set'.
--
-- Note that unlike @Data.Set.deleteFindMin@ for 'Set', this cannot ever
-- fail, and so is a total function. However, the result 'Set' is
-- potentially empty, since the original set might have contained just
-- a single item.
--
-- > deleteFindMin (fromList (5 :| [3, 10])) == (3, Data.Set.fromList [5, 10])
deleteFindMin :: NESet a -> (a, Set a)
deleteFindMin (NESet x s) = (x, s)
{-# INLINE deleteFindMin #-}

-- | /O(log n)/. Delete and find the minimal element.
--
-- Note that unlike @Data.Set.deleteFindMax@ for 'Set', this cannot ever
-- fail, and so is a total function. However, the result 'Set' is
-- potentially empty, since the original set might have contained just
-- a single item.
--
-- > deleteFindMax (fromList (5 :| [3, 10])) == (10, Data.Set.fromList [3, 5])
deleteFindMax :: NESet a -> (a, Set a)
deleteFindMax (NESet x s) = maybe (x, S.empty) (second (insertMinSet x))
                          . S.maxView
                          $ s
{-# INLINE deleteFindMax #-}

-- | /O(n)/. An alias of 'toAscList'. The elements of a set in ascending
-- order.
elems :: NESet a -> NonEmpty a
elems = toList
{-# INLINE elems #-}

-- | /O(n)/. Convert the set to an ascending non-empty list of elements.
toAscList :: NESet a -> NonEmpty a
toAscList = toList
{-# INLINE toAscList #-}

-- | /O(n)/. Convert the set to a descending non-empty list of elements.
toDescList :: NESet a -> NonEmpty a
toDescList (NESet x s) = S.foldl' (flip (NE.<|)) (x :| []) s
{-# INLINE toDescList #-}

-- ---------------------------
-- Combining functions
-- ---------------------------
--
-- Code comes from "Data.Set.Internal" from containers, modified slightly
-- to work with NonEmpty
--
-- Copyright   :  (c) Daan Leijen 2002

combineEq :: Eq a => NonEmpty a -> NonEmpty a
combineEq (x :| xs) = go x xs
  where
    go z [] = z :| []
    go z (y:ys)
      | z == y    = go z ys
      | otherwise = z NE.<| go y ys