{-# LANGUAGE UnboxedTuples, MultiWayIf, BangPatterns, Trustworthy #-}
module Data.RangeSet.Internal (
    module Data.RangeSet.Internal,
    RangeSet(..), E, SRangeList(..), StrictMaybeE(..),
    size, height, foldE,
    module Data.RangeSet.Internal.Enum,
    module Data.RangeSet.Internal.SmartConstructors,
    module Data.RangeSet.Internal.Inserters,
    module Data.RangeSet.Internal.Extractors,
    module Data.RangeSet.Internal.Lumpers,
    module Data.RangeSet.Internal.Splitters,
    module Data.RangeSet.Internal.Heuristics
  ) where

import Prelude

import Data.RangeSet.Internal.Types
import Data.RangeSet.Internal.Enum
import Data.RangeSet.Internal.SmartConstructors
import Data.RangeSet.Internal.Inserters
import Data.RangeSet.Internal.Extractors
import Data.RangeSet.Internal.Lumpers
import Data.RangeSet.Internal.Splitters
import Data.RangeSet.Internal.Heuristics
import Data.Bits (shiftR)

{-# INLINEABLE insertE #-}
insertE :: E -> RangeSet a -> RangeSet a
insertE :: forall a. E -> RangeSet a -> RangeSet a
insertE !E
x RangeSet a
Tip = forall a. E -> E -> RangeSet a
single E
x E
x
insertE E
x t :: RangeSet a
t@(Fork H
h E
l E
u RangeSet a
lt RangeSet a
rt)
  -- Nothing happens when it's already in range
  | E
l forall a. Ord a => a -> a -> Bool
<= E
x = if
    | E
x forall a. Ord a => a -> a -> Bool
<= E
u -> RangeSet a
t
  -- If it is adjacent to the upper range, it may fuse
    | E
x forall a. Eq a => a -> a -> Bool
== forall a. Enum a => a -> a
succ E
u -> forall {a}. H -> E -> E -> RangeSet a -> RangeSet a -> RangeSet a
fuseRight H
h E
l E
x RangeSet a
lt RangeSet a
rt                                         -- we know x > u since x <= l && not x <= u
  -- Otherwise, insert and balance for right
    | Bool
otherwise -> forall a.
RangeSet a
-> RangeSet a
-> RangeSet a
-> (RangeSet a -> RangeSet a)
-> RangeSet a
ifStayedSame RangeSet a
rt (forall a. E -> RangeSet a -> RangeSet a
insertE E
x RangeSet a
rt) RangeSet a
t (forall a. E -> E -> RangeSet a -> RangeSet a -> RangeSet a
balance E
l E
u RangeSet a
lt)               -- cannot be biased, because fusion can shrink a tree
  | {- x < l -} Bool
otherwise = if
  -- If it is adjacent to the lower, it may fuse
    E
x forall a. Eq a => a -> a -> Bool
== forall a. Enum a => a -> a
pred E
l then forall {a}. H -> E -> E -> RangeSet a -> RangeSet a -> RangeSet a
fuseLeft H
h E
x E
u RangeSet a
lt RangeSet a
rt                                          -- the equality must be guarded by an existence check
  -- Otherwise, insert and balance for left
                else forall a.
RangeSet a
-> RangeSet a
-> RangeSet a
-> (RangeSet a -> RangeSet a)
-> RangeSet a
ifStayedSame RangeSet a
lt (forall a. E -> RangeSet a -> RangeSet a
insertE E
x RangeSet a
lt) RangeSet a
t forall a b. (a -> b) -> a -> b
$ \RangeSet a
lt' -> forall a. E -> E -> RangeSet a -> RangeSet a -> RangeSet a
balance E
l E
u RangeSet a
lt' RangeSet a
rt -- cannot be biased, because fusion can shrink a tree
  where
    {-# INLINE fuseLeft #-}
    fuseLeft :: H -> E -> E -> RangeSet a -> RangeSet a -> RangeSet a
fuseLeft !H
h !E
x !E
u RangeSet a
Tip !RangeSet a
rt = forall {a}. H -> E -> E -> RangeSet a -> RangeSet a -> RangeSet a
Fork H
h E
x E
u forall a. RangeSet a
Tip RangeSet a
rt
    fuseLeft H
h E
x E
u lt :: RangeSet a
lt@(Fork H
_ E
ll E
lu RangeSet a
llt RangeSet a
lrt) RangeSet a
rt
      | (# !E
l, !E
x', RangeSet a
lt' #) <- forall a.
E -> E -> RangeSet a -> RangeSet a -> (# E, E, RangeSet a #)
maxDelete E
ll E
lu RangeSet a
llt RangeSet a
lrt
      -- we know there exists an element larger than x'
      -- if x == x' + 1, we fuse (x != x' since that breaks disjointness, x == pred l)
      , E
x forall a. Eq a => a -> a -> Bool
== forall a. Enum a => a -> a
succ E
x' = forall a. E -> E -> RangeSet a -> RangeSet a -> RangeSet a
balanceR E
l E
u RangeSet a
lt' RangeSet a
rt
      | Bool
otherwise    = forall {a}. H -> E -> E -> RangeSet a -> RangeSet a -> RangeSet a
Fork H
h E
x E
u RangeSet a
lt RangeSet a
rt
    {-# INLINE fuseRight #-}
    fuseRight :: H -> E -> E -> RangeSet a -> RangeSet a -> RangeSet a
fuseRight !H
h !E
l !E
x !RangeSet a
lt RangeSet a
Tip = forall {a}. H -> E -> E -> RangeSet a -> RangeSet a -> RangeSet a
Fork H
h E
l E
x RangeSet a
lt forall a. RangeSet a
Tip
    fuseRight H
h E
l E
x RangeSet a
lt rt :: RangeSet a
rt@(Fork H
_ E
rl E
ru RangeSet a
rlt RangeSet a
rrt)
      | (# !E
x', !E
u, RangeSet a
rt' #) <- forall a.
E -> E -> RangeSet a -> RangeSet a -> (# E, E, RangeSet a #)
minDelete E
rl E
ru RangeSet a
rlt RangeSet a
rrt
      -- we know there exists an element smaller than x'
      -- if x == x' - 1, we fuse (x != x' since that breaks disjointness, as x == succ u)
      , E
x forall a. Eq a => a -> a -> Bool
== forall a. Enum a => a -> a
pred E
x' = forall a. E -> E -> RangeSet a -> RangeSet a -> RangeSet a
balanceL E
l E
u RangeSet a
lt RangeSet a
rt'
      | Bool
otherwise    = forall {a}. H -> E -> E -> RangeSet a -> RangeSet a -> RangeSet a
Fork H
h E
l E
x RangeSet a
lt RangeSet a
rt

{-# INLINEABLE deleteE #-}
deleteE :: E -> RangeSet a -> RangeSet a
deleteE :: forall a. E -> RangeSet a -> RangeSet a
deleteE !E
_ RangeSet a
Tip = forall a. RangeSet a
Tip
deleteE E
x t :: RangeSet a
t@(Fork H
h E
l E
u RangeSet a
lt RangeSet a
rt) =
  case forall a. Ord a => a -> a -> Ordering
compare E
l E
x of
    -- If its the only part of the range, the node is removed
    Ordering
EQ | E
x forall a. Eq a => a -> a -> Bool
== E
u    -> forall a. RangeSet a -> RangeSet a -> RangeSet a
glue RangeSet a
lt RangeSet a
rt
    -- If it's at an extreme, it shrinks the range
       | Bool
otherwise -> forall {a}. H -> E -> E -> RangeSet a -> RangeSet a -> RangeSet a
Fork H
h (forall a. Enum a => a -> a
succ E
l) E
u RangeSet a
lt RangeSet a
rt
    Ordering
LT -> case forall a. Ord a => a -> a -> Ordering
compare E
x E
u of
    -- If it's at an extreme, it shrinks the range
       Ordering
EQ          -> forall {a}. H -> E -> E -> RangeSet a -> RangeSet a -> RangeSet a
Fork H
h E
l (forall a. Enum a => a -> a
pred E
u) RangeSet a
lt RangeSet a
rt
    -- Otherwise, if it's still in range, the range undergoes fission
       Ordering
LT          -> forall a. E -> E -> E -> RangeSet a -> RangeSet a -> RangeSet a
fission E
l E
x E
u RangeSet a
lt RangeSet a
rt
    -- Otherwise delete and balance for one of the left or right
       Ordering
GT          -> forall a.
RangeSet a
-> RangeSet a
-> RangeSet a
-> (RangeSet a -> RangeSet a)
-> RangeSet a
ifStayedSame RangeSet a
rt (forall a. E -> RangeSet a -> RangeSet a
deleteE E
x RangeSet a
rt) RangeSet a
t forall a b. (a -> b) -> a -> b
$ forall a. E -> E -> RangeSet a -> RangeSet a -> RangeSet a
balance E
l E
u RangeSet a
lt             -- cannot be biased, because fisson can grow a tree
    Ordering
GT             -> forall a.
RangeSet a
-> RangeSet a
-> RangeSet a
-> (RangeSet a -> RangeSet a)
-> RangeSet a
ifStayedSame RangeSet a
lt (forall a. E -> RangeSet a -> RangeSet a
deleteE E
x RangeSet a
lt) RangeSet a
t forall a b. (a -> b) -> a -> b
$ \RangeSet a
lt' -> forall a. E -> E -> RangeSet a -> RangeSet a -> RangeSet a
balance E
l E
u RangeSet a
lt' RangeSet a
rt -- cannot be biased, because fisson can grow a tree
  where
    {- Fission breaks a node into two new ranges
       we'll push the range down into the smallest child, ensuring it's balanced -}
    {-# INLINE fission #-}
    fission :: E -> E -> E -> RangeSet a -> RangeSet a -> RangeSet a
    fission :: forall a. E -> E -> E -> RangeSet a -> RangeSet a -> RangeSet a
fission !E
l1 !E
x !E
u2 !RangeSet a
lt !RangeSet a
rt
      | forall a. RangeSet a -> H
height RangeSet a
lt forall a. Ord a => a -> a -> Bool
> forall a. RangeSet a -> H
height RangeSet a
rt = forall a. E -> E -> RangeSet a -> RangeSet a -> RangeSet a
fork E
l1 E
u1 RangeSet a
lt (forall a. E -> E -> RangeSet a -> RangeSet a
unsafeInsertL E
l2 E
u2 RangeSet a
rt)
      | Bool
otherwise = forall a. E -> E -> RangeSet a -> RangeSet a -> RangeSet a
fork E
l1 E
u1 (forall a. E -> E -> RangeSet a -> RangeSet a
unsafeInsertR E
l2 E
u2 RangeSet a
lt) RangeSet a
rt
      where
        !u1 :: E
u1 = forall a. Enum a => a -> a
pred E
x
        !l2 :: E
l2 = forall a. Enum a => a -> a
succ E
x

uncheckedSubsetOf :: RangeSet a -> RangeSet a -> Bool
uncheckedSubsetOf :: forall a. RangeSet a -> RangeSet a -> Bool
uncheckedSubsetOf RangeSet a
Tip RangeSet a
_ = Bool
True
uncheckedSubsetOf RangeSet a
_ RangeSet a
Tip = Bool
False
uncheckedSubsetOf (Fork H
_ E
ll E
lu RangeSet a
llt RangeSet a
lrt) (Fork H
_ E
rl E
ru RangeSet a
rlt RangeSet a
rrt) = case forall a.
E
-> E
-> E
-> E
-> RangeSet a
-> RangeSet a
-> (# RangeSet a, RangeSet a, RangeSet a #)
splitOverlapFork E
ll E
lu E
rl E
ru RangeSet a
rlt RangeSet a
rrt of
  (# RangeSet a
lt', Fork H
1 E
x E
y RangeSet a
_ RangeSet a
_, RangeSet a
rt' #) ->
       E
x forall a. Eq a => a -> a -> Bool
== E
ll Bool -> Bool -> Bool
&& E
y forall a. Eq a => a -> a -> Bool
== E
lu
    Bool -> Bool -> Bool
&& forall a. RangeSet a -> RangeSet a -> Bool
uncheckedSubsetOf RangeSet a
llt RangeSet a
lt' Bool -> Bool -> Bool
&& forall a. RangeSet a -> RangeSet a -> Bool
uncheckedSubsetOf RangeSet a
lrt RangeSet a
rt'
  (# RangeSet a, RangeSet a, RangeSet a #)
_                              -> Bool
False

{-# INLINEABLE fromDistinctAscRangesSz #-}
fromDistinctAscRangesSz :: SRangeList -> Int -> RangeSet a
fromDistinctAscRangesSz :: forall a. SRangeList -> E -> RangeSet a
fromDistinctAscRangesSz SRangeList
rs !E
n = case forall a. SRangeList -> E -> E -> (# H, RangeSet a, SRangeList #)
go SRangeList
rs E
0 (E
n forall a. Num a => a -> a -> a
- E
1) of (# H
_, RangeSet a
t, SRangeList
_ #) -> RangeSet a
t
  where
    go :: SRangeList -> Int -> Int -> (# H, RangeSet a, SRangeList #)
    go :: forall a. SRangeList -> E -> E -> (# H, RangeSet a, SRangeList #)
go SRangeList
rs !E
i !E
j
      | E
i forall a. Ord a => a -> a -> Bool
> E
j     = (# H
1, forall a. RangeSet a
Tip, SRangeList
rs #)
      | Bool
otherwise =
        let !mid :: E
mid = (E
i forall a. Num a => a -> a -> a
+ E
j) forall a. Bits a => a -> E -> a
`shiftR` E
1
        in case forall a. SRangeList -> E -> E -> (# H, RangeSet a, SRangeList #)
go SRangeList
rs E
i (E
mid forall a. Num a => a -> a -> a
- E
1) of
             (# H
_, RangeSet a
lt, SRangeList
rs' #) ->
                let !(SRangeCons E
l E
u SRangeList
rs'') = SRangeList
rs'
                in case forall a. SRangeList -> E -> E -> (# H, RangeSet a, SRangeList #)
go SRangeList
rs'' (E
mid forall a. Num a => a -> a -> a
+ E
1) E
j of
                      -- there is a height bias to the right, so the height of the right tree is all we need
                      -- perhaps this can be computed though from mid somehow instead of passing back?
                      (# !H
h, RangeSet a
rt, SRangeList
rs''' #) -> (# H
h forall a. Num a => a -> a -> a
+ H
1, forall {a}. H -> E -> E -> RangeSet a -> RangeSet a -> RangeSet a
Fork H
h E
l E
u RangeSet a
lt RangeSet a
rt, SRangeList
rs''' #)

{-# INLINE insertRangeE #-}
-- This could be improved, but is OK
insertRangeE :: E -> E -> RangeSet a -> RangeSet a
insertRangeE :: forall a. E -> E -> RangeSet a -> RangeSet a
insertRangeE !E
l !E
u RangeSet a
Tip = forall a. E -> E -> RangeSet a
single E
l E
u
insertRangeE E
l E
u (Fork H
_ E
l' E
u' RangeSet a
lt RangeSet a
rt) = let (# RangeSet a
lt', RangeSet a
rt' #) = forall a.
E
-> E
-> E
-> E
-> RangeSet a
-> RangeSet a
-> (# RangeSet a, RangeSet a #)
splitFork E
l E
u E
l' E
u' RangeSet a
lt RangeSet a
rt in forall a. E -> E -> RangeSet a -> RangeSet a -> RangeSet a
link E
l E
u RangeSet a
lt' RangeSet a
rt'