(******************************************************************************)
(* Introduction ***************************************************************)
(******************************************************************************)

module type SET = sig
  type t 
  type set
  val empty : set 
  val length : set -> int 
  val mem : t -> set -> bool
  val add : t -> set -> set
  val remove : t -> set -> set
  val to_list : set -> t list
end;;



(******************************************************************************)
(* 1. Structures simples ******************************************************)
(******************************************************************************)

(* 1.1. Un exemple ************************************************************)

module IntSetByFun : (SET with type t = int) =
struct
  type t = int
  type set = (int -> bool) * int

  let empty = ((fun x -> false), 0)
  
  let length (f,l) = l 
  
  let mem x (f,l) = (f x)
  
  let add  y (f,l) = 
    if (f y) then (f,l) 
    else ((fun x -> if x = y then true else (f x)) , l+1)

  let remove y (f,l) = 
    if (f y) then ((fun x -> if x = y then false else (f x)) , l-1)
    else (f,l)
        
  let to_list (f,l) =  
    (* n : nombre d'éléments à trouver            *)
    (* i : élément à tester pour son appartenance *)
    let rec to_list_aux n i =
      if n = 0 then
	[]
      else if (f i) then
	i::(to_list_aux (n-1) (i+1)) 
      else
	(to_list_aux n (i+1))
    in to_list_aux l 0
end;;



(* 1.2. Listes ****************************************************************)

(* Question 1. *)

module IntSetByList : (SET with type t = int) =
struct
  type t = int
  type set = t list
  let empty = []
  let length = List.length
  let mem = List.mem
  let rec add x l = match l with
    | [] -> [x]
    | y::q -> if x = y then l
	      else y::(add x q)
  let rec remove x l = match l with
    | [] -> []
    | y::q -> if x = y then q
	      else y::(remove x q)
  let rec to_list l = l
end;;


(* Question 2. *)

module type ELEMENT =
sig
  type t
end;;

module SetByList (A : ELEMENT) : (SET with type t = A.t) =
struct
  type t = A.t
  type set = t list
  let empty = []
  let length = List.length
  let mem = List.mem
  let add x l = if List.mem x l then l else x::l
  let rec remove x l = match l with
    | [] -> []
    | y::q -> if x = y then q
	      else y::(remove x q)
  let rec to_list l = l
end;;

module Integer : (ELEMENT with type t = int) =
struct
  type t = int
end;;



(* 1.3. Vers une nouvelle structure de données ********************************)

module type ORDERED =
sig
  type t
  val compare : t -> t -> int
    (* (compare x y) < 0  si x < y *)
    (* (compare x y) = 0  si x = y *)
    (* (compare x y) > 0  si x > y *)
end;;

module Integer : (ORDERED with type t = int) = 
struct
  type t = int
  let compare x y = (x - y)
end;;



(******************************************************************************)
(* 2. Arbres binaires de recherche ********************************************)
(******************************************************************************)

(* 2.3. À l'attaque ***********************************************************)

(* Question 1. *)

module SetByTree (A : ORDERED) : (SET with type t = A.t) =
struct
  type t = A.t
  type set = Node of (t * set * set) | Leaf
  let empty = Leaf
  let rec length = function
    | Leaf -> 0
    | Node (_,l,r) -> 1 + (length l) + (length r)
  let rec mem x = function
    | Leaf -> false
    | Node (y,l,r) ->
	if (compare x y) < 0 then mem x l
	else if (compare x y) > 0 then mem x r
	else true
  let rec add x = function
    | Leaf -> Node (x, Leaf, Leaf)
    | Node (y,l,r) ->
	if (compare x y) < 0 then Node (y,(add x l),r)
	else if (compare x y) > 0 then Node (y,l,(add x r))
	else Node (y,l,r)
  let rec remove x =
    let rec aux = function
      | Leaf -> failwith "Apocalypse has occured."
      | Node (y,l,r) -> if r = Leaf then (y,l)
                        else let (z,t) = aux r in (z, Node(y,l,t))
    in function
      | Leaf -> Leaf
      | Node (y,l,r) ->
	  if (compare x y) < 0 then Node (y,(remove x l),r)
	  else if (compare x y) > 0 then Node (y,l,(remove x r))
	  else if r = Leaf then l
	  else if l = Leaf then r
	  else
	    let (z,t) = aux l in Node (z,t,r)
  let rec to_list = function
    | Leaf -> []
    | Node (y,l,r) -> (to_list l)@(y::(to_list r))
end;;



(******************************************************************************)
(* 3. Arbres AVL **************************************************************)
(******************************************************************************)

(* 3.3. Insertions et suppressions ********************************************)

(* Question 1. *)

(*
  8

  8
   \
    9

    9
   / \
  8   11

    9
   / \
  8   11
       \
        13

      9
     / \
    8   11
   /     \
  2       13

      9
     / \
    8   13
   /   / \
  2   11  32

      9
     / \
    8   13
   /   / \
  2   11  32
         /
        18

      9
     / \
    8   13
   /   / \
  2   11  32
         / \
        18  52

      9
     / \
    8   18
   /   / \
  2   13  32
     /   / \
    11  19  52

*)

(* Question 2. *)

(*
      9
     / \
    8   18
   /   / \
  2   13  32
     /   /
    11  19

      11
     / \
    8   18
   /   / \
  2   13  32
         /
        19

      11
     / \
    8   19
   /   / \
  2   18  32

*)



(* 3.4. La vraie vie **********************************************************)

(* Questions 3., 4., 5., 6. *)

module SetByAVL (A : ORDERED) =
struct
  type t = A.t
  type tree = Node of (A.t * avl * avl) | Leaf
  and  avl = tree * int
  type set = avl
  let empty = (Leaf,0)
  let rec length t = match (fst t) with
    | Leaf -> 0
    | Node (_,l,r) -> 1 + (length l) + (length r)
  let rec mem x t = match (fst t) with
    | Leaf -> false
    | Node (y,l,r) ->
	if (compare x y) < 0 then mem x l
	else if (compare x y) > 0 then mem x r
	else true

  let r_rot  = function
    | (Node(x,(Node(y,tA,tB),_),tC),_) ->
	(Node (y,
	       tA,
	       (Node (x,tB,tC), 1 + (max (snd tB) (snd tC)))
	      ), 1 + (max (snd tA) (1 + (max (snd tB) (snd tC)))))
    | _ -> failwith "Incorrect AVL structure for left rotation."

  let l_rot = function
    | (Node(x,tA,(Node(y,tB,tC),_)),_) ->
	(Node (y,
	       (Node (x,tA,tB), 1 + (max (snd tA) (snd tB))),
	       tC
	      ), 1 + (max (1 + (max (snd tA) (snd tB))) (snd tC)))
    | _ -> failwith "Incorrect AVL structure for right rotation."

  let lr_rot = function
    | (Node(x,tA,tB),_) -> r_rot (Node (x,(l_rot tA),tB), 0)
    | _ -> failwith "Incorrect AVL structure for left-right rotation."

  let rl_rot = function
    | (Node(x,tA,tB),_) -> l_rot (Node (x,tA,(r_rot tB)), 0)
    | _ -> failwith "Incorrect AVL structure for right-left rotation."

  let balance (t,_) = match t with
    | Leaf -> (Leaf,0)
    | Node (x,tA,tB) ->
	if (snd tA) = (snd tB)+2 then
	  match (fst tA) with
	    | Leaf -> failwith "Apocalypse has occured."
	    | Node (y,tAl,tAr) ->
		if (snd tAl) = (snd tB)+1 then
		  r_rot (t,0)
		else
		  lr_rot (t,0)
	else if (snd tB) = (snd tA)+2 then
	  match (fst tB) with
	    | Leaf -> failwith "Apocalypse has occured."
	    | Node (y,tBl,tBr) ->
		if (snd tBr) = (snd tA)+1 then
		  l_rot (t,0)
		else
		  rl_rot (t,0)
	else
	  (Node(x,tA,tB), 1 + (max (snd tA) (snd tB)))

  let rec add x t = match (fst t) with
    | Leaf -> (Node (x,empty,empty), 1)
    | Node (y,l,r) ->
	if (compare x y) < 0 then
	  balance (Node (y, (add x l), r), 0)
	else if (compare x y) > 0 then
	  balance (Node (y, l, (add x r)), 0)
	else t

  let rec remove x t =
    let rec aux t = match (fst t) with
      | Leaf -> failwith "Apocalypse has occured."
      | Node (y,l,r) -> if (fst r) = Leaf then y else aux r
    in match (fst t) with
      | Leaf -> empty
      | Node (y,l,r) ->
	  if (compare x y) < 0 then
	    balance (Node (y, (remove x l), r), 0)
	  else if (compare x y) > 0 then
	    balance (Node (y, l, (remove x r)), 0)
	  else if (fst r) = Leaf then l
	  else if (fst l) = Leaf then r
	  else
	    let z = aux l in balance (Node (z,(remove z l),r), 0)

  let rec to_list t = match (fst t) with
    | Leaf -> []
    | Node (y,l,r) -> (to_list l)@(y::(to_list r))

  let rec pop t = match (fst t) with
    | Leaf -> failwith "Empty tree."
    | Node (y,l,r) -> if (fst l) = Leaf then y
                      else pop l

  let sort l =
    let rec aux1 = function
      | [] -> empty
      | x::q -> add x (aux1 q)
    and aux2 t = match (fst t) with
      | Leaf -> []
      | _ -> let x = pop t in x::(aux2 (remove x t))
    in aux2 (aux1 l)
end;;



module S = SetByAVL(Integer);;

let t = S.empty;;
let s = S.add 19 (S.add 52 (S.add 18 (S.add 32 (S.add 2 (S.add 13 (S.add 11 (S.add 9 (S.add 8 t))))))));;
let r = S.remove 13 (S.remove 9 (S.remove 52 s));;

let l = [19;52;18;32;2;13;11;9;8];;
S.sort l;;



(******************************************************************************)
(* 4. Arbres rouge-noir *******************************************************)
(******************************************************************************)

(* 4.4. Allons-y gaiement *****************************************************)

module SetByRBTree (A : ORDERED) =
struct
  type t = A.t
  type color = Red | Black
  type tree = Node of (A.t * set * set) | Leaf
  and  set = tree * color
  let empty = (Leaf,Black)
  let rec length t = match (fst t) with
    | Leaf -> 0
    | Node (_,l,r) -> 1 + (length l) + (length r)
  let rec mem x t = match (fst t) with
    | Leaf -> false
    | Node (y,l,r) ->
	if (compare x y) < 0 then mem x l
	else if (compare x y) > 0 then mem x r
	else true

  (* On a besoin de savoir si le noeud courant est la racine de l'arbre, car *)
  (* il faut respecter la propriété 1. On utilise pour cela is_root.         *)
  let balance_add (t,c) is_root = match t with
    | Leaf -> (Leaf,Black)
    | Node (g,tA,tB) ->
	if (c = Black) && ((snd tA) = Red) then
	  match (fst tA) with
	    | Node (p,tAl,tAr) ->
		if (snd tB) = Red then
		  (Node (g,(fst tA, Black),(fst tB, Black)), if is_root then Black else Red)
		else if ((snd tAl) = Black) && ((snd tAr) = Black) then
		  (t, if is_root then Black else c)
		else if ((snd tAl) = Red) then
		  (Node (p,
			 tAl,
			 (Node (g,tAr,tB), Red)
			), Black)
		else
		  (match (fst tAr) with
		     | Node (n,tArl,tArr) ->
			 (Node (n,
				(Node (p,tAl,tArl), Red),
				(Node (g,tArr,tB), Red)
			       ), Black)
		     | _ -> failwith "Apocalypse has occured.")
	    | _ -> failwith "Apocalypse has occured."
	else if (c = Black) && ((snd tB) = Red) then
	  match (fst tB) with
	    | Node (p,tBl,tBr) ->
		if (snd tA) = Red then
		  (Node (g,(fst tA, Black),(fst tB, Black)), if is_root then Black else Red)
		else if ((snd tBl) = Black) && ((snd tBr) = Black) then
		  (t, if is_root then Black else c)
		else if ((snd tBr) = Red) then
		  (Node (p,
			 (Node (g,tA,tBl), Red),
			 tBr
			), Black)
		else
		  (match (fst tBl) with
		     | Node (n,tBll,tBlr) ->
			 (Node (n,
				(Node (g,tA,tBll), Red),
				(Node (p,tBlr,tBr), Red)
			       ), Black)
		     | _ -> failwith "Apocalypse has occured.")
	    | _ -> failwith "Apocalypse has occured."
	else
	  (t, if is_root then Black else c)

  let add x t =
    let rec add_aux t is_root = match (fst t) with
      | Leaf -> (Node (x,empty,empty), Red)
      | Node (y,l,r) ->
	  if (compare x y) < 0 then
	    balance_add (Node (y, (add_aux l false), r), snd t) is_root
	  else if (compare x y) > 0 then
	    balance_add (Node (y, l, (add_aux r false)), snd t) is_root
	  else t
    in add_aux t true

  (* Cette fonction est assez tordue, car un simple pattern-matching sur les  *)
  (* couleurs ne nous permet pas de décider quel est le sous-arbre            *)
  (* déséquilibré. On utilise donc le booléen left qui est vrai si le         *)
  (* sous-arbre de gauche est déséquilibré, et faux si c'est celui de droite. *)
  (* Il faut aussi savoir si l'opération a complètement rééquilibré l'arbre,  *)
  (* ou s'il faut encore l'appliquer lors de la reconstruction de l'arbre.    *)
  (* La fonction renvoie pour cela un couple (arbre, booléen), où le booléen  *)
  (* est mis à vrai si l'arbre ne risque pas de provoquer de déséquilibre,    *)
  (* faux sinon.                                                              *)
  let rec balance_rem (t,c) left = match t with
    | Leaf -> ((Leaf,Black),true)
    | Node (p,tA,tB) ->
	let (tN,tS) = if left then (tA,tB) else (tB,tA) in
	  match (fst tS) with
	    | Node (s,tSl,tSr) ->
		if (c = Black) && ((snd tS) = Black) && ((snd tSl) = Black) && ((snd tSr) = Black) then
		  if left then
		    ((Node(p,tN,(fst tS,Red)),Black),false)
		  else
		    ((Node(p,(fst tS,Red),tN),Black),false)
		else if (c = Red) && ((snd tS) = Black) && ((snd tSl) = Black) && ((snd tSr) = Black) then
		  if left then
		    ((Node(p,tN,(fst tS,Red)),Black),true)
		  else
		    ((Node(p,(fst tS,Red),tN),Black),true)
		else if ((snd tS) = Black) then
		  if left then
		    if ((snd tSr) = Red) then
		      ((Node(s,(Node(p,tN,tSl),Black),(fst tSr,Black)),c),true)
		    else
		      (match (fst tSl) with
			 | Node (sl,tSll,tSlr) ->
			     ((Node(sl,
				    (Node(p,tN,tSll),Black),
				    (Node(s,tSlr,tSr),Black)
				   ),c),true)
			 | _ -> failwith "Apocalypse has occured.")
		  else
		    if ((snd tSl) = Red) then
		      ((Node(s,(fst tSl,Black),(Node(p,tSr,tN),Black)),c),true)
		    else
		      (match (fst tSr) with
			 | Node (sr,tSrl,tSrr) ->
			     ((Node(sr,
				    (Node(s,tSl,tSrl),Black),
				    (Node(p,tSrr,tN),Black)
				   ),c),true)
			 | _ -> failwith "Apocalypse has occured.")
		else
		  if left then
		    ((Node(s,
			   (fst (balance_rem (Node(p,tN,tSl),Red) left)),
			   tSr
			  ),Black),true)
		  else
		    ((Node(s,
			   tSl,
			   (fst (balance_rem (Node(p,tSr,tN),Red) left))
			  ),Black),true)
	    | _ -> failwith "Apocalypse has occured."

  let remove x t =
    let rec aux t = match (fst t) with
      | Leaf -> failwith "Apocalypse has occured."
      | Node (y,l,r) -> if (fst r) = Leaf then y else aux r
    and rem_aux x t = match (fst t) with
      | Leaf -> (empty,true)
      | Node (y,l,r) ->
	  if (compare x y) < 0 then
	    let (ln,b) = rem_aux x l in
	      if not b then
		balance_rem (Node (y,ln,r), snd t) true
	      else
		((Node (y,ln,r), snd t), true)
	  else if (compare x y) > 0 then
	    let (rn,b) = rem_aux x r in
	      if not b then
		balance_rem (Node (y,l,rn), snd t) false
	      else
		((Node (y,l,rn), snd t), true)
	  else if (fst r) = Leaf then
	    if ((snd t) = Red) or ((snd l) = Red) then
	      ((fst l,Black),true)
	    else
	      ((fst l,Black),false)
	  else if (fst l) = Leaf then
	    if ((snd t) = Red) or ((snd r) = Red) then
	      ((fst r,Black),true)
	    else
	      ((fst r,Black),false)
	  else
	    let z = aux l in
	    let (ln,b) = rem_aux z l in
	      if not b then
		balance_rem (Node (z,ln,r), snd t) true
	      else
		((Node (z,ln,r), snd t), true)
    in fst (rem_aux x t)

  let rec to_list t = match (fst t) with
    | Leaf -> []
    | Node (y,l,r) -> (to_list l)@(y::(to_list r))

  let rec pop t = match (fst t) with
    | Leaf -> failwith "Empty tree."
    | Node (y,l,r) -> if (fst l) = Leaf then y
                      else pop l

  let sort l =
    let rec aux1 = function
      | [] -> empty
      | x::q -> add x (aux1 q)
    and aux2 t = match (fst t) with
      | Leaf -> []
      | _ -> let x = pop t in x::(aux2 (remove x t))
    in aux2 (aux1 l)
end;;

module S = SetByRBTree(Integer);;

let t = S.empty;;
let s = S.add 19 (S.add 52 (S.add 18 (S.add 32 (S.add 2 (S.add 13 (S.add 11 (S.add 9 (S.add 8 t))))))));;
let r = S.remove 13 (S.remove 9 (S.remove 52 s));;

let l = [19;52;18;32;2;13;11;9;8];;
S.sort l;;