(***************************************************************
 Sigmex adds the sigma construct to Bindex.
 ****************************************************************)

module Sigmex = struct

  open Sexp
  open List

  exception SyntaxError of string
  exception Unbound of string list 

  (************************************************************
   Abstract Syntax
   ************************************************************)

  type var = string

  type pgm = Pgm of string list * exp (* param names, body *)

  and exp =
      Lit of int (* integer literal with value *)
    | Var of var (* variable reference *)
    | BinApp of binop * exp * exp (* primitive application with rator, rands *)
    | Bind of var * exp * exp (* bind name to value of defn in body *)
    | Sigma of var * exp * exp * exp (* name * lo * hi * body *)

  and binop = 
    | Add | Sub | Mul | Div | Rem (* binary arithmetic ops *)


  (************************************************************
   Free Variables 
   ************************************************************)

  module StrSet = Set.Make(String) (* String Sets *)

  (* val addElts : StrSet.elt list -> StrSet.t -> StrSet.t *)
  let addElts strs set = (* add a list of elements to a set *)
    List.fold_right StrSet.add strs set

  (* val listToSet : StrSet.elt list -> StrSet.t *)
  let listToSet strs = addElts strs StrSet.empty

  (* val setToList : StrSet.t -> StrSet.elt list *)
  let setToList set = StrSet.elements set 

  (* val varCheck : pgm -> unit *)
  let rec varCheck pgm = 
    let unbounds = freeVarsPgm pgm 
     in if StrSet.is_empty unbounds then
          () (* OK *)
        else 
          raise (Unbound (StrSet.elements unbounds))

  (* val freeVarsPgm : pgm -> StrSet.t *)
  (* Returns the free variables of a program *)
  and freeVarsPgm (Pgm(fmls,body)) = 
    StrSet.diff (freeVarsExp body) (listToSet fmls)

  (* val freeVarsExp : exp -> StrSet.t *)
  (* Returns the free variables of an expression *)
  (* direct version *)
  and freeVarsExp e = 
    match e with 
      Lit i -> StrSet.empty
    | Var s -> StrSet.singleton s
    | BinApp(_,r1,r2) -> freeVarsExps [r1;r2]
    | Bind(name,defn,body) -> 
        StrSet.union (freeVarsExp defn)
                     (StrSet.diff (freeVarsExp body)
                                  (StrSet.singleton name))
    | Sigma(name,lo,hi,body) -> StrSet.empty (* replace this stub for 2a *)

  (* val freeVarsExps : exp list -> StrSet.t *)
  (* Returns the free variables of a list of expressions *)
  (* direct version *)
  and freeVarsExps es = 
    List.fold_right StrSet.union (List.map freeVarsExp es) StrSet.empty


  (************************************************************
   Substitution & Renaming
   ************************************************************)

  (* val subst : exp -> exp Env.env -> exp *)
  let rec subst exp env = 
    match exp with 
      Lit i -> exp
    | Var v -> (match Env.lookup v env with 
	          Some e -> e
                | None -> exp)
    | BinApp(op,r1,r2) -> BinApp(op, subst r1 env, subst r2 env)
    | Bind(name,defn,body) -> 
        (* Take the simple approach of renaming every name. 
           With more work, could avoid renaming unless absolutely necessary. *)
	let name' = StringUtils.fresh name in 
          Bind(name', subst defn env, subst (rename1 name name' body) env)
	  (* note: could be cleverer and do a single substitution/renaming *)
    | Sigma(name,lo,hi,body) -> exp (* replace this stub for 2b *)

  (* val subst1 : var -> exp -> exp -> exp *)
  and subst1 name newexp exp = subst exp (Env.make [name] [newexp])

  (* val subst_all : string list -> exp list -> exp -> exp *)
  and subst_all names newexps exp = subst exp (Env.make names newexps)

  (* val rename1 : var -> var -> exp -> exp *)
  and rename1 oldname newname exp = subst1 oldname (Var newname) exp 

  (* val rename_all : string list -> var list -> exp -> exp *)
  and rename_all oldnames newnames exp = 
    subst_all oldnames (List.map (fun s -> Var s) newnames) exp 

  (************************************************************
   Parsing from S-Expressions 
   ************************************************************)

  (* val sexpToPgm : Sexp.sexp -> pgm *)
  let rec sexpToPgm sexp = 
    match sexp with 
      Seq [Sym "sigmex"; Seq formals; body] -> 
            Pgm(List.map symToString formals, sexpToExp body)
    | _ -> raise (SyntaxError ("invalid Sigmex program: " ^ (sexpToString sexp)))

  (* val symToString : Sexp.sexp -> string *)
  and symToString sexp =
    match sexp with 
      Sym s -> s
    | _ -> raise (SyntaxError ("symToString: not a string -- " ^ (sexpToString sexp)))

  (* val sexpToExp : Sexp.sexp -> exp *)
  and sexpToExp sexp = 
    match sexp with 
      Int i -> Lit i
    | Sym s -> Var s
    | Seq [Sym p; rand1; rand2] -> 
        BinApp(stringToBinop p, sexpToExp rand1, sexpToExp rand2)
    | Seq [Sym "bind"; Sym name; defn; body] -> 
        Bind (name, sexpToExp defn, sexpToExp body)
    (* add a clause here for 2c *)
    | _ ->  raise (SyntaxError ("invalid Sigmex expression: "  
                                ^ (sexpToString sexp)))

  (* val stringToBinop : string -> binop *)
  and stringToBinop s = 
    match s with 
    | "+" -> Add
    | "-" -> Sub
    | "*" -> Mul
    | "/" -> Div
    | "%" -> Rem
    |  _ -> raise (SyntaxError ("invalid Sigmex primop: " ^ s))

  (* val stringToExp : string -> exp *)
  and stringToExp s = sexpToExp (stringToSexp s) 

  (* val stringToPgm : string -> pgm *)
  and stringToPgm s = sexpToPgm (stringToSexp s)

  (************************************************************
   Unparsing to S-Expressions
   ************************************************************)

  (* val pgmToSexp : pgm -> Sexp.sexp *)
  let rec pgmToSexp p = 
    match p with 
      Pgm (fmls, e) -> 
        Seq [Sym "sigmex"; Seq(List.map (fun s -> Sym s) fmls); expToSexp e]

  (* val expToSexp : exp -> Sexp.sexp *)
  and expToSexp e = 
    match e with 
      Lit i -> Int i
    | Var s -> Sym s
    | BinApp (rator, rand1, rand2) -> 
        Seq [Sym (binopToString rator); expToSexp rand1; expToSexp rand2]
    | Bind(n,d,b) -> Seq [Sym "bind"; Sym n; expToSexp d; expToSexp b]
    | Sigma(name,lo,hi,body) -> Int 17 (* replace this stub for 2d *)

  (* val binopToString : binop -> string *)
  and binopToString p = 
    match p with 
    | Add -> "+"
    | Sub -> "-" 
    | Mul -> "*" 
    | Div -> "/" 
    | Rem -> "%" 

  (* val expToString : exp -> string *)
  and expToString s = sexpToString (expToSexp s)

  (* val pgmToString : pgm -> string *)
  and pgmToString s = sexpToString (pgmToSexp s)

end