(*[
 datatype dict with nat
 datasort dict : 
        badLeft < dict; badRoot < dict; badRight < dict;
   rbt < badLeft;   rbt < badRoot;    rbt < badRight;
                  nonempty < rbt;   black < rbt;
   red < nonempty;  nonemptyBlack < nonempty; nonemptyBlack < black

 datacon Empty : black(0)

 datacon Black : -all h : nat-
        int * dict(h) * dict(h) -> dict(h+1)
     &  int * rbt(h) * rbt(h) -> nonemptyBlack(h+1)
     &  int * badRoot(h) * rbt(h) -> badLeft(h+1)
     &  int * rbt(h) * badRoot(h) -> badRight(h+1)

 datacon Red : -all h : nat-
        int * dict(h) * dict(h) -> dict(h)
     &  int * black(h) * black(h) -> red(h)
     &  int * rbt(h) * black(h) -> badRoot(h)
     &  int * black(h) * rbt(h) -> badRoot(h)
]*)
datatype dict = Empty
             | Black of int * dict * dict
             | Red of int * dict * dict
(*[
  (* z : zipper(h, hz) if  zip(z, t) : rbt(hz), where t : rbt(h). *)
  datatype zipper with nat * nat

 (* The datasort refinement below captures two distinct properties:
   - the color of the hole's parent:
     If blackZipper, the root (of the zipper, that is, the parent of the hole)
     is black and can therefore tolerate black trees as well as red;
   - the color of the root _of the corresponding zipped tree_.   *)

 datasort zipper :
   topOrBR < zipper;  blackZipper < zipper;
   BRzipper < topOrBR;
   topZipper < topOrBR;  topZipper < blackZipper;
   blackBRzipper < topZipper; blackBRzipper < BRzipper

 datacon TOP : -all h : nat- topZipper(h, h)

 datacon LEFTB : -all h, hz : nat-
      int*rbt(h)*zipper(h+1,hz) -> blackZipper(h,hz)
   &  int*rbt(h)*topOrBR(h+1,hz) -> blackBRzipper(h,hz)

 datacon RIGHTB : -all h, hz : nat-
      rbt(h)*int*zipper(h+1,hz) -> blackZipper(h,hz)
   &  rbt(h)*int*topOrBR(h+1,hz) -> blackBRzipper(h,hz)

 datacon LEFTR : -all h, hz : nat-
      int*black(h)*blackZipper(h,hz) -> zipper(h,hz)
   &  int*black(h)*blackBRzipper(h,hz) -> BRzipper(h,hz)

 datacon RIGHTR : -all h, hz : nat-
      black(h)*int*blackZipper(h,hz) -> zipper(h,hz)
   &  black(h)*int*blackBRzipper(h,hz) -> BRzipper(h,hz)
]*)
datatype zipper = TOP
               | LEFTB of int * dict * zipper
               | LEFTR of int * dict * zipper
               | RIGHTB of dict * int * zipper
               | RIGHTR of dict * int * zipper ;

(*[ val zip : -all h, hz : nat-
         blackZipper(h, hz) * rbt(h) -> rbt(hz)
       & zipper(h, hz) * black(h) -> rbt(hz)
       & blackBRzipper(h, hz) * rbt(h) -> black(hz)
       & topOrBR(h, hz) * black(h) -> black(hz)
]*)
fun zip arg = case arg of
   (TOP, t) => t
 | (LEFTB (x, b, z as _), a) => zip(z, Black(x, a, b))
 | (RIGHTB(a, x, z as _), b) => zip(z, Black(x, a, b))
 | (LEFTR (x, b, z), a)      => zip(z, Red(x, a, b))
 | (RIGHTR(a, x, z), b)      => zip(z, Red(x, a, b))

(* bbZip propagates a black deficit up the tree until either the top
* is reached, or the deficit can be covered.  It returns a boolean
* that is true if there is still a deficit and the zipped tree. *)

(*[ val bbZip : -all h,hz : nat-
     zipper(h+1,hz)*rbt(h) -> ((bool(true)*rbt(hz-1))
                            \/ (bool(false)*rbt(hz)))
 & BRzipper(h+1,hz)*rbt(h) -> ((bool(true)*black(hz-1))
                            \/ (bool(false)*black(hz)))
 & topOrBR(h+1,hz)*black(h) -> ((bool(true)*black(hz-1))
                             \/ (bool(false)*black(hz)))
]*)
fun bbZip arg = case arg of
  (TOP, t) => (true, t)

| (LEFTB(x, Red(y,c,d), z), a) => (* case 1L-Black *)
    bbZip(LEFTR(x, c, LEFTB(y, d, z)), a)

| (LEFTB(x, Black(w, Red(y,c,d), e), z), a) => (* case 3L-Black *)
   (false, zip(z, Black(y, Black(x,a,c), Black(w,d,e))))
| (LEFTR(x, Black(w, Red(y,c,d), e), z), a) => (* case 3L-Red *)
   (false, zip(z, Red(y, Black(x,a,c), Black(w,d,e))))

| (LEFTB(x, Black(y, c, Red(w,d,e)), z), a) => (* case 4L-Black *)
   (false, zip(z, Black(y, Black(x,a,c), Black(w,d,e))))
| (LEFTR(x, Black(y, c, Red(w,d,e)), z), a) => (* case 4L-Red *)
   (false, zip(z, Red(y, Black(x,a,c), Black(w,d,e))))

| (LEFTR(x, Black(y,c,d), z), a) => (* case 2L-Red *)
   (false, zip(z, Black(x, a, Red(y,c,d))))
| (LEFTB(x, Black(y,c,d), z), a) => (* case 2L-Black *)
   bbZip(z, Black(x, a, Red(y,c,d)))

| (RIGHTB(Red(y,c,d), x, z), b) => (* case 1R-Black *)
   bbZip(RIGHTR(d, x, RIGHTB(c,y,z)), b)

(*(NJ library:)   | (RIGHTR(Red(y,c,d), x, z), b) => (* case 1R-Red *)   (* ?!?  This is a color violation: RIGHTRed - Red! *)
   bbZip(RIGHTR(d, x, RIGHTB(c, y, z)), b)             (* ...and there is no corresponding arm for 1L *)
*)

| (RIGHTB(Black(y, Red(w,c,d), e), x, z), b) => (* case 3R-Black *)
   (false, zip(z, Black(y, Black(w,c,d), Black(x,e,b))))
| (RIGHTR(Black(y, Red(w,c,d), e), x, z), b) => (* case 3R-Red *)
   (false, zip(z, Red(y, Black(w,c,d), Black(x,e,b))))

(* This 4R is correct -- unlike the buggy NJ library *)
| (RIGHTB(Black(y, c, Red(w,d,e)), x, z), b) => (* case 4R-Black *)
   (false, zip(z, Black(w, Black(y,c,d), Black(x,e,b))))
| (RIGHTR(Black(y, c, Red(w,d,e)), x, z), b) => (* case 4R-Red *)
   (false, zip(z, Red(w, Black(y,c,d), Black(x,e,b))))

| (RIGHTR(Black(y,c,d), x, z), b) => (* case 2R-Red *)
   (false, zip(z, Black(x, Red(y,c,d), b)))
| (RIGHTB(Black(y,c,d), x, z), b) => (* case 2R-Black *)
   bbZip(z, Black(x, Red(y,c,d), b))


(*[ val delMin : -all h, hz : nat-

      nonempty(h) * blackZipper(h, hz)
   -> int * ((bool(false)*rbt(hz))
           \/(bool(true)*rbt(hz-1)))

 &    nonemptyBlack(h) * zipper(h, hz)
   -> int * ((bool(false)*rbt(hz))
           \/(bool(true)*rbt(hz-1)))

 &    nonempty(h) * blackBRzipper(h, hz)
   -> int * ((bool(false)*black(hz))
           \/(bool(true)*black(hz-1)))

 &    nonemptyBlack(h) * BRzipper(h, hz) 
   -> int * ((bool(false)*black(hz))
           \/(bool(true)*black(hz-1)))    ]*)

fun delMin arg = case arg of
  (Red(y, Empty, b), z) => 
    (y, (false (* i.e., no deficit, black height unchanged *), zip(z,b)))

|  (Black(y, Empty, b), z) =>
    (* This black node is the minimum: deleting it yields a black deficit. *)
    (y, bbZip(z,b))
     (* Call bbZip; the flag is important, since it tells the caller whether the
      resulting tree has an internal black deficit (if true) or not (if false).
     (For example, if this is the only non-Empty node in the tree originally
     passed to delMin, there is no way to fix the deficit (nowhere to put
     the "extra blackness") in this function, and the flag will be `true'.)  *)

|  (Black(y, a, b), z) => delMin(a, LEFTB(y, b, z))  
|  (Red(y, a, b), z) => delMin(a, LEFTR(y, b, z))

(*[ val joinRed : -all h,hz:nat-
        black(h)*black(h)*blackZipper(h, hz) -> rbt  ]*)
fun joinRed arg = case arg of
  (Empty, Empty, z) => zip(z, Empty)
| (a, Empty, z) => #2(bbZip(z, a))
| (Empty, b, z) => #2(bbZip(z, b))
| (a, Black(x,Empty,bb), z)=>#2(bbZip(RIGHTR(a,x,z),bb))
| (a, Black(y,aa,bb), z) =>
   let in case delMin(aa, LEFTB(y, bb, TOP)) of
       (x, (needB as false, b')) =>
         zip(z, Red(x,a,b'))
     | (x, (needB as true, b')) =>
         #2(bbZip(RIGHTR(a,x,z), b'))
   end

(*[ val joinBlack : -all h,hz:nat-
        rbt(h)*rbt(h)*zipper(h+1, hz) -> rbt  ]*)
fun joinBlack arg = case arg of
  (a, Empty, z) => #2(bbZip(z, a))
| (Empty, b, z) => #2(bbZip(z, b))
| (a, b, z) =>
   let in case delMin(b, TOP) of
     (x, (needB as false, b')) => zip(z, Black(x,a,b'))
   | (x, (needB as true, b')) =>
         #2(bbZip(RIGHTB(a,x,z), b'))
   end

(*[ val delete : -all h : nat- rbt(h) -> int -> rbt ]*)
fun delete t key =
  let
   (*[ val del : -all h, hz : nat-
                   rbt(h) * blackZipper(h, hz) -> rbt
                &  black(h) * zipper(h, hz) -> rbt  ]*)
   fun del arg = case arg of
      (Empty, z) => raise NotFound
    | (Black(key1, a, b), z) =>
       if key = key1 then joinBlack (a, b, z)
       else if key < key1 then del (a, LEFTB(key1, b, z))
       else del (b, RIGHTB(a, key1, z))
    | (Red(key1, a, b), z) => 
       if key = key1 then joinRed (a, b, z)
       else if key < key1 then del (a, LEFTR(key1, b, z))
       else del (b, RIGHTR(a, key1, z))
  in
    del(t, TOP)
  end