RedBlack.idr

||| Red-black trees that intrinsically verifies it preserves the color invariants.
||| Originally written by Dan Licata in Agda.
module RedBlack
%default total
%hide Prelude.List.toList

data Color = Red | Black

mutual
  ||| Idris has induction-recursion but not induction-induction
  data WellRedTree : Type -> Type where
    Empty : WellRedTree a
    RedNode : (l : WellRedTree a)
           -> a
           -> (r : WellRedTree a)
           -> (rl : l `hasRoot` Black)
           -> (rr : r `hasRoot` Black)
           -> WellRedTree a
    BlackNode : (l : WellRedTree a) -> a -> (r : WellRedTree a) -> WellRedTree a

  ||| A predicate to check what color the root of a WellRedTree is.
  ||| Licata's version has a data type instead of a predicate,
  ||| but Idris has totality issues with that.
  hasRoot : WellRedTree a -> Color -> Type
  hasRoot Empty Black = Unit
  hasRoot (RedNode _ _ _ _ _) Red = Unit
  hasRoot (BlackNode _ _ _) Black = Unit
  hasRoot _ _ = Void

||| In an almost well red tree, no red node has a red child, except possibly the root
data AlmostWellRedTree : Type -> Type where
  AEmpty : AlmostWellRedTree a
  ANode  : (l : WellRedTree a) -> Color -> a -> (r : WellRedTree a) -> AlmostWellRedTree a

||| A function that casts a well-red tree to an almost-well-red tree.
promote : WellRedTree a -> AlmostWellRedTree a
promote Empty = AEmpty
promote (RedNode l kv r _ _) = ANode l Red kv r
promote (BlackNode l kv r)   = ANode l Black kv r

||| Every tree has either a black root or a red root.
decideRoot : (t : WellRedTree a) -> Either (t `hasRoot` Black) (t `hasRoot` Red)
decideRoot Empty = Left ()
decideRoot (RedNode _ _ _ _ _) = Right ()
decideRoot (BlackNode _ _ _) = Left ()

|||(1) For all l:tree, kv:(k, v), r:tree
|||    if l is almost well red
|||      r is well-red
|||    then balance l kv Black r is well-red.
balanceLeft : AlmostWellRedTree a -> a -> WellRedTree a -> WellRedTree a
-- these are the two rotation cases
balanceLeft (ANode (RedNode a x b _ _) Red y c) z d = RedNode (BlackNode a x b) y (BlackNode c z d) () ()
balanceLeft (ANode a Red x (RedNode b y c _ _)) z d = RedNode (BlackNode a x b) y (BlackNode c z d) () ()
-- need to expand the catch-all, because the *proof* is different in each case.
balanceLeft AEmpty kv r = BlackNode Empty kv r
balanceLeft (ANode a Black x b) kv r = BlackNode (BlackNode a x b) kv r
balanceLeft (ANode Empty Red x Empty) kv r = BlackNode (RedNode Empty x Empty () ()) kv r
balanceLeft (ANode Empty Red x (BlackNode l kv r)) kv' r' = BlackNode (RedNode Empty x (BlackNode l kv r) () ()) kv r'
balanceLeft (ANode (BlackNode a1 x1 a2) Red x Empty) y c = BlackNode (RedNode (BlackNode a1 x1 a2) x Empty () ()) y c
balanceLeft (ANode (BlackNode a1 x1 a2) Red x (BlackNode b1 y1 b2)) y c = BlackNode (RedNode (BlackNode a1 x1 a2) x (BlackNode b1 y1 b2) () ()) y c

|||(2) For all l:tree, kv:(k, v), r:tree
|||    if l is well-red
|||      r is almost-well-red
|||    then balance l kv Black r is well-red.
balanceRight : WellRedTree a -> a -> AlmostWellRedTree a -> WellRedTree a
-- these are the two rotation cases
balanceRight a x (ANode (RedNode b y c _ _) Red z d) = RedNode (BlackNode a x b) y (BlackNode c z d) () ()
balanceRight a x (ANode b Red y (RedNode c z d _ _)) = RedNode (BlackNode a x b) y (BlackNode c z d) () ()
-- catch-all
balanceRight a x AEmpty = BlackNode a x Empty
balanceRight a x (ANode l Black kv r) = BlackNode a x (BlackNode l kv r)
balanceRight a x (ANode Empty Red kv Empty) = BlackNode a x (RedNode Empty kv Empty () ())
balanceRight a x (ANode Empty Red kv (BlackNode l kv' r)) = BlackNode a x (RedNode Empty kv (BlackNode l kv' r) () ())
balanceRight a x (ANode (BlackNode l kv r) Red kv' Empty) = BlackNode a x (RedNode (BlackNode l kv r) kv' Empty () ())
balanceRight a x (ANode (BlackNode l kv r) Red kv' (BlackNode l' kv0 r')) = BlackNode a x (RedNode (BlackNode l kv r) kv' (BlackNode l' kv0 r') () ())

|||(3) For all l:tree, kv:(k, v), r:tree,
|||    if l is well-red
|||      r is well-red
|||    then balance l kv Red r is almost-well-red
balanceBreak : WellRedTree a -> a -> WellRedTree a -> AlmostWellRedTree a
balanceBreak l kv r = ANode l Red kv r

mutual
  |||(1) For all t, if t is well-red and the root of t is black, then
  |||    ins t (k,v) is well-red.
  insBlack : Ord a => (t : WellRedTree a) -> (kv : a) -> t `hasRoot` Black -> WellRedTree a
  insBlack Empty kv rt = RedNode Empty kv Empty () ()
  insBlack (RedNode l kv r _ _) kv' _ impossible
  insBlack (BlackNode l x' r) x rt with (compare x x')
    | LT = balanceLeft (insAny l x) x' r
    | GT = balanceRight l x (insAny r x)
    | EQ = BlackNode l x r

  |||(2) For all t, if t is well-red and the root of t is red, then
  |||    ins t (k,v) is almost-well-red.
  insRed : Ord a => (t : WellRedTree a) -> (kv : a) -> t `hasRoot` Red -> AlmostWellRedTree a
  insRed Empty kv _ impossible
  insRed (RedNode l x' r rl rr) x rc with (compare x x')
    | LT = balanceBreak (insBlack l x rl) x' r
    | GT = balanceBreak l x' (insBlack r x rr)
    | EQ = ANode l Red x r
  insRed (BlackNode l kv r) kv' _ impossible

  |||(3) Therefore, for all t, if t is well-red, then ins t (k,v) is almost-well-red.
  insAny : Ord a => (t : WellRedTree a) -> (kv : a) -> AlmostWellRedTree a
  insAny t kv with (decideRoot t)
    | Left rt = promote (insBlack t kv rt)
    | Right rt = insRed t kv rt

mutual
  |||(1) For all t, if t is well-red and the root of t is black, then
  |||    ins t (k,v) is well-red.
  insKeyBlack : Ord k => (t : WellRedTree (k, v)) -> (kv : (k, v)) -> t `hasRoot` Black -> WellRedTree (k, v)
  insKeyBlack Empty kv rt = RedNode Empty kv Empty () ()
  insKeyBlack (RedNode l kv r _ _) kv' _ impossible
  insKeyBlack (BlackNode l (k' , v') r) (k , v) rt with (compare k k')
    | LT = balanceLeft (insKeyAny l (k , v)) (k' , v') r
    | GT = balanceRight l (k , v) (insKeyAny r (k , v))
    | EQ = BlackNode l (k , v) r

  |||(2) For all t, if t is well-red and the root of t is red, then
  |||    ins t (k,v) is almost-well-red.
  insKeyRed : Ord k => (t : WellRedTree (k, v)) -> (kv : (k, v)) -> t `hasRoot` Red -> AlmostWellRedTree (k, v)
  insKeyRed Empty kv _ impossible
  insKeyRed (RedNode l (k' , v') r rl rr) (k , v) rc with (compare k k')
    | LT = balanceBreak (insKeyBlack l (k , v) rl) (k' , v') r
    | GT = balanceBreak l (k' , v') (insKeyBlack r (k , v) rr)
    | EQ = ANode l Red (k , v) r
  insKeyRed (BlackNode l kv r) kv' _ impossible

  |||(3) Therefore, for all t, if t is well-red, then ins t (k,v) is almost-well-red.
  insKeyAny : Ord k => (t : WellRedTree (k, v)) -> (kv : (k, v)) -> AlmostWellRedTree (k, v)
  insKeyAny t kv with (decideRoot t)
    | Left rt = promote (insKeyBlack t kv rt)
    | Right rt = insKeyRed t kv rt

||| For all t, if t is almost-well-red, then blackenRoot t is well-red.
blackenRoot : AlmostWellRedTree a -> WellRedTree a
blackenRoot AEmpty = Empty
blackenRoot (ANode l _ kv r) = BlackNode l kv r

||| An intrinsically verified insertion function
insert : Ord a => WellRedTree a -> a -> WellRedTree a
insert t kv = blackenRoot (insAny t kv)

||| An intrinsically verified insertion function
insertWithKey : Ord k => WellRedTree (k, v) -> (k, v) -> WellRedTree (k, v)
insertWithKey t kv = blackenRoot (insKeyAny t kv)

fromList : Ord k => List (k, v) -> WellRedTree (k, v)
fromList = foldl insertWithKey Empty

toList : WellRedTree a -> List a
toList Empty = []
toList (RedNode l kv r _ _) = toList l ++ kv :: toList r
toList (BlackNode l kv r) = toList l ++ kv :: toList r

lookup : Ord k => k -> WellRedTree (k, v) -> Maybe v
lookup _ Empty = Nothing
lookup k1 (RedNode l (k2, v) r _ _) with (compare k1 k2)
  | LT = lookup k1 l
  | GT = lookup k1 r
  | EQ = Just v
lookup k1 (BlackNode l (k2, v) r) with (compare k1 k2)
  | LT = lookup k1 l
  | GT = lookup k1 r
  | EQ = Just v

elem : Ord a => a -> WellRedTree a -> Bool
elem _ Empty = False
elem x (RedNode l x' r _ _) = x == x' || elem x l || elem x r
elem x (BlackNode l x' r) = x == x' || elem x l || elem x r

Foldable WellRedTree where
  foldr f x Empty = x
  foldr f x (RedNode l kv r _ _) = foldr f (f kv (foldr f x r)) l
  foldr f x (BlackNode l kv r) = foldr f (f kv (foldr f x r)) l

Ord a => Semigroup (WellRedTree a) where
  t1 <+> t2 = foldl insert t2 t1

Ord a => Monoid (WellRedTree a) where
  neutral = Empty

||| A different `Semigroup` implementation in case your values do not have
||| an `Ord` implementation.
[keySemigroup] Ord k => Semigroup (WellRedTree (k, v)) where
  t1 <+> t2 = foldl insertWithKey t2 t1

||| A different `Monoid` implementation in case your values do not have
||| an `Ord` implementation.
[keyMonoid] Ord a => Monoid (WellRedTree a) where
  neutral = Empty