(***************************************************************
 HOFL adds abstractions, applications, recursion, pairs, and lists to Ibex. 
 Binds are no longer necessary as kernel constructs
   (they desugar into applications of abstractions.)

(* still to do: 
   + change read_line() in other repls to readSexp
   + abstract over repl
   + better repl (with top-level definitions and ability to read more than one line)
 *)
 ****************************************************************)

module Hofl = struct

  open Sexp
  open List
  open FunUtils
  open StringUtils

  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 *)
    | If of exp * exp * exp (* conditional with test, then, else *)
    | Abs of var * exp (* function abstraction *)
    | App of exp * exp (* function application *)
    | Bindrec of var list * exp list * exp (* recursive bindings *)

  and valu = 
      Int of int 
    | Bool of bool
    | Fun of var * exp * valu Env.env
    | Pair of valu * valu
    | List of valu list

  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 *)
    | PairFn | Left | Right (* pairs *)
    | Prepend | Head | Tail | Null | IsNull (* lists *)
    | IsInt  | IsBool | IsFun | IsPair | IsList (* predicates *)

  (************************************************************
   Folding over Hofl Expressions
   ************************************************************)
(*
  (* Not updated yet *)
  (* 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
      | 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
    | If(tst,thn,els) -> freeVarsExps [tst;thn;els]
    | Abs(fml,body) -> S.diff (freeVarsExp body) (S.singleton fml)
    | App(rator,rand) -> freeVarsExps [rator;rand]
    | Bindrec(names,defns,body) -> 
	S.diff (S.union (freeVarsExps defns) (freeVarsExp body))
               (listToSet names)

  (* 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)
    | If(tst,thn,els) -> If(subst tst env, subst thn env, subst els env)
    | Abs(fml,body) -> 
	let fml' = fresh fml in 
	  Abs(fml', subst (rename1 fml fml' body) env)
    | App(rator,rand) -> App(subst rator env, subst rand env)
    | Bindrec(names,defns,body) -> 
	let names' = map fresh names in 
	  Bindrec(names', 
		  map (rename_all names names') defns, 
		  rename_all names names' body)

  (* 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 
	  (* Special handling for constructs with specially interpreted parens *)
	  Seq [Sym "bindrec"; Seq bindings; body] -> 
	    let (names,defns) = 
	      ListUtils.unzip (map (fun binding -> 
		                      match binding with 
					Seq [Sym name; defn] -> (name,defn)
				      |	_ -> raise (SyntaxError ("ill-formed bindrec binding")))
				   bindings)
	    in Seq [Sym "bindrec"; 
		     Seq (map2 (fun n d -> Seq [Sym n; desugar d]) names defns);
		     desugar body]
        | 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]
    (* All forms of bind desugar into other HOFL forms *)
    | Seq [Sym "bind"; Sym name; defn; body] -> 
	Seq [Seq[Sym "abs"; Sym name; body]; defn]
    | 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]]
    (* In Hofl, *can* handle bindpar as an s-expression desugaring *)
    | Seq [Sym "bindpar"; Seq bindings; body] -> 
	let (names, defns) = 
          ListUtils.unzip (map (fun binding -> 
	                         (match binding with 
                                   Seq[Sym name; defn] -> (name, defn)
				 | _ -> raise (SyntaxError ("ill-formed bindpar binding"
							    ^ (sexpToString binding)))))
                             bindings) in 
	Seq (Seq[Sym "fun"; Seq(map (fun n -> Sym n)  names); body] :: defns)
    (* LET and LET* are synonyms for BINDPAR and BINDSEQ *)
    | Seq [Sym "let"; bindings; body] -> Seq [Sym "bindpar"; bindings; body] 
    | Seq [Sym "let*"; bindings; body] -> Seq [Sym "bindseq"; bindings; body] 
    (* COND is a Scheme-like multi-branch conditional *)
    | Seq [Sym "cond"; Seq [Sym "else"; default]] -> default
    | Seq (Sym "cond" :: Seq [test; body] :: clauses) -> 
        Seq [Sym "if"; test; body; Seq(Sym "cond" :: clauses)]
    (* Handle Intex arg refs as var refs *)
    | Seq [Sym "$"; Sexp.Int i] -> Sym ("$" ^ (string_of_int i)) 
    (* Sugar for lists *) 
    | Seq [Sym "list"] -> Seq [Sym "null"]
    | Seq (Sym "list" :: x :: xs) -> Seq [Sym "prepend"; x; Seq (Sym "list" :: xs)]
    (* Use the "fun" syntax for multiple argument functions (just curried functions) *)
    | Seq [Sym "fun"; Seq []; body] -> Seq [Sym "abs"; Sym (fresh "ignore"); body]
    | Seq [Sym "fun"; Seq [Sym fml]; body] -> Seq [Sym "abs"; Sym fml; body]
    | Seq [Sym "fun"; Seq (Sym fml :: formals) ; body] -> 
	Seq [Sym "abs"; Sym fml; Seq [Sym "fun"; Seq formals ; body]]
    (* Desugar multiple argument function calls into nested single-argument calls *)
    | Seq (sexp1 :: sexp2 :: sexp3 :: rest) when (not (isKeyword sexp1)) -> 
       Seq(Seq[sexp1; sexp2] :: sexp3 :: rest)
    (* Desugar nullary function call into application to false *)
    | Seq [sexp1] when (not (isKeyword sexp1)) -> Seq[sexp1; Sym "false"]
    | _ -> sexp

  and isSpecial s = 
    List.mem s ["if"; "abs"; "&&"; "||"; "fun"; "bind"; "bindseq"; "bindpar"; 
		 "bindrec"; "let"; "let*"; "cond"; "fun"]

  and isKeyword sexp = 
    match sexp with 
      Sym s -> isSpecial s || isPrimop s
    | _ -> false

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

  (* val sexpToPgm : Sexp.sexp -> pgm *)
  and sexpToPgm sexp = 
    match sexp with 
      Seq (Sym "hofl" :: Seq formals :: body :: defs) -> 
	let (defNames,defExps) = ListUtils.unzip (map defToBinding defs)
	in Pgm(map symToString formals, 
	       Bindrec(defNames, defExps, sexpToExp body))
    (* Handle Ibex programs as well *)
    | Seq [Sym "ibex"; Seq formals; body] -> 
            Pgm(map symToString formals, sexpToExp body)
    (* Handle Bindex 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 Hofl program: " ^ (sexpToString sexp)))

  (* val defToBinding: Sexp.sexp -> string * Hofl.exp *)
  (* Defs have the form (def I_name I_defn). 
     Defs of the form (def (I_0 I_1 .. I_n) I_body) are treated
       as sugar for (def I_0 (fun (I_1 .. I_n) I_body))
   *)
  and defToBinding def = 
    match desugarDef def with 
      Seq [Sym "def"; Sym name; defn] -> (name, sexpToExp defn)
    | _ -> raise (SyntaxError 
		    ("ill-formed def: "
		     ^ (sexpToString def)))

  and desugarDef def = 
    match def with 
      Seq [Sym "def"; Sym name; defn] -> Seq [Sym "def"; Sym name; desugar defn]
    | Seq [Sym "def"; Seq (Sym name :: formals); body] -> 
	Seq [Sym "def"; Sym name; desugar (Seq [Sym "fun"; Seq formals; body])]
    | _ -> raise (SyntaxError 
		    ("ill-formed def: "
		     ^ (sexpToString def)))

  (* 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 "if"; tst; thn; els] -> 
        If(sexpToExp' tst, sexpToExp' thn, sexpToExp' els)
    | Seq [Sym "abs"; Sym fml; body] -> Abs(fml, sexpToExp' body)
    | Seq [Sym "bindrec"; 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 bindrec binding"
							    ^ (sexpToString binding)))))
                                bindings) in 
	Bindrec(names,defns,sexpToExp' body)
    (* This clause must be last! *)
    | Seq (rator :: rands) -> 
        (match (rator,rands) with 
           (Sym p, _) when isPrimop p -> 
	     PrimApp(stringToPrimop p, map sexpToExp' rands)
	 | (_, [rand]) -> App(sexpToExp' rator, sexpToExp' rand)
	 | _ -> raise (SyntaxError ("invalid Hofl application: "
                                    ^ (sexpToString sexp)))
	  )
    | _ ->  raise (SyntaxError ("invalid Hofl expression: "  
                                ^ (sexpToString sexp)))

  (* 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 
    | "pair" -> PairFn
    | "left" -> Left 
    | "right" -> Right 
    | "prepend" -> Prepend 
    | "head" -> Head 
    | "tail" -> Tail 
    | "null" -> Null 
    | "null?" -> IsNull 
    | "int?" -> IsInt  
    | "bool?" -> IsBool 
    | "fun?" -> IsFun
    | "pair?" -> IsPair 
    | "list?" -> IsList 
    |  _ -> raise (SyntaxError ("invalid Hofl primop: " ^ s))

  and isPrimop s = try let _ = stringToPrimop s in true with SyntaxError _ -> false 

  (* 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 *)
  and pgmToSexp p = 
    match p with 
      Pgm (fmls, e) -> 
        Seq [Sym "hofl"; Seq(map (fun s -> Sym s) fmls); expToSexp e]

  (* val expToSexp : exp -> Sexp.sexp *)
  and expToSexp e = 
    match e with 
      Lit v -> valuToSexp v 
    | Var s -> Sym s
    | PrimApp (rator, rands) -> 
        Seq (Sym (primopToString rator) :: map expToSexp rands)
    | If(tst,thn,els) -> Seq [Sym "if"; expToSexp tst; expToSexp thn; expToSexp els]
    | Abs(fml,body) -> Seq [Sym "abs"; Sym fml; expToSexp body]
    | App(rator,rand) -> Seq [expToSexp rator; expToSexp rand]
    | Bindrec(names,defns,body) -> 
	Seq [Sym "bindrec"; 
	     Seq (map2 (fun name defn -> Seq[Sym name; expToSexp defn]) names defns);
	     expToSexp body]

  (* 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" 
    | PairFn -> "pair"
    | Left -> "left"
    | Right -> "right"
    | Prepend -> "prepend"
    | Head -> "head"
    | Tail -> "tail" 
    | Null -> "null" 
    | IsNull -> "null?" 
    | IsInt -> "int?" 
    | IsBool -> "bool?" 
    | IsFun -> "fun?" 
    | IsPair -> "pair?" 
    | IsList -> "list?" 

  (* val valuToSexp : valu -> sexp *)
  and valuToSexp valu = 
    match valu with 
      Int i -> Sexp.Int i
    | Bool b -> Sym (if b then "true" else "false")
    | Fun _ -> Sym "<fun>"
    | Pair(l,r) -> Seq [Sym "pair"; valuToSexp l; valuToSexp r]
    | List xs -> Seq (Sym "list" :: (map valuToSexp xs))

  (* val valuToString : valu -> string *)
  and valuToString valu = 
    match valu with 
      Int i -> string_of_int i
    | Bool b -> string_of_bool b
    | Fun _ -> "<fun>"
    | Pair(l,r) -> "<" ^ (valuToString l) ^ "," ^ (valuToString r) ^ ">"
    | List xs -> "{" ^ (String.concat "," (map valuToString xs)) ^ "}"

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

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

end