--------------------------------------------------------------------------------
-- |
-- 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 Show
cmpS :: S -> S -> Ordering
cmpS = comparing (\(S s) -> length 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 f x s = let (l,s') = Set.spanAntitone (g LT . f) s
(m,r) = Set.spanAntitone (g EQ . f) s'
g c y = y `compare` x == c
in (l,m,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 f s = let (l,s') = Set.spanAntitone ((==) LT . f) s
(m,r) = Set.spanAntitone ((==) EQ . f) s'
in (l,m,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 cmp xs = withOrd cmp (extractOrd1 . Set.fromList . map O $ 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 = 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 cmp x s = withOrd cmp $ liftOrd1 (Set.insert $ O x) 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 cmp x s = withOrd cmp $ liftOrd1 (Set.delete $ O x) 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 cmp fs q s = withOrd cmp $ liftOrd1 (fs $ O q) 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 = queryBy cmpS Set.lookupGE (S "22") $ fromListBy cmpS [S "a" , S "bbb" , S "ddddddd"]
-- test = succBy cmpS (S "22") $ fromListBy cmpS [S "a" , S "bbb" , S "ddddddd"]