(****************************************************************************)
(*                                                                          *)
(*                                   OCaml                                  *)
(*                                                                          *)
(*                            INRIA Rocquencourt                            *)
(*                                                                          *)
(*  Copyright  2007  Institut  National  de  Recherche en Informatique et   *)
(*  en Automatique.  All rights reserved.  This file is distributed under   *)
(*  the terms of the GNU Library General Public License, with the special   *)
(*  exception on linking described in LICENSE at the top of the Camlp4      *)
(*  source tree.                                                            *)
(*                                                                          *)
(****************************************************************************)

type variable = string
 and term =
  | Var   of variable
  | Lam   of variable * term
  | App   of term * term
  | Const of constant
 and constant =
  | CInt    of int
  | CString of string
;;

class fold = Camlp4Filters.GenerateFold.generated;;
(* class fold = Camlp4FoldGenerator.generated;; *)

module VarSet = Set.Make(String);;

(* Compute free variables with the fold class *)
let free_variables_v1 =
  let o =
    object (self)
      inherit fold as super
      val fv = VarSet.empty
      method fv = fv
      method empty_fv = {< fv = VarSet.empty >}
      method term t =
        match t with
        | Var(v) -> {< fv = VarSet.add v fv >}
        | Lam(v, t) ->
            let fv1 = VarSet.remove v (self#empty_fv#term t)#fv in
            {< fv = VarSet.union fv fv1 >}
        | _ -> super#term t
    end
  in fun t -> VarSet.elements (o#term t)#fv
;;

(* Let's try to abstract that a little *)

let fold_term f t init =
  let o =
    object (self)
      inherit fold as super
      val acc = init
      method get = acc
      method reset = {< acc = init >}
      method term t =
        {< acc = f t acc (fun t -> (self#reset#term t)#get)
                         (fun t -> (super#term t)#get) >}
    end
  in
  (o#term t)#get
;;

(* A nicer version of free_variables *)
let free_variables_v2 t =
  VarSet.elements begin
    fold_term begin fun t fv self next ->
      match t with
      | Var(v)    -> VarSet.add v fv
      | Lam(v, t) -> VarSet.union fv (VarSet.remove v (self t))
      | _         -> next t
    end t VarSet.empty
  end
;;

let term1 =
  App(
    App(Var"x1",
      Lam("x",
        App(Var"x", App(Var"y", (Lam("y", Lam("z", (App(Var"y", App(Var"x4",Var"z")))))))))),
    Var"x3")

;;

let fv1 = free_variables_v1 term1;;
let fv2 = free_variables_v2 term1;;

(* Low cost syntax *)
let ( ^-> ) v t = Lam(v, t)
let ( @ ) t1 t2 = App(t1, t2)
let ( ! ) s = Var s

let term2 =
  !"x1" @
  ("x" ^-> !"x" @ !"y" @ ("y" ^-> ("z" ^-> !"y" @ !"x4" @ !"z"))) @
  !"x3"

;;

let fv1' = free_variables_v1 term2;;
let fv2' = free_variables_v2 term2;;