(******************************************************************************
 * code for sumultaneously performing substitutions and maintaining sets
 * of Free and Bound variables for each sub-expression of a lambda term
 *****************************************************************************)

(* FV.set is the type sets of free variables, BV.set is the type of sets of
 * bound variables.  
 * We use the smlnj-lib ListSetFn implementation of sets as lists.
 *
 * all constants and routines are prefixed with "FV." or "BV."
 * constant:    empty             the empty set.
 * functions:   singleton(x)      make singleton set
 *              add(set,item)     make a new set like "set" but containing item
 *              union(s1, s2)     set union
 *              difference(s1,s2) make a new set containing s1 - s2
 *)

use "list-set-fn.sml";

structure FV = ListSetFn(struct 
                           type ord_key = string 
			   val compare = String.compare 
			 end)

structure BV = ListSetFn(struct 
                            type ord_key = string 
			    val compare = String.compare 
			 end)


datatype lam = var of string
             | apl of (lamFB * lamFB)
             | abs of (string * lamFB)
withtype lamFB = lam * FV.set * BV.set

fun termOf ((term,_,_)  : lamFB) = term
fun fvsOf  ((_, fvs, _) : lamFB) = fvs
fun bvsOf  ((_, _, bvs) : lamFB) = bvs

fun lamFBofVar x = (var x, FV.singleton(x), BV.empty)

(* 
Substitution rules:
1. x [M/x]        = M                    
2. y [M/x]        = y                    
3. (YZ) [M/x]     = (Y [M/x]) (Z [M/x])  
4. (\x.Z) [M/x]   = (\x.Z)               
5. (\y.Z) [M/x]   = (\y.(Z [M/x])), y not free in M
6. (\y.Z) [M/x]   = \w.((Z [w/y]) [M/x]), w not free in Z or M.
*)

(* freshVarString : FV.set -> string
 * guarantee: returned string is not in the argument set of strings
 *
 * Use (lamFBofVar o freshVarString) to get back a lamFB
 *)
val freshVarString =
  let val n = ref 1 
      fun fresh (fvs) =
	  let val w = "w" ^ Int.toString(! n) before n := ! n + 1
	  in 
	      if FV.member(fvs, w) then fresh (fvs)
	      else w
	  end
  in fresh
  end



(* subst : lamFB * (lamFB * string) -> lamFB
 * performs substitution and keeps free variables and bound variables 
 * up to date in the process.
 * N [M/x] is calculated by   subst (N, (M, x))
 *)
fun subst (t as (var v, _, _), (M, x)) = 
    raise Fail "Rules 1 and 2 not implemented"
  | subst _ =
    raise Fail "clauses of subst for rules 3 -- 6 not implemented"
	   
(* The following will catch a type error, but for debugging purposes it is 
 * often better to ascribe types ro the arguments and return values
 *  subst (Z : lamFB, (M : lamFB, x : string)) : lamFB = ...
 *)
val _ = op subst : lamFB * (lamFB * string) -> lamFB



(****************************************************************************
 * The following are for ease of writing test cases: use lam0 types and
 * convert them to / from lamFB types
 ***************************************************************************)

(* for checking a result. Also as a pattern for your code 
 * take a lamFB to lamFB with correct calculation of free and bound variables 
 *)
fun correctFvBvs (var x, _, _) = lamFBofVar x
  | correctFvBvs (apl (M,N), _, _) =
    let val M' = correctFvBvs M
	val N' = correctFvBvs N
    in
	(apl (M',N'), 
	 FV.union(fvsOf M', fvsOf N'), 
	 BV.union(bvsOf M', bvsOf N'))
    end
  | correctFvBvs (abs (x,M), _, _) =
    let val M' = correctFvBvs M in
	(abs (x,M'), 
	 FV.difference(fvsOf M', FV.singleton(x)),
	 BV.add(bvsOf M', x))
    end


datatype lam0 = va  of string
             | ap of (lam0 * lam0)
             | ab of (string * lam0)

(* from lam0 to lamFB for ease in writing test cases *)
fun buildFvBvs ((va x) : lam0) : lamFB = lamFBofVar x
  | buildFvBvs (ap (M,N)) =
    let val M' = buildFvBvs M
	val N' = buildFvBvs N
    in
	(apl (M',N'), 
	 FV.union(fvsOf M', fvsOf N'), 
	 BV.union(bvsOf M', bvsOf N'))
    end
  | buildFvBvs (ab (x,M)) =
    let val M' = buildFvBvs M in
	(abs (x,M'), 
	 FV.difference(fvsOf M', FV.singleton(x)),
	 BV.add(bvsOf M', x))
    end

(* from lamFB to lam0 for ease in reading back results without free/bound *)
fun stripFvBvs (var x, _, _) = va x
  | stripFvBvs (apl (M,N), _, _) = ap (stripFvBvs M, stripFvBvs N)
  | stripFvBvs (abs (x,M), _, _) = ab (x, stripFvBvs M)


(*****************************************************************************
 *
 * Some test cases
 * You should write at least one other test case.

val test1term = buildFvBvs 
  (ap (ab ("w", ab ("x", va "x")), ab ("x", ap (va "x", va "w"))));

val test1subst = buildFvBvs
 (ab("y", ap(va "y", va "w")));

val test1Result = subst(test1term, (test1subst, "x"));

val test1stripped = stripFvBvs test1Result;

test1Result = correctFvBvs test1Result;



val test2term = buildFvBvs 
  (ab ("y", ap (va "x", ab ("x", va "x"))));

val test2subst = buildFvBvs (ap (va "y", va "x"));

val test2Result = subst(test2term, (test2subst, "x"));

val test2stripped = stripFvBvs test2Result;

test2Result = correctFvBvs test2Result;
******************************************************************************)