{-# 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
data MaybeS v
= JustS !v
| NothingS
instance Functor MaybeS where
fmap _ NothingS = NothingS
fmap f (JustS v) = JustS (f v)
instance Alternative MaybeS where
empty = NothingS
mv@JustS{} <|> _ = mv
NothingS <|> v = v
instance Applicative MaybeS where
pure = JustS
NothingS <*> _ = NothingS
JustS{} <*> NothingS = NothingS
JustS f <*> JustS x = JustS (f x)
fromMaybeS :: a -> MaybeS a -> a
fromMaybeS r NothingS = r
fromMaybeS _ (JustS v) = v
data Updated a
= Updated !a
| Unchanged !a
updatedValue :: Updated a -> a
updatedValue (Updated a) = a
updatedValue (Unchanged a) = a
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 = 3
ratio = 2
balanceL :: (IsBinTree c e) => e -> c -> c -> c
balanceL p l r = do
case asBin l of
BinTree l_pair ll lr | size l > max 1 (delta*size r) ->
case asBin lr of
BinTree lr_pair lrl lrr | size lr >= max 2 (ratio*size ll) ->
bin lr_pair (bin l_pair ll lrl) (bin p lrr r)
_ -> bin l_pair ll (bin p lr r)
_ -> bin p l r
{-# INLINE balanceL #-}
balanceR :: (IsBinTree c e) => e -> c -> c -> c
balanceR p l r = do
case asBin r of
BinTree r_pair rl rr | size r > max 1 (delta*size l) ->
case asBin rl of
BinTree rl_pair rll rlr | size rl >= max 2 (ratio*size rr) ->
(bin rl_pair $! bin p l rll) $! bin r_pair rlr rr
_ -> bin r_pair (bin p l rl) rr
_ -> bin p l r
{-# INLINE balanceR #-}
insertMax :: IsBinTree c e => e -> c -> c
insertMax p t =
case asBin t of
TipTree -> bin p tip tip
BinTree q l r -> balanceR q l (insertMax p r)
insertMin :: IsBinTree c e => e -> c -> c
insertMin p t =
case asBin t of
TipTree -> bin p tip tip
BinTree q l r -> balanceL q (insertMin p l) r
link :: IsBinTree c e => e -> c -> c -> c
link p l r =
case (asBin l, asBin r) of
(TipTree, _) -> insertMin p r
(_, TipTree) -> insertMax p l
(BinTree py ly ry, BinTree pz lz rz)
| delta*size l < size r -> balanceL pz (link p l lz) rz
| delta*size r < size l -> balanceR py ly (link p ry r)
| otherwise -> bin p l r
{-# INLINE link #-}
data PairS f s = PairS !f !s
deleteFindMin :: IsBinTree c e => e -> c -> c -> PairS e c
deleteFindMin p l r =
case asBin l of
TipTree -> PairS p r
BinTree lp ll lr ->
case deleteFindMin lp ll lr of
PairS q l' -> PairS q (balanceR p l' r)
{-# INLINABLE deleteFindMin #-}
deleteFindMax :: IsBinTree c e => e -> c -> c -> PairS e c
deleteFindMax p l r =
case asBin r of
TipTree -> PairS p l
BinTree rp rl rr ->
case deleteFindMax rp rl rr of
PairS q r' -> PairS q (balanceL p l r')
{-# INLINABLE deleteFindMax #-}
merge :: IsBinTree c e => c -> c -> c
merge l r =
case (asBin l, asBin r) of
(TipTree, _) -> r
(_, TipTree) -> l
(BinTree x lx rx, BinTree y ly ry)
| delta*size l < size r -> balanceL y (merge l ly) ry
| delta*size r < size l -> balanceR x lx (merge rx r)
| size l > size r ->
case deleteFindMax x lx rx of
PairS q l' -> balanceR q l' r
| otherwise ->
case deleteFindMin y ly ry of
PairS q r' -> balanceL q l r'
{-# INLINABLE merge #-}
insert :: (IsBinTree c e) => (e -> e -> Ordering) -> e -> c -> Updated c
insert comp x t =
case asBin t of
TipTree -> Updated (bin x tip tip)
BinTree y l r ->
case comp x y of
LT ->
case insert comp x l of
Updated l' -> Updated (balanceL y l' r)
Unchanged l' -> Unchanged (bin y l' r)
GT ->
case insert comp x r of
Updated r' -> Updated (balanceR y l r')
Unchanged r' -> Unchanged (bin y l r')
EQ -> Unchanged (bin x l r)
{-# INLINABLE insert #-}
glue :: IsBinTree c e => c -> c -> c
glue l r =
case (asBin l, asBin r) of
(TipTree, _) -> r
(_, TipTree) -> l
(BinTree x lx rx, BinTree y ly ry)
| size l > size r ->
case deleteFindMax x lx rx of
PairS q l' -> balanceR q l' r
| otherwise ->
case deleteFindMin y ly ry of
PairS q r' -> balanceL q l r'
{-# INLINABLE glue #-}
delete :: IsBinTree c e
=> (e -> Ordering)
-> c
-> MaybeS c
delete k t =
case asBin t of
TipTree -> NothingS
BinTree p l r ->
case k p of
LT -> (\l' -> balanceR p l' r) <$> delete k l
GT -> (\r' -> balanceL p l r') <$> delete k r
EQ -> JustS (glue l r)
{-# INLINABLE delete #-}
filterGt :: IsBinTree c e => (e -> Ordering) -> c -> MaybeS c
filterGt k t =
case asBin t of
TipTree -> NothingS
BinTree x l r ->
case k x of
LT -> (\l' -> link x l' r) <$> filterGt k l
GT -> filterGt k r <|> JustS r
EQ -> JustS r
{-# INLINABLE filterGt #-}
filterLt :: IsBinTree c e => (e -> Ordering) -> c -> MaybeS c
filterLt k t =
case asBin t of
TipTree -> NothingS
BinTree x l r ->
case k x of
LT -> filterLt k l <|> JustS l
GT -> (\r' -> link x l r') <$> filterLt k r
EQ -> JustS l
{-# INLINABLE filterLt #-}
insertR :: forall c e . (IsBinTree c e) => (e -> e -> Ordering) -> e -> c -> c
insertR comp e m = fromMaybeS m (go e m)
where
go :: e -> c -> MaybeS c
go x t =
case asBin t of
TipTree -> JustS (bin x tip tip)
BinTree y l r ->
case comp x y of
LT -> (\l' -> balanceL y l' r) <$> go x l
GT -> (\r' -> balanceR y l r') <$> go x r
EQ -> NothingS
{-# INLINABLE insertR #-}
union :: (IsBinTree c e) => (e -> e -> Ordering) -> c -> c -> c
union comp t1 t2 =
case (asBin t1, asBin t2) of
(TipTree, _) -> t2
(_, TipTree) -> t1
(_, BinTree p (asBin -> TipTree) (asBin -> TipTree)) -> insertR comp p t1
(BinTree x l r, _) ->
link x
(hedgeUnion_UB comp x l t2)
(hedgeUnion_LB comp x r t2)
{-# INLINABLE union #-}
hedgeUnion_LB :: (IsBinTree c e) => (e -> e -> Ordering) -> e -> c -> c -> c
hedgeUnion_LB comp lo t1 t2 =
case (asBin t1, asBin t2) of
(_, TipTree) -> t1
(TipTree, _) -> fromMaybeS t2 (filterGt (comp lo) t2)
(_, BinTree k _ r) | comp k lo <= EQ -> hedgeUnion_LB comp lo t1 r
(_, BinTree x (asBin -> TipTree) (asBin -> TipTree)) -> insertR comp x t1
(BinTree x l r, _) ->
link x
(hedgeUnion_LB_UB comp lo x l t2)
(hedgeUnion_LB comp x r t2)
{-# INLINABLE hedgeUnion_LB #-}
hedgeUnion_UB :: (IsBinTree c e) => (e -> e -> Ordering) -> e -> c -> c -> c
hedgeUnion_UB comp hi t1 t2 =
case (asBin t1, asBin t2) of
(_, TipTree) -> t1
(TipTree, _) -> fromMaybeS t2 (filterLt (comp hi) t2)
(_, BinTree x l _) | comp x hi >= EQ -> hedgeUnion_UB comp hi t1 l
(_, BinTree x (asBin -> TipTree) (asBin -> TipTree)) -> insertR comp x t1
(BinTree x l r, _) ->
link x
(hedgeUnion_UB comp x l t2)
(hedgeUnion_LB_UB comp x hi r t2)
{-# INLINABLE hedgeUnion_UB #-}
hedgeUnion_LB_UB :: (IsBinTree c e) => (e -> e -> Ordering) -> e -> e -> c -> c -> c
hedgeUnion_LB_UB comp lo hi t1 t2 =
case (asBin t1, asBin t2) of
(_, TipTree) -> t1
(_, BinTree k _ r) | comp k lo <= EQ -> hedgeUnion_LB_UB comp lo hi t1 r
(_, BinTree k l _) | comp k hi >= EQ -> hedgeUnion_LB_UB comp lo hi t1 l
(TipTree, BinTree x l r) ->
case (filterGt (comp lo) l, filterLt (comp hi) r) of
(NothingS, NothingS) -> t2
(l',r') -> link x (fromMaybeS l l') (fromMaybeS r r')
(_, BinTree x (asBin -> TipTree) (asBin -> TipTree)) -> insertR comp x t1
(BinTree x l r, _) ->
link x
(hedgeUnion_LB_UB comp lo x l t2)
(hedgeUnion_LB_UB comp x hi r t2)
{-# INLINABLE hedgeUnion_LB_UB #-}