{-|
Description      : Utilities for balanced binary trees.
Copyright        : (c) Galois, Inc 2014-2019
Maintainer       : Joe Hendrix <jhendrix@galois.com>
-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE Safe #-}
module Data.Parameterized.Utils.BinTree
  ( MaybeS(..)
  , fromMaybeS
  , Updated(..)
  , updatedValue
  , TreeApp(..)
  , IsBinTree(..)
  , balanceL
  , balanceR
  , glue
  , merge
  , filterGt
  , filterLt
  , insert
  , delete
  , union
  , link
  , PairS(..)
  ) where

import Control.Applicative

------------------------------------------------------------------------
-- MaybeS

-- | A strict version of 'Maybe'
data MaybeS v
   = JustS !v
   | NothingS

instance Functor MaybeS where
  fmap :: (a -> b) -> MaybeS a -> MaybeS b
fmap a -> b
_ MaybeS a
NothingS = MaybeS b
forall v. MaybeS v
NothingS
  fmap a -> b
f (JustS a
v) = b -> MaybeS b
forall v. v -> MaybeS v
JustS (a -> b
f a
v)

instance Alternative MaybeS where
  empty :: MaybeS a
empty = MaybeS a
forall v. MaybeS v
NothingS
  mv :: MaybeS a
mv@JustS{} <|> :: MaybeS a -> MaybeS a -> MaybeS a
<|> MaybeS a
_ = MaybeS a
mv
  MaybeS a
NothingS <|> MaybeS a
v = MaybeS a
v

instance Applicative MaybeS where
  pure :: a -> MaybeS a
pure = a -> MaybeS a
forall v. v -> MaybeS v
JustS

  MaybeS (a -> b)
NothingS <*> :: MaybeS (a -> b) -> MaybeS a -> MaybeS b
<*> MaybeS a
_ = MaybeS b
forall v. MaybeS v
NothingS
  JustS{} <*> MaybeS a
NothingS = MaybeS b
forall v. MaybeS v
NothingS
  JustS a -> b
f <*> JustS a
x = b -> MaybeS b
forall v. v -> MaybeS v
JustS (a -> b
f a
x)

fromMaybeS :: a -> MaybeS a -> a
fromMaybeS :: a -> MaybeS a -> a
fromMaybeS a
r MaybeS a
NothingS = a
r
fromMaybeS a
_ (JustS a
v) = a
v

------------------------------------------------------------------------
-- Updated

-- | @Updated a@ contains a value that has been flagged on whether it was
-- modified by an operation.
data Updated a
   = Updated   !a
   | Unchanged !a

updatedValue :: Updated a -> a
updatedValue :: Updated a -> a
updatedValue (Updated a
a) = a
a
updatedValue (Unchanged a
a) = a
a

------------------------------------------------------------------------
-- IsBinTree

data TreeApp e t
   = BinTree !e !t !t
   | TipTree

class IsBinTree t e | t -> e where
  asBin :: t -> TreeApp e t
  tip :: t

  bin :: e -> t -> t -> t
  size :: t -> Int

delta,ratio :: Int
delta :: Int
delta = Int
3
ratio :: Int
ratio = Int
2

-- | @balanceL p l r@ returns a balanced tree for the sequence @l ++ [p] ++ r@.
--
-- It assumes that @l@ and @r@ are close to being balanced, and that only
-- @l@ may contain too many elements.
balanceL :: (IsBinTree c e) => e -> c -> c -> c
balanceL :: e -> c -> c -> c
balanceL e
p c
l c
r = do
  case c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin c
l of
    BinTree e
l_pair c
ll c
lr | c -> Int
forall t e. IsBinTree t e => t -> Int
size c
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 (Int
deltaInt -> Int -> Int
forall a. Num a => a -> a -> a
*c -> Int
forall t e. IsBinTree t e => t -> Int
size c
r) ->
      case c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin c
lr of
        BinTree e
lr_pair c
lrl c
lrr | c -> Int
forall t e. IsBinTree t e => t -> Int
size c
lr Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
2 (Int
ratioInt -> Int -> Int
forall a. Num a => a -> a -> a
*c -> Int
forall t e. IsBinTree t e => t -> Int
size c
ll) ->
          e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
bin e
lr_pair (e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
bin e
l_pair c
ll c
lrl) (e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
bin e
p c
lrr c
r)
        TreeApp e c
_ -> e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
bin e
l_pair c
ll (e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
bin e
p c
lr c
r)

    TreeApp e c
_ -> e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
bin e
p c
l c
r
{-# INLINE balanceL #-}

-- | @balanceR p l r@ returns a balanced tree for the sequence @l ++ [p] ++ r@.
--
-- It assumes that @l@ and @r@ are close to being balanced, and that only
-- @r@ may contain too many elements.
balanceR :: (IsBinTree c e) => e -> c -> c -> c
balanceR :: e -> c -> c -> c
balanceR e
p c
l c
r = do
  case c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin c
r of
    BinTree e
r_pair c
rl c
rr | c -> Int
forall t e. IsBinTree t e => t -> Int
size c
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 (Int
deltaInt -> Int -> Int
forall a. Num a => a -> a -> a
*c -> Int
forall t e. IsBinTree t e => t -> Int
size c
l) ->
      case c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin c
rl of
        BinTree e
rl_pair c
rll c
rlr | c -> Int
forall t e. IsBinTree t e => t -> Int
size c
rl Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
2 (Int
ratioInt -> Int -> Int
forall a. Num a => a -> a -> a
*c -> Int
forall t e. IsBinTree t e => t -> Int
size c
rr) ->
          (e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
bin e
rl_pair (c -> c -> c) -> c -> c -> c
forall a b. (a -> b) -> a -> b
$! e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
bin e
p c
l c
rll) (c -> c) -> c -> c
forall a b. (a -> b) -> a -> b
$! e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
bin e
r_pair c
rlr c
rr
        TreeApp e c
_ -> e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
bin e
r_pair (e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
bin e
p c
l c
rl) c
rr
    TreeApp e c
_ -> e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
bin e
p c
l c
r
{-# INLINE balanceR #-}

-- | Insert a new maximal element.
insertMax :: IsBinTree c e => e -> c -> c
insertMax :: e -> c -> c
insertMax e
p c
t =
  case c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin c
t of
    TreeApp e c
TipTree -> e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
bin e
p c
forall t e. IsBinTree t e => t
tip c
forall t e. IsBinTree t e => t
tip
    BinTree e
q c
l c
r -> e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
balanceR e
q c
l (e -> c -> c
forall c e. IsBinTree c e => e -> c -> c
insertMax e
p c
r)

-- | Insert a new minimal element.
insertMin :: IsBinTree c e => e -> c -> c
insertMin :: e -> c -> c
insertMin e
p c
t =
  case c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin c
t of
    TreeApp e c
TipTree -> e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
bin e
p c
forall t e. IsBinTree t e => t
tip c
forall t e. IsBinTree t e => t
tip
    BinTree e
q c
l c
r -> e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
balanceL e
q (e -> c -> c
forall c e. IsBinTree c e => e -> c -> c
insertMin e
p c
l) c
r

-- | @link@ is called to insert a key and value between two disjoint subtrees.
link :: IsBinTree c e => e -> c -> c -> c
link :: e -> c -> c -> c
link e
p c
l c
r =
  case (c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin c
l, c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin c
r) of
    (TreeApp e c
TipTree, TreeApp e c
_) -> e -> c -> c
forall c e. IsBinTree c e => e -> c -> c
insertMin e
p c
r
    (TreeApp e c
_, TreeApp e c
TipTree) -> e -> c -> c
forall c e. IsBinTree c e => e -> c -> c
insertMax e
p c
l
    (BinTree e
py c
ly c
ry, BinTree e
pz c
lz c
rz)
     | Int
deltaInt -> Int -> Int
forall a. Num a => a -> a -> a
*c -> Int
forall t e. IsBinTree t e => t -> Int
size c
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< c -> Int
forall t e. IsBinTree t e => t -> Int
size c
r -> e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
balanceL e
pz (e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
link e
p c
l c
lz) c
rz
     | Int
deltaInt -> Int -> Int
forall a. Num a => a -> a -> a
*c -> Int
forall t e. IsBinTree t e => t -> Int
size c
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< c -> Int
forall t e. IsBinTree t e => t -> Int
size c
l -> e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
balanceR e
py c
ly (e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
link e
p c
ry c
r)
     | Bool
otherwise             -> e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
bin e
p c
l c
r
{-# INLINE link #-}

-- | A Strict pair
data PairS f s = PairS !f !s

deleteFindMin :: IsBinTree c e => e -> c -> c -> PairS e c
deleteFindMin :: e -> c -> c -> PairS e c
deleteFindMin e
p c
l c
r =
  case c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin c
l of
    TreeApp e c
TipTree -> e -> c -> PairS e c
forall f s. f -> s -> PairS f s
PairS e
p c
r
    BinTree e
lp c
ll c
lr ->
      case e -> c -> c -> PairS e c
forall c e. IsBinTree c e => e -> c -> c -> PairS e c
deleteFindMin e
lp c
ll c
lr of
        PairS e
q c
l' -> e -> c -> PairS e c
forall f s. f -> s -> PairS f s
PairS e
q (e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
balanceR e
p c
l' c
r)
{-# INLINABLE deleteFindMin #-}

deleteFindMax :: IsBinTree c e => e -> c -> c -> PairS e c
deleteFindMax :: e -> c -> c -> PairS e c
deleteFindMax e
p c
l c
r =
  case c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin c
r of
    TreeApp e c
TipTree -> e -> c -> PairS e c
forall f s. f -> s -> PairS f s
PairS e
p c
l
    BinTree e
rp c
rl c
rr ->
      case e -> c -> c -> PairS e c
forall c e. IsBinTree c e => e -> c -> c -> PairS e c
deleteFindMax e
rp c
rl c
rr of
        PairS e
q c
r' -> e -> c -> PairS e c
forall f s. f -> s -> PairS f s
PairS e
q (e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
balanceL e
p c
l c
r')
{-# INLINABLE deleteFindMax #-}

-- | Concatenate two trees that are ordered with respect to each other.
merge :: IsBinTree c e => c -> c -> c
merge :: c -> c -> c
merge c
l c
r =
  case (c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin c
l, c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin c
r) of
    (TreeApp e c
TipTree, TreeApp e c
_) -> c
r
    (TreeApp e c
_, TreeApp e c
TipTree) -> c
l
    (BinTree e
x c
lx c
rx, BinTree e
y c
ly c
ry)
      | Int
deltaInt -> Int -> Int
forall a. Num a => a -> a -> a
*c -> Int
forall t e. IsBinTree t e => t -> Int
size c
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< c -> Int
forall t e. IsBinTree t e => t -> Int
size c
r -> e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
balanceL e
y (c -> c -> c
forall c e. IsBinTree c e => c -> c -> c
merge c
l c
ly) c
ry
      | Int
deltaInt -> Int -> Int
forall a. Num a => a -> a -> a
*c -> Int
forall t e. IsBinTree t e => t -> Int
size c
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< c -> Int
forall t e. IsBinTree t e => t -> Int
size c
l -> e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
balanceR e
x c
lx (c -> c -> c
forall c e. IsBinTree c e => c -> c -> c
merge c
rx c
r)
      | c -> Int
forall t e. IsBinTree t e => t -> Int
size c
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> c -> Int
forall t e. IsBinTree t e => t -> Int
size c
r ->
        case e -> c -> c -> PairS e c
forall c e. IsBinTree c e => e -> c -> c -> PairS e c
deleteFindMax e
x c
lx c
rx of
          PairS e
q c
l' -> e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
balanceR e
q c
l' c
r
      | Bool
otherwise ->
        case e -> c -> c -> PairS e c
forall c e. IsBinTree c e => e -> c -> c -> PairS e c
deleteFindMin e
y c
ly c
ry of
          PairS e
q c
r' -> e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
balanceL e
q c
l c
r'
{-# INLINABLE merge #-}

------------------------------------------------------------------------
-- Ordered operations

-- | @insert p m@ inserts the binding into @m@.  It returns
-- an Unchanged value if the map stays the same size and an updated
-- value if a new entry was inserted.
insert :: (IsBinTree c e) => (e -> e -> Ordering) -> e -> c -> Updated c
insert :: (e -> e -> Ordering) -> e -> c -> Updated c
insert e -> e -> Ordering
comp e
x c
t =
  case c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin c
t of
    TreeApp e c
TipTree -> c -> Updated c
forall a. a -> Updated a
Updated (e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
bin e
x c
forall t e. IsBinTree t e => t
tip c
forall t e. IsBinTree t e => t
tip)
    BinTree e
y c
l c
r ->
      case e -> e -> Ordering
comp e
x e
y of
        Ordering
LT ->
          case (e -> e -> Ordering) -> e -> c -> Updated c
forall c e.
IsBinTree c e =>
(e -> e -> Ordering) -> e -> c -> Updated c
insert e -> e -> Ordering
comp e
x c
l of
            Updated c
l'   -> c -> Updated c
forall a. a -> Updated a
Updated   (e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
balanceL e
y c
l' c
r)
            Unchanged c
l' -> c -> Updated c
forall a. a -> Updated a
Unchanged (e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
bin       e
y c
l' c
r)
        Ordering
GT ->
          case (e -> e -> Ordering) -> e -> c -> Updated c
forall c e.
IsBinTree c e =>
(e -> e -> Ordering) -> e -> c -> Updated c
insert e -> e -> Ordering
comp e
x c
r of
            Updated c
r'   -> c -> Updated c
forall a. a -> Updated a
Updated   (e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
balanceR e
y c
l c
r')
            Unchanged c
r' -> c -> Updated c
forall a. a -> Updated a
Unchanged (e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
bin       e
y c
l c
r')
        Ordering
EQ -> c -> Updated c
forall a. a -> Updated a
Unchanged (e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
bin e
x c
l c
r)
{-# INLINABLE insert #-}

-- | @glue l r@ concatenates @l@ and @r@.
--
-- It assumes that @l@ and @r@ are already balanced with respect to each other.
glue :: IsBinTree c e => c -> c -> c
glue :: c -> c -> c
glue c
l c
r =
  case (c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin c
l, c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin c
r) of
    (TreeApp e c
TipTree, TreeApp e c
_) -> c
r
    (TreeApp e c
_, TreeApp e c
TipTree) -> c
l
    (BinTree e
x c
lx c
rx, BinTree e
y c
ly c
ry)
     | c -> Int
forall t e. IsBinTree t e => t -> Int
size c
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> c -> Int
forall t e. IsBinTree t e => t -> Int
size c
r ->
       case e -> c -> c -> PairS e c
forall c e. IsBinTree c e => e -> c -> c -> PairS e c
deleteFindMax e
x c
lx c
rx of
         PairS e
q c
l' -> e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
balanceR e
q c
l' c
r
     | Bool
otherwise ->
       case e -> c -> c -> PairS e c
forall c e. IsBinTree c e => e -> c -> c -> PairS e c
deleteFindMin e
y c
ly c
ry of
         PairS e
q c
r' -> e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
balanceL e
q c
l c
r'
{-# INLINABLE glue #-}

delete :: IsBinTree c e
       => (e -> Ordering)
          -- ^ Predicate that returns whether the entry is less than, greater than, or equal
          -- to the key we are entry that we are looking for.
       -> c
       -> MaybeS c
delete :: (e -> Ordering) -> c -> MaybeS c
delete e -> Ordering
k c
t =
  case c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin c
t of
    TreeApp e c
TipTree -> MaybeS c
forall v. MaybeS v
NothingS
    BinTree e
p c
l c
r ->
      case e -> Ordering
k e
p of
        Ordering
LT -> (\c
l' -> e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
balanceR e
p c
l' c
r) (c -> c) -> MaybeS c -> MaybeS c
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (e -> Ordering) -> c -> MaybeS c
forall c e. IsBinTree c e => (e -> Ordering) -> c -> MaybeS c
delete e -> Ordering
k c
l
        Ordering
GT -> (\c
r' -> e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
balanceL e
p c
l c
r') (c -> c) -> MaybeS c -> MaybeS c
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (e -> Ordering) -> c -> MaybeS c
forall c e. IsBinTree c e => (e -> Ordering) -> c -> MaybeS c
delete e -> Ordering
k c
r
        Ordering
EQ -> c -> MaybeS c
forall v. v -> MaybeS v
JustS (c -> c -> c
forall c e. IsBinTree c e => c -> c -> c
glue c
l c
r)
{-# INLINABLE delete #-}

------------------------------------------------------------------------
-- filter

-- | Returns only entries that are less than predicate with respect to the ordering
-- and Nothing if no elements are discarded.
filterGt :: IsBinTree c e => (e -> Ordering) -> c -> MaybeS c
filterGt :: (e -> Ordering) -> c -> MaybeS c
filterGt e -> Ordering
k c
t =
  case c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin c
t of
    TreeApp e c
TipTree -> MaybeS c
forall v. MaybeS v
NothingS
    BinTree e
x c
l c
r ->
      case e -> Ordering
k e
x of
        Ordering
LT -> (\c
l' -> e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
link e
x c
l' c
r) (c -> c) -> MaybeS c -> MaybeS c
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (e -> Ordering) -> c -> MaybeS c
forall c e. IsBinTree c e => (e -> Ordering) -> c -> MaybeS c
filterGt e -> Ordering
k c
l
        Ordering
GT -> (e -> Ordering) -> c -> MaybeS c
forall c e. IsBinTree c e => (e -> Ordering) -> c -> MaybeS c
filterGt e -> Ordering
k c
r MaybeS c -> MaybeS c -> MaybeS c
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> c -> MaybeS c
forall v. v -> MaybeS v
JustS c
r
        Ordering
EQ -> c -> MaybeS c
forall v. v -> MaybeS v
JustS c
r
{-# INLINABLE filterGt #-}


-- | @filterLt k m@ returns submap of @m@ that only contains entries
-- that are smaller than @k@.  If no entries are deleted then return Nothing.
filterLt :: IsBinTree c e => (e -> Ordering) -> c -> MaybeS c
filterLt :: (e -> Ordering) -> c -> MaybeS c
filterLt e -> Ordering
k c
t =
  case c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin c
t of
    TreeApp e c
TipTree -> MaybeS c
forall v. MaybeS v
NothingS
    BinTree e
x c
l c
r ->
      case e -> Ordering
k e
x of
        Ordering
LT -> (e -> Ordering) -> c -> MaybeS c
forall c e. IsBinTree c e => (e -> Ordering) -> c -> MaybeS c
filterLt e -> Ordering
k c
l MaybeS c -> MaybeS c -> MaybeS c
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> c -> MaybeS c
forall v. v -> MaybeS v
JustS c
l
        Ordering
GT -> (\c
r' -> e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
link e
x c
l c
r') (c -> c) -> MaybeS c -> MaybeS c
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (e -> Ordering) -> c -> MaybeS c
forall c e. IsBinTree c e => (e -> Ordering) -> c -> MaybeS c
filterLt e -> Ordering
k c
r
        Ordering
EQ -> c -> MaybeS c
forall v. v -> MaybeS v
JustS c
l
{-# INLINABLE filterLt #-}

------------------------------------------------------------------------
-- Union

-- | Insert a new key and value in the map if it is not already present.
-- Used by 'union'.
insertR :: forall c e . (IsBinTree c e) => (e -> e -> Ordering) -> e -> c -> c
insertR :: (e -> e -> Ordering) -> e -> c -> c
insertR e -> e -> Ordering
comp e
e c
m = c -> MaybeS c -> c
forall a. a -> MaybeS a -> a
fromMaybeS c
m (e -> c -> MaybeS c
go e
e c
m)
  where
    go :: e -> c -> MaybeS c
    go :: e -> c -> MaybeS c
go e
x c
t =
      case c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin c
t of
        TreeApp e c
TipTree -> c -> MaybeS c
forall v. v -> MaybeS v
JustS (e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
bin e
x c
forall t e. IsBinTree t e => t
tip c
forall t e. IsBinTree t e => t
tip)
        BinTree e
y c
l c
r ->
          case e -> e -> Ordering
comp e
x e
y of
            Ordering
LT -> (\c
l' -> e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
balanceL e
y c
l' c
r) (c -> c) -> MaybeS c -> MaybeS c
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> e -> c -> MaybeS c
go e
x c
l
            Ordering
GT -> (\c
r' -> e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
balanceR e
y c
l c
r') (c -> c) -> MaybeS c -> MaybeS c
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> e -> c -> MaybeS c
go e
x c
r
            Ordering
EQ -> MaybeS c
forall v. MaybeS v
NothingS
{-# INLINABLE insertR #-}

-- | Union two sets
union :: (IsBinTree c e) => (e -> e -> Ordering) -> c -> c -> c
union :: (e -> e -> Ordering) -> c -> c -> c
union e -> e -> Ordering
comp c
t1 c
t2 =
  case (c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin c
t1, c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin c
t2) of
    (TreeApp e c
TipTree, TreeApp e c
_) -> c
t2
    (TreeApp e c
_, TreeApp e c
TipTree) -> c
t1
    (TreeApp e c
_, BinTree e
p (c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin -> TreeApp e c
TipTree) (c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin -> TreeApp e c
TipTree)) -> (e -> e -> Ordering) -> e -> c -> c
forall c e. IsBinTree c e => (e -> e -> Ordering) -> e -> c -> c
insertR e -> e -> Ordering
comp e
p c
t1
    (BinTree e
x c
l c
r, TreeApp e c
_) ->
      e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
link e
x
           ((e -> e -> Ordering) -> e -> c -> c -> c
forall c e.
IsBinTree c e =>
(e -> e -> Ordering) -> e -> c -> c -> c
hedgeUnion_UB e -> e -> Ordering
comp e
x   c
l c
t2)
           ((e -> e -> Ordering) -> e -> c -> c -> c
forall c e.
IsBinTree c e =>
(e -> e -> Ordering) -> e -> c -> c -> c
hedgeUnion_LB e -> e -> Ordering
comp e
x c
r   c
t2)
{-# INLINABLE union #-}

-- | Hedge union where we only add elements in second map if key is
-- strictly above a lower bound.
hedgeUnion_LB :: (IsBinTree c e) => (e -> e -> Ordering) -> e -> c -> c -> c
hedgeUnion_LB :: (e -> e -> Ordering) -> e -> c -> c -> c
hedgeUnion_LB e -> e -> Ordering
comp e
lo c
t1 c
t2 =
  case (c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin c
t1, c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin c
t2) of
    (TreeApp e c
_, TreeApp e c
TipTree) -> c
t1
    (TreeApp e c
TipTree, TreeApp e c
_) -> c -> MaybeS c -> c
forall a. a -> MaybeS a -> a
fromMaybeS c
t2 ((e -> Ordering) -> c -> MaybeS c
forall c e. IsBinTree c e => (e -> Ordering) -> c -> MaybeS c
filterGt (e -> e -> Ordering
comp e
lo) c
t2)
    -- Prune left tree.
    (TreeApp e c
_, BinTree e
k c
_ c
r) | e -> e -> Ordering
comp e
k e
lo Ordering -> Ordering -> Bool
forall a. Ord a => a -> a -> Bool
<= Ordering
EQ -> (e -> e -> Ordering) -> e -> c -> c -> c
forall c e.
IsBinTree c e =>
(e -> e -> Ordering) -> e -> c -> c -> c
hedgeUnion_LB e -> e -> Ordering
comp e
lo c
t1 c
r
    -- Special case when t2 is a single element.
    (TreeApp e c
_, BinTree e
x (c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin -> TreeApp e c
TipTree) (c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin -> TreeApp e c
TipTree)) -> (e -> e -> Ordering) -> e -> c -> c
forall c e. IsBinTree c e => (e -> e -> Ordering) -> e -> c -> c
insertR e -> e -> Ordering
comp e
x c
t1
    -- Split on left-and-right subtrees of t1.
    (BinTree e
x c
l c
r, TreeApp e c
_) ->
      e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
link e
x
           ((e -> e -> Ordering) -> e -> e -> c -> c -> c
forall c e.
IsBinTree c e =>
(e -> e -> Ordering) -> e -> e -> c -> c -> c
hedgeUnion_LB_UB e -> e -> Ordering
comp e
lo e
x  c
l c
t2)
           ((e -> e -> Ordering) -> e -> c -> c -> c
forall c e.
IsBinTree c e =>
(e -> e -> Ordering) -> e -> c -> c -> c
hedgeUnion_LB    e -> e -> Ordering
comp e
x     c
r c
t2)
{-# INLINABLE hedgeUnion_LB #-}

-- | Hedge union where we only add elements in second map if key is
-- strictly below a upper bound.
hedgeUnion_UB :: (IsBinTree c e) => (e -> e -> Ordering) -> e -> c -> c -> c
hedgeUnion_UB :: (e -> e -> Ordering) -> e -> c -> c -> c
hedgeUnion_UB e -> e -> Ordering
comp e
hi c
t1 c
t2 =
  case (c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin c
t1, c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin c
t2) of
    (TreeApp e c
_, TreeApp e c
TipTree) -> c
t1
    (TreeApp e c
TipTree, TreeApp e c
_) -> c -> MaybeS c -> c
forall a. a -> MaybeS a -> a
fromMaybeS c
t2 ((e -> Ordering) -> c -> MaybeS c
forall c e. IsBinTree c e => (e -> Ordering) -> c -> MaybeS c
filterLt (e -> e -> Ordering
comp e
hi) c
t2)
    -- Prune right tree.
    (TreeApp e c
_, BinTree e
x c
l c
_) | e -> e -> Ordering
comp e
x e
hi Ordering -> Ordering -> Bool
forall a. Ord a => a -> a -> Bool
>= Ordering
EQ -> (e -> e -> Ordering) -> e -> c -> c -> c
forall c e.
IsBinTree c e =>
(e -> e -> Ordering) -> e -> c -> c -> c
hedgeUnion_UB e -> e -> Ordering
comp e
hi c
t1 c
l
    -- Special case when t2 is a single element.
    (TreeApp e c
_, BinTree e
x (c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin -> TreeApp e c
TipTree) (c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin -> TreeApp e c
TipTree))  -> (e -> e -> Ordering) -> e -> c -> c
forall c e. IsBinTree c e => (e -> e -> Ordering) -> e -> c -> c
insertR e -> e -> Ordering
comp e
x c
t1
    -- Split on left-and-right subtrees of t1.
    (BinTree e
x c
l c
r, TreeApp e c
_) ->
      e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
link e
x
           ((e -> e -> Ordering) -> e -> c -> c -> c
forall c e.
IsBinTree c e =>
(e -> e -> Ordering) -> e -> c -> c -> c
hedgeUnion_UB    e -> e -> Ordering
comp e
x      c
l c
t2)
           ((e -> e -> Ordering) -> e -> e -> c -> c -> c
forall c e.
IsBinTree c e =>
(e -> e -> Ordering) -> e -> e -> c -> c -> c
hedgeUnion_LB_UB e -> e -> Ordering
comp e
x  e
hi  c
r c
t2)
{-# INLINABLE hedgeUnion_UB #-}

-- | Hedge union where we only add elements in second map if key is
-- strictly between a lower and upper bound.
hedgeUnion_LB_UB :: (IsBinTree c e) => (e -> e -> Ordering) -> e -> e -> c -> c -> c
hedgeUnion_LB_UB :: (e -> e -> Ordering) -> e -> e -> c -> c -> c
hedgeUnion_LB_UB e -> e -> Ordering
comp e
lo e
hi c
t1 c
t2 =
  case (c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin c
t1, c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin c
t2) of
    (TreeApp e c
_, TreeApp e c
TipTree) -> c
t1
    -- Prune left tree.
    (TreeApp e c
_,   BinTree e
k c
_ c
r) | e -> e -> Ordering
comp e
k e
lo Ordering -> Ordering -> Bool
forall a. Ord a => a -> a -> Bool
<= Ordering
EQ -> (e -> e -> Ordering) -> e -> e -> c -> c -> c
forall c e.
IsBinTree c e =>
(e -> e -> Ordering) -> e -> e -> c -> c -> c
hedgeUnion_LB_UB e -> e -> Ordering
comp e
lo e
hi c
t1 c
r
    -- Prune right tree.
    (TreeApp e c
_,   BinTree e
k c
l c
_) | e -> e -> Ordering
comp e
k e
hi Ordering -> Ordering -> Bool
forall a. Ord a => a -> a -> Bool
>= Ordering
EQ -> (e -> e -> Ordering) -> e -> e -> c -> c -> c
forall c e.
IsBinTree c e =>
(e -> e -> Ordering) -> e -> e -> c -> c -> c
hedgeUnion_LB_UB e -> e -> Ordering
comp e
lo e
hi c
t1 c
l
    -- When t1 becomes empty (assumes lo <= k <= hi)
    (TreeApp e c
TipTree, BinTree e
x c
l c
r) ->
      case ((e -> Ordering) -> c -> MaybeS c
forall c e. IsBinTree c e => (e -> Ordering) -> c -> MaybeS c
filterGt (e -> e -> Ordering
comp e
lo) c
l, (e -> Ordering) -> c -> MaybeS c
forall c e. IsBinTree c e => (e -> Ordering) -> c -> MaybeS c
filterLt (e -> e -> Ordering
comp e
hi) c
r) of
        -- No variables in t2 were eliminated.
        (MaybeS c
NothingS, MaybeS c
NothingS) -> c
t2
        -- Relink t2 with filtered elements removed.
        (MaybeS c
l',MaybeS c
r') -> e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
link e
x (c -> MaybeS c -> c
forall a. a -> MaybeS a -> a
fromMaybeS c
l MaybeS c
l') (c -> MaybeS c -> c
forall a. a -> MaybeS a -> a
fromMaybeS c
r MaybeS c
r')
    -- Special case when t2 is a single element.
    (TreeApp e c
_, BinTree e
x (c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin -> TreeApp e c
TipTree) (c -> TreeApp e c
forall t e. IsBinTree t e => t -> TreeApp e t
asBin -> TreeApp e c
TipTree)) -> (e -> e -> Ordering) -> e -> c -> c
forall c e. IsBinTree c e => (e -> e -> Ordering) -> e -> c -> c
insertR e -> e -> Ordering
comp e
x c
t1
    -- Split on left-and-right subtrees of t1.
    (BinTree e
x c
l c
r, TreeApp e c
_) ->
      e -> c -> c -> c
forall t e. IsBinTree t e => e -> t -> t -> t
link e
x
           ((e -> e -> Ordering) -> e -> e -> c -> c -> c
forall c e.
IsBinTree c e =>
(e -> e -> Ordering) -> e -> e -> c -> c -> c
hedgeUnion_LB_UB e -> e -> Ordering
comp e
lo e
x  c
l c
t2)
           ((e -> e -> Ordering) -> e -> e -> c -> c -> c
forall c e.
IsBinTree c e =>
(e -> e -> Ordering) -> e -> e -> c -> c -> c
hedgeUnion_LB_UB e -> e -> Ordering
comp e
x  e
hi c
r c
t2)
{-# INLINABLE hedgeUnion_LB_UB #-}