Local Notation int := I.t.Definition t := tree.(** ** Height of trees *)

Definition height (s : t) : int :=
  match s with
  | Leaf => 0
  | Node h _ _ _ => h
  end.(** ** Singleton set *)

Definition singleton x := Node 1 Leaf x Leaf.(** ** Helper functions *)

(** [create l x r] creates a node, assuming [l] and [r]
    to be balanced and [|height l - height r| <= 2]. *)

Definition create l x r :=
   Node (max (height l) (height r) + 1) l x r.(** [bal l x r] acts as [create], but performs one step of
    rebalancing if necessary, i.e. assumes [|height l - height r| <= 3]. *)

Definition assert_false := create.Definition bal l x r :=
  let hl := height l in
  let hr := height r in
  if (hr+2) <? hl then
    match l with
     | Leaf => assert_false l x r
     | Node _ ll lx lr =>
       if (height lr) <=? (height ll) then
         create ll lx (create lr x r)
       else
         match lr with
          | Leaf => assert_false l x r
          | Node _ lrl lrx lrr =>
              create (create ll lx lrl) lrx (create lrr x r)
         end
    end
  else
    if (hl+2) <? hr then
      match r with
       | Leaf => assert_false l x r
       | Node _ rl rx rr =>
         if (height rl) <=? (height rr) then
            create (create l x rl) rx rr
         else
           match rl with
            | Leaf => assert_false l x r
            | Node _ rll rlx rlr =>
                create (create l x rll) rlx (create rlr rx rr)
           end
      end
    else
      create l x r.(** ** Insertion *)

Fixpoint add x s := match s with
   | Leaf => Node 1 Leaf x Leaf
   | Node h l y r =>
      match X.compare x y with
         | Lt => bal (add x l) y r
         | Eq => Node h l y r
         | Gt => bal l y (add x r)
      end
  end.(** ** Join

    Same as [bal] but does not assume anything regarding heights
    of [l] and [r].
*)

Fixpoint join l : elt -> t -> t :=
  match l with
    | Leaf => add
    | Node lh ll lx lr => fun x =>
       fix join_aux (r:t) : t := match r with
          | Leaf => add x l
          | Node rh rl rx rr =>
               if (rh+2) <? lh then bal ll lx (join lr x r)
               else if (lh+2) <? rh then bal (join_aux rl) rx rr
               else create l x r
          end
  end.(** ** Extraction of minimum element

  Morally, [remove_min] is to be applied to a non-empty tree
  [t = Node h l x r]. Since we can't deal here with [assert false]
  for [t=Leaf], we pre-unpack [t] (and forget about [h]).
*)

Fixpoint remove_min l x r : t*elt :=
  match l with
    | Leaf => (r,x)
    | Node lh ll lx lr =>
       let (l',m) := remove_min ll lx lr in (bal l' x r, m)
  end.(** ** Merging two trees

  [merge t1 t2] builds the union of [t1] and [t2] assuming all elements
  of [t1] to be smaller than all elements of [t2], and
  [|height t1 - height t2| <= 2].
*)

Definition merge s1 s2 :=  match s1,s2 with
  | Leaf, _ => s2
  | _, Leaf => s1
  | _, Node _ l2 x2 r2 =>
        let (s2',m) := remove_min l2 x2 r2 in bal s1 m s2'
end.(** ** Deletion *)

Fixpoint remove x s := match s with
  | Leaf => Leaf
  | Node _ l y r =>
      match X.compare x y with
         | Lt => bal (remove x l) y r
         | Eq => merge l r
         | Gt => bal l y (remove x r)
      end
   end.(** ** Concatenation

    Same as [merge] but does not assume anything about heights.
*)

Definition concat s1 s2 :=
   match s1, s2 with
      | Leaf, _ => s2
      | _, Leaf => s1
      | _, Node _ l2 x2 r2 =>
            let (s2',m) := remove_min l2 x2 r2 in
            join s1 m s2'
   end.(** ** Splitting

    [split x s] returns a triple [(l, present, r)] where
    - [l] is the set of elements of [s] that are [< x]
    - [r] is the set of elements of [s] that are [> x]
    - [present] is [true] if and only if [s] contains  [x].
*)

Record triple := mktriple { t_left:t; t_in:bool; t_right:t }.(** ** Splitting

    [split x s] returns a triple [(l, present, r)] where
    - [l] is the set of elements of [s] that are [< x]
    - [r] is the set of elements of [s] that are [> x]
    - [present] is [true] if and only if [s] contains  [x].
*)

Record triple := mktriple { t_left:t; t_in:bool; t_right:t }.(** ** Splitting

    [split x s] returns a triple [(l, present, r)] where
    - [l] is the set of elements of [s] that are [< x]
    - [r] is the set of elements of [s] that are [> x]
    - [present] is [true] if and only if [s] contains  [x].
*)

Record triple := mktriple { t_left:t; t_in:bool; t_right:t }.(** ** Splitting

    [split x s] returns a triple [(l, present, r)] where
    - [l] is the set of elements of [s] that are [< x]
    - [r] is the set of elements of [s] that are [> x]
    - [present] is [true] if and only if [s] contains  [x].
*)

Record triple := mktriple { t_left:t; t_in:bool; t_right:t }.(** ** Splitting

    [split x s] returns a triple [(l, present, r)] where
    - [l] is the set of elements of [s] that are [< x]
    - [r] is the set of elements of [s] that are [> x]
    - [present] is [true] if and only if [s] contains  [x].
*)

Record triple := mktriple { t_left:t; t_in:bool; t_right:t }.Notation "<< l , b , r >>" := (mktriple l b r) (at level 9).Fixpoint split x s : triple := match s with
  | Leaf => << Leaf, false, Leaf >>
  | Node _ l y r =>
     match X.compare x y with
      | Lt => let (ll,b,rl) := split x l in << ll, b, join rl y r >>
      | Eq => << l, true, r >>
      | Gt => let (rl,b,rr) := split x r in << join l y rl, b, rr >>
     end
 end.(** ** Intersection *)

Fixpoint inter s1 s2 := match s1, s2 with
    | Leaf, _ => Leaf
    | _, Leaf => Leaf
    | Node _ l1 x1 r1, _ =>
            let (l2',pres,r2') := split x1 s2 in
            if pres then join (inter l1 l2') x1 (inter r1 r2')
            else concat (inter l1 l2') (inter r1 r2')
    end.(** ** Difference *)

Fixpoint diff s1 s2 := match s1, s2 with
 | Leaf, _ => Leaf
 | _, Leaf => s1
 | Node _ l1 x1 r1, _ =>
    let (l2',pres,r2') := split x1 s2 in
    if pres then concat (diff l1 l2') (diff r1 r2')
    else join (diff l1 l2') x1 (diff r1 r2')
end.(** ** Union *)

(** In ocaml, heights of [s1] and [s2] are compared each time in order
   to recursively perform the split on the smaller set.
   Unfortunately, this leads to a non-structural algorithm. The
   following code is a simplification of the ocaml version: no
   comparison of heights. It might be slightly slower, but
   experimentally all the tests I've made in ocaml have shown this
   potential slowdown to be non-significant. Anyway, the exact code
   of ocaml has also been formalized thanks to Function+measure, see
   [ocaml_union] in [MSetFullAVL].
*)

Fixpoint union s1 s2 :=
 match s1, s2 with
  | Leaf, _ => s2
  | _, Leaf => s1
  | Node _ l1 x1 r1, _ =>
     let (l2',_,r2') := split x1 s2 in
     join (union l1 l2') x1 (union r1 r2')
 end.(** ** Filter *)

Fixpoint filter (f:elt->bool) s := match s with
  | Leaf => Leaf
  | Node _ l x r =>
    let l' := filter f l in
    let r' := filter f r in
    if f x then join l' x r' else concat l' r'
 end.(** ** Partition *)

Fixpoint partition (f:elt->bool)(s : t) : t*t :=
  match s with
   | Leaf => (Leaf, Leaf)
   | Node _ l x r =>
      let (l1,l2) := partition f l in
      let (r1,r2) := partition f r in
      if f x then (join l1 x r1, concat l2 r2)
      else (concat l1 r1, join l2 x r2)
  end.(** Inductions principles for some of the set operators *)

#[local] Ltac caseq :=
match goal with [ |- context [match ?t with _ => _ end] ] =>
  let cmp := fresh in
  let H := fresh in
  remember t as cmp eqn:H; symmetry in H; destruct cmp
end.

Lemma bal_ind [P : t -> X.t -> t -> tree -> Prop] : (forall (l : t) (x : X.t) (r : t), let hl := height l in let hr := height r in (hr + 2 <? hl) = true -> l = Leaf -> P Leaf x r (assert_false l x r)) -> (forall (l : t) (x : X.t) (r : t), let hl := height l in let hr := height r in (hr + 2 <? hl) = true -> forall (_x : I.t) (ll : tree) (lx : X.t) (lr : tree), l = Node _x ll lx lr -> (height lr <=? height ll) = true -> P (Node _x ll lx lr) x r (create ll lx (create lr x r))) -> (forall (l : t) (x : X.t) (r : t), let hl := height l in let hr := height r in (hr + 2 <? hl) = true -> forall (_x : I.t) (ll : tree) (lx : X.t) (lr : tree), l = Node _x ll lx lr -> (height lr <=? height ll) = false -> lr = Leaf -> P (Node _x ll lx Leaf) x r (assert_false l x r)) -> (forall (l : t) (x : X.t) (r : t), let hl := height l in let hr := height r in (hr + 2 <? hl) = true -> forall (_x : I.t) (ll : tree) (lx : X.t) (lr : tree), l = Node _x ll lx lr -> (height lr <=? height ll) = false -> forall (_x0 : I.t) (lrl : tree) (lrx : X.t) (lrr : tree), lr = Node _x0 lrl lrx lrr -> P (Node _x ll lx (Node _x0 lrl lrx lrr)) x r (create (create ll lx lrl) lrx (create lrr x r))) -> (forall (l : t) (x : X.t) (r : t), let hl := height l in let hr := height r in (hr + 2 <? hl) = false -> (hl + 2 <? hr) = true -> r = Leaf -> P l x Leaf (assert_false l x r)) -> (forall (l : t) (x : X.t) (r : t), let hl := height l in let hr := height r in (hr + 2 <? hl) = false -> (hl + 2 <? hr) = true -> forall (_x : I.t) (rl : tree) (rx : X.t) (rr : tree), r = Node _x rl rx rr -> (height rl <=? height rr) = true -> P l x (Node _x rl rx rr) (create (create l x rl) rx rr)) -> (forall (l : t) (x : X.t) (r : t), let hl := height l in let hr := height r in (hr + 2 <? hl) = false -> (hl + 2 <? hr) = true -> forall (_x : I.t) (rl : tree) (rx : X.t) (rr : tree), r = Node _x rl rx rr -> (height rl <=? height rr) = false -> rl = Leaf -> P l x (Node _x Leaf rx rr) (assert_false l x r)) -> (forall (l : t) (x : X.t) (r : t), let hl := height l in let hr := height r in (hr + 2 <? hl) = false -> (hl + 2 <? hr) = true -> forall (_x : I.t) (rl : tree) (rx : X.t) (rr : tree), r = Node _x rl rx rr -> (height rl <=? height rr) = false -> forall (_x0 : I.t) (rll : tree) (rlx : X.t) (rlr : tree), rl = Node _x0 rll rlx rlr -> P l x (Node _x (Node _x0 rll rlx rlr) rx rr) (create (create l x rll) rlx (create rlr rx rr))) -> (forall (l : t) (x : X.t) (r : t), let hl := height l in let hr := height r in (hr + 2 <? hl) = false -> (hl + 2 <? hr) = false -> P l x r (create l x r)) -> forall (l : t) (x : X.t) (r : t), P l x r (bal l x r).