open Support.Pervasive
open Support.Error
open Cpp_ast
open Printf

let rec lower_fun_exp (e : exp) : (exp * decl list)  = 
  match e with
    NullE | IntE _ | FloatE _ | DoubleE _
  | BoolE _ | StringE _ | CharE _ | VarE _ -> 
      (e, [])

  | CastE (t, e) ->
      let (e,ds) = lower_fun_exp e in
	(CastE (t, e), ds)
  | AnyCastE (t, e) ->
      let (e,ds) = lower_fun_exp e in
	(AnyCastE (t, e), ds)
	  
  | DestroyE (t, e) -> 
      let e_ds = lower_fun_exp e in
      (DestroyE (t, fst e_ds), snd e_ds)
  
  | NewE (a, t, args) -> 
      let args_ds = List.map lower_fun_exp args in
      let args = List.map fst args_ds in
      let ds = List.concat (List.map snd args_ds) in
      (NewE (a, t, args), ds)

  | PlacementNewE (p, t, args) -> 
      let p_ds = lower_fun_exp p in
      let args_ds = List.map lower_fun_exp args in
      let args = List.map fst args_ds in
      let ds = List.concat (List.map snd args_ds) in
      (PlacementNewE (fst p_ds, t, args), (snd p_ds) @ ds)

  | NewArrayE (a, t, e) -> 
      let (e,ds) = lower_fun_exp e in
      (NewArrayE (a, t, e), ds)

  | SizeofE t ->  (SizeofE t, [])

  | FunE (ps, rt, inits, body) ->
      let (body, ds1) = lower_fun_stmt body in
      let inits_ds = List.map (fun (n,(e,t)) -> (n,(lower_fun_exp e, t))) inits in
      let inits = List.map (fun (n,((e,ds),t)) -> (n,(e,t))) inits_ds in
      let ds2 = List.concat (ds1::(List.map (fun (n,((e,ds),t)) -> ds) inits_ds)) in
      
      (* create a functor class *)
      let fn = sprintf "__functor_%d" (make_id()) in
      let ms = List.map (fun (n,(e,t)) -> Field (n, t)) inits in
      let psc = List.map (fun (n,(e,t)) -> (n,t)) inits in
      let cons = Constructor (0, psc, List.map (fun (n,_) -> (n,[VarE n])) inits, 
			    CompoundS []) in
      let call_op = Method ("operator()", false, ps, rt, body) in
      let dest = Destructor (CompoundS []) in
      let ds3 = ds2@[LocalD [ClassD (fn, [], ms @ cons::dest::call_op::[])]] in
      (* leave behind a call to the constructor *)
      (ClassE (Prim.Stack, fn, 0, List.map fst (List.map snd inits)), 
       ds3)

  | ApplyE (f, args) ->
      let (f,ds) = lower_fun_exp f in
      let args_ds = List.map lower_fun_exp args in
      let args = List.map fst args_ds in
      let ds = List.concat (ds::(List.map snd args_ds)) in
      (ApplyE (f, args), ds)

  | PrimE p -> 
      (PrimE p, [])

  | ClassE (a, cn, k, args) ->
      let args_ds = List.map lower_fun_exp args in      
      let args = List.map fst args_ds in
      let ds = List.concat (List.map snd args_ds) in
      (ClassE (a, cn, k, args), ds)

  | PlacementClassE (p, cn, k, args) ->
      let p_ds = lower_fun_exp p in
      let args_ds = List.map lower_fun_exp args in      
      let args = List.map fst args_ds in
      let ds = List.concat (List.map snd args_ds) in
      (PlacementClassE (fst p_ds, cn, k, args), (snd p_ds) @ ds)

  | MemE (e, f) ->
      let (e,ds) = lower_fun_exp e in
	(MemE (e,f), ds)

  | ScopedE ns -> 
      (ScopedE ns, [])

  | IfE (c, els, thn) ->
      let (c,ds1) = lower_fun_exp c in
      let (els,ds2) = lower_fun_exp els in
      let (thn,ds3) = lower_fun_exp thn in
      (IfE (c, els, thn), ds1@ds2@ds3)

  | SeqE args ->
      let args_ds = List.map lower_fun_exp args in
      let args = List.map fst args_ds in
      let ds = List.concat (List.map snd args_ds) in
	(SeqE args, ds)

and lower_fun_stmt (s : stmt) : (stmt * decl list) =
  match s with
    VarS (t, n, e) ->
      let (e,ds) = lower_fun_exp e in
      (VarS (t, n, e), ds)
  | ExprS e ->
      let (e,ds) = lower_fun_exp e in
      (ExprS e, ds)
  | ReturnS None ->
      (ReturnS None, [])
  | ReturnS (Some e) ->
      let (e,ds) = lower_fun_exp e in
      (ReturnS (Some e), ds)
  | IfS (c, els, thn) ->
      let (c,ds1) = lower_fun_exp c in
      let (els,ds2) = lower_fun_stmt els in
      let (thn,ds3) = lower_fun_stmt thn in
      (IfS (c, els, thn), ds1@ds2@ds3)
  | WhileS (c, body) ->
      let (c,ds1) = lower_fun_exp c in
      let (body,ds2) = lower_fun_stmt body in
      (WhileS (c, body), ds1@ds2)
  | CompoundS (ss) ->
      let ss_ds = List.map lower_fun_stmt ss in      
      let ss = List.map fst ss_ds in
      let ds = List.map snd ss_ds in
      (CompoundS ss, List.concat ds)
  | SwitchS (c, s) ->
      let (c,ds1) = lower_fun_exp c in
      let (s,ds2) = lower_fun_stmt s in
      (SwitchS (c, s), ds1@ds2)
  | CaseS (e,s) ->
      let (e,ds1) = lower_fun_exp e in
      let (s,ds2) = lower_fun_stmt s in
      (CaseS (e, s), ds1@ds2)
  | DefaultS s ->
      let (s,ds) = lower_fun_stmt s in
      (DefaultS s, ds)
  | BreakS -> (BreakS, [])
  | PureS -> (PureS, [])

and lower_fun_class_mem m =
  match m with
      Field (n,t) ->
	(Field (n,t), [])
    | Constructor (k, ps, inits, body) ->
	let inits_ds = 
	  List.map
	    (fun (n,es) -> 
	       let es_dss = List.map lower_fun_exp es in
	       let es = List.map fst es_dss 
	       and dss = List.map snd es_dss in
		 ((n,es),List.concat dss)) inits in
	let inits = List.map fst inits_ds in
	let ds1 = List.map snd inits_ds in
	let (body,ds2) = lower_fun_stmt body in
	  (Constructor (k, ps, inits, body), (List.concat ds1)@ds2)
    | Destructor body ->
	let (body, ds) = lower_fun_stmt body in
	  (Destructor body, ds)
    | Method (mn, virt, ps, rt, body) ->
	let (body,ds) = lower_fun_stmt body in
	  (Method (mn,virt,ps,rt,body), ds)
    | MethodFwd (mn, virt, ps, rt) ->
	(MethodFwd (mn,virt,ps,rt), [])
    | Enum (n,ns) ->
	(Enum (n,ns), [])

and lower_fun_decl (d : decl) : decl list =
  match d with
    VarFwdD (t, n) ->
      [VarFwdD (t,n)]
    | VarD (t, n, e) ->
	let (e,ds) = lower_fun_exp e in
      ds@[VarD (t,n, e)]
  | FunD (fn, ps, rt, body) ->
      let (body, ds) = lower_fun_stmt body in
	ds@[FunD (fn, ps, rt, body)]
  | FunFwdD (fn, ps, rt) ->
      [FunFwdD (fn, ps, rt)]

  | ExternFunFwdD (l, fn, ps, rt) ->
      [ExternFunFwdD (l, fn, ps, rt)]

  | ExternScopeD (lang, ds) ->
      [ExternScopeD (lang, ds)]

  | ClassFwdD sn ->
      [ClassFwdD sn]

  | ClassD (cn, bs, ms) ->
      let ms_ds = map lower_fun_class_mem ms in
	(concat (map snd ms_ds))@[ClassD (cn, bs, map fst ms_ds)]

  | MethodD (cn, mn, ps, rt, body) ->
      let (body,ds) = lower_fun_stmt body in
	ds @ [MethodD (cn,mn,ps,rt,body)]

  | IncludeD fn ->
      [IncludeD fn]

  | LocalD ds ->
      [LocalD (concat (map lower_fun_decl ds))]

and lower_fun_decls (ds : decl list) : decl list =
  List.concat (List.map lower_fun_decl ds)

and lower_fun_typ (t : typ) : typ = t