{-# 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 (0, 1)
[(0,1)]
>>> intervalIS (1, 0)
[]
-}
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 2 $ fromListIS [(0,1),(3,4)]
[(0,4)]
>>> insertIS 1 $ fromListIS [(0,1)]
[(0,1)]
-}
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 (2,4) $ fromListIS [(1,1),(3,3),(5,5)]
[(1,5)]
-}
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 1 $ fromListIS [0..2]
[(0,0),(2,2)]
-}
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 (1,7) $ fromListIS [(0,1),(3,5),(7,8)]
[(0,0),(8,8)]
>>> deleteIntervalIS (1,1) $ fromListIS [(0,0)]
[(0,0)]
-}
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 0 (fromListIS [(-1,1)])
True
>>> memberIS 1 (fromListIS [(-1,1)])
True
>>> memberIS 2 (fromListIS [(-1,1)])
False
-}
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 (-2) $ intervalIS (-1,1)
Nothing
>>> lookupIS (-1) $ intervalIS (-1,1)
Just (-1,1)
>>> lookupIS 0 $ intervalIS (-1,1)
Just (-1,1)
>>> lookupIS 1 $ intervalIS (-1,1)
Just (-1,1)
>>> lookupIS 2 $ intervalIS (-1,1)
Nothing
-}
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 (1,8) $ fromList [(0,2),(4,5),(7,9)]
[(0,2),(4,5),(7,9)]
>>> lookupIntervalIS (3,6) $ fromList [(0,2),(4,5),(7,9)]
[(4,5)]
>>> lookupIntervalIS (3,3) $ fromList [(0,2),(4,5),(7,9)]
[]
>>> lookupIntervalIS (0, 1) $ fromList [(1, 2)]
[(1,2)]
>>> lookupIntervalIS (0, 1) $ fromList [(-1, 0)]
[(-1,0)]
-}
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)
_ -> []

{- | minimum excluded value

>>> mex $ fromListIS [(0,1),(3,5)]
2
>>> mex $ intervalIS (1, 5)
0
-}
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