(* ssa-pp-fn.sml
 *
 * COPYRIGHT (c) 2010 The Diderot Project (http://diderot-language.cs.uchicago.edu)
 * All rights reserved.
 *
 * Pretty printing for SSA representations
 *)

functor SSAPPFn (IL : SSA) : sig

    val assignToString : IL.assign -> string

    val output : TextIO.outstream * string * IL.program -> unit

  end = struct

    structure Op = IL.Op
    structure Var = IL.Var
    structure Ty = IL.Ty

    local
      val {getFn, setFn} = IL.Node.newFlag()
    in
    val isMarked = getFn
    fun mark nd = setFn(nd, true)
    fun clear nd = setFn(nd, false)
    end

  (* if true, the output is in compact form *)
    val compact = ref true

    fun indent (outS, i) = TextIO.output(outS, StringCvt.padLeft #" " i "")
    fun incIndent (outS, i) = (outS, i+2)
    fun pr ((outS, _), s) = TextIO.output(outS, s)
    fun prl (out, l) = pr(out, concat l)
    fun prln (out, l) = (indent out; prl(out, l))

    fun typedVar x = String.concat [
	    Ty.toString(Var.ty x), " ", Var.toString x, "#", Int.toString(Var.useCount x)
	  ]

    fun assignToString (y, rhs) = let
	  val rhs = (case rhs
		 of IL.VAR x => [Var.toString x, ";"]
		  | IL.LIT lit => [Literal.toString lit, ";"]
		  | IL.OP(rator, []) => [Op.toString rator, ";"]
		  | IL.OP(rator, args) => [
			Op.toString rator, "(",
			String.concatWith "," (List.map Var.toString args), ");"
		      ]
		  | IL.CONS xs => ["[", String.concatWith "," (List.map Var.toString xs), "];"]
		(* end case *))
	  in
	    String.concat(typedVar y :: " = " :: rhs)
	  end

    fun labelOf (IL.ND{id, ...}) = "L"^Stamp.toString id

    fun ppCFG (out, cfg as IL.CFG{entry, exit}) = let
	  fun goto (out, nd) = if (List.length(IL.Node.preds nd) > 1)
		then (
		  prln(incIndent out, ["goto ", IL.Node.toString nd, "\n"]);
		  ppNd (out, false, nd))
		else ppNd (out, true, nd)
	  and ppNd (out, noLabel, nd) = let
		val out1 = incIndent out
		fun prPhi (y, xs) =
		      prln (out1, [
			  typedVar y, " = phi(",
			  String.concatWith "," (List.map Var.toString xs), ")\n"
			])
		in
		  if isMarked nd
		    then ()
		    else (
		      mark nd;
		      if noLabel andalso (! compact)
			then ()
			else prln (out, [
			    IL.Node.toString nd, ":  preds = [",
			    String.concatWith "," (List.map IL.Node.toString (IL.Node.preds nd)),
			    "]\n"
			  ]);
		      case IL.Node.kind nd
		       of IL.NULL => ()
			| IL.ENTRY{succ} => goto (out, !succ)
			| IL.JOIN{phis, succ, ...} => (
			    List.app prPhi (!phis);
			    goto (out, !succ))
			| IL.COND{cond, trueBranch, falseBranch, ...} => (
			    prln (out1, [
				"if ", Var.toString cond,
				" then goto ", IL.Node.toString(!trueBranch),
				" else goto ", IL.Node.toString(!falseBranch), "\n"
			      ]);
			    ppNd (out, false, !trueBranch);
			    ppNd (out, false, !falseBranch))
			| IL.COM{text, succ, ...} => (
			    List.app (fn s => prln (out1, ["//", s, "\n"])) text;
			    goto (out, !succ))
			| IL.ASSIGN{stm, succ, ...} => (
			    prln (out1, [assignToString stm, "\n"]);
			    goto (out, !succ))
			| IL.NEW{strand, args, succ, ...} => (
			    prln (out1, [
				"new ", Atom.toString strand, "(",
				String.concatWith "," (List.map Var.toString args), ");\n"
			      ]);
			    goto (out, !succ))
			| IL.EXIT{kind, live, ...} => let
			    val live = (case live
				   of [] => "()"
				    | [x] => String.concat["(", Var.toString x, ")"]
				    | xs => String.concat[
					  "(",
					  String.concatWith "," (List.map Var.toString xs),
					  ")"
					]
				  (* end case *))
			    in
			      case kind
			       of ExitKind.FRAGMENT =>
				    prln (out1, ["live vars = ", live, "\n"])
				| ExitKind.RETURN =>
				    prln (out1, ["return ", live, "\n"])
				| ExitKind.ACTIVE => 
				    prln (out1, ["active ", live, "\n"])
				| ExitKind.STABILIZE => 
				    prln (out1, ["stabilize ", live, "\n"])
				| ExitKind.DIE => 
				    prln (out1, ["die ", live, "\n"])
			      (* end case *)
			    end
		      (* end case *))
		end
	  in
	    ppNd (out, false, entry);
	  (* clear marks *)
	    IL.CFG.apply clear cfg
	  end

    fun ppInitially (out, IL.Initially{isArray, rangeInit, iters, create}) = let
	  val (initArg, strand, args) = create
	  fun ppIters (out, []) = (
		ppCFG (out, initArg);
		prln (out, [
		    "new ", Atom.toString strand, "(",
		    String.concatWith "," (List.map Var.toString args), ");\n"
		  ]))
	    | ppIters (out, (param, lo, hi)::iters) = (
		indent out; prl(out, [
		    "for ", typedVar param, " = ", Var.toString lo, " .. ", Var.toString hi, "\n"
		  ]);
		ppIters (incIndent out, iters))
	  in
	    indent out; prl(out, if isArray then ["ARRAY\n"] else ["COLLECTION\n"]);
	    ppCFG (incIndent out, rangeInit);
	    ppIters (incIndent out, iters)
	  end

    fun ppMethod (out, IL.Method{name, stateIn, stateOut, body}) = let
	  val out1 = incIndent out
	  fun prVars xs = List.app (fn x => prl(out, [" ", typedVar x, ";"])) xs
	  in
	    indent out; prl(out, ["method ", Atom.toString name, "\n"]);
	    indent out1; pr(out1, "state in: "); prVars stateIn; pr(out1, "\n");
	    ppCFG (incIndent out1, body);
	    indent out1; pr(out1, "state out:"); prVars stateOut; pr(out1, "\n");
	    indent out; prl(out, ["end ", Atom.toString name, "\n"])
	  end

    and ppStrand (out, IL.Strand{name, params, state, stateInit, methods}) = let
	  val out1 = incIndent out
	  fun prVars xs = List.app (fn x => prl(out, [" ", typedVar x, ";"])) xs
	  in
	    indent out;
	    prl(out, [
		"strand ", Atom.toString name, " (",
		String.concatWith ", " (List.map typedVar params), ")\n"
	      ]);
	    indent out1; pr(out1, "state: "); prVars state; pr(out1, "\n");
	    ppCFG (incIndent out1, stateInit);
	    List.app (fn m => ppMethod(out1, m)) methods;
	    indent out; prl(out, ["end ", Atom.toString name, "\n"])
	  end

    fun output (outS, msg, IL.Program{globals, globalInit, initially, strands}) = let
	  val out = (outS, 0)
	  val out1 = incIndent out
	  in
	    pr (out, concat["##### ", msg, " ####\n"]);
	    pr (out, "## globals\n");
	    List.app
	      (fn x => (indent out1; prl(out1, ["global ", typedVar x, "\n"]))) 
		globals;
	    pr (out, "## global initialization\n");
	    ppCFG (out1, globalInit);
	    pr (out, "## initially\n");
	    ppInitially (out1, initially);
	    pr (out, "## strands\n");
	    List.app (fn strand => ppStrand(out1, strand)) strands;
	    pr (out, "#### end program ####\n")
	  end

  end