inch-0.1.0: examples/RedBlackCost.hs
{-# OPTIONS_GHC -F -pgmF inch #-}
{-# LANGUAGE RankNTypes, GADTs, KindSignatures, ScopedTypeVariables,
NPlusKPatterns #-}
{-
An implementation of red-black tree insertion and deletion that uses
Nils Anders Danielsson's technique for semiformal time complexity
analysis to show that these operations are linear in black
height. See the RedBlack module for an implementation of the tree
operations without time complexity annotations, and the Cost module
for the definitions of the library primitives used in the analysis.
-}
module RedBlackCost where
import Cost
type Black = 0
type Red = 1
data Colour :: Integer -> * where
Black :: Colour Black
Red :: Colour Red
deriving Show
data Tree :: Integer -> Integer -> Integer -> Nat -> * where
E :: forall (lo hi :: Integer) . lo < hi => Tree lo hi Black 0
TR :: forall (lo hi :: Integer)(n :: Nat) . pi (x :: Integer) .
Tree lo x Black n -> Tree x hi Black n -> Tree lo hi Red n
TB :: forall (lo hi cl cr :: Integer)(n :: Nat) . pi (x :: Integer) .
Tree lo x cl n -> Tree x hi cr n -> Tree lo hi Black (n+1)
deriving Show
data RBT :: Integer -> Integer -> * where
RBT :: forall (lo hi :: Integer)(n :: Nat) . Tree lo hi Black n -> RBT lo hi
deriving Show
empty = RBT E
data TreeZip :: Integer -> Integer -> Integer -> Nat ->
Integer -> Integer -> Integer -> Nat ->
Nat -> * where
Root :: forall (lo hi c :: Integer)(n :: Nat) . TreeZip lo hi c n lo hi c n 0
ZRL :: forall (rlo rhi lo hi rc :: Integer)(rn n d :: Nat) . pi (x :: Integer) .
TreeZip rlo rhi rc rn lo hi Red n d -> Tree x hi Black n ->
TreeZip rlo rhi rc rn lo x Black n (d + 1)
ZRR :: forall (rlo rhi lo hi rc :: Integer)(rn n d :: Nat) . pi (x :: Integer) .
Tree lo x Black n -> TreeZip rlo rhi rc rn lo hi Red n d ->
TreeZip rlo rhi rc rn x hi Black n (d + 1)
ZBL :: forall (rlo rhi lo hi rc c hc :: Integer)(rn n d :: Nat) . pi (x :: Integer) .
TreeZip rlo rhi rc rn lo hi Black (n+1) d -> Tree x hi c n ->
TreeZip rlo rhi rc rn lo x hc n (d + 1)
ZBR :: forall (rlo rhi lo hi rc c hc :: Integer)(rn n d :: Nat) . pi (x :: Integer) .
Tree lo x c n -> TreeZip rlo rhi rc rn lo hi Black (n+1) d ->
TreeZip rlo rhi rc rn x hi hc n (d + 1)
deriving Show
plug :: forall (rlo rhi lo hi rc rn c n d :: Integer) . Tree lo hi c n ->
TreeZip rlo rhi rc rn lo hi c n d -> Cost (d + 1) (Tree rlo rhi rc rn)
plug t Root = tick (returnCost t)
plug t (ZRL {x} z r) = tick (plug (TR {x} t r) z)
plug t (ZRR {x} l z) = tick (plug (TR {x} l t) z)
plug t (ZBL {x} z r) = tick (plug (TB {x} t r) z)
plug t (ZBR {x} l z) = tick (plug (TB {x} l t) z)
plugBR :: forall (rlo rhi lo hi n rn d :: Integer) . Tree lo hi Black n ->
TreeZip rlo rhi Black rn lo hi Red n d -> Cost (d + 1) (Tree rlo rhi Black rn)
plugBR t (ZBL {x} z r) = plug t (ZBL {x} z r)
plugBR t (ZBR {x} l z) = plug t (ZBR {x} l z)
data SearchResult :: Integer -> Integer -> Integer -> Integer -> * where
Nope :: forall (x rlo rhi lo hi :: Integer)(rn d :: Nat) .
(d <= (2 * rn), lo < x, x < hi) =>
TreeZip rlo rhi Black rn lo hi Black 0 d -> SearchResult x rlo rhi rn
Yep :: forall (x rlo rhi lo hi c :: Integer)(rn n d :: Nat) .
((2 * n) + d) <= (2 * rn) =>
TreeZip rlo rhi Black rn lo hi c n d -> Tree lo hi c n ->
SearchResult x rlo rhi rn
search :: forall (rlo rhi :: Integer)(rn :: Nat) .
pi (x :: Integer) . (rlo < x, x < rhi) =>
Tree rlo rhi Black rn -> Cost (2 * rn + 1) (SearchResult x rlo rhi rn)
search {x} = helpB Root
where
help :: forall (lo hi c :: Integer)(n d :: Nat) .
((1 + (2 * n) + d) <= (2 * rn), lo < x, x < hi) =>
TreeZip rlo rhi Black rn lo hi c n d -> Tree lo hi c n ->
Cost (2 + 2 * n) (SearchResult x rlo rhi rn)
help z E = tick (returnW (Nope z))
help z (TR {y} l r) | {x < y} = tick (helpB (ZRL {y} z r) l)
help z (TR {y} l r) | {x ~ y} = tick (returnW (Yep z (TR {y} l r)))
help z (TR {y} l r) | {x > y} = tick (helpB (ZRR {y} l z) r)
help z (TB {y} l r) | {x < y} = tick (weakenBy {1} (help (ZBL {y} z r) l))
help z (TB {y} l r) | {x ~ y} = tick (returnW (Yep z (TB {y} l r)))
help z (TB {y} l r) | {x > y} = tick (weakenBy {1} (help (ZBR {y} l z) r))
helpB :: forall (lo hi :: Integer)(n d :: Nat) .
(((2 * n) + d) <= (2 * rn), lo < x, x < hi) =>
TreeZip rlo rhi Black rn lo hi Black n d -> Tree lo hi Black n ->
Cost (2 * n + 1) (SearchResult x rlo rhi rn)
helpB z E = tick (returnW (Nope z))
helpB z (TB {y} l r) | {x < y} = tick (help (ZBL {y} z r) l)
helpB z (TB {y} l r) | {x ~ y} = tick (returnW (Yep z (TB {y} l r)))
helpB z (TB {y} l r) | {x > y} = tick (help (ZBR {y} l z) r)
member :: forall (lo hi :: Integer)(n :: Nat) .
pi (x :: Integer) . (lo < x, x < hi) =>
Tree lo hi Black n -> Cost (2 * n + 3) Bool
member {x} t = tick (bindCost (search {x} t) f)
where
f :: SearchResult x lo hi n -> Cost 1 Bool
f (Nope _) = tick (returnCost False)
f (Yep _ _) = tick (returnCost True)
data InsProb :: Integer -> Integer -> Integer -> Integer -> * where
Level :: forall (lo hi c ci :: Integer)( n :: Nat) .
Colour ci -> Tree lo hi ci n -> InsProb lo hi c n
PanicRB :: forall (lo hi :: Integer)(n :: Nat) . pi (x :: Integer) .
Tree lo x Red n -> Tree x hi Black n -> InsProb lo hi Red n
PanicBR :: forall (lo hi :: Integer)(n :: Nat) . pi (x :: Integer) .
Tree lo x Black n -> Tree x hi Red n -> InsProb lo hi Red n
deriving Show
solveIns :: forall (rlo rhi lo hi c rc :: Integer)(rn n d :: Nat) .
InsProb lo hi c n -> TreeZip rlo rhi rc rn lo hi c n d ->
Cost (d + 1) (RBT rlo rhi)
solveIns (Level c t) Root = tick (returnCost (rbt c t))
solveIns (Level Red t) (ZRL {x} z r) = tick (solveIns (PanicRB {x} t r) z)
solveIns (Level Red t) (ZRR {x} l z) = tick (solveIns (PanicBR {x} l t) z)
solveIns (Level Black t) (ZRL {x} z r) = tick (solveIns (Level Red (TR {x} t r)) z)
solveIns (Level Black t) (ZRR {x} l z) = tick (solveIns (Level Red (TR {x} l t)) z)
solveIns (Level col t) (ZBL {x} z r) = tick (solveIns (Level Black (TB {x} t r)) z)
solveIns (Level col t) (ZBR {x} l z) = tick (solveIns (Level Black (TB {x} l t)) z)
solveIns (PanicRB {xi} (TR {xil} lil ril) ri) (ZBL {x} z r) =
tick (solveIns (Level Red (TR {xi} (TB {xil} lil ril) (TB {x} ri r))) z)
solveIns (PanicBR {xi} li (TR {xir} lir rir)) (ZBL {x} z r) =
tick (solveIns (Level Red (TR {xir} (TB {xi} li lir) (TB {x} rir r))) z)
solveIns (PanicRB {xi} (TR {xil} lil ril) ri) (ZBR {x} l z) =
tick (solveIns (Level Red (TR {xil} (TB {x} l lil) (TB {xi} ril ri))) z)
solveIns (PanicBR {xi} li (TR {xir} lir rir)) (ZBR {x} l z) =
tick (solveIns (Level Red (TR {xi} (TB {x} l li) (TB {xir} lir rir))) z)
insert :: forall (lo hi :: Integer)(n :: Nat) . pi (x :: Integer) . (lo < x, x < hi) =>
Tree lo hi Black n -> Cost (4 * n + 6) (RBT lo hi)
insert {x} t = tick (bindCost (search {x} t) f)
where
f :: SearchResult x lo hi n -> Cost (2 * n + 4) (RBT lo hi)
f (Nope z) = tick (weaken (solveIns (Level Red (TR {x} E E)) z))
f (Yep _ _) = tick (returnW (RBT t))
r2b :: forall (lo hi n :: Integer) . Tree lo hi Red n -> Tree lo hi Black (n+1)
r2b (TR {x} l r) = TB {x} l r
rbt :: forall (lo hi c :: Integer)(n :: Nat) . Colour c -> Tree lo hi c n -> RBT lo hi
rbt Black t = RBT t
rbt Red t = RBT (r2b t)
solveDel :: forall (rlo rhi lo hi :: Integer)(rn n d :: Nat) . Tree lo hi Black n ->
TreeZip rlo rhi Black rn lo hi Black (n+1) d ->
Cost (d + 1) (RBT rlo rhi)
solveDel t Root = tick (returnW (RBT t))
solveDel t (ZRL {x} z (TB {y} (TR {lx} ll lr) r)) = tick (mapCost RBT (plug (TR {lx} (TB {x} t ll) (TB {y} lr r)) z))
solveDel t (ZRL {x} z (TB {y} l (TR {rx} rl rr))) = tick (mapCost RBT (plug (TR {y} (TB {x} t l) (TB {rx} rl rr)) z))
-- Arrgh: these are one line in Agda because we can pattern match on the colours being black
solveDel t (ZRL {x} z (TB {y} E E)) = tick (mapCost RBT (plugBR (TB {x} t (TR {y} E E)) z))
solveDel t (ZRL {x} z (TB {y} (TB {lx} ll lr) (TB {rx} rl rr))) = tick (mapCost RBT (plugBR (TB {x} t (TR {y} (TB {lx} ll lr) (TB {rx} rl rr))) z))
solveDel t (ZRR {x} (TB {y} (TR {lx} ll lr) r) z) = tick (mapCost RBT (plug (TR {y} (TB {lx} ll lr) (TB {x} r t)) z))
solveDel t (ZRR {x} (TB {y} l (TR {rx} rl rr)) z) = tick (mapCost RBT (plug (TR {rx} (TB {y} l rl) (TB {x} rr t)) z))
-- Arrgh
solveDel t (ZRR {x} (TB {y} E E) z) = tick (mapCost RBT (plugBR (TB {y} E (TR {x} E t)) z))
solveDel t (ZRR {x} (TB {y} (TB {lx} ll lr) (TB {rx} rl rr)) z) = tick (mapCost RBT (plugBR (TB {y} (TB {lx} ll lr) (TR {x} (TB {rx} rl rr) t)) z))
-- Arrgh
solveDel t (ZBL {x} z (TR {y} (TB {lx} E lr) r)) = tick (mapCost RBT (plug (TB {y} (TB {lx} (TR {x} t E) lr) r) z))
solveDel t (ZBL {x} z (TR {y} (TB {lx} (TB {llx} lll llr) lr) r)) = tick (mapCost RBT (plug (TB {y} (TB {lx} (TR {x} t (TB {llx} lll llr)) lr) r) z))
solveDel t (ZBL {x} z (TR {y} (TB {lx} (TR {llx} lll llr) lr) r)) = tick (mapCost RBT (plug (TB {llx} (TB {x} t lll) (TR {y} (TB {lx} llr lr) r)) z))
-- Arrgh
solveDel t (ZBL {x} z (TB {y} E r)) = tick (solveDel (TB {y} (TR {x} t E) r) z)
solveDel t (ZBL {x} z (TB {y} (TB {lx} ll lr) r)) = tick (solveDel (TB {y} (TR {x} t (TB {lx} ll lr)) r) z)
-- Arrgh
solveDel t (ZBL {x} z (TB {y} (TR {lx} ll lr) E)) = tick (solveDel (TB {lx} (TR {x} t ll) (TR {y} lr E)) z)
solveDel t (ZBL {x} z (TB {y} (TR {lx} ll lr) (TB {rx} rl rr))) = tick (solveDel (TB {lx} (TR {x} t ll) (TR {y} lr (TB {rx} rl rr))) z)
solveDel t (ZBL {x} z (TB {y} (TR {lx} ll lr) (TR {rx} rl rr))) = tick (mapCost RBT (plug (TB {lx} (TB {x} t ll) (TB {y} lr (TR {rx} rl rr))) z))
-- Arrgh
solveDel t (ZBR {x} (TR {y} l (TB {rx} rl E)) z) = tick (mapCost RBT (plug (TB {y} l (TB {rx} rl (TR {x} E t))) z))
solveDel t (ZBR {x} (TR {y} l (TB {rx} rl (TB {rrx} rrl rrr))) z) = tick (mapCost RBT (plug (TB {y} l (TB {rx} rl (TR {x} (TB {rrx} rrl rrr) t))) z))
solveDel t (ZBR {x} (TR {y} l (TB {rx} rl (TR {rrx} rrl rrr))) z) = tick (mapCost RBT (plug (TB {rrx} (TR {y} l (TB {rx} rl rrl)) (TB {x} rrr t)) z))
-- Arrgh
solveDel t (ZBR {x} (TB {y} l E) z) = tick (solveDel (TB {y} l (TR {x} E t)) z)
solveDel t (ZBR {x} (TB {y} l (TB {lx} ll lr)) z) = tick (solveDel (TB {y} l (TR {x} (TB {lx} ll lr) t)) z)
-- Arrgh
solveDel t (ZBR {x} (TB {y} E (TR {rx} rl rr)) z) = tick (solveDel (TB {rx} (TR {y} E rl) (TR {x} rr t)) z)
solveDel t (ZBR {x} (TB {y} (TB {lx} ll lr) (TR {rx} rl rr)) z) = tick (solveDel (TB {rx} (TR {y} (TB {lx} ll lr) rl) (TR {x} rr t)) z)
solveDel t (ZBR {x} (TB {y} (TR {lx} ll lr) (TR {rx} rl rr)) z) = tick (mapCost RBT (plug (TB {y} (TB {lx} ll lr) (TB {rx} rl (TR {x} rr t))) z))
findMin :: forall (rlo rhi lo hi c :: Integer)(rn n d :: Nat) . Tree lo hi c (n+1) ->
(pi (k :: Integer) . lo < k => TreeZip rlo rhi Black rn k hi c (n+1) d) ->
Cost (3 * n + d + 4) (RBT rlo rhi)
findMin (TR {x} (TB {y} E E) r) f = tick (weaken (solveDel E (ZRL {x} (f {y}) r)))
findMin (TR {x} (TB {y} E (TR {lx} ll lr)) r) f = tick (weaken (mapCost RBT (plug (TB {lx} ll lr) (ZRL {x} (f {y}) r))))
findMin (TR {x} (TB {y} (TR {k} E E) lr) r) f = tick (weaken (mapCost RBT (plug E (ZBL {y} (ZRL {x} (f {k}) r) lr))))
findMin (TB {x} (TR {y} E E) r) f = tick (weaken (mapCost RBT (plug E (ZBL {x} (f {y}) r))))
findMin (TB {x} E (TR {lx} ll lr)) f = tick (weaken (mapCost RBT (plug (TB {lx} ll lr) (f {x}))))
findMin (TB {x} E E) f = tick (weaken (solveDel E (f {x})))
findMin (TR {x} (TB {y} (TB {llx} lll llr) lr) r) f = tick (findMin (TB {llx} lll llr) (\ {k} -> ZBL {y} (ZRL {x} (f {k}) r) lr))
findMin (TB {x} (TB {lx} ll lr) r) f = tick (weakenBy {1} (findMin (TB {lx} ll lr) (\ {k} -> ZBL {x} (f {k}) r)))
wkTree :: forall (lo hi ha c n :: Integer) . hi < ha => Tree lo hi c n -> Tree lo ha c n
wkTree E = E
wkTree (TR {x} l r) = TR {x} l (wkTree r)
wkTree (TB {x} l r) = TB {x} l (wkTree r)
delFocus :: forall (rlo rhi lo hi c :: Integer)(rn n d :: Nat) . Tree lo hi c n ->
TreeZip rlo rhi Black rn lo hi c n d ->
Cost (3 * n + d + 3) (RBT rlo rhi)
delFocus (TR {x} E E) z = tick (weakenBy {1} (mapCost RBT (plugBR E z)))
delFocus (TR {x} l (TB {rx} rl rr)) z = tick (findMin (TB {rx} rl rr) (\ {k} -> ZRR {k} (wkTree l) z))
delFocus E z = tick (weaken (mapCost RBT (plug E z)))
delFocus (TB {x} E E) z = tick (weaken (solveDel E z))
delFocus (TB {x} (TR {y} E E) E) z = tick (weaken (mapCost RBT (plug (TB {y} E E) z)))
delFocus (TB {x} E (TR {y} E E)) z = tick (weaken (mapCost RBT (plug (TB {y} E E) z)))
delFocus (TB {x} (TR {k} E E) (TR {y} E E)) z = tick (weaken (mapCost RBT (plug (TB {k} E (TR {y} E E)) z)))
delFocus (TB {x} l (TB {rx} rl rr)) z = tick (weakenBy {3} (findMin (TB {rx} rl rr) (\ {k} -> ZBR {k} (wkTree l) z)))
delFocus (TB {x} (TB {lx} ll lr) r) z = tick (weakenBy {3} (findMin r (\ {k} -> ZBR {k} (wkTree (TB {lx} ll lr)) z)))
delete :: forall (lo hi :: Integer)(n :: Nat) . pi (x :: Integer) . (lo < x, x < hi) =>
Tree lo hi Black n -> Cost (5 * n + 6) (RBT lo hi)
delete {x} t = tick (bindCost (search {x} t) f)
where
f :: SearchResult x lo hi n -> Cost (3 * n + 4) (RBT lo hi)
f (Nope _) = tick (returnW (RBT t))
f (Yep z t) = tick (weaken (delFocus t z))