
module RedBlackTree

  (* Red-black trees (data type) *)

  type key = int
  type value = int

  type color = Red | Black

  type tree =
    | Leaf
    | Node color tree key value tree

  (* occurrence of a key/value pair in a tree *)

  predicate memt (t : tree) (k : key) (v : value) =
    match t with
    | Leaf -> false
    | Node _ l k' v' r -> (k = k' /\ v = v') \/ memt l k v \/ memt r k v
    end

  lemma memt_color :
    forall l r : tree, k k' : key, v v' : value, c c' : color.
    memt (Node c l k v r) k' v' -> memt (Node c' l k v r) k' v'

  (* binary search tree *)

  use int.Int

  predicate lt_tree (x : key) (t : tree) =
    forall k : key. forall v : value. memt t k v -> k < x

  predicate gt_tree (x : key) (t : tree) =
    forall k : key. forall v : value. memt t k v -> x < k

  lemma lt_leaf: forall x: key. lt_tree x Leaf

  lemma gt_leaf: forall x: key. gt_tree x Leaf

  lemma lt_tree_node:
    forall x y: key, v: value, l r: tree, c: color.
    lt_tree x l -> lt_tree x r -> y < x -> lt_tree x (Node c l y v r)

  lemma gt_tree_node:
    forall x y: key, v: value, l r: tree, c: color.
    gt_tree x l -> gt_tree x r -> x < y -> gt_tree x (Node c l y v r)

  lemma lt_node_lt:
    forall x y: key, v: value, l r: tree, c: color.
    lt_tree x (Node c l y v r) -> y < x

  lemma gt_node_gt:
    forall x y: key, v: value, l r: tree, c: color.
    gt_tree x (Node c l y v r) -> x < y

  lemma lt_left:
    forall x y: key, v: value, l r: tree, c : color.
    lt_tree x (Node c l y v r) -> lt_tree x l

  lemma lt_right:
    forall x y: key, v: value,  l r: tree, c: color.
    lt_tree x (Node c l y v r) -> lt_tree x r

  lemma gt_left:
    forall x y: key, v: value, l r: tree, c: color.
    gt_tree x (Node c l y v r) -> gt_tree x l

  lemma gt_right:
    forall x y: key, v: value, l r: tree, c: color.
    gt_tree x (Node c l y v r) -> gt_tree x r

  lemma lt_tree_not_in:
    forall x: key, t: tree. lt_tree x t -> forall v: value. not (memt t x v)

  lemma lt_tree_trans:
    forall x y: key. x < y -> forall t: tree. lt_tree x t -> lt_tree y t

  lemma gt_tree_not_in :
    forall x: key, t: tree. gt_tree x t -> forall v: value. not (memt t x v)

  lemma gt_tree_trans :
    forall x y: key. y < x -> forall t: tree. gt_tree x t -> gt_tree y t

  predicate bst (t : tree) =
    match t with
    | Leaf -> true
    | Node _ l k _ r -> bst l /\ bst r /\ lt_tree k l /\ gt_tree k r
    end

  lemma bst_Leaf : bst Leaf

  lemma bst_left:
    forall k: key, v: value, l r: tree, c: color. bst (Node c l k v r) -> bst l

  lemma bst_right:
    forall k: key, v: value, l r: tree, c: color. bst (Node c l k v r) -> bst r

  lemma bst_color:
    forall c c': color, k: key, v: value, l r: tree.
    bst (Node c l k v r) -> bst (Node c' l k v r)

  lemma rotate_left:
    forall kx ky: key, vx vy: value, a b c: tree, c1 c2 c3 c4: color.
    bst (Node c1 a kx vx (Node c2 b ky vy c)) ->
    bst (Node c3 (Node c4 a kx vx b) ky vy c)

  lemma rotate_right:
    forall kx ky: key, vx vy: value, a b c: tree, c1 c2 c3 c4: color.
    bst (Node c3 (Node c4 a kx vx b) ky vy c) ->
    bst (Node c1 a kx vx (Node c2 b ky vy c))

  (* [rbtree n t]: red black tree invariant
     [t] is a properly balanced red-black tree, i.e. it
      satisfies the following two invariants:
      - a red node has no red son
      - any path from the root to a leaf has exactly [n] black nodes
   *)

  predicate is_not_red (t : tree) =
    match t with
    | Node Red _ _ _ _ -> false
    | Leaf | Node Black _ _ _ _ -> true
    end

  predicate rbtree (n : int) (t : tree) =
    match t with
    | Leaf ->
        n = 0
    | Node Red   l _ _ r ->
        rbtree n l /\ rbtree n r /\ is_not_red l /\ is_not_red r
    | Node Black l _ _ r ->
        rbtree (n-1) l /\ rbtree (n-1) r
    end

  lemma rbtree_Leaf:
    rbtree 0 Leaf

  lemma rbtree_Node1:
    forall k:key, v:value. rbtree 0 (Node Red Leaf k v Leaf)

  lemma rbtree_left:
    forall x: key, v: value, l r: tree, c: color.
    (exists n: int. rbtree n (Node c l x v r)) -> exists n: int. rbtree n l

  lemma rbtree_right:
    forall x: key, v: value, l r: tree, c: color.
    (exists n: int. rbtree n (Node c l x v r)) -> exists n: int. rbtree n r

  (* lookup *)

  exception Not_found

  let rec find (t : tree) (k : key) : value
    requires { bst t }
    variant { t }
    ensures { memt t k result }
    raises { Not_found -> forall v : value. not (memt t k v) }
  = match t with
    | Leaf -> raise Not_found
    | Node _ l k' v r ->
        if k = k' then v
        else if k < k' then find l k
        else (* k > k' *) find r k
    end

  (* insertion *)

  (** `almost_rbtree n t`: `t` may have one red-red conflict at its root;
      it satisfies `rbtree n` everywhere else *)

  predicate almost_rbtree (n : int) (t : tree) =
    match t with
    | Leaf ->
        n = 0
    | Node Red   l _ _ r ->
        rbtree n l /\ rbtree n r
    | Node Black l _ _ r ->
        rbtree (n-1) l /\ rbtree (n-1) r
    end

  lemma rbtree_almost_rbtree:
    forall n: int, t: tree. rbtree n t -> almost_rbtree n t

  lemma rbtree_almost_rbtree_ex:
    forall s: tree.
    (exists n: int. rbtree n s) -> exists n: int. almost_rbtree n s

  lemma almost_rbtree_rbtree_black:
    forall x: key, v: value, l r: tree, n: int.
    almost_rbtree n (Node Black l x v r) -> rbtree n (Node Black l x v r)

  (** `lbalance c x l r` acts as a black node constructor,
      solving a possible red-red conflict on `l`, using the following
      schema:

      B (R (R a x b) y c) z d
      B (R a x (R b y c)) z d -> R (B a x b) y (R c z d)
      B a x b -> B a x b

      The result is not necessarily a black node. *)

  let lbalance (l : tree) (k : key) (v : value) (r : tree)
    requires { lt_tree k l /\ gt_tree k r /\ bst l /\ bst r }
    ensures { bst result /\
      (forall n : int. almost_rbtree n l -> rbtree n r -> rbtree (n+1) result)
      /\
      forall k':key, v':value.
        memt result k' v' <->
        if k' = k then v' = v else (memt l k' v' \/ memt r k' v') }
  = match l with
    | Node Red (Node Red a kx vx b) ky vy c
    | Node Red a kx vx (Node Red b ky vy c) ->
        Node Red (Node Black a kx vx b) ky vy (Node Black c k v r)
    | _ ->
        Node Black l k v r
    end

  (** `rbalance l x r` is similar to `lbalance`, solving a possible red-red
      conflict on `r`. The balancing schema is similar:

      B a x (R (R b y c) z d)
      B a x (R b y (R c z d)) -> R (B a x b) y (R c z d)
      B a x b -> B a x b
   *)

  let rbalance (l : tree) (k : key) (v : value) (r : tree)
    requires { lt_tree k l /\ gt_tree k r /\ bst l /\ bst r }
    ensures { bst result /\
      (forall n : int. almost_rbtree n r -> rbtree n l -> rbtree (n+1) result)
      /\
      forall k':key, v':value.
        memt result k' v' <->
        if k' = k then v' = v else (memt l k' v' \/ memt r k' v') }
  = match r with
    | Node Red (Node Red b ky vy c) kz vz d
    | Node Red b ky vy (Node Red c kz vz d) ->
        Node Red (Node Black l k v b) ky vy (Node Black c kz vz d)
    | _ ->
        Node Black l k v r
    end

  (* `insert x s` inserts `x` in tree `s`, resulting in a possible top red-red
     conflict when `s` is red. *)

  let rec insert (t : tree) (k : key) (v : value) : tree
    requires { bst t /\ exists n: int. rbtree n t }
    variant { t }
    ensures { bst result /\
      (forall n : int. rbtree n t -> almost_rbtree n result /\
         (is_not_red t -> rbtree n result)) /\
      memt result k v /\
      forall k':key, v':value.
      memt result k' v' <-> if k' = k then v' = v else memt t k' v' }
  = match t with
    | Leaf ->
        Node Red Leaf k v Leaf
    | Node Red l k' v' r ->
        if k < k' then Node Red (insert l k v) k' v' r
        else if k' < k then Node Red l k' v' (insert r k v)
        else (* k = k' *) Node Red l k' v r
    | Node Black l k' v' r ->
        if k < k' then lbalance (insert l k v) k' v' r
        else if k' < k then rbalance l k' v' (insert r k v)
        else (* k = k' *) Node Black l k' v r
    end

  (* finally `add x s` calls `insert` and recolors the root as black *)

  let add (t : tree) (k : key) (v : value) : tree
    requires { bst t /\ exists n:int. rbtree n t }
    ensures { bst result /\ (exists n:int. rbtree n result) /\
      memt result k v /\
      forall k':key, v':value.
      memt result k' v' <-> if k' = k then v' = v else memt t k' v' }
  = match insert t k v with
    | Node _ l k' v' r -> Node Black l k' v' r
    | Leaf -> absurd
    end

end

module Vacid0

  (* the VACID-0 interface = mutable map with default value*)

  use ref.Ref
  use RedBlackTree

  type rbt = (value, tree)

  predicate inv (r : rbt) =
    let (_, t) = r in bst t /\ exists n : int. rbtree n t

  let function default (r : rbt) : value =
    let (d, _) = r in d

  predicate mem (r : rbt) (k : key) (v : value) =
    let (d, t) = r in
    memt t k v \/ (v = d /\ forall v':value. not (memt t k v'))

  let create (d : int)
    ensures { inv !result /\
      default !result = d /\
      forall k:key, v:value. mem !result k v <-> v = d }
  = let x = (d, Leaf) in ref x (* BUG: ref (d, Leaf) *)

  let replace (m : ref rbt) k v
    requires { inv !m }
    ensures { inv !m /\
      default !m = default (old !m) /\
      forall k':key, v':value.
      mem !m k' v' <-> if k' = k then v' = v else mem (old !m) k' v' }
  = let (d, t) = !m in
    m := (d, add t k v)

  let lookup (m : ref rbt) k
    requires { inv !m }
    ensures { mem !m k result }
  = let (d, t) = !m in
    try find t k with Not_found -> d end

  (* the easy way: implements `remove` using `replace` *)
  let remove (m : ref rbt) k
    requires { inv !m }
    ensures { inv !m /\
      default !m = default (old !m) /\
      forall k':key, v':value.
      mem !m k' v' <-> if k' = k then v' = default !m else mem (old !m) k' v' }
  = replace m k (default !m)

end
