(* simple-pp.sml
 *
 * COPYRIGHT (c) 2010 The Diderot Project (http://diderot.cs.uchicago.edu)
 * All rights reserved.
 *
 * Pretty printing for the Simple-AST representation.
 *)

structure SimplePP : sig

    val output : TextIO.outstream * Simple.program -> unit

  end = struct

    structure PP = TextIOPP
    structure TU = TypeUtil
    structure S = Simple

    val indent = PP.Abs 2

    fun ppList ppFn (left, sep, right) (ppStrm, list) = let
	  fun sp () = PP.space ppStrm 1
	  val string = PP.string ppStrm
	  fun pp [] = string right
	    | pp [x] = (ppFn(ppStrm, x); string right)
	    | pp (x::xs) = (ppFn(ppStrm, x); string sep; sp(); pp xs)
	  in
	    string left; pp list
	  end

  (* print type arguments; we use "#" to denote differentiation arguments, "$" to denote
   * shape arguments, and "%" to denote dimension arguments.
   *)
    fun ppTyArgs (ppStrm, mvs) = let
	  val string = PP.string ppStrm
	  fun ppTyArg (_, mv) = (case mv
		 of Types.TYPE tv => string(TU.toString(TU.resolve tv))
		  | Types.DIFF dv => string("#"^TU.diffToString(TU.resolveDiff dv))
		  | Types.SHAPE sv => string("$"^TU.shapeToString(TU.resolveShape sv))
		  | Types.DIM dv => string("%"^TU.dimToString(TU.resolveDim dv))
		(* end case *))
	  in
	    ppList ppTyArg ("<", ";", ">") (ppStrm, mvs)
	  end

    fun ppVar (ppStrm, x) = PP.string ppStrm (Var.uniqueNameOf x)

    fun ppVarDecl ppStrm = let
	  fun sp () = PP.space ppStrm 1
	  val string = PP.string ppStrm
	  in
	    fn x => (
	      PP.openHBox ppStrm;
		case Var.kindOf x
		 of S.InputVar => (string "input"; sp())
		  | S.ActorOutputVar => (string "output"; sp())
		  | _ => ()
		(* end case *);
		string(TU.toString(#2(Var.typeOf x))); sp(); string(Var.uniqueNameOf x); string ";";
	      PP.closeBox ppStrm)
	  end

    fun ppExp (ppStrm, e) = let
	  fun sp () = PP.space ppStrm 1
	  val string = PP.string ppStrm
	  fun var x = string(Var.nameOf x)
	  fun pp e = (case e
		 of S.E_Var x => string(Var.uniqueNameOf x)
		  | S.E_Lit lit => string (Literal.toString lit)
		  | S.E_Tuple es => ppArgs (ppStrm, es)
		  | S.E_Apply(f, [], args, _) => (var f; sp(); ppArgs (ppStrm, args))
		  | S.E_Apply(f, mvs, args, _) => (
		      var f; ppTyArgs (ppStrm, mvs); sp(); ppArgs (ppStrm, args))
		  | S.E_Cons es => (
		      ppList ppVar ("[", ",", "]") (ppStrm, es))
		(* end case *))
	  in
	    pp e
	  end

    and ppArgs (ppStrm, args) = ppList ppVar ("(", ",", ")") (ppStrm, args)

    fun ppBlock (ppStrm, vars, stms) = let
	  fun sp () = PP.space ppStrm 1
	  fun nl () = PP.newline ppStrm
	  val string = PP.string ppStrm
	  fun var x = string(Var.uniqueNameOf x)
	  fun ppStmt stmt = (case stmt
		 of S.S_Block stms => ppBlock (ppStrm, [], stms)
		  | S.S_Assign(x, e) => (
		      PP.openHBox ppStrm;
			var x; sp(); string "="; sp(); ppExp(ppStrm, e); string ";";
		      PP.closeBox ppStrm;
		      nl())
		  | S.S_IfThenElse(x, S.S_Block stms, S.S_Block[]) => (
		      PP.openHBox ppStrm;
			string "if"; sp(); ppVar(ppStrm, x);
			sp(); ppBlock (ppStrm, [], stms);
		      PP.closeBox ppStrm)
		  | S.S_IfThenElse(x, s1, S.S_Block[]) => (
		      PP.openVBox ppStrm indent;
			PP.openHBox ppStrm;
			  string "if"; sp(); ppVar(ppStrm, x);
			PP.closeBox ppStrm;
			nl();
			ppStmt s1;
		      PP.closeBox ppStrm;
		      nl())
		  | S.S_IfThenElse(x, S.S_Block stms1, S.S_Block stms2) => (
		      PP.openHBox ppStrm;
			string "if"; sp(); ppVar(ppStrm, x);
			sp(); ppBlock (ppStrm, [], stms1);
		      PP.closeBox ppStrm;
		      PP.openHBox ppStrm;
			string "else"; sp(); ppBlock (ppStrm, [], stms2);
		      PP.closeBox ppStrm)
		  | S.S_IfThenElse(x, S.S_Block stms1, s2) => raise Fail "FIXME"
		  | S.S_IfThenElse(x, s1, S.S_Block stms2) => raise Fail "FIXME"
		  | S.S_IfThenElse(x, s1, s2) => (
		      PP.openVBox ppStrm indent;
			PP.openHBox ppStrm;
			  string "if"; sp(); ppVar(ppStrm, x);
			PP.closeBox ppStrm;
			nl();
			ppStmt s1;
		      PP.closeBox ppStrm;
		      nl();
		      PP.openVBox ppStrm indent;
			string "else"; nl();
			ppStmt s2;
		      PP.closeBox ppStrm;
		      nl())
		  | S.S_New(actor, args) => (
		      PP.openHBox ppStrm;
			string "new"; sp(); string(Atom.toString actor); sp();
			ppArgs (ppStrm, args); string ";";
		      PP.closeBox ppStrm;
		      nl())
		  | S.S_Die => (string "die;"; nl())
		  | S.S_Stabilize => (string "stabilize;"; nl())
		(* end case *))
	  in
	    PP.openVBox ppStrm (PP.Abs 0);
	      string "{"; nl();
	      PP.openVBox ppStrm indent;
	        List.app (fn vdcl => (ppVarDecl ppStrm vdcl; nl())) vars;
		List.app ppStmt stms;
	      PP.closeBox ppStrm;
	      string "}"; nl();
	    PP.closeBox ppStrm
	  end

    fun ppVarsAndStmt (ppStrm, vars, S.S_Block stms) = ppBlock(ppStrm, vars, stms)
      | ppVarsAndStmt (ppStrm, vars, stm) = ppBlock(ppStrm, vars, [stm])

    fun ppActor ppStrm (S.Actor{name, params, state, stateInit, methods}) = let
	  fun sp () = PP.space ppStrm 1
	  fun nl () = PP.newline ppStrm
	  val string = PP.string ppStrm
	  fun var x = string(Var.nameOf x)
	  fun ppMethod (S.M_Method(name, body)) = (
		nl(); string(Atom.toString name); nl(); ppVarsAndStmt (ppStrm, [], body))
	  in
	    PP.openHBox ppStrm;
	      string "actor"; sp(); string(Atom.toString name); sp();
	      ppList (fn (_, x) => (string(TU.toString(#2(Var.typeOf x))); sp(); var x))
		("(", ",", ")") (ppStrm, params);
	    PP.closeBox ppStrm;
	    nl();
	    PP.openVBox ppStrm indent;
	      string "{";
	      ppVarsAndStmt (ppStrm, state, stateInit);
	      List.app ppMethod methods;
	    PP.closeBox ppStrm;
	    nl();
	    string "}";  nl()
	  end

    fun output (outS, S.Program{globals, globalInit, actors}) = let
	  val ppStrm = PP.openOut {dst = outS, wid = 120}
	  in
	    PP.openVBox ppStrm (PP.Abs 0);
	      PP.string ppStrm "/* Simplified Program start */"; PP.newline ppStrm;
	      ppVarsAndStmt (ppStrm, globals, globalInit);
	      List.app (ppActor ppStrm) actors;
	      PP.string ppStrm "/* Program end */"; PP.newline ppStrm;
	    PP.closeBox ppStrm;
	    PP.closeStream ppStrm
	  end

  end