rbtree2.hs

import Gibbon.Maybe
import Gibbon.Prelude
import Gibbon.PList

data B = B Bool

-- packed primitives

tru :: B
tru = B True

fal :: B
fal = B False

unB :: B -> Bool
unB bPacked = case bPacked of
  B b -> b

data I = I Int

unI :: I -> Int
unI iPacked = case iPacked of
  I i -> i

gt :: I -> I -> Bool
gt i1 i2 = unI i1 > unI i2

lt :: I -> I -> Bool
lt i1 i2 = unI i1 < unI i2

-- tree

data RBT
  = Empty
  | Node B I RBT RBT

-- printer

fold_plist :: (a -> b -> a) -> a -> PList b -> a
fold_plist f init lst = case lst of
  Cons h t -> 
    let x = f init h in
    fold_plist f x t
  Nil -> init

cat_plist :: PList a -> PList a -> PList a
cat_plist l1 l2 = case l1 of
  Cons x xs -> Cons x (cat_plist xs l2)
  Nil -> l2

join_plist :: PList (PList a) -> PList a
join_plist lst = case lst of
  Cons x xs -> cat_plist x (join_plist xs)
  Nil -> Nil

bind_plist :: (a -> PList b) -> PList a -> PList b
bind_plist f xs = join_plist (map_plist f xs)

any_plist :: PList Bool -> Bool
any_plist lst = case lst of
  Nil -> False
  Cons h t -> h || any_plist t

is_pop :: Maybe a -> Bool
is_pop m = case m of
  Just z -> True
  Nothing -> False

print_rbt :: RBT -> ()
print_rbt t = print_rbt_h1 (Cons (Just t) Nil)

print_rbt_h1 :: PList (Maybe RBT) -> ()
print_rbt_h1 lst = 
  let x = bind_plist print_rbt_h2 lst in
  if any_plist (map_plist is_pop x) then
      let _ = print_newline () in
      print_rbt_h1 x
  else ()

print_rbt_h2 :: Maybe RBT -> PList (Maybe RBT)
print_rbt_h2 t_opt = let
    xs = case t_opt of
      Nothing -> print_empty (quote "..")
      Just t -> case t of
        Empty -> print_empty (quote "**")
        Node bp ip l r -> 
          let
            _ = printsym (if unB bp then quote "R" else quote "B")
            _ = printint (unI ip)
          in
          Cons (Just l) (Cons (Just r) Nil)
    _ = print_space ()
  in xs

print_empty :: Sym -> PList (Maybe RBT)
print_empty s = 
  let _ = printsym s in
  Cons Nothing (Cons Nothing Nil)

-- impl

empty :: RBT
empty = Empty

singleton :: I -> RBT
singleton x = Node tru x Empty Empty

insert :: I -> RBT -> RBT
insert e1 t = case t of
  Empty -> singleton e1
  Node colorx x left right -> 
    flipColors (rotateRight (rotateLeft (
      if lt x e1 then Node colorx x left (insert e1 right)
      else if gt x e1 then Node colorx x (insert e1 left) right
      else t
    )))

rotateLeft :: RBT -> RBT
rotateLeft t = case t of
  Empty -> Empty
  Node colorx x leftx rightx -> case rightx of
    Empty -> t
    Node c z leftz rightz -> 
      if isRed rightx && isBlack leftx
      then Node colorx z (Node tru x leftx leftz) rightz
      else t

rotateRight :: RBT -> RBT
rotateRight t = case t of
  Empty -> Empty
  Node colorx x leftx rightx -> case leftx of
    Empty -> t
    Node c y lefty righty -> 
      if isRed leftx && isRed lefty
      then Node colorx y lefty (Node tru x righty rightx)
      else t

flipColors :: RBT -> RBT
flipColors t = case t of
  Empty -> Empty
  Node c x leftx rightx -> case leftx of
    Empty -> t
    Node c1 y lefty righty -> case rightx of
      Empty -> t
      Node c2 z leftz rightz -> 
        if isRed leftx && isRed rightx
        then Node tru x (Node fal y lefty righty) (Node fal z leftz rightz)
        else t

isRed :: RBT -> Bool
isRed t = case t of
  Empty -> False
  Node c1 x l r -> unB c1

isBlack :: RBT -> Bool
isBlack t = case t of
  Empty -> True
  Node c1 x l r -> if unB c1 then False else True

ins :: Int -> RBT -> RBT
ins x t = insert (I x) t

mini :: RBT -> I
mini t = case t of
  Empty -> I (0 - 1)
  Node c x l r -> case l of
    Empty -> x
    Node cl xl ll rl -> mini l

gibbon_main = 
  let
    t1 = ins 2 empty
    t2 = ins 1 t1
    t3 = ins 3 t2
    t4 = ins 4 t3
    _ = print_rbt t4
  in ()