module Data.NonEmpty.Set (
   T,
   insert,
   singleton,
   member,
   size,
   fromList,
   fromAscList,
   toAscList,
   fetch,
   flatten,
   union,
   unionLeft,
   unionRight,

   findMin,
   findMax,
   delete,
   deleteMin,
   deleteMax,
   deleteFindMin,
   deleteFindMax,
   minView,
   maxView,
   ) where

import qualified Data.NonEmpty.Class as C
import qualified Data.NonEmpty as NonEmpty

import qualified Data.Set as Set
import Data.Set (Set, )

import Control.Monad (mzero, )
import Control.Applicative (liftA2)
import Control.DeepSeq (NFData, rnf, )
import Data.Maybe (fromMaybe, mapMaybe)
import Data.Tuple.HT (forcePair, mapSnd, )

import qualified Test.QuickCheck as QC


{-
The first field will always contain the smallest element.
We do not use the NonEmpty data type here
since it is easy to break this invariant using NonEmpty.!:.
The custom type is also consistent with Map.
-}
data T a = Cons a (Set a)
   deriving (T a -> T a -> Bool
(T a -> T a -> Bool) -> (T a -> T a -> Bool) -> Eq (T a)
forall a. Eq a => T a -> T a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: T a -> T a -> Bool
$c/= :: forall a. Eq a => T a -> T a -> Bool
== :: T a -> T a -> Bool
$c== :: forall a. Eq a => T a -> T a -> Bool
Eq, Eq (T a)
Eq (T a)
-> (T a -> T a -> Ordering)
-> (T a -> T a -> Bool)
-> (T a -> T a -> Bool)
-> (T a -> T a -> Bool)
-> (T a -> T a -> Bool)
-> (T a -> T a -> T a)
-> (T a -> T a -> T a)
-> Ord (T a)
T a -> T a -> Bool
T a -> T a -> Ordering
T a -> T a -> T a
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall a. Ord a => Eq (T a)
forall a. Ord a => T a -> T a -> Bool
forall a. Ord a => T a -> T a -> Ordering
forall a. Ord a => T a -> T a -> T a
min :: T a -> T a -> T a
$cmin :: forall a. Ord a => T a -> T a -> T a
max :: T a -> T a -> T a
$cmax :: forall a. Ord a => T a -> T a -> T a
>= :: T a -> T a -> Bool
$c>= :: forall a. Ord a => T a -> T a -> Bool
> :: T a -> T a -> Bool
$c> :: forall a. Ord a => T a -> T a -> Bool
<= :: T a -> T a -> Bool
$c<= :: forall a. Ord a => T a -> T a -> Bool
< :: T a -> T a -> Bool
$c< :: forall a. Ord a => T a -> T a -> Bool
compare :: T a -> T a -> Ordering
$ccompare :: forall a. Ord a => T a -> T a -> Ordering
$cp1Ord :: forall a. Ord a => Eq (T a)
Ord)

instance (Show a) => Show (T a) where
   showsPrec :: Int -> T a -> ShowS
showsPrec Int
p T a
xs =
      Bool -> ShowS -> ShowS
showParen (Int
pInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>Int
10) (ShowS -> ShowS) -> ShowS -> ShowS
forall a b. (a -> b) -> a -> b
$
         String -> ShowS
showString String
"NonEmptySet.fromList " ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
         Int -> T [] a -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec Int
11 (T a -> T [] a
forall a. T a -> T [] a
toAscList T a
xs)


instance (NFData a) => NFData (T a) where
   rnf :: T a -> ()
rnf = T a -> ()
forall (f :: * -> *) a. (NFData f, NFData a) => f a -> ()
C.rnf

instance C.NFData T where
   rnf :: T a -> ()
rnf (Cons a
x Set a
xs) = (a, ()) -> ()
forall a. NFData a => a -> ()
rnf (a
x, Set a -> ()
forall (f :: * -> *) a. (NFData f, NFData a) => f a -> ()
C.rnf Set a
xs)


instance (QC.Arbitrary a, Ord a) => QC.Arbitrary (T a) where
   arbitrary :: Gen (T a)
arbitrary = (a -> Set a -> T a) -> Gen a -> Gen (Set a) -> Gen (T a)
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 a -> Set a -> T a
forall a. Ord a => a -> Set a -> T a
insert Gen a
forall a. Arbitrary a => Gen a
QC.arbitrary Gen (Set a)
forall a. Arbitrary a => Gen a
QC.arbitrary
   shrink :: T a -> [T a]
shrink = (Set a -> Maybe (T a)) -> [Set a] -> [T a]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe Set a -> Maybe (T a)
forall a. Ord a => Set a -> Maybe (T a)
fetch ([Set a] -> [T a]) -> (T a -> [Set a]) -> T a -> [T a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set a -> [Set a]
forall a. Arbitrary a => a -> [a]
QC.shrink (Set a -> [Set a]) -> (T a -> Set a) -> T a -> [Set a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. T a -> Set a
forall a. Ord a => T a -> Set a
flatten


{- |
We cannot have a reasonable @instance Insert Set@,
since the @instance Insert (NonEmpty Set)@
would preserve duplicate leading elements, whereas 'Set' does not.

However, the @instance Insert NonEmpty@ is not the problem.
A general type like

> insertSet :: (Insert f, Ord a) => a -> f a -> NonEmpty f a

cannot work, since it can be instantiated to

> insertSet :: (Ord a) => a -> NonEmpty Set a -> NonEmpty (NonEmpty Set) a

and this is obviously wrong:
@insertSet x (singleton x)@ has only one element, not two.
-}
insert :: Ord a => a -> Set a -> T a
insert :: a -> Set a -> T a
insert = ((a, a) -> a) -> a -> Set a -> T a
forall a. Ord a => ((a, a) -> a) -> a -> Set a -> T a
insertGen (a, a) -> a
forall a b. (a, b) -> a
fst

insertGen :: Ord a => ((a,a) -> a) -> a -> Set a -> T a
insertGen :: ((a, a) -> a) -> a -> Set a -> T a
insertGen (a, a) -> a
select a
y Set a
xt =
   (a -> Set a -> T a) -> (a, Set a) -> T a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> Set a -> T a
forall a. a -> Set a -> T a
Cons ((a, Set a) -> T a) -> (a, Set a) -> T a
forall a b. (a -> b) -> a -> b
$
   (a, Set a) -> Maybe (a, Set a) -> (a, Set a)
forall a. a -> Maybe a -> a
fromMaybe (a
y, Set a
xt) (Maybe (a, Set a) -> (a, Set a)) -> Maybe (a, Set a) -> (a, Set a)
forall a b. (a -> b) -> a -> b
$ do
      (a
x,Set a
xs) <- Set a -> Maybe (a, Set a)
forall a. Set a -> Maybe (a, Set a)
Set.minView Set a
xt
      case a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare a
y a
x of
         Ordering
GT -> (a, Set a) -> Maybe (a, Set a)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
x, a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
Set.insert a
y Set a
xs)
         Ordering
EQ -> (a, Set a) -> Maybe (a, Set a)
forall (m :: * -> *) a. Monad m => a -> m a
return ((a, a) -> a
select (a
y,a
x), Set a
xs)
         Ordering
LT -> Maybe (a, Set a)
forall (m :: * -> *) a. MonadPlus m => m a
mzero

singleton :: a -> T a
singleton :: a -> T a
singleton a
a = a -> Set a -> T a
forall a. a -> Set a -> T a
Cons a
a Set a
forall a. Set a
Set.empty

member :: (Ord a) => a -> T a -> Bool
member :: a -> T a -> Bool
member a
y (Cons a
x Set a
xs) =
   a
ya -> a -> Bool
forall a. Eq a => a -> a -> Bool
==a
x Bool -> Bool -> Bool
|| a -> Set a -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member a
y Set a
xs

size :: T a -> Int
size :: T a -> Int
size (Cons a
_ Set a
xs) = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Set a -> Int
forall a. Set a -> Int
Set.size Set a
xs


findMin :: T a -> a
findMin :: T a -> a
findMin (Cons a
x Set a
_) = a
x

findMax :: T a -> a
findMax :: T a -> a
findMax (Cons a
x Set a
xs) =
   if Set a -> Bool
forall a. Set a -> Bool
Set.null Set a
xs then a
x else Set a -> a
forall a. Set a -> a
Set.findMax Set a
xs


delete :: (Ord k) => k -> T k -> Set k
delete :: k -> T k -> Set k
delete k
y (Cons k
x Set k
xs) =
   if k
y k -> k -> Bool
forall a. Eq a => a -> a -> Bool
== k
x then Set k
xs else k -> Set k -> Set k
forall a. Ord a => a -> Set a -> Set a
Set.insert k
x (Set k -> Set k) -> Set k -> Set k
forall a b. (a -> b) -> a -> b
$ k -> Set k -> Set k
forall a. Ord a => a -> Set a -> Set a
Set.delete k
y Set k
xs

deleteMin :: T a -> Set a
deleteMin :: T a -> Set a
deleteMin (Cons a
_ Set a
xs) = Set a
xs

deleteMax :: (Ord a) => T a -> Set a
deleteMax :: T a -> Set a
deleteMax (Cons a
x Set a
xs) =
   if Set a -> Bool
forall a. Set a -> Bool
Set.null Set a
xs then Set a
forall a. Set a
Set.empty else a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
Set.insert a
x (Set a -> Set a) -> Set a -> Set a
forall a b. (a -> b) -> a -> b
$ Set a -> Set a
forall a. Set a -> Set a
Set.deleteMax Set a
xs


deleteFindMin :: T a -> (a, Set a)
deleteFindMin :: T a -> (a, Set a)
deleteFindMin (Cons a
x Set a
xs) = (a
x, Set a
xs)

deleteFindMax :: (Ord a) => T a -> (a, Set a)
deleteFindMax :: T a -> (a, Set a)
deleteFindMax (Cons a
x Set a
xs) =
   if Set a -> Bool
forall a. Set a -> Bool
Set.null Set a
xs
     then (a
x, Set a
forall a. Set a
Set.empty)
     else (Set a -> Set a) -> (a, Set a) -> (a, Set a)
forall b c a. (b -> c) -> (a, b) -> (a, c)
mapSnd (a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
Set.insert a
x) ((a, Set a) -> (a, Set a)) -> (a, Set a) -> (a, Set a)
forall a b. (a -> b) -> a -> b
$ Set a -> (a, Set a)
forall a. Set a -> (a, Set a)
Set.deleteFindMax Set a
xs


minView :: T a -> (a, Set a)
minView :: T a -> (a, Set a)
minView (Cons a
x Set a
xs) = (a
x,Set a
xs)

maxView :: (Ord a) => T a -> (a, Set a)
maxView :: T a -> (a, Set a)
maxView (Cons a
x Set a
xs) =
   (a, Set a) -> (a, Set a)
forall a b. (a, b) -> (a, b)
forcePair ((a, Set a) -> (a, Set a)) -> (a, Set a) -> (a, Set a)
forall a b. (a -> b) -> a -> b
$
   case Set a -> Maybe (a, Set a)
forall a. Set a -> Maybe (a, Set a)
Set.maxView Set a
xs of
      Maybe (a, Set a)
Nothing -> (a
x,Set a
xs)
      Just (a
y,Set a
ys) -> (a
y, a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
Set.insert a
x Set a
ys)


fromList :: (Ord a) => NonEmpty.T [] a -> T a
fromList :: T [] a -> T a
fromList (NonEmpty.Cons a
x [a]
xs) = a -> Set a -> T a
forall a. Ord a => a -> Set a -> T a
insert a
x (Set a -> T a) -> Set a -> T a
forall a b. (a -> b) -> a -> b
$ [a] -> Set a
forall a. Ord a => [a] -> Set a
Set.fromList [a]
xs

fromAscList :: (Ord a) => NonEmpty.T [] a -> T a
fromAscList :: T [] a -> T a
fromAscList (NonEmpty.Cons a
x [a]
xs) = a -> Set a -> T a
forall a. a -> Set a -> T a
Cons a
x (Set a -> T a) -> Set a -> T a
forall a b. (a -> b) -> a -> b
$ [a] -> Set a
forall a. Eq a => [a] -> Set a
Set.fromAscList [a]
xs

toAscList :: T a -> NonEmpty.T [] a
toAscList :: T a -> T [] a
toAscList (Cons a
x Set a
xs) = a -> [a] -> T [] a
forall (f :: * -> *) a. a -> f a -> T f a
NonEmpty.Cons a
x ([a] -> T [] a) -> [a] -> T [] a
forall a b. (a -> b) -> a -> b
$ Set a -> [a]
forall a. Set a -> [a]
Set.toAscList Set a
xs

fetch :: (Ord a) => Set a -> Maybe (T a)
fetch :: Set a -> Maybe (T a)
fetch  =  ((a, Set a) -> T a) -> Maybe (a, Set a) -> Maybe (T a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((a -> Set a -> T a) -> (a, Set a) -> T a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> Set a -> T a
forall a. a -> Set a -> T a
Cons) (Maybe (a, Set a) -> Maybe (T a))
-> (Set a -> Maybe (a, Set a)) -> Set a -> Maybe (T a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set a -> Maybe (a, Set a)
forall a. Set a -> Maybe (a, Set a)
Set.minView

flatten :: (Ord a) => T a -> Set a
flatten :: T a -> Set a
flatten (Cons a
x Set a
xs) = a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
Set.insert a
x Set a
xs

union :: (Ord a) => T a -> T a -> T a
union :: T a -> T a -> T a
union (Cons a
x Set a
xs) (Cons a
y Set a
ys) =
   (a -> Set a -> T a) -> (a, Set a) -> T a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> Set a -> T a
forall a. a -> Set a -> T a
Cons ((a, Set a) -> T a) -> (a, Set a) -> T a
forall a b. (a -> b) -> a -> b
$
   case Set a -> Set a -> Set a
forall a. Ord a => Set a -> Set a -> Set a
Set.union Set a
xs Set a
ys of
      Set a
zs ->
         case a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare a
x a
y of
            Ordering
LT -> (a
x, Set a -> Set a -> Set a
forall a. Ord a => Set a -> Set a -> Set a
Set.union Set a
zs (Set a -> Set a) -> Set a -> Set a
forall a b. (a -> b) -> a -> b
$ a -> Set a
forall a. a -> Set a
Set.singleton a
y)
            Ordering
GT -> (a
y, a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
Set.insert a
x Set a
zs)
            Ordering
EQ -> (a
x, Set a
zs)

unionLeft :: (Ord a) => Set a -> T a -> T a
unionLeft :: Set a -> T a -> T a
unionLeft Set a
xs (Cons a
y Set a
ys) =
   ((a, a) -> a) -> a -> Set a -> T a
forall a. Ord a => ((a, a) -> a) -> a -> Set a -> T a
insertGen (a, a) -> a
forall a b. (a, b) -> b
snd a
y (Set a -> T a) -> Set a -> T a
forall a b. (a -> b) -> a -> b
$ Set a -> Set a -> Set a
forall a. Ord a => Set a -> Set a -> Set a
Set.union Set a
xs Set a
ys

unionRight :: (Ord a) => T a -> Set a -> T a
unionRight :: T a -> Set a -> T a
unionRight (Cons a
x Set a
xs) Set a
ys =
   ((a, a) -> a) -> a -> Set a -> T a
forall a. Ord a => ((a, a) -> a) -> a -> Set a -> T a
insertGen (a, a) -> a
forall a b. (a, b) -> a
fst a
x (Set a -> T a) -> Set a -> T a
forall a b. (a -> b) -> a -> b
$ Set a -> Set a -> Set a
forall a. Ord a => Set a -> Set a -> Set a
Set.union Set a
xs Set a
ys


{-
According Set functions are only available since containers-0.5.2, i.e. GHC-7.8.

elemAt :: Int -> T a -> a
elemAt k (Cons x xs) =
   if k==0 then x else Set.elemAt (pred k) xs

deleteAt :: Int -> T a -> Set a
deleteAt k (Cons _ xs) =
   if k==0 then xs else Set.deleteAt (pred k) xs
-}