(* Convert ordinary Pnf code to "extra-normalized" Pnf. 
   In this form, Let expressions bind only 
   Triv, App, Record, Select, or Variant expressions
   (collectively the "bindable" expressions); 
   moreover, these bindable expressions can *only* appear 
   within a Let binding.
   This implementation assumes that in the input Pnf term,
   all Switch arms are manifest Abs terms.
   Also, this implementation will duplicate any common code
   following a Switch.
*)


module Pnfnorm:
sig 
  val norm : Pnf.exp -> Pnf.exp
end = 
struct
  open Pnf
  open Pnfcheck

  exception Non_bindable
  exception Bad_switch_arm

  (* Normalize e *)
  let rec norm (e:exp) : exp =
    match e with
    | Let(vt,e1,e2) -> 
	normalize e1 (fun bexp -> Let(vt,bexp,norm e2))
    | Fix(fs,b) ->  
	Fix(List.map (fun (ft,tre) -> (ft,norm_triv tre)) fs,
	    norm b) 
    | Switch(tre,r) -> Switch(norm_triv tre,norm_row r)
    | Triv trexp -> Triv(norm_triv trexp)
    | _ -> 
	let xt = gensym(check e)  in
	Let(xt,norm_bindable e,Triv(Var xt))
  
  and norm_bindable (e:exp) : exp =
    match e with
      Triv trexp -> Triv(norm_triv trexp)
    | App(tre1,tre2) -> App(norm_triv tre1,norm_triv tre2)
    | Record r -> Record(norm_row r)
    | Select(l,tre) -> Select(l,norm_triv tre)
    | Variant(l,tre) -> Variant(l,norm_triv tre)
    | _ -> raise Non_bindable

  and norm_triv (tre:trexp) : trexp =
    match tre with
      Var vt -> Var vt
    | Abs(vt,e) -> Abs(vt,norm e)
    | Int i -> Int i

  and norm_row r = List.map (fun (l,tre) -> (l,norm_triv tre)) r

  (* Convert e to bindable form and pass it to k *)
  and normalize (e:exp) (k:exp -> exp) : exp =
    match e with
      Let(vt,e1,e2) -> 
	normalize e1 (fun bexp ->
	  Let(vt,bexp,normalize e2 k)) 
    | Fix(fs,b) -> 
	Fix(List.map (fun (ft,tre) -> (ft,norm_triv tre)) fs,
	    normalize b k)
    | Switch(tre,r) ->
	let norm_arm (l,tre) =
	  match tre with
	    Abs(vt,e) -> 
	      let e' = normalize e k in
	      (l,Abs(vt,e'))
	  | _ -> raise Bad_switch_arm in
	Switch(norm_triv tre, List.map norm_arm r)  (* may duplicate code !! *)
    | _ -> k(norm_bindable e)
	  
end