(* inliner.sml
 *
 * COPYRIGHT (c) 2013 The Diderot Project (http://diderot-language.cs.uchicago.edu)
 * All rights reserved.
 *
 * This pass eliminates the function definitions by inlining them.
 *)

structure Inliner : sig

    val transform : Simple.program -> Simple.program

  end = struct

    structure S = Simple
    structure V = SimpleVar

  (* beta reduce the application "lhs = f(args)" by creating a fresh copy of f's body
   * while mapping the parameters to arguments.
   *)
    fun beta (lhs, S.Func{f, params, body}, args) = let
	  val needsLHSPreDecl = ref false (* set to true if the lhs needs to be declared before the body *)
	  fun rename env x = (case V.Map.find(env, x)
		 of SOME x' => x'
		  | NONE => if SimpleVar.isGlobal x
		      then x
		      else raise Fail("unknown variable " ^ V.uniqueNameOf x)
		(* end case *))
          fun doBlock (env, isTop, S.Block stms) = let
                fun f (stm, (env, stms)) = let
                        val (env, stm) = doStmt (env, isTop, stm)
                        in
                          (env, stm::stms)
                        end
                val (_, stms) = List.foldl f (env, []) stms
                in
                  S.Block(List.rev stms)
                end
          and doStmt (env, isTop, stm) = (case stm
                 of S.S_Var x => let
		      val x' = V.copy x
		      in
			(V.Map.insert(env, x, x'), S.S_Var x')
		      end
                  | S.S_Assign(x, e) => (case V.Map.find(env, x)
		       of SOME x' => (env, S.S_Assign(x', doExp env e))
			| NONE => let
			    val x' = V.copy x
			    in
			      (V.Map.insert(env, x, x'), S.S_Assign(x', doExp env e))
			    end
		      (* end case *))
                  | S.S_IfThenElse(x, b1, b2) =>
		      (env, S.S_IfThenElse(rename env x, doBlock(env, false, b1), doBlock(env, false, b2)))
                  | S.S_New(strnd, xs) => (env, S.S_New(strnd, List.map (rename env) xs))
                  | S.S_Die => (env, stm)
                  | S.S_Stabilize => (env, stm)
                  | S.S_Return x => (
		      if not isTop then needsLHSPreDecl := true else ();
		      (env, S.S_Assign(lhs, S.E_Var(rename env x))))
                  | S.S_Print xs => (env, S.S_Print(List.map (rename env) xs))
                (* end case *))
          and doExp env exp = (case exp
                 of S.E_Var x => S.E_Var(rename env x)
                  | S.E_Lit _ => exp
                  | S.E_Tuple xs => S.E_Tuple(List.map (rename env) xs)
                  | S.E_Apply(f, xs, ty) => S.E_Apply(f, List.map (rename env) xs, ty)
                  | S.E_Prim(f, tys, xs, ty) =>
		      S.E_Prim(f, tys, List.map (rename env) xs, ty)
                  | S.E_Cons xs => S.E_Cons(List.map (rename env) xs)
                  | S.E_Slice(x, xs, ty) =>
		      S.E_Slice(rename env x, List.map (Option.map (rename env)) xs, ty)
                  | S.E_Coerce{srcTy, dstTy, x} =>
		      S.E_Coerce{srcTy=srcTy, dstTy=dstTy, x=rename env x}
		  | S.E_Input(_, _, _, NONE) => exp
		  | S.E_Input(ty, name, desc, SOME x) =>
		      S.E_Input(ty, name, desc, SOME(rename env x))
                  | S.E_LoadImage _ => exp
                (* end case *))
	(* build the initial environment by mapping parameters to arguments *)
	  val env = ListPair.foldlEq
		(fn (x, x', env) => V.Map.insert(env, x, x'))
		  V.Map.empty (params, args)
	  val blk as S.Block stms = doBlock (env, true, body)
          in
	    if !needsLHSPreDecl
	      then S.Block(S.S_Var lhs :: stms)
	      else blk
          end

  (* inline expand user-function calls in a block *)
    fun expandBlock funcTbl = let
	  val findFunc = V.Tbl.find funcTbl
	  fun expandBlk (S.Block stms) =
		S.Block(List.foldr expandStm [] stms)
	  and expandStm (stm, stms') = (case stm
		 of S.S_Assign(x, S.E_Apply(f, xs, _)) => (case findFunc f
		       of NONE => stm :: stms'
			| SOME func => let
			    val S.Block stms = beta(x, func, xs)
			    in
			      stms @ stms'
			    end
		      (* end case *))
		  | S.S_IfThenElse(x, b1, b2) =>
		      S.S_IfThenElse(x, expandBlk b1, expandBlk b2) :: stms'
		  | _ => stm :: stms'
		(* end case *))
	  in
	    expandBlk
	  end

    fun expandFunc funcTbl (S.Func{f, params, body}) = let
	  val body' = expandBlock funcTbl body
	  val func' = S.Func{f=f, params=params, body=body'}
	  in
	    V.Tbl.insert funcTbl (f, func')
	  end

    fun expandStrand funcTbl = let
	  val expandBlock = expandBlock funcTbl
	  fun expandMeth (S.Method(m, blk)) = S.Method(m, expandBlock blk)
	  fun expand (S.Strand{name, params, state, stateInit, methods}) = S.Strand{
		  name = name,
		  params = params,
		  state = state,
		  stateInit = expandBlock stateInit,
		  methods = List.map expandMeth methods
		}
	  in
	    expand
	  end

    fun expandInit funcTbl (S.Initially{isArray, rangeInit, iters, create}) = let
	  val expandBlock = expandBlock funcTbl
	  fun expandCreate (S.C_Create{argInit, name, args}) = S.C_Create{
		  argInit = expandBlock argInit,
		  name = name, args = args
		}
	  in
	    S.Initially{
		isArray = isArray,
		rangeInit = expandBlock rangeInit,
		iters = iters,
		create = expandCreate create
	      }
	  end

    fun transform (prog as S.Program{funcs=[], ...}) = prog
      | transform (S.Program{globals, globalInit, funcs, strands, init}) = let
	(* a table that maps function names to their definitions *)
	  val funcTbl = V.Tbl.mkTable (List.length funcs, Fail "funcTbl")
	(* first we inline expand the function bodies in definition order *)
	  val _ = List.app (expandFunc funcTbl) funcs
	  val expandBlock = expandBlock funcTbl
	  in
	    S.Program{
		globals = globals,
		globalInit = expandBlock globalInit,
		funcs = [],
		strands = List.map (expandStrand funcTbl) strands,
		init = expandInit funcTbl init
	      }
	  end

  end