inch-0.1.0: examples/RedBlack.hs

{-# OPTIONS_GHC -F -pgmF inch #-}
{-# LANGUAGE RankNTypes, GADTs, KindSignatures, ScopedTypeVariables,
             NPlusKPatterns #-}

{-
  An implementation of red-black tree insertion and deletion using an
  indexed zipper. The type indices guarantee that the ordering, colour
  and height invariants are preserved. 
-}

module RedBlack where

-- We can't (yet) lift types to kinds automatically, but we can
-- represent finite enumerations using numbers. Here we use 0 for
-- black and 1 for red, and use a singleton type to fake pi-types for
-- colours. Proper lifting of algebraic data types to kinds would be
-- better.

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 -> * where
    Root  :: forall (lo hi c :: Integer)(n :: Nat) . TreeZip lo hi c n lo hi c n
    ZRL   :: forall (rlo rhi lo hi rc :: Integer)(rn n :: Nat) . pi (x :: Integer) .
                 TreeZip rlo rhi rc rn lo hi Red n -> Tree x hi Black n ->
                     TreeZip rlo rhi rc rn lo x Black n
    ZRR   :: forall (rlo rhi lo hi rc :: Integer)(rn n :: Nat) . pi (x :: Integer) .
                 Tree lo x Black n -> TreeZip rlo rhi rc rn lo hi Red n  ->
                     TreeZip rlo rhi rc rn x hi Black n
    ZBL   :: forall (rlo rhi lo hi rc c hc :: Integer)(rn n  :: Nat) . pi (x :: Integer) . 
                 TreeZip rlo rhi rc rn lo hi Black (n+1)  -> Tree x hi c n ->
                     TreeZip rlo rhi rc rn lo x hc n
    ZBR   :: forall (rlo rhi lo hi rc c hc :: Integer)(rn n  :: Nat) . pi (x :: Integer) .
                 Tree lo x c n -> TreeZip rlo rhi rc rn lo hi Black (n+1) ->
                     TreeZip rlo rhi rc rn x hi hc n
  deriving Show

plug ::  forall (rlo rhi lo hi rc rn c n :: Integer) . Tree lo hi c n ->
             TreeZip rlo rhi rc rn lo hi c n -> Tree rlo rhi rc rn
plug t Root           = t
plug t (ZRL {x} z r)  = plug (TR {x} t r) z
plug t (ZRR {x} l z)  = plug (TR {x} l t) z
plug t (ZBL {x} z r)  = plug (TB {x} t r) z
plug t (ZBR {x} l z)  = plug (TB {x} l t) z

plugBR :: forall (rlo rhi lo hi n rn :: Integer) . Tree lo hi Black n ->
              TreeZip rlo rhi Black rn lo hi Red n -> 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 :: Nat) . (lo < x, x < hi) =>
                TreeZip rlo rhi Black rn lo hi Black 0 -> SearchResult x rlo rhi rn
  Yep   ::  forall (x rlo rhi lo hi c :: Integer)(rn n :: Nat) .
                TreeZip rlo rhi Black rn lo hi c n -> 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 -> SearchResult x rlo rhi rn
search {x} = help Root
  where
    help ::  forall (lo hi c :: Integer)(n :: Nat) . (lo < x, x < hi) =>
                 TreeZip rlo rhi Black rn lo hi c n -> Tree lo hi c n ->
                     SearchResult x rlo rhi rn
    help z E                       = Nope z
    help z (TR {y} l r) | {x < y}  = help (ZRL {y} z r) l
    help z (TR {y} l r) | {x ~ y}  = Yep z (TR {y} l r)
    help z (TR {y} l r) | {x > y}  = help (ZRR {y} l z) r
    help z (TB {y} l r) | {x < y}  = help (ZBL {y} z r) l
    help z (TB {y} l r) | {x ~ y}  = Yep z (TB {y} l r)
    help z (TB {y} l r) | {x > y}  = help (ZBR {y} l z) r

member ::  forall (lo hi :: Integer) . pi (x :: Integer) . (lo < x, x < hi) =>
               RBT lo hi -> Bool
member {x} (RBT t) = case search {x} t of
                       Nope _   -> False
                       Yep _ _  -> 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

solveIns ::  forall (rlo rhi lo hi c rc :: Integer)(rn n :: Nat) . 
                InsProb lo hi c n -> TreeZip rlo rhi rc rn lo hi c n ->
                    RBT rlo rhi
solveIns (Level c t)      Root           = rbt c t

solveIns (Level Red t)    (ZRL {x} z r)  = solveIns (PanicRB {x} t r) z
solveIns (Level Red t)    (ZRR {x} l z)  = solveIns (PanicBR {x} l t) z
solveIns (Level Black t)  (ZRL {x} z r)  = solveIns (Level Red (TR {x} t r)) z
solveIns (Level Black t)  (ZRR {x} l z)  = solveIns (Level Red (TR {x} l t)) z
solveIns (Level col t)    (ZBL {x} z r)  = solveIns (Level Black (TB {x} t r)) z
solveIns (Level col t)    (ZBR {x} l z)  = solveIns (Level Black (TB {x} l t)) z

solveIns (PanicRB {xi} (TR {xil} lil ril) ri)  (ZBL {x} z r)  =
    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)  =
    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)  =
    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)  =
    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 -> RBT lo hi
insert {x} t = case search {x} t :: SearchResult x lo hi n of
    Nope z   -> solveIns (Level Red (TR {x} E E)) z
    Yep _ _  -> 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 :: Nat) . Tree lo hi Black n ->
                 TreeZip rlo rhi Black rn lo hi Black (n+1) -> RBT rlo rhi
solveDel t Root = RBT t

solveDel t (ZRL {x} z (TB {y} (TR {lx} ll lr) r)) = 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))) = 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))              = 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)))  = 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)  = 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)  = RBT (plug (TR {rx} (TB {y} l rl) (TB {x} rr t)) z)

-- Arrgh
solveDel t (ZRR {x} (TB {y} E E) z)              = 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)  = 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))  = 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))  = 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))  = 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)) = solveDel (TB {y} (TR {x} t E) r) z
solveDel t (ZBL {x} z (TB {y} (TB {lx} ll lr) r))  = 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))  = 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)))  = 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))) = 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) = 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)  = 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)  = 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)  = solveDel (TB {y} l (TR {x} E t)) z
solveDel t (ZBR {x} (TB {y} l (TB {lx} ll lr)) z)  = 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)  = 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)  = 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) = 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 :: Nat) . Tree lo hi c (n+1) ->
                (pi (k :: Integer) . lo < k => TreeZip rlo rhi Black rn k hi c (n+1)) ->
                    RBT rlo rhi
findMin (TR {x} (TB {y} E E) r)                    f = solveDel E (ZRL {x} (f {y}) r)
findMin (TR {x} (TB {y} E (TR {lx} ll lr)) r)      f = RBT (plug (TB {lx} ll lr) (ZRL {x} (f {y}) r))

findMin (TR {x} (TB {y} (TR {k} E E) lr) r)        f = RBT (plug E (ZBL {y} (ZRL {x} (f {k}) r) lr))

findMin (TB {x} (TR {y} E E) r)                    f = RBT (plug E (ZBL {x} (f {y}) r))
findMin (TB {x} E (TR {lx} ll lr))                 f = RBT (plug (TB {lx} ll lr) (f {x}))
findMin (TB {x} E E)                               f = solveDel E (f {x})

findMin (TR {x} (TB {y} (TB {llx} lll llr) lr) r)  f = findMin (TB {llx} lll llr) (\ {k} -> ZBL {y} (ZRL {x} (f {k}) r) lr)
findMin (TB {x} (TB {lx} ll lr) r)                 f = 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 :: Nat) . Tree lo hi c n ->
                 TreeZip rlo rhi Black rn lo hi c n -> RBT rlo rhi
delFocus E                                   z = RBT (plug E z)
delFocus (TR {x} E E)                        z = RBT (plugBR E z)
delFocus (TR {x} l (TB {rx} rl rr))          z = findMin (TB {rx} rl rr) (\ {k} -> ZRR {k} (wkTree l) z)
delFocus (TB {x} E E)                        z = solveDel E z
delFocus (TB {x} (TR {y} E E) E)             z = RBT (plug (TB {y} E E) z)
delFocus (TB {x} E (TR {y} E E))             z = RBT (plug (TB {y} E E) z)
delFocus (TB {x} (TR {k} E E) (TR {y} E E))  z = RBT (plug (TB {k} E (TR {y} E E)) z)
delFocus (TB {x} l (TB {rx} rl rr))          z = findMin (TB {rx} rl rr) (\ {k} -> ZBR {k} (wkTree l) z)
delFocus (TB {x} (TB {lx} ll lr)  r)         z = findMin r (\ {k} -> ZBR {k} (wkTree (TB {lx} ll lr)) z)

 
delete ::  forall (lo hi :: Integer) . pi (x :: Integer) . (lo < x, x < hi) =>
               RBT lo hi -> RBT lo hi
delete {x} (RBT t) = f (search {x} t)
  where
    f :: forall (n :: Nat) . SearchResult x lo hi n -> RBT lo hi
    f (Nope _)   = RBT t
    f (Yep z t)  = delFocus t z



-- Suppose we want to hide the bounds from the user of our red-black
-- tree library. In a dependently typed language, we could add top and
-- bottom elements to the order, but we can't do so here for the
-- integers. Instead, here's a solution that weakens the bounds on the
-- tree as necessary. Note that wkTree2 could safely be implemented
-- using unsafeCoerce. 

data T where
    T :: forall (n :: Nat)(lo hi :: Num) . Tree lo hi Black n -> T
  deriving Show

emptyT = T E

rbtToT :: forall (lo hi :: Num) . RBT lo hi -> T
rbtToT (RBT t) = T t

insertT :: pi (x :: Num) . T -> T
insertT {x} (T t) = rbtToT (insert {x} (weakling {x} t))

deleteT :: pi (x :: Num) . T -> T
deleteT {x} (T t) = rbtToT (delete {x} (RBT (weakling {x} t)))

weakling :: forall (lo hi c n :: Num) . pi (x :: Num) . Tree lo hi c n ->
              Tree (min lo (x-1)) (max hi (x+1)) c n
weakling {x} t = wkTree2 t

wkTree2 :: forall (lo lo' hi hi' c n :: Num) . (lo' <= lo, hi <= hi') =>
               Tree lo hi c n -> Tree lo' hi' c n
wkTree2 E            = E
wkTree2 (TB {x} l r) = TB {x} (wkTree2 l) (wkTree2 r)
wkTree2 (TR {x} l r) = TR {x} (wkTree2 l) (wkTree2 r)