(* ssa-pp-fn.sml
 *
 * This code is part of the Diderot Project (http://diderot-language.cs.uchicago.edu)
 *
 * COPYRIGHT (c) 2015 The University of Chicago
 * 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.STATE x => [IL.StateVar.toString x, ";"]
                  | 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.APPLY(f, args) => [
                        MathFuns.toString f, "(",
                        String.concatWith "," (List.map Var.toString args), ");"
                      ]
                  | IL.CONS(ty, xs) => [
                        "<", Ty.toString ty, ">[",
                        String.concatWith "," (List.map Var.toString xs), "];"
                      ]
                (* end case *))
          in
            String.concat(typedVar y :: " = " :: rhs)
          end

    fun massignToString ([], rator, xs) = String.concat [
            Op.toString rator, "(", String.concatWith "," (List.map Var.toString xs), ");"
          ]
      | massignToString (ys, rator, xs) = String.concat [
            "(", String.concatWith "," (List.map typedVar ys), ") = ",
            Op.toString rator, "(", String.concatWith "," (List.map Var.toString xs), ");"
          ]

    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.MASSIGN{stm, succ, ...} => (
                            prln (out1, [massignToString 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.SAVE{lhs, rhs, succ, ...} => (
                            prln (out1, [
                                IL.StateVar.toString lhs, " = ", Var.toString rhs, ";\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.SINIT =>
                                    prln (out1, ["strand_init ", 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, body}) = let
          val out1 = incIndent out
          fun prVars xs = List.app (fn x => prl(out, [" ", typedVar x, ";"])) xs
          in
            indent out; prl(out, ["method ", StrandUtil.nameToString name, "\n"]);
            ppCFG (incIndent out1, body);
            indent out; prl(out, ["end ", StrandUtil.nameToString name, "\n"])
          end

    and ppStrand (out, IL.Strand{name, params, state, stateInit, methods}) = let
          val out1 = incIndent out
          fun stateVar x = let
                val l = [Ty.toString(IL.StateVar.ty x), " ", IL.StateVar.toString x, ";"]
                val l = if IL.StateVar.isOutput x then "output " :: l else l
                in
                  prl (out1, l)
                end
          in
            indent out;
            prl(out, [
                "strand ", Atom.toString name, " (",
                String.concatWith ", " (List.map typedVar params), ")\n"
              ]);
            indent out1; pr(out1, "state: "); List.app stateVar 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{props, globalInit, initially, strands}) = let
          val out = (outS, 0)
          val out1 = incIndent out
          in
            pr (out, concat["##### ", msg, " ####\n"]);
            pr (out, "## properties\n");
	    prln (out1, [StrandUtil.propsToString props, "\n"]);
            pr (out, "## globals\n");
            List.app
              (fn x => (indent out1; prl(out1, ["global ", typedVar x, "\n"]))) 
                (IL.CFG.liveAtExit globalInit);
            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