(*
copyright (c) 2013, simon cruanes
all rights reserved.

redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

redistributions of source code must retain the above copyright notice, this
list of conditions and the following disclaimer.  redistributions in binary
form must reproduce the above copyright notice, this list of conditions and the
following disclaimer in the documentation and/or other materials provided with
the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*)

(** {1 Rational Terms} *)

module type SYMBOL = sig
  type t
  val compare : t -> t -> int
  val to_string : t -> string
end

module type S = sig
  module Symbol : SYMBOL

  type t = private
    | Var of int
    | Ref of int
    | App of Symbol.t * int option * t list

  type term = t
  
  (** Structural equality and comparisons. Two terms being different
      for {!eq} may still be equal, but with distinct representations.
      For instance [r:f(f(r))] and [r:f(r)] are the same term but they
      are not equal structurally. *)

  val eq : t -> t -> bool
  val cmp : t -> t -> int

  val var : unit -> t
    (** free variable, with a fresh name *)

  val app : Symbol.t -> t list -> t
    (** Application of a symbol to a list, possibly with a unique label *)

  val app_fix : Symbol.t -> (t -> t list) -> t
    (** Same as {!app}, but the arguments are generated by a function
        that takes a reference to the result as a parameter. *)

  val pp : Buffer.t -> t -> unit
  val fmt : Format.formatter -> t -> unit
  val to_string : t -> string

  val rename : t -> t
    (** Rename all variables and references to fresh ones *)

  module Subst : sig
    type t
    val empty : t
    val bind : t -> int -> term -> t
    val deref : t -> term -> term
    val apply : t -> term -> term

    val pp : Buffer.t -> t -> unit
    val fmt : Format.formatter -> t -> unit
    val to_string : t -> string
  end

  val matching : ?subst:Subst.t -> term -> term -> Subst.t option
  val unify : ?subst:Subst.t -> term -> term -> Subst.t option
end

module Make(Symbol : SYMBOL) = struct
  module Symbol = Symbol

  type t =
    | Var of int
    | Ref of int
    | App of Symbol.t * int option * t list

  type term = t

  module IMap = Map.Make(struct
    type t = int
    let compare i j = i-j
  end)
  module IHTbl = Hashtbl.Make(struct
    type t = int
    let equal i j = i=j
    let hash i = i land max_int
  end)

  let _add_ref map t = match t with
    | App (_, Some i, _) -> IMap.add i t map
    | _ -> map

  let _to_int = function
    | Var _ -> 1
    | Ref _ -> 2
    | App _ -> 3

  let rec cmp t1 t2 = match t1, t2 with
    | Var i1, Var i2 -> i1 - i2
    | Ref i1, Ref i2 -> i1 - i2
    | App (f1, r1, l1), App (f2, r2, l2) ->
        let c = Symbol.compare f1 f2 in
        if c <> 0 then c
        else cmp_list l1 l2
    | _ -> _to_int t1 - _to_int t2
  and cmp_list l1 l2 = match l1, l2 with
    | [], [] -> 0
    | [], _ -> -1
    | _, [] -> 1
    | t1::l1', t2::l2' ->
        let c = cmp t1 t2 in
        if c <> 0 then c else cmp_list l1' l2'

  let eq t1 t2 = cmp t1 t2 = 0

  let _count = ref 0

  let var () =
    let v = Var !_count in
    incr _count;
    v

  let app s l = App (s, None, l)

  let app_fix s f =
    let i = !_count in
    incr _count;
    let l = f (Ref i) in
    match l with
      | [] -> App (s, None, [])
      | _ -> App (s, Some i, l)

  let rec pp buf t = match t with
    | Var i -> Printf.bprintf buf "X%d" i
    | Ref i -> Printf.bprintf buf "*%d" i
    | App (_, Some _, []) -> assert false
    | App (s, None, []) ->
        Buffer.add_string buf (Symbol.to_string s)
    | App (s, Some r, l) ->
        Printf.bprintf buf "%d: %s(%a)" r (Symbol.to_string s) pp_list l
    | App (s, None, l) ->
        Printf.bprintf buf "%s(%a)" (Symbol.to_string s) pp_list l
  and pp_list buf l = match l with
    | [] -> ()
    | [x] -> pp buf x
    | x::((_::_) as l') -> 
      pp buf x; Buffer.add_string buf ", "; pp_list buf l'

  let to_string t =
    let b = Buffer.create 16 in
    pp b t;
    Buffer.contents b

  let fmt fmt t = Format.pp_print_string fmt (to_string t)

  let rename t =
    let names = IHTbl.create 16 in
    let rec rename back t = match t with
      | Var i ->
          begin try IHTbl.find names i
          with Not_found ->
            (* rename variable into a fresh one *)
            let v = var() in
            IHTbl.add names i v;
            v
          end
      | Ref i ->
          begin try IMap.find i back
          with Not_found -> failwith "ill-formed term"
          end
      | App (s, None, l) ->
          app s (List.map (rename back) l)
      | App (s, Some r, l) ->
          app_fix s
            (fun r' ->
              let back' = IMap.add r r' back in
              List.map (rename back') l)
    in rename IMap.empty t

  module Subst = struct
    type t = term IMap.t

    let empty = IMap.empty

    let bind s i t =
      match t with
        | _ when IMap.mem i s -> failwith "Subst.bind"
        | Var j when i=j -> s  (* id *)
        | _ -> IMap.add i t s

    let rec deref s t = match t with
      | Var i ->
          begin try deref s (IMap.find i s)
          with Not_found -> t
          end
      | Ref _
      | App _ -> t

    (* does the variable [v] occur in [subst(t)]? *)
    let rec _occur subst ~var t =
      match deref subst t with
      | Var _ -> eq var t
      | Ref _
      | App (_, _, []) -> false
      | App (_, _, l) -> List.exists (_occur subst ~var) l

    let apply subst t =
      let names = IHTbl.create 16 in
      let rec apply subst t = match t with
        | Ref i ->
            begin try
              IHTbl.find names i
            with Not_found -> failwith "Subst.apply: invalid term"
            end
        | Var i ->
            let t' = deref subst t in
            (* interesting case. Either [t] is bound to a term [t']
               that contains it, which makes a cyclic term, or it's
               not in which case it's easy. *)
            begin match t' with
              | Ref _ -> t'
              | App (s, Some r, l) ->
                if _occur subst ~var:t t'
                then
                  (* in any case we are possibly going to modify [r']
                     by replacing [x] by a backref. Therefore we need a
                     new backref to designate the new term. *)
                  app_fix s
                    (fun r' ->
                      (* alias r -> r' *)
                      IHTbl.add names r r';
                      (* now [x] designates the cyclic term, ie [r'] *)
                      let subst' = IMap.add i r' subst in
                      List.map (apply subst') l)
                else
                  (* simply evaluate [t'] *)
                  apply subst t'
              | App (s, None, l) ->
                (* same problem: does [x] introduce a cycle? *)
                if _occur subst ~var:t t'
                then
                  app_fix s
                    (fun r' ->
                      let subst' = IMap.add i r' subst in
                      List.map (apply subst') l)
                else apply subst t'
              | Var j ->
                assert (not (IMap.mem j subst));
                t'
            end
        | App (s, None, l) ->
            app s (List.map (apply subst) l)
        | App (s, Some r, l) ->
            begin try
              (* alias? *)
              IHTbl.find names r
            with Not_found ->
              app_fix s
                (fun r' ->
                  IHTbl.add names r r';
                  List.map (apply subst) l)
            end
      in apply subst t

    let pp buf subst =
      Buffer.add_string buf "{";
      let first = ref true in
      IMap.iter
        (fun i t ->
          if !first then first:= false else Buffer.add_string buf ", ";
          Printf.bprintf buf "X%d → %a" i pp t)
        subst;
      Buffer.add_string buf "}";
      ()

    let to_string t =
      let b = Buffer.create 16 in
      pp b t;
      Buffer.contents b

    let fmt fmt t = Format.pp_print_string fmt (to_string t)
  end

  exception Fail

  let _add_tbl tbl r t = match r with
    | None -> ()
    | Some r -> IHTbl.replace tbl r t

  let matching ?(subst=Subst.empty) t1 t2 =
    assert false  (* TODO (need to gather variables of [t2]... *)

  let unify ?(subst=Subst.empty) t1 t2 =
    let names = IHTbl.create 16 in
    let rec unif subst t1 t2 =
      match Subst.deref subst t1, Subst.deref subst t2 with
      | Var i, Var j when i=j -> subst
      | Var i, _ -> Subst.bind subst i t2
      | _, Var j -> Subst.bind subst j t1
      | Ref i1, _ -> unif subst (IHTbl.find names i1) t2
      | _, Ref i2 -> unif subst t1 (IHTbl.find names i2)
      | App (f1, r1, l1), App (f2, r2, l2) ->
        if Symbol.compare f1 f2 <> 0 then raise Fail;
        _add_tbl names r1 t1;
        _add_tbl names r2 t2;
        try
          List.fold_left2 unif subst l1 l2
        with Invalid_argument _ -> raise Fail
    in
    try Some (unif subst t1 t2)
    with Fail -> None
end

module Str = struct
  type t = string
  let compare = String.compare
  let to_string s = s
end

module Default = Make(Str)