(***************************************************************
 Ibex adds booleans, conditionals, and generalized primitive 
 application to Bindex. 
 ****************************************************************)

module Ibex = struct

  open Sexp
  open List
  open FunUtils

  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 valu (* integer & boolean literals *)
    | Var of var (* variable reference *)
    | PrimApp of primop * exp list (* primitive application with rator, rands *)
    | Bind of var * exp * exp (* bind name to value of defn in body *)
    | If of exp * exp * exp (* conditional with test, then, else *)

  and valu = Int of int | Bool of bool

  and primop = 
    | Add | Sub | Mul | Div | Rem (* binary arithmetic ops *)
    | LT | LE | EQ | NE | GE | GT (* binary relational ops *)
    | And | Or (* binary logical ops ( *not* short-circuit!) *)
    | Not (* unary logical negation *)

  (************************************************************
   Folding over Ibex Expressions
   ************************************************************)

  (* val fold :
        (int -> 'a) ->
        (var -> 'a) ->
        (primop -> 'a list -> 'a)
        (var -> 'a -> 'a -> 'a) -> 
        ('a -> 'a -> 'a -> 'a) -> 
        -> exp -> 'a *)
  let rec fold litfun varfun appfun bindfun iffun exp = 
    let fold' e = fold litfun varfun appfun bindfun iffun e in 
      match exp with 
        Lit i -> litfun i 
      | Var s -> varfun s
      | Bind(name,defn,body) -> bindfun name (fold' defn) (fold' body)
      | PrimApp(op, rands) -> appfun op (map fold' rands)
      | If(test,thn,els) -> iffun (fold' test) (fold' thn) (fold' els)

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

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

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

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

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

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

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

  (* val freeVarsExp : exp -> S.t *)
  (* Returns the free variables of an expression *)
  (* direct version *)
  and freeVarsExp e = 
    match e with 
      Lit i -> S.empty
    | Var s -> S.singleton s
    | PrimApp(_,rands) -> freeVarsExps rands
    | Bind(name,defn,body) -> 
        S.union (freeVarsExp defn)
                     (S.diff (freeVarsExp body)
                                  (S.singleton name))
    | If(tst,thn,els) -> freeVarsExps [tst;thn;els]

  (* val freeVarsExps : exp list -> S.t *)
  (* Returns the free variables of a list of expressions *)
  (* direct version *)
  and freeVarsExps es = 
    fold_right S.union (map freeVarsExp es) S.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)
    | PrimApp(op,rands) -> PrimApp(op, map (flip subst env) rands)
    | 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 *)
    | If(tst,thn,els) -> If(subst tst env, subst thn env, subst els env)

  (* 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 (map (fun s -> Var s) newnames) exp 

  (************************************************************
   Desugaring (val desugar : sexp -> sexp)
   ************************************************************)

  let rec desugar sexp = 
    let sexp' = desugarRules sexp in 
      if sexp' = sexp then (* efficient in OCAML if they're pointer equivalent *)
        match sexp with 
          Seq sexps -> Seq (map desugar sexps)
        | _ -> sexp
      else desugar sexp'

  and desugarRules sexp = 
    match sexp with 
    (* Note: the following desugarings for && and || allow 
       non-boolean expressions for second argument! *)
      Seq [Sym "&&"; x; y] -> Seq [Sym "if"; x; y; Sym "false"]
    | Seq [Sym "||"; x; y] -> Seq [Sym "if"; x; Sym "true"; y]
    | Seq [Sym "bindseq"; Seq[]; body] -> body
    | Seq [Sym "bindseq"; Seq ((Seq[Sym name;defn])::bindings); body]
        -> Seq[Sym "bind"; Sym name; defn; Seq[Sym "bindseq"; Seq bindings; body]]
    (* Note: can't handle bindpar here, because it requires renaming *)
    (* See sexpToExp' below for handling bindpar *)
    | Seq [Sym "cond"; Seq [Sym "else"; default]] -> default
    | Seq (Sym "cond" :: Seq [test; body] :: clauses) -> 
        Seq [Sym "if"; test; body; Seq(Sym "cond" :: clauses)]
    | Seq [Sym "$"; Sexp.Int i] -> Sym ("$" ^ (string_of_int i)) 
       (* Handle Intex arg refs as var refs *)
    | _ -> sexp

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

  (* val sexpToPgm : Sexp.sexp -> pgm *)
  let rec sexpToPgm sexp = 
    match sexp with 
      Seq [Sym "ibex"; Seq formals; body] -> 
            Pgm(map symToString formals, sexpToExp body)
    (* Handle Ibex programs as well *)
    | Seq [Sym "bindex"; Seq formals; body] -> 
            Pgm(map symToString formals, sexpToExp body)
    (* Handle Intex programs as well *)
    | Seq [Sym "intex"; Sexp.Int n; body] -> 
            Pgm(map 
                  (fun i -> "$" ^ (string_of_int i))
                  (ListUtils.range 1 n),
                sexpToExp body)
    | _ -> raise (SyntaxError ("invalid Ibex 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)))

  and sexpToExp sexp = sexpToExp' (desugar sexp)

  (* val sexpToExp : Sexp.sexp -> exp *)
  and sexpToExp' sexp = 
    match sexp with 
      Sexp.Int i -> Lit (Int i)
    | Sym "true" -> Lit (Bool true)    (* true and false are keywords *)
    | Sym "false" -> Lit (Bool false)  (* for literals, not variables *)
    | Sym s -> Var s
    | Seq [Sym "bind"; Sym(name); defn; body] -> 
        Bind (name, sexpToExp' defn, sexpToExp' body)
    | Seq [Sym "if"; tst; thn; els] -> 
        If(sexpToExp' tst, sexpToExp' thn, sexpToExp' els)
    | Seq [Sym "bindpar"; Seq bindings; body] -> 
	let (names, defns) = 
          ListUtils.unzip (map (fun binding -> 
	                         (match binding with 
                                   Seq[Sym name; defn] -> (name, sexpToExp' defn)
				 | _ -> raise (SyntaxError ("ill-formed bindpar binding"
							    ^ (sexpToString binding)))))
                                bindings) in 
	desugarBindpar names defns (sexpToExp' body)
    (* This clause must be last! *)
    | Seq (Sym p :: rands) -> 
        PrimApp(stringToPrimop p, map sexpToExp' rands)

    | _ ->  raise (SyntaxError ("invalid Ibex expression: "  
                                ^ (sexpToString sexp)))

  (* Can't handle desugarings with renamings in desugar *)
  and desugarBindpar names defns body = 
    let freshNames = map StringUtils.fresh names in 
      fold_right2 (fun n d b -> Bind(n,d,b))
                  freshNames
                  defns
                  (rename_all names freshNames body)

  (* val stringToPrimop : string -> primop *)
  and stringToPrimop s = 
    match s with 
    | "+" -> Add
    | "-" -> Sub
    | "*" -> Mul
    | "/" -> Div
    | "%" -> Rem
    | "<" -> LT
    | "<=" -> LE
    | "=" -> EQ
    | "!=" -> NE
    | ">=" -> GE
    | ">" -> GT
    | "and" -> And (* non-short-circuit *)
    | "or" -> Or (* non-short-circuit *)
    | "!" -> Not 
    |  _ -> raise (SyntaxError ("invalid Ibexx primop: " ^ s))

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

  (* 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 "ibex"; Seq(map (fun s -> Sym s) fmls); expToSexp e]

  (* val expToSexp : exp -> Sexp.sexp *)
  and expToSexp e = 
    match e with 
      Lit (Int i) -> Sexp.Int i
    | Lit (Bool b) -> Sym (if b then "true" else "false")
    | Var s -> Sym s
    | PrimApp (rator, rands) -> 
        Seq (Sym (primopToString rator) :: map expToSexp rands)
    | Bind(n,d,b) -> Seq [Sym "bind"; Sym n; expToSexp d; expToSexp b]
    | If(tst,thn,els) -> Seq [Sym "if"; expToSexp tst; expToSexp thn; expToSexp els]

  (* val primopToString : primop -> string *)
  and primopToString p = 
    match p with 
    | Add -> "+"
    | Sub -> "-" 
    | Mul -> "*" 
    | Div -> "/" 
    | Rem -> "%" 
    | LT -> "<" 
    | LE -> "<=" 
    | EQ -> "=" 
    | NE -> "!=" 
    | GE -> ">=" 
    | GT -> ">" 
    | Not -> "!" 
    | And -> "and" 
    | Or -> "or" 

  (* val valuToString : valu -> string *)
  and valuToString valu = 
    match valu with 
      Int i -> string_of_int i
    | Bool b -> string_of_bool b

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

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

end