--------------------------------------------------------------------------------
-- |
-- Module      :  Data.Set.Util
-- Copyright   :  (C) Frank Staals
-- License     :  see the LICENSE file
-- Maintainer  :  Frank Staals
--------------------------------------------------------------------------------
module Data.Set.Util where

import           Data.DynamicOrd
import           Data.Set (Set)
import qualified Data.Set as Set
import qualified Data.Set.Internal as Internal


import           Data.Ord (comparing)

data S = S String deriving Int -> S -> ShowS
[S] -> ShowS
S -> String
(Int -> S -> ShowS) -> (S -> String) -> ([S] -> ShowS) -> Show S
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [S] -> ShowS
$cshowList :: [S] -> ShowS
show :: S -> String
$cshow :: S -> String
showsPrec :: Int -> S -> ShowS
$cshowsPrec :: Int -> S -> ShowS
Show
cmpS :: S -> S -> Ordering
cmpS :: S -> S -> Ordering
cmpS = (S -> Int) -> S -> S -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing (\(S String
s) -> String -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length String
s)


-- $setup
-- >>> import Data.Ord(comparing)
-- >>> data S = S String deriving Show
-- >>> cmpS = comparing (\(S s) -> length s)
--

-- | Given a monotonic function f that maps a to b, split the sequence s
-- depending on the b values. I.e. the result (l,m,r) is such that
-- * all (< x) . fmap f $ l
-- * all (== x) . fmap f $ m
-- * all (> x) . fmap f $ r
--
-- running time: \(O(\log n)\)
splitOn       :: Ord b => (a -> b) -> b -> Set a -> (Set a, Set a, Set a)
splitOn :: (a -> b) -> b -> Set a -> (Set a, Set a, Set a)
splitOn a -> b
f b
x Set a
s = let (Set a
l,Set a
s') = (a -> Bool) -> Set a -> (Set a, Set a)
forall a. (a -> Bool) -> Set a -> (Set a, Set a)
Set.spanAntitone (Ordering -> b -> Bool
g Ordering
LT (b -> Bool) -> (a -> b) -> a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b
f) Set a
s
                    (Set a
m,Set a
r)  = (a -> Bool) -> Set a -> (Set a, Set a)
forall a. (a -> Bool) -> Set a -> (Set a, Set a)
Set.spanAntitone (Ordering -> b -> Bool
g Ordering
EQ (b -> Bool) -> (a -> b) -> a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b
f) Set a
s'
                    g :: Ordering -> b -> Bool
g Ordering
c b
y  = b
y b -> b -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` b
x Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
== Ordering
c
                in (Set a
l,Set a
m,Set a
r)

-- | Given a monotonic function f that orders @a@, split the sequence @s@
-- into three parts. I.e. the result (lt,eq,gt) is such that
-- * all (\x -> f x == LT) . fmap f $ lt
-- * all (\x -> f x == EQ) . fmap f $ eq
-- * all (\x -> f x == GT) . fmap f $ gt
--
-- running time: \(O(\log n)\)
splitBy       :: (a -> Ordering) -> Set a -> (Set a, Set a, Set a)
splitBy :: (a -> Ordering) -> Set a -> (Set a, Set a, Set a)
splitBy a -> Ordering
f Set a
s = let (Set a
l,Set a
s') = (a -> Bool) -> Set a -> (Set a, Set a)
forall a. (a -> Bool) -> Set a -> (Set a, Set a)
Set.spanAntitone (Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
(==) Ordering
LT (Ordering -> Bool) -> (a -> Ordering) -> a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Ordering
f) Set a
s
                  (Set a
m,Set a
r)  = (a -> Bool) -> Set a -> (Set a, Set a)
forall a. (a -> Bool) -> Set a -> (Set a, Set a)
Set.spanAntitone (Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
(==) Ordering
EQ (Ordering -> Bool) -> (a -> Ordering) -> a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Ordering
f) Set a
s'
              in (Set a
l,Set a
m,Set a
r)

-- | Constructs a Set using the given Order.
--
-- Note that this is dangerous as the resulting set may not abide the
-- ordering expected of such sets.
--
-- running time: \(O(n\log n)\)
fromListBy        :: (a -> a -> Ordering) -> [a] -> Set a
fromListBy :: (a -> a -> Ordering) -> [a] -> Set a
fromListBy a -> a -> Ordering
cmp [a]
xs = (a -> a -> Ordering)
-> (forall s. Reifies s (OrdDict a) => O s (Set a)) -> Set a
forall a b.
(a -> a -> Ordering)
-> (forall s. Reifies s (OrdDict a) => O s b) -> b
withOrd a -> a -> Ordering
cmp (Set (O s a) -> O s (Set a)
forall (f :: * -> *) s a. f (O s a) -> O s (f a)
extractOrd1 (Set (O s a) -> O s (Set a))
-> ([a] -> Set (O s a)) -> [a] -> O s (Set a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [O s a] -> Set (O s a)
forall a. Ord a => [a] -> Set a
Set.fromList ([O s a] -> Set (O s a)) -> ([a] -> [O s a]) -> [a] -> Set (O s a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> O s a) -> [a] -> [O s a]
forall a b. (a -> b) -> [a] -> [b]
map a -> O s a
forall s a. a -> O s a
O ([a] -> O s (Set a)) -> [a] -> O s (Set a)
forall a b. (a -> b) -> a -> b
$ [a]
xs)

-- | Given two sets l and r, such that all elements of l occur before
-- r, join the two sets into a combined set.
--
-- running time: \(O(\log n)\)
join :: Set a -> Set a -> Set a
join :: Set a -> Set a -> Set a
join = Set a -> Set a -> Set a
forall a. Set a -> Set a -> Set a
Internal.merge


-- | Inserts an element into the set, assuming that the set is ordered
-- by the given order.
--
-- >>> insertBy cmpS (S "ccc") $ fromListBy cmpS [S "a" , S "bb" , S "dddd"]
-- fromList [S "a",S "bb",S "ccc",S "dddd"]
--
-- When trying to insert an element that equals an element already in
-- the set (according to the given comparator), this function replaces
-- the old element by the new one:
--
-- >>> insertBy cmpS (S "cc") $ fromListBy cmpS [S "a" , S "bb" , S "dddd"]
-- fromList [S "a",S "cc",S "dddd"]
--
-- running time: \(O(\log n)\)
insertBy         :: (a -> a -> Ordering) -> a -> Set a -> Set a
insertBy :: (a -> a -> Ordering) -> a -> Set a -> Set a
insertBy a -> a -> Ordering
cmp a
x Set a
s = (a -> a -> Ordering)
-> (forall s. Reifies s (OrdDict a) => O s (Set a)) -> Set a
forall a b.
(a -> a -> Ordering)
-> (forall s. Reifies s (OrdDict a) => O s b) -> b
withOrd a -> a -> Ordering
cmp ((forall s. Reifies s (OrdDict a) => O s (Set a)) -> Set a)
-> (forall s. Reifies s (OrdDict a) => O s (Set a)) -> Set a
forall a b. (a -> b) -> a -> b
$ (Set (O s a) -> Set (O s a)) -> Set a -> O s (Set a)
forall (f :: * -> *) s a (g :: * -> *).
(f (O s a) -> g (O s a)) -> f a -> O s (g a)
liftOrd1 (O s a -> Set (O s a) -> Set (O s a)
forall a. Ord a => a -> Set a -> Set a
Set.insert (O s a -> Set (O s a) -> Set (O s a))
-> O s a -> Set (O s a) -> Set (O s a)
forall a b. (a -> b) -> a -> b
$ a -> O s a
forall s a. a -> O s a
O a
x) Set a
s


-- | Deletes an element from the set, assuming the set is ordered by
-- the given ordering.
--
-- >>> deleteAllBy cmpS (S "bb") $ fromListBy cmpS [S "a" , S "bb" , S "dddd"]
-- fromList [S "a",S "dddd"]
-- >>> deleteAllBy cmpS (S "bb") $ fromListBy cmpS [S "a" , S "bb" , S "cc", S "dd", S "ee", S "ff", S "dddd"]
-- fromList [S "a",S "dddd"]
--
-- running time: \(O(\log n)\)
deleteAllBy         :: (a -> a -> Ordering) -> a -> Set a -> Set a
deleteAllBy :: (a -> a -> Ordering) -> a -> Set a -> Set a
deleteAllBy a -> a -> Ordering
cmp a
x Set a
s = (a -> a -> Ordering)
-> (forall s. Reifies s (OrdDict a) => O s (Set a)) -> Set a
forall a b.
(a -> a -> Ordering)
-> (forall s. Reifies s (OrdDict a) => O s b) -> b
withOrd a -> a -> Ordering
cmp ((forall s. Reifies s (OrdDict a) => O s (Set a)) -> Set a)
-> (forall s. Reifies s (OrdDict a) => O s (Set a)) -> Set a
forall a b. (a -> b) -> a -> b
$ (Set (O s a) -> Set (O s a)) -> Set a -> O s (Set a)
forall (f :: * -> *) s a (g :: * -> *).
(f (O s a) -> g (O s a)) -> f a -> O s (g a)
liftOrd1 (O s a -> Set (O s a) -> Set (O s a)
forall a. Ord a => a -> Set a -> Set a
Set.delete (O s a -> Set (O s a) -> Set (O s a))
-> O s a -> Set (O s a) -> Set (O s a)
forall a b. (a -> b) -> a -> b
$ a -> O s a
forall s a. a -> O s a
O a
x) Set a
s

-- | Run a query, eg. lookupGE, on the set with the given ordering.
--
-- Note: The 'Algorithms.BinarySearch.binarySearchIn' function may be
-- a useful alternative to 'queryBy'
--
-- >>> queryBy cmpS Set.lookupGE (S "22") $ fromListBy cmpS [S "a" , S "bbb" , S "ddddddd"]
-- Just (S "bbb")
-- >>> queryBy cmpS Set.lookupLE (S "22") $ fromListBy cmpS [S "a" , S "bbb" , S "ddddddd"]
-- Just (S "a")
-- >>> queryBy cmpS Set.lookupGE (S "333") $ fromListBy cmpS [S "a" , S "bbb" , S "ddddddd"]
-- Just (S "bbb")
queryBy           :: (a -> a -> Ordering)
                  -> (forall b. Ord b => b -> Set b -> t b)
                  -> a -> Set a -> t a
queryBy :: (a -> a -> Ordering)
-> (forall b. Ord b => b -> Set b -> t b) -> a -> Set a -> t a
queryBy a -> a -> Ordering
cmp forall b. Ord b => b -> Set b -> t b
fs a
q Set a
s = (a -> a -> Ordering)
-> (forall s. Reifies s (OrdDict a) => O s (t a)) -> t a
forall a b.
(a -> a -> Ordering)
-> (forall s. Reifies s (OrdDict a) => O s b) -> b
withOrd a -> a -> Ordering
cmp ((forall s. Reifies s (OrdDict a) => O s (t a)) -> t a)
-> (forall s. Reifies s (OrdDict a) => O s (t a)) -> t a
forall a b. (a -> b) -> a -> b
$ (Set (O s a) -> t (O s a)) -> Set a -> O s (t a)
forall (f :: * -> *) s a (g :: * -> *).
(f (O s a) -> g (O s a)) -> f a -> O s (g a)
liftOrd1 (O s a -> Set (O s a) -> t (O s a)
forall b. Ord b => b -> Set b -> t b
fs (O s a -> Set (O s a) -> t (O s a))
-> O s a -> Set (O s a) -> t (O s a)
forall a b. (a -> b) -> a -> b
$ a -> O s a
forall s a. a -> O s a
O a
q) Set a
s




-- queryBy'           :: Ord r
--                    => (a -> r)
--                    -> r
--                   -> (forall b. Ord b => b -> Set b -> t b)
--                   -> a -> Set a -> t a
-- queryBy' g fs q s = queryBy
--   where



--   withOrd cmp $ liftOrd1 (fs $ O q) s



  -- withOrd cmp $ liftOrd1 (Set.lookupGE $ O q) s




test :: Maybe S
test = (S -> S -> Ordering)
-> (forall b. Ord b => b -> Set b -> Maybe b)
-> S
-> Set S
-> Maybe S
forall a (t :: * -> *).
(a -> a -> Ordering)
-> (forall b. Ord b => b -> Set b -> t b) -> a -> Set a -> t a
queryBy S -> S -> Ordering
cmpS forall b. Ord b => b -> Set b -> Maybe b
Set.lookupGE (String -> S
S String
"22") (Set S -> Maybe S) -> Set S -> Maybe S
forall a b. (a -> b) -> a -> b
$ (S -> S -> Ordering) -> [S] -> Set S
forall a. (a -> a -> Ordering) -> [a] -> Set a
fromListBy S -> S -> Ordering
cmpS [String -> S
S String
"a" , String -> S
S String
"bbb" , String -> S
S String
"ddddddd"]
-- test = succBy cmpS (S "22") $ fromListBy cmpS [S "a" , S "bbb" , S "ddddddd"]