{-# LANGUAGE TypeFamilies #-}
module Data.IntervalSet where
import qualified Data.Foldable as F
import qualified Data.IntMap.Strict as IM
import GHC.Exts
newtype IntervalSet = IntervalSet (IM.IntMap Int)
deriving (IntervalSet -> IntervalSet -> Bool
(IntervalSet -> IntervalSet -> Bool)
-> (IntervalSet -> IntervalSet -> Bool) -> Eq IntervalSet
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: IntervalSet -> IntervalSet -> Bool
== :: IntervalSet -> IntervalSet -> Bool
$c/= :: IntervalSet -> IntervalSet -> Bool
/= :: IntervalSet -> IntervalSet -> Bool
Eq)
instance Show IntervalSet where
show :: IntervalSet -> String
show (IntervalSet IntMap Int
s) = [(Int, Int)] -> String
forall a. Show a => a -> String
show ([(Int, Int)] -> String) -> [(Int, Int)] -> String
forall a b. (a -> b) -> a -> b
$ IntMap Int -> [(Int, Int)]
forall a. IntMap a -> [(Int, a)]
IM.toAscList IntMap Int
s
emptyIS :: IntervalSet
emptyIS :: IntervalSet
emptyIS = IntMap Int -> IntervalSet
IntervalSet IntMap Int
forall a. IntMap a
IM.empty
singletonIS :: Int -> IntervalSet
singletonIS :: Int -> IntervalSet
singletonIS Int
x = IntMap Int -> IntervalSet
IntervalSet (Int -> Int -> IntMap Int
forall a. Int -> a -> IntMap a
IM.singleton Int
x Int
x)
intervalIS :: (Int, Int) -> IntervalSet
intervalIS :: (Int, Int) -> IntervalSet
intervalIS (Int
l, Int
r)
| Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
r = IntMap Int -> IntervalSet
IntervalSet (Int -> Int -> IntMap Int
forall a. Int -> a -> IntMap a
IM.singleton Int
l Int
r)
| Bool
otherwise = IntervalSet
emptyIS
insertIS :: Int -> IntervalSet -> IntervalSet
insertIS :: Int -> IntervalSet -> IntervalSet
insertIS Int
x (IntervalSet IntMap Int
s) = IntMap Int -> IntervalSet
IntervalSet (IntMap Int -> IntervalSet) -> IntMap Int -> IntervalSet
forall a b. (a -> b) -> a -> b
$ case Int -> IntMap Int -> Maybe (Int, Int)
forall a. Int -> IntMap a -> Maybe (Int, a)
IM.lookupLE Int
x IntMap Int
s of
Just (Int
l, Int
r)
| Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
r -> IntMap Int
s
| Bool
otherwise -> case Int -> IntMap Int -> Maybe Int
forall a. Int -> IntMap a -> Maybe a
IM.lookup (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) IntMap Int
s of
Just Int
r'
| Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 -> Int -> IntMap Int -> IntMap Int
forall a. Int -> IntMap a -> IntMap a
IM.delete (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (IntMap Int -> IntMap Int) -> IntMap Int -> IntMap Int
forall a b. (a -> b) -> a -> b
$ Int -> Int -> IntMap Int -> IntMap Int
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert Int
l Int
r' IntMap Int
s
| Bool
otherwise -> Int -> IntMap Int -> IntMap Int
forall a. Int -> IntMap a -> IntMap a
IM.delete (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (IntMap Int -> IntMap Int) -> IntMap Int -> IntMap Int
forall a b. (a -> b) -> a -> b
$ Int -> Int -> IntMap Int -> IntMap Int
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert Int
x Int
r' IntMap Int
s
Maybe Int
Nothing
| Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 -> Int -> Int -> IntMap Int -> IntMap Int
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert Int
l Int
x IntMap Int
s
| Bool
otherwise -> Int -> Int -> IntMap Int -> IntMap Int
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert Int
x Int
x IntMap Int
s
Maybe (Int, Int)
Nothing -> case Int -> IntMap Int -> Maybe Int
forall a. Int -> IntMap a -> Maybe a
IM.lookup (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) IntMap Int
s of
Just Int
r' -> Int -> IntMap Int -> IntMap Int
forall a. Int -> IntMap a -> IntMap a
IM.delete (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (IntMap Int -> IntMap Int) -> IntMap Int -> IntMap Int
forall a b. (a -> b) -> a -> b
$ Int -> Int -> IntMap Int -> IntMap Int
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert Int
x Int
r' IntMap Int
s
Maybe Int
Nothing -> Int -> Int -> IntMap Int -> IntMap Int
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert Int
x Int
x IntMap Int
s
insertIntervalIS :: (Int, Int) -> IntervalSet -> IntervalSet
insertIntervalIS :: (Int, Int) -> IntervalSet -> IntervalSet
insertIntervalIS (Int
l, Int
r) (IntervalSet IntMap Int
s)
| Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
r = IntMap Int -> IntervalSet
IntervalSet (IntMap Int -> IntMap Int
go IntMap Int
s)
| Bool
otherwise = IntMap Int -> IntervalSet
IntervalSet IntMap Int
s
where
!l' :: Int
l' = case Int -> IntMap Int -> Maybe (Int, Int)
forall a. Int -> IntMap a -> Maybe (Int, a)
IM.lookupLE Int
l IntMap Int
s of
Just (Int
ll, Int
lr) | Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
lr Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 -> Int
ll
Maybe (Int, Int)
_ -> Int
l
!r' :: Int
r' = case Int -> IntMap Int -> Maybe (Int, Int)
forall a. Int -> IntMap a -> Maybe (Int, a)
IM.lookupLE Int
r IntMap Int
s of
Just (Int
_, Int
rr) | Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
rr -> Int
rr
Maybe (Int, Int)
_ -> Int -> Int -> IntMap Int -> Int
forall a. a -> Int -> IntMap a -> a
IM.findWithDefault Int
r (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) IntMap Int
s
go :: IntMap Int -> IntMap Int
go IntMap Int
im = case Int -> IntMap Int -> Maybe (Int, Int)
forall a. Int -> IntMap a -> Maybe (Int, a)
IM.lookupLE Int
r' IntMap Int
im of
Just (Int
kl, Int
_)
| Int
l' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
kl -> IntMap Int -> IntMap Int
go (IntMap Int -> IntMap Int) -> IntMap Int -> IntMap Int
forall a b. (a -> b) -> a -> b
$ Int -> IntMap Int -> IntMap Int
forall a. Int -> IntMap a -> IntMap a
IM.delete Int
kl IntMap Int
im
Maybe (Int, Int)
_ -> Int -> Int -> IntMap Int -> IntMap Int
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert Int
l' Int
r' IntMap Int
im
deleteIS :: Int -> IntervalSet -> IntervalSet
deleteIS :: Int -> IntervalSet -> IntervalSet
deleteIS Int
x (IntervalSet IntMap Int
s) = IntMap Int -> IntervalSet
IntervalSet (IntMap Int -> IntervalSet) -> IntMap Int -> IntervalSet
forall a b. (a -> b) -> a -> b
$ case Int -> IntMap Int -> Maybe (Int, Int)
forall a. Int -> IntMap a -> Maybe (Int, a)
IM.lookupLE Int
x IntMap Int
s of
Just (Int
l, Int
r) | Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
r -> Int -> IntMap Int -> IntMap Int
fl Int
l (IntMap Int -> IntMap Int) -> IntMap Int -> IntMap Int
forall a b. (a -> b) -> a -> b
$ Int -> IntMap Int -> IntMap Int
fr Int
r IntMap Int
s
Maybe (Int, Int)
_ -> IntMap Int
s
where
fl :: Int -> IntMap Int -> IntMap Int
fl Int
l IntMap Int
im
| Int
l Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
x = Int -> Int -> IntMap Int -> IntMap Int
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert Int
l (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) IntMap Int
im
| Bool
otherwise = Int -> IntMap Int -> IntMap Int
forall a. Int -> IntMap a -> IntMap a
IM.delete Int
l IntMap Int
im
fr :: Int -> IntMap Int -> IntMap Int
fr Int
r IntMap Int
im
| Int
r Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
x = Int -> Int -> IntMap Int -> IntMap Int
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
r IntMap Int
im
| Bool
otherwise = IntMap Int
im
deleteIntervalIS :: (Int, Int) -> IntervalSet -> IntervalSet
deleteIntervalIS :: (Int, Int) -> IntervalSet -> IntervalSet
deleteIntervalIS (Int
l, Int
r) (IntervalSet IntMap Int
s)
| Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
r = IntMap Int -> IntervalSet
IntervalSet (IntMap Int -> IntervalSet) -> IntMap Int -> IntervalSet
forall a b. (a -> b) -> a -> b
$ IntMap Int -> IntMap Int
go IntMap Int
s
| Bool
otherwise = (IntMap Int -> IntervalSet
IntervalSet IntMap Int
s)
where
go :: IntMap Int -> IntMap Int
go !IntMap Int
im = case Int -> IntMap Int -> Maybe (Int, Int)
forall a. Int -> IntMap a -> Maybe (Int, a)
IM.lookupLE Int
r IntMap Int
im of
Just (Int
l', Int
r')
| Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
l' ->
if Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
r'
then IntMap Int -> IntMap Int
go (IntMap Int -> IntMap Int)
-> (IntMap Int -> IntMap Int) -> IntMap Int -> IntMap Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> IntMap Int -> IntMap Int
forall a. Int -> IntMap a -> IntMap a
IM.delete Int
l' (IntMap Int -> IntMap Int) -> IntMap Int -> IntMap Int
forall a b. (a -> b) -> a -> b
$ Int -> Int -> IntMap Int -> IntMap Int
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
r' IntMap Int
im
else IntMap Int -> IntMap Int
go (IntMap Int -> IntMap Int) -> IntMap Int -> IntMap Int
forall a b. (a -> b) -> a -> b
$ Int -> IntMap Int -> IntMap Int
forall a. Int -> IntMap a -> IntMap a
IM.delete Int
l' IntMap Int
im
| Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
r' -> Int -> Int -> IntMap Int -> IntMap Int
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
r' (IntMap Int -> IntMap Int) -> IntMap Int -> IntMap Int
forall a b. (a -> b) -> a -> b
$ Int -> Int -> IntMap Int -> IntMap Int
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert Int
l' (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) IntMap Int
im
| Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
r' -> Int -> Int -> IntMap Int -> IntMap Int
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert Int
l' (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) IntMap Int
im
| Bool
otherwise -> IntMap Int
im
Maybe (Int, Int)
Nothing -> IntMap Int
im
memberIS :: Int -> IntervalSet -> Bool
memberIS :: Int -> IntervalSet -> Bool
memberIS Int
x (IntervalSet IntMap Int
s) = case Int -> IntMap Int -> Maybe (Int, Int)
forall a. Int -> IntMap a -> Maybe (Int, a)
IM.lookupLE Int
x IntMap Int
s of
Just (Int
_, Int
r) -> Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
r
Maybe (Int, Int)
Nothing -> Bool
False
lookupIS :: Int -> IntervalSet -> Maybe (Int, Int)
lookupIS :: Int -> IntervalSet -> Maybe (Int, Int)
lookupIS Int
x (IntervalSet IntMap Int
s) = case Int -> IntMap Int -> Maybe (Int, Int)
forall a. Int -> IntMap a -> Maybe (Int, a)
IM.lookupLE Int
x IntMap Int
s of
Just (Int
l, Int
r) | Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
r -> (Int, Int) -> Maybe (Int, Int)
forall a. a -> Maybe a
Just (Int
l, Int
r)
Maybe (Int, Int)
_ -> Maybe (Int, Int)
forall a. Maybe a
Nothing
lookupIntervalIS :: (Int, Int) -> IntervalSet -> [(Int, Int)]
lookupIntervalIS :: (Int, Int) -> IntervalSet -> [(Int, Int)]
lookupIntervalIS (Int
l0, Int
r0) (IntervalSet IntMap Int
s)
| Int
l0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
r0 = case Int -> IntMap Int -> Maybe (Int, Int)
forall a. Int -> IntMap a -> Maybe (Int, a)
IM.lookupLE Int
l0 IntMap Int
s of
Just (Int
l, Int
r) | Int
l0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
r -> (Int
l, Int
r) (Int, Int) -> [(Int, Int)] -> [(Int, Int)]
forall a. a -> [a] -> [a]
: Int -> [(Int, Int)]
go Int
r
Maybe (Int, Int)
_ -> Int -> [(Int, Int)]
go Int
l0
| Bool
otherwise = []
where
go :: Int -> [(Int, Int)]
go Int
r = case Int -> IntMap Int -> Maybe (Int, Int)
forall a. Int -> IntMap a -> Maybe (Int, a)
IM.lookupGT Int
r IntMap Int
s of
Just (Int
l', Int
r')
| Int
l' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
r0 -> (Int
l', Int
r') (Int, Int) -> [(Int, Int)] -> [(Int, Int)]
forall a. a -> [a] -> [a]
: Int -> [(Int, Int)]
go Int
r'
Maybe (Int, Int)
_ -> []
mex :: IntervalSet -> Int
mex :: IntervalSet -> Int
mex (IntervalSet IntMap Int
s) = case Int -> IntMap Int -> Maybe (Int, Int)
forall a. Int -> IntMap a -> Maybe (Int, a)
IM.lookupLE Int
0 IntMap Int
s of
Just (Int
_, Int
r) | Int
0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
r -> Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
Maybe (Int, Int)
_ -> Int
0
fromListIS :: [(Int, Int)] -> IntervalSet
fromListIS :: [(Int, Int)] -> IntervalSet
fromListIS = (IntervalSet -> (Int, Int) -> IntervalSet)
-> IntervalSet -> [(Int, Int)] -> IntervalSet
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
F.foldl' (((Int, Int) -> IntervalSet -> IntervalSet)
-> IntervalSet -> (Int, Int) -> IntervalSet
forall a b c. (a -> b -> c) -> b -> a -> c
flip (Int, Int) -> IntervalSet -> IntervalSet
insertIntervalIS) IntervalSet
emptyIS
toListIS :: IntervalSet -> [(Int, Int)]
toListIS :: IntervalSet -> [(Int, Int)]
toListIS (IntervalSet IntMap Int
s) = IntMap Int -> [(Int, Int)]
forall a. IntMap a -> [(Int, a)]
IM.toAscList IntMap Int
s
instance IsList IntervalSet where
type Item IntervalSet = (Int, Int)
fromList :: [Item IntervalSet] -> IntervalSet
fromList = [(Int, Int)] -> IntervalSet
[Item IntervalSet] -> IntervalSet
fromListIS
toList :: IntervalSet -> [Item IntervalSet]
toList = IntervalSet -> [(Int, Int)]
IntervalSet -> [Item IntervalSet]
toListIS