(*
This file is part of licht.

licht is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

licht is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with licht.  If not, see <http://www.gnu.org/licenses/>.
*)

module type KEY = sig

  type 'a t

  (** Parametrized comparator *)
  val comp: 'a t -> 'b t -> ('a, 'b) Tools.cmp

  val repr: Format.formatter -> 'a t -> unit

end

module Make (El : KEY) = struct

  type container = C : ('a El.t * 'a) -> container [@@unboxed]

  type leaf (** Fantom type for typing the tree *)
  type node (** Fantom type for typing the tree *)

  type 'a branch =
    | Leaf : leaf branch
    | Node : _ branch * ('a El.t * 'a) * _ branch -> node branch

  type t = T : 'a branch ref -> t [@@unboxed]

  let empty = T (ref Leaf)

  let isEmpty (T tree) = match !tree with
    | Leaf -> true
    | _ -> false

  let rec splay : type a. a El.t -> node branch -> node branch = fun x t -> begin
    let Node (l, y, r) = t in
    begin match El.comp x (fst y)  with
      | Tools.Eq -> t
      | Tools.Lt ->
        begin match l with
          | Leaf -> t
          | Node (ll, z, rr) ->
            begin match El.comp x (fst z)  with
              | Tools.Eq -> Node (ll, z, Node (rr, y, r))
              | Tools.Lt ->
                begin match ll with
                  | Leaf -> Node (ll, z, Node (rr, y, r))
                  | Node _ as ll ->
                    let Node (newL, newV, newR) = splay x ll
                    in Node (newL, newV, Node (newR, z, Node (rr, y, r)))
                end
              | Tools.Gt ->
                begin match rr with
                  | Leaf -> Node (ll, z, Node (rr, y, r))
                  | Node _ as rr ->
                    let Node (newL, newV, newR) = splay x rr
                    in Node (Node (ll, z, newL), newV, Node (newR, y, r))
                end
            end
        end
      | Tools.Gt ->
        begin match r with
          | Leaf -> t
          | Node (ll, z, rr) ->
            begin match El.comp x (fst z)  with
              | Tools.Eq -> Node (Node (l, y, ll), z, rr)
              | Tools.Lt ->
                begin match ll with
                  | Leaf -> Node (Node (l, y, ll), z, rr)
                  | Node _ as ll ->
                    let Node (newL, newV, newR) = splay x ll
                    in Node (Node (l, y, newL), newV, Node (newR, z, rr))
                end
              | Tools.Gt ->
                begin match rr with
                  | Leaf -> Node (Node (l, y, ll), z, rr)
                  | Node _ as rr ->
                    let Node (newL, newV, newR) = splay x rr
                    in Node (Node (Node(l, y, ll), z, newL), newV, newR)
                end
            end
        end
    end
  end

  let member: type a. a El.t -> t -> bool = fun x (T t) -> match !t with
    | Leaf -> false
    | Node _ as root ->
        let root' = splay x root in
        t := root';
        let Node (_, c', _) = root' in
        begin match El.comp (fst c') x  with
          | Tools.Eq -> true
          | _ -> false
        end

  let find: type a. a El.t -> t -> a = fun x (T t) -> match !t with
    | Leaf -> raise Not_found
    | Node _ as root ->
        let root' = splay x root in
        t := root';
        let Node (_, c', _) = root' in
        begin match El.comp (fst c') x with
          | Tools.Eq -> snd c'
          | _ -> raise Not_found
        end

  let add: type a. a El.t -> a -> t -> t = fun key value (T t) -> match !t with
    | Leaf -> T (ref (Node (Leaf, (key, value), Leaf)))
    | Node _ as root ->
        let root' = splay key root in
        let Node (l, y, r) = root' in
        begin match El.comp key (fst y)  with
          | Tools.Eq -> T (ref (Node(l, (key, value), r)))
          | Tools.Lt -> T (ref (Node (l, (key, value), Node (Leaf, y, r))))
          | Tools.Gt -> T (ref (Node (Node (l, y, Leaf), (key, value), r)))
        end

  let rec _subtree_maximum:type a. a branch -> a branch = fun t -> begin match t with
  | Leaf -> Leaf
  | Node (_, _, (Node (_, _, _) as x)) -> _subtree_maximum x
  | Node (_, (key, value), Leaf) -> splay key t
  end

  let rec _subtree_minimum: type a. a branch -> a branch = fun t -> begin match t with
  | Leaf -> Leaf
  | Node ((Node (_, _, _) as x), _, _) -> _subtree_minimum x
  | Node (Leaf, (key, value), _) -> splay key t
  end

  let remove: type a. a El.t -> t -> t = fun key (T t) -> begin match !t with
    | Leaf -> empty
    | Node _ as root ->
      let root' = splay key root in
      let Node (l, c', r) = root' in
      begin match El.comp (fst c') key with
        | Tools.Eq -> begin match _subtree_maximum l with
          | Node(l, c, Leaf) -> T (ref (Node(l, c, r)))
          | Node(l, c, _) -> raise Not_found
          | Leaf -> begin match _subtree_minimum r with
            | Leaf -> empty
            | Node(Leaf, c, r) -> T (ref (Node(l, c, r)))
            | Node(_, c, r) -> raise Not_found
          end
        end
        (* The key is not present, return the splayed tree *)
        | _ -> T (ref root')
      end
  end

  (** Existencial type for the branches *)
  type exBranch = Branch : _ branch -> exBranch [@@unboxed]

  let fold f init (T t) = begin
    let rec _fold : type b. (container * exBranch) list -> 'a -> b branch -> 'a = begin
      fun acc v -> function
    (* We have a node : we accumulate the right part, and process the left branch *)
    | Node (left, (key, value), right) ->
        let c = C (key, value) in
        (_fold [@tailcall]) ((c, Branch right)::acc) v left
    (* We have nothing left, we process the values delayed *)
    | Leaf -> begin match acc with
      | [] -> v
      | (c, (Branch right))::tl -> (_fold [@tailcall]) tl (f v c) right
      end
    end in
    _fold [] init !t
  end

  let choose (T tree) = begin match (_subtree_minimum !tree) with
  | Leaf -> raise Not_found
  | Node (left, (key, value), right) -> C (key, value)
  end

  let repr formatter (T t) = begin

    let repr_edge from formatter dest = begin
      Format.fprintf formatter "\"%a\" -> \"%a\"\n"
        El.repr from
        El.repr dest
    end in

    let rec repr': type a b. a El.t -> Format.formatter -> b branch -> unit = fun parent formatter -> function
    | Leaf -> ()
    | Node (l, c, r) ->
        let key = fst c in
        Format.fprintf formatter "%a%a%a"
        (repr_edge parent) key
        (repr' key) l
        (repr' key) r in

    begin match !t with
    | Leaf -> Format.fprintf formatter "digraph G {}"
    | Node (l, c, r) ->
      let key = fst c in
      Format.fprintf formatter "digraph G {\n%a%a}"
        (repr' key) l
        (repr' key) r
    end

  end

end