rbt.timl

(* Red-black tree with invariants for well-red-ness and black-height-blancedness *)

(* The main functions are [insert] and [lookup]. *)

structure LogFacts = struct
datatype lemma1 {a b : Nat} =
         Lemma1 {2 ** a <= b + 1 -> a <= ceil (log2 $(b + 1))} of lemma1 {a} {b}

fun lemma1 {a b : Nat} () return lemma1 {a} {b} =
    @Lemma1 {a} {b} {admit} (* assume this fact as axiom *)
end

structure Compare = struct
datatype cmp_result = Equal | Less | Greater
end

(* rbt is implemented as a functor parametrized on the key type *)
functor Rbt (K : sig
                           (* signatures can have both type parameter and index parameter *)
                           type t
                           idx m : Time
                           val cmp : t * t -- m --> Compare.cmp_result
              end) =
struct
        
open Basic
open Compare
(* open LogFacts *)
(* open K *)

type key = K.t

datatype color : {Bool} =
         Black of color {true}
         | Red of color {false}
                          
datatype rbt 'a : {Nat(*size*)} {Bool} {Nat(*black-height*)} =
         Leaf of rbt 'a {0} {true} {0}
       | Node {lcolor color rcolor : Bool}
              {lsize rsize bh : Nat}
              {color = false -> lcolor = true /\ rcolor = true }
              (* We have to add the following two extra invariants. These two invariants can be derived from the other invariants, but because in TiML lemmas, lemma invocations, and inductions must be written as ordinary functions (like in Dafny) which increase a program's running time, deriving these invariants on the fly will increase asymptotic complexity. *)
              {2 ** (bh + b2n (not lcolor) + 1) <= 2 * (lsize + 1) /\ 2 * (lsize + 1) <= 2 ** (2 * bh + b2n (not lcolor) + 1)}
              {2 ** (bh + b2n (not rcolor) + 1) <= 2 * (rsize + 1) /\ 2 * (rsize + 1) <= 2 ** (2 * bh + b2n (not rcolor) + 1)}
         of color {color} * rbt 'a {lsize} {lcolor} {bh} * (key * 'a) * rbt 'a {rsize} {rcolor} {bh} --> rbt 'a {lsize + 1 + rsize} {color} {bh + b2n color}

datatype size_good {color : Bool} {size bh : Nat} =
         SizeGood {2 ** (bh + b2n (not color) + 1) <= 2 * (size + 1) /\ 2 * (size + 1) <= 2 ** (2 * bh + b2n (not color) + 1)}
         of size_good {color} {size} {bh}
                 
fun rbt_size_good ['a] {color : Bool} {size bh : Nat} (tr : rbt 'a {size} {color} {bh}) return size_good {color} {size} {bh} =
    case tr of
        Leaf => SizeGood
      | Node (color, _, _, _) =>
        case color of
            Black => SizeGood
          | Red => SizeGood
                  
datatype violation 'a : {Nat} {Nat} =
         ViolateLeft {lsize rsize bh : Nat}
         of rbt 'a {lsize} {false} {bh} * (key * 'a) * rbt 'a {rsize} {true} {bh} --> violation 'a {lsize + 1 + rsize} {bh}
       | ViolateRight {lsize rsize bh : Nat}
         of rbt 'a {lsize} {true} {bh} * (key * 'a) * rbt 'a {rsize} {false} {bh} --> violation 'a {lsize + 1 + rsize} {bh}

fun balance_left ['a] {rcolor : Bool} {lsize rsize bh : Nat}
                 (left : violation 'a {lsize} {bh})
                 (center as z : key * 'a)
                 (right as d : rbt 'a {rsize} {rcolor} {bh})
                 return rbt 'a {lsize + 1 + rsize} {false} {bh + 1} =
    case left of
        ViolateLeft (Node (Red, a, x, b), y, c) =>
        let
          val SizeGood _ = rbt_size_good c
          val SizeGood _ = rbt_size_good d
        in
          Node (Red, Node (Black, a, x, b), y, Node (Black, c, z, d))
        end
      | ViolateRight (a, x, Node (Red, b, y, c)) =>
        let
          val SizeGood _ = rbt_size_good a
          val SizeGood _ = rbt_size_good d
        in
          Node (Red, Node (Black, a, x, b), y, Node (Black, c, z, d))
        end
      | _ => never

fun balance_right ['a] {lcolor : Bool} {lsize rsize bh : Nat}
                  (left as a : rbt 'a {lsize} {lcolor} {bh})
                  (center as x : key * 'a)
                  (right : violation 'a {rsize} {bh})
                  return rbt 'a {lsize + 1 + rsize} {false} {bh + 1} =
    case right of
        ViolateLeft (Node (Red, b, y, c), z, d) =>
        let
          val SizeGood _ = rbt_size_good a
          val SizeGood _ = rbt_size_good d
        in
          Node (Red, Node (Black, a, x, b), y, Node (Black, c, z, d))
        end
      | ViolateRight (b, y, Node (Red, c, z, d)) =>
        let
          val SizeGood _ = rbt_size_good a
          val SizeGood _ = rbt_size_good b
        in
          Node (Red, Node (Black, a, x, b), y, Node (Black, c, z, d))
        end
      | _ => never
               
(* arbt: 'almost' red black tree, except that wellredness may be violated between root and one of its children *)
datatype arbt 'a {size bh : Nat} : {Bool (*whether the rbt is already good*)} {Bool} =
         Good {color : Bool}
         of color {color} * rbt 'a {size} {color} {bh} --> arbt 'a {size} {bh} {true} {color}
       | Bad {size bh : Nat}
         of violation 'a {size} {bh} --> arbt 'a {size} {bh} {false} {false}

datatype ins_result 'a {input_color : Bool} {input_size bh : Nat} =
         InsResult {output_color is_good : Bool}
                   {output_size : Nat}
                   {input_color = true -> is_good = true}
                   {output_size = input_size \/ output_size = input_size + 1}
         of arbt 'a {output_size} {bh} {is_good} {output_color} --> ins_result 'a {input_color} {input_size} {bh}

absidx T_ins : BigO (fn n => $n) with                                   
fun ins ['v] {input_color : Bool} {input_size bh : Nat}
        (tr : rbt 'v {input_size} {input_color} {bh}) (new as (k, _))
        return ins_result 'v {input_color} {input_size} {bh} using T_ins (2 * bh + b2n (not input_color)) =
    case tr of
        Leaf =>
        let
          val tr = Node (Red, Leaf, new, Leaf)
          val tr = Good (Red, tr)
          val tr = InsResult tr
        in
          tr
        end
      | Node (Red, left, center as (k', _), right) =>
        (case K.cmp (k, k') of
             Compare.Equal => InsResult (Good (Red, Node (Red, left, new, right)))
           | Compare.Less =>
             let
               val (InsResult left) = ins left new
             in
               case left of
                   Good (color, left) =>
                   (case color of
                        Red =>
                        InsResult (Bad (ViolateLeft (left, center, right)))
                      | Black =>
                        let
                          val SizeGood _ = rbt_size_good left
                          val SizeGood _ = rbt_size_good right
                        in
                          InsResult (Good (Red, Node (Red, left, center, right)))
                        end
                   )
                 | _ => never
             end
           | Compare.Greater =>
             let
               val (InsResult right) = ins right new
             in
               case right of
                   Good (color, right) =>
                   (case color of
                        Red =>
                        InsResult (Bad (ViolateRight (left, center, right)))
                      | Black =>
                        let
                          val SizeGood _ = rbt_size_good left
                          val SizeGood _ = rbt_size_good right
                        in
                          InsResult (Good (Red, Node (Red, left, center, right)))
                        end
                   )
                 | _ => never
             end
        )
      | Node (Black, left, center as (k', _), right) =>
        (case K.cmp (k, k') of
             Compare.Equal =>
             InsResult (Good (Black, Node (Black, left, new, right)))
           | Compare.Less => 
             let
               val (InsResult left) = ins left new
             in
               case left of
                   Good (_, left) =>
                   let
                     val SizeGood _ = rbt_size_good left
                     val SizeGood _ = rbt_size_good right
                   in
                     InsResult (Good (Black, Node (Black, left, center, right)))
                   end
                 | Bad left =>
                   InsResult (Good (Red, balance_left left center right))
             end
           | Greater => 
             let
               val (InsResult right) = ins right new
             in
               case right of
                   Good (_, right) =>
                   let
                     val SizeGood _ = rbt_size_good left
                     val SizeGood _ = rbt_size_good right
                   in
                     InsResult (Good (Black, Node (Black, left, center, right)))
                   end
                 | Bad right =>
                   InsResult (Good (Red, balance_right left center right))
             end
        )
end

(* simplify time complexity *)
absidx T_insert_rbt' : BigO (fn n => $n) with                                   
fun insert_rbt' ['v] {color : Bool} {size bh : Nat}
        (tr : rbt 'v {size} {color} {bh}) new
        return using T_insert_rbt' bh =
    ins tr new
end

absidx T_insert_rbt : BigO (fn n => log2 $n) with                                   
fun insert_rbt ['v] {color : Bool} {size bh : Nat}
        (tr : rbt 'v {size} {color} {bh}) new
        return using T_insert_rbt size =
    let
      val SizeGood _ = rbt_size_good tr
      val LogFacts.Lemma1 _ = @LogFacts.lemma1 {bh} {size} ()
    in
      insert_rbt' tr new using 2.0 + T_insert_rbt' (ceil (log2 $(size + 1)))
    end
end

fun blacken_root ['a] {size bh : Nat} (tr : rbt 'a {size} {false} {bh}) return rbt 'a {size} {true} {bh + 1} =
    case tr of
        Node (Red, l, c, r) => Node (Black, l, c, r)
      | _ => never

(* final packaging: hide color and black-height; root must be black. *)
datatype rb_tree 'a {size : Nat} =
         RBTree {bh : Nat} of rbt 'a {size} {true} {bh} --> rb_tree 'a {size}

val empty = RBTree Leaf

(* for insert, the size of the resulting tree may be one node larger than the input *) 
datatype may_grow_one 'a {size : Nat} =
         MayGrowOne {size' : Nat | size' = size \/ size' = size + 1} of rb_tree 'a {size'} --> may_grow_one 'a {size}

absidx T_insert : BigO (fn n => log2 $n) with
fun insert ['v] {size : Nat} (tr : rb_tree 'v {size}) (new: key * 'v) return may_grow_one 'v {size} using T_insert size =
    case tr of
        RBTree tr =>
        case insert_rbt tr new of
            InsResult (Good (color, tr)) =>
            (case color of
                 Black => MayGrowOne (RBTree tr)
               | Red => MayGrowOne (RBTree (blacken_root tr))
            )
          | _ => never
end

absidx T_lookup_rbt' : BigO (fn n => $n) with
fun lookup_rbt' ['v] {color : Bool} {size bh : Nat} (tr :rbt 'v {size} {color} {bh}) k return option 'v using T_lookup_rbt' (2 * bh + b2n (not color)) =
    case tr of
        Leaf => NONE
      | Node (_, left, (k', v), right) =>
        case K.cmp (k, k') of
            Equal => SOME v
          | Less => lookup_rbt' left k
          | Greater => lookup_rbt' right k
end

(* simplify time complexity *)
absidx T_lookup_rbt : BigO (fn n => $n) with                                   
fun lookup_rbt ['v] {color : Bool} {size bh : Nat}
        (tr : rbt 'v {size} {color} {bh}) k
        return using T_lookup_rbt bh =
    lookup_rbt' tr k
end

absidx T_lookup : BigO (fn n => log2 $n) with
fun lookup ['v] {size : Nat} (tr : rb_tree 'v {size}) (k: key) return option 'v using T_lookup size =
    case tr of
        @RBTree {_ bh} tr =>
        let
          open LogFacts
          val SizeGood _ = rbt_size_good tr
          val Lemma1 _ = @lemma1 {bh} {size} ()
        in
          lookup_rbt tr k using 2.0 + T_lookup_rbt (ceil (log2 $(size + 1)))
        end
end

end

(* test the functor *)

structure IntKey = struct

type t = int
           
absidx m = 7.0 with
fun cmp (a : int, b : int) return using m = Compare.Equal
end
                     
end

structure IntRbt = Rbt (IntKey)

structure TestIntRbt = struct
open IntRbt

fun test1 () =
  let
    val m = empty
    val MayGrowOne m = insert m (77, true)
    val r1 = lookup m 77 using 2.0 + T_lookup 1
    val r1 = case r1 of
                 SOME true => true
               | _ => false
    val r2 = lookup m 88 using 2.0 + T_lookup 1
    val r2 = case r2 of
                 NONE => true
               | _ => false
  in
    ()
  end

end