(* Copyright (C) 1999-2002 Henry Cejtin, Matthew Fluet, Suresh
 *    Jagannathan, and Stephen Weeks.
 * Copyright (C) 1997-1999 NEC Research Institute.
 *
 * MLton is released under the GNU General Public License (GPL).
 * Please see the file MLton-LICENSE for license information.
 *)
structure Main =
struct

type int = Int.t

val fail = Process.fail

fun usage msg =
   Process.usage {usage = "[-mlkit] [-mosml] [-smlnj] bench1 bench2 ...",
		  msg = msg}

val doHtml = ref false
val doOnce = ref false
val runArgs : string list ref = ref []
   
fun withInput (file, f: unit -> 'a): 'a =
   let
      open FileDesc
      val inFd =
	 let
	    open Pervasive.Posix.FileSys
	 in
	    openf (file, O_RDONLY, O.flags [])
	 end
   in
      DynamicWind.wind
      (fn () => FileDesc.fluidLet (FileDesc.stdin, inFd, f),
       fn () => FileDesc.close inFd)
   end

fun ignoreOutput f =
   let
      val nullFd =
	 let
	    open Pervasive.Posix.FileSys
	 in
	    openf ("/dev/null", O_WRONLY, O.flags [])
	 end
      open FileDesc
   in
      DynamicWind.wind
      (fn () => fluidLet (stderr, nullFd, fn () =>
			  fluidLet (stdout, nullFd, f)),
       fn () => close nullFd)
   end

datatype command =
   Explicit of {args: string list,
		com: string}
  | Shell of string

fun timeIt ca =
   Process.time
   (fn () =>
    case ca of
       Explicit {args, com} =>
	  Process.wait (Process.spawnp {file = com, args = com :: args})
     | Shell s => Process.system s)
   
local
   val trialTime = Time.seconds (IntInf.fromInt 60)
in
   fun timeCall (com, args): real =
      let 
	 fun doit ac =
	    let
	       val {user, system} = timeIt (Explicit {args = args, com = com})
	       val op + = Time.+
	    in ac + user + system
	    end
	 fun loop (n, ac: Time.t): real =
	    if Time.> (ac, trialTime)
	       then Time.toReal ac / Real.fromInt n
	    else loop (n + 1, doit ac)
      in 
	 if !doOnce
	    then Time.toReal (doit Time.zero)
	 else loop (0, Time.zero)
      end
end

val benchCounts: (string * int * int) list =
   [("barnes-hut", 4096, 1024),
    ("boyer", 1500, 500),
    ("checksum", 1500, 150),
    ("count-graphs", 3, 1),
    ("DLXSimulator", 50, 15),
    ("fft", 256, 128),
    ("fib", 6, 1),
    ("flat-array", 2000, 400),
    ("hamlet", 50, 5),
    ("imp-for", 1000, 300),
    ("knuth-bendix", 500, 100),
    ("lexgen", 300, 50),
    ("life", 3, 1),
    ("logic", 40, 7),
    ("mandelbrot", 2, 1),
    ("matrix-multiply", 2, 2),
    ("md5", 3, 1),
    ("merge", 2000, 500),
    ("mlyacc", 500, 150),
    ("model-elimination", 0, 0),
    ("mpuz", 20, 5),
    ("nucleic", 500, 150),
    ("output1", 1, 1),
    ("peek", 1000, 100),
    ("psdes-random", 3, 1),
    ("ratio-regions", 512, 256),
    ("ray", 100, 30),
    ("raytrace", 10, 3),
    ("simple", 100, 20),
    ("smith-normal-form", 6, 1),
    ("tailfib", 200, 60),
    ("tak", 2, 1),
    ("tensor", 3, 1),
    ("tsp", 4, 1),
    ("tyan", 80, 13),
    ("vector-concat", 10, 2),
    ("vector-rev", 10, 2),
    ("vliw", 150, 30),
    ("wc-input1", 4000, 1000),
    ("wc-scanStream", 3000, 1000),
    ("zebra", 15, 3),
    ("zern", 2000, 500)]

val benchCount =
   String.memoize
   (fn s =>
    case List.peek (benchCounts, fn (b, _, _) => b = s) of
       NONE => Error.bug (concat ["no benchCount for ", s])
     | SOME (_, x86, sparc) =>
	  Int.toString
	  let
	     open MLton.Platform.Arch
	  in
	     case host of
		PowerPC => x86
	      | Sparc => sparc
	      | X86 => x86
	  end)

fun compileSizeRun {command, exe, doTextPlusData: bool} =
   Escape.new
   (fn e =>
    let
       val exe = "./" ^ exe
       val {system, user} = timeIt command
	  handle _ => Escape.escape (e, {compile = NONE,
					 run = NONE,
					 size = NONE})
       val compile = SOME (Time.toReal (Time.+ (system, user)))
       val size =
	  if doTextPlusData
	     then
		let 
		   val {text, data, ...} = Process.size exe
		in SOME (Position.fromInt (text + data))
		end
	  else SOME (File.size exe)
       val run =
	  timeCall (exe, !runArgs)
	  handle _ => Escape.escape (e, {compile = compile,
					 run = NONE,
					 size = size})
    in {compile = compile,
	run = SOME run,
	size = size}
    end)

fun batch bench = concat [bench, ".batch.sml"]

local
   val n = Counter.new 0
   val exe = "a.out"
in
   fun makeMLton commandPattern =
      case ChoicePattern.expand commandPattern of
	 Result.No m => usage m
       | Result.Yes coms => 
	    List.map
	    (coms, fn com =>
	     {name = com,
	      abbrv = "MLton" ^ (Int.toString (Counter.next n)),
	      test = (fn {bench} =>
		      compileSizeRun
		      {command = Shell (concat [com, " -output ", exe, " ", batch bench]),
		       exe = exe,
		       doTextPlusData = true})})
end

fun kitCompile {bench} =
   compileSizeRun {command = Explicit {args = [batch bench],
				       com = "mlkit"},
		   exe = "run",
		   doTextPlusData = true}
   
fun mosmlCompile {bench} =
   compileSizeRun
   {command = Explicit {args = ["-orthodox", "-standalone", "-toplevel",
				batch bench],
			com = "mosmlc"},
    exe = "a.out",
    doTextPlusData = false}

fun njCompile {bench} =
   Escape.new
   (fn e =>
    let
       (* sml should start SML/NJ *)
       val sml = "sml"
       val {system, user} =
	  File.withTempOut
	  (fn out =>
	   (Out.output
	    (out, "local\nval _ = SMLofNJ.Internals.GC.messages false\n")
	    ; File.outputContents (concat [bench, ".sml"], out)
	    ; (Out.output
	       (out,
		concat
		["in val _ = SMLofNJ.exportFn (\"", bench,
		 "\", fn _ => (Main.doit ", benchCount bench,
		 "; OS.Process.success))\nend\n"]
		 ))),
           fn input => withInput (input, fn () => timeIt (Explicit {args = [],
								    com = sml})))
         handle _ => Escape.escape (e, {compile = NONE,
					run = NONE,
					size = NONE})
       val suffix =
		 let
		   datatype z = datatype MLton.Platform.Arch.t
		   datatype z = datatype MLton.Platform.OS.t
		 in
		   case (MLton.Platform.Arch.host, MLton.Platform.OS.host) of
		     (X86, Linux) => ".x86-linux"
		   | (Sparc, Solaris) => ".sparc-solaris"
		   | _ => raise Fail "don't know SML/NJ suffix for host type"
		 end
       val heap = concat [bench, suffix]
    in
       if not (File.doesExist heap)
	  then {compile = NONE,
		run = NONE,
		size = NONE}
       else
          let
	     val compile = Time.toReal (Time.+ (user, system))
	     val size = SOME (File.size heap)
	     val run =
		  timeCall (sml, [concat ["@SMLload=", heap]])
		  handle _ => Escape.escape (e, {compile = SOME compile,
						 run = NONE,
						 size = size})
	  in {compile = SOME compile,
	      run = SOME run,
	      size = size}
	  end
    end)
		
fun polyCompile {bench} =
   Escape.new
   (fn e =>
    let
       val originalDbase = "/usr/lib/poly/ML_dbase"
       val poly = "poly"
    in File.withTemp
       (fn dbase =>
	let
	   val _ = File.copy (originalDbase, dbase)
	   val original = File.size dbase
	   val {system, user} =
	      File.withTempOut
	      (fn out =>
	       Out.output
	       (out,
		concat ["use \"", bench, ".sml\" handle _ => PolyML.quit ();\n",
			"if PolyML.commit() then () else ",
			"(Main.doit ", benchCount bench, "; ());\n",
			"PolyML.quit();\n"]),
	       fn input =>
	       withInput
	       (input, fn () =>
		timeIt (Explicit {args = [dbase],
				  com = "poly"})))
	   val after = File.size dbase
	in
	   if original = after
	      then {compile = NONE,
		    run = NONE,
		    size = NONE}
	   else
	       let
		  val compile = SOME (Time.toReal (Time.+ (user, system)))
		  val size = SOME (after - original)
		  val run =
		     timeCall (poly, [dbase])
		     handle _ => Escape.escape (e, {compile = compile,
						    run = NONE,
						    size = size})
	       in
		  {compile = compile,
		   run = SOME run,
		   size = size}
	       end
	end)
    end)

type 'a data = {bench: string,
		compiler: string,
		value: 'a} list

fun main args =
   let
      val compilers: {name: string,
		      abbrv: string,
		      test: {bench: File.t} -> {compile: real option,
						run: real option,
						size: Position.int option}} list ref 
	= ref []
      fun pushCompiler compiler = List.push(compilers, compiler)
      fun pushCompilers compilers' = compilers := (List.rev compilers') @ (!compilers)

      fun setData (switch, data, str) =
	 let
	    fun die () = usage (concat ["invalid -", switch, " argument: ", str])
	    open Regexp
	    val numSave = Save.new ()
	    val regexpSave = Save.new ()
	    val re = seq [save (star digit, numSave),
			  char #",",
			  save (star any, regexpSave)]
	    val reC = compileDFA re
	 in
	    case Compiled.matchAll (reC, str) of
	       NONE => die ()
	     | SOME match => 
		  let
		     val num = Match.lookupString (match, numSave)
		     val num = case Int.fromString num of
		                  NONE => die ()
				| SOME num => num
		     val regexp = Match.lookupString (match, regexpSave)
		     val (regexp, saves) = 
		        case Regexp.fromString regexp of
			   NONE => die ()
			 | SOME regexp => regexp
		     val save = if 0 <= num andalso num < Vector.length saves
				   then Vector.sub (saves, num)
				else die ()
		     val regexpC = compileDFA regexp
		     fun doit s =
		         Option.map
			 (Compiled.matchAll (regexpC, s),
			  fn match => Match.lookupString (match, save))
		  in
		    data := SOME (str, doit)
		  end
	 end
      val outData : (string * (string -> string option)) option ref = ref NONE
      val setOutData = fn str => setData ("out", outData, str)
      val errData : (string * (string -> string option)) option ref = ref NONE
      val setErrData = fn str => setData ("err", errData, str)
      (* Set the stack limit to its max, since mlkit segfaults on some benchmarks
       * otherwise.
       *)
       val _ =
	  let
	     open MLton.Platform.OS
	  in
	     if host = Linux
		then
		   let
		      open MLton.Rlimit
		      val {hard, ...} = get stackSize
		   in
		      set (stackSize, {hard = hard, soft = hard})
		   end
	     else ()
	  end
      local
	 open Popt
      in
	 val res =
	    parse
	    {switches = args,
	     opts = [("args",
		      SpaceString
		      (fn args =>
		       runArgs := String.tokens (args, Char.isSpace))),
		     ("err", SpaceString setErrData),
		     ("html", trueRef doHtml),
		     ("mlkit", 
		      None (fn () => pushCompiler
			    {name = "ML-Kit",
			     abbrv = "ML-Kit",
			     test = kitCompile})),
		     ("mosml",
		      None (fn () => pushCompiler
			    {name = "Moscow-ML",
			     abbrv = "Moscow-ML",
			     test = mosmlCompile})),
		     ("mlton",
		      SpaceString (fn arg => pushCompilers
				   (makeMLton arg))),
		     ("once", trueRef doOnce),
		     ("out", SpaceString setOutData),
		     ("poly",
		      None (fn () => pushCompiler
			    {name = "Poly/ML",
			     abbrv = "Poly/ML",
			     test = polyCompile})),
		     ("smlnj",
		      None (fn () => pushCompiler
			    {name = "SML/NJ",
			     abbrv = "SML/NJ",
			     test = njCompile})),
		     trace]}
      end
   in
      case res of
	 Result.No msg => usage msg
       | Result.Yes benchmarks =>
	    let
	       val compilers = List.rev (!compilers)
	       val base = #name (hd compilers)
	       val _ =
		  let
		     open MLton.Signal
		  in
		     setHandler (Posix.Signal.pipe, Handler.ignore)
		  end
	       fun r2s n r = Real.format (r, Real.Format.fix (SOME n))
	       val i2s = Int.toCommaString
	       val p2s = i2s o Position.toInt
	       val s2s = fn s => s
	       val failures = ref []
	       fun show ({compiles, runs, sizes, errs, outs}, {showAll}) =
		  let
		     val out = Out.standard
		     val _ =
			List.foreach
			(compilers, fn {name, abbrv, ...} =>
			 Out.output (out, concat [abbrv, " -- ", name, "\n"]))
		     val _ =
			case !failures of
			   [] => ()
			 | fs =>
			    Out.output
			    (out,
			     concat ["WARNING: ", base, " failed on: ",
				     concat (List.separate (fs, ", ")),
				     "\n"])
		     fun show (title, data: 'a data, toString, toStringHtml) =
			let
			   val _ = Out.output (out, concat [title, "\n"])
			   val compilers =
			      List.fold
			      (compilers, [],
			       fn ({name = n, abbrv = n', ...}, ac) =>
			       if showAll
				  orelse (List.exists
					  (data, fn {compiler = c, ...} =>
					   n = c))
				  then (n, n') :: ac
			       else ac)
			   val benchmarks =
			      List.fold
			      (benchmarks, [], fn (b, ac) =>
			       if showAll
				  orelse List.exists (data, fn {bench, ...} =>
						      bench = b)
				  then b :: ac
			       else ac)
			   fun rows toString =
			      ("benchmark"
			       :: List.revMap (compilers, fn (_, n') => n'))
			      :: (List.revMap
				  (benchmarks, fn b =>
				   b :: (List.revMap
					 (compilers, fn (n, _) =>
					  case (List.peek
						(data, fn {bench = b',
							   compiler = c', ...} =>
						 b = b' andalso n = c')) of
					     NONE => "*"
					   | SOME {value = v, ...} =>
						toString v))))
			   open Justify
			   val () =
			      outputTable
			      (table {columnHeads = NONE,
				      justs = (Left ::
					       List.revMap (compilers,
							    fn _ => Right)),
				      rows = rows toString},
			       out)
			   fun p ss =
			      (Out.output (out, concat ss)
			       ; Out.newline out)
			   fun prow (ns, c) =
			      case ns of
				 [] => raise Fail "bug"
			       | b :: ns =>
				    (p ["<tr>"]
				     ; p ["<t", c, " align = left>", b,
					  "</t", c, ">"]
				     ; (List.foreach
					(ns, fn n =>
					 p ["<t", c, " align = right>", n,
					    "</t", c, ">"]))
				     ; p ["</tr>"])
			   val _ =
			      if !doHtml
				 then
				    let
				       val rows = rows toStringHtml
				    in					  
				       prow (hd rows, "h")
				       ; (List.foreach
					  (tl rows,
					   fn [] => raise Fail "bug"
					    | b :: r =>
						 let
						    val b = 
						       concat
						       ["<a href = \"benchmarks/",
							b, ".sml\">", b, "</a>"]
						 in
							prow (b :: r, "d")
						 end))
				    end
			      else ()
			in
			   ()
			end
		     val bases = List.keepAll (runs, fn {compiler, ...} =>
					       compiler = base)
		     val ratios =
			List.fold
			(runs, [], fn ({bench, compiler, value}, ac) =>
			 if compiler = base andalso not showAll
			    then ac
			 else
			    {bench = bench,
			     compiler = compiler,
			     value =
			     case List.peek (bases, fn {bench = b, ...} =>
					     bench = b) of
				NONE => ~1.0
			      | SOME {value = v, ...} => value / v} :: ac)
		     val _ = show ("run time ratio", ratios, r2s 2, r2s 1)
		     val _ = show ("size", sizes, p2s, p2s)
		     val _ = show ("compile time", compiles, r2s 2, r2s 2)
		     val _ = show ("run time", runs, r2s 2, r2s 2)
		     val _ = case !outData of
                        NONE => ()
                      | SOME (out, _) =>
                           show (concat ["out: ", out], outs, s2s, s2s)
		     val _ = case !errData of
			NONE => ()
		      | SOME (err, _) =>
			   show (concat ["err: ", err], errs, s2s, s2s)
		  in ()
		  end
	       val totalFailures = ref []
	       val data = 
		  List.fold
		  (benchmarks, {compiles = [], runs = [], sizes = [],
				outs = [], errs = []},
		   fn (bench, ac) =>
		   let
		      val _ =
			 File.withOut
			 (batch bench, fn out =>
			  (File.outputContents (concat [bench, ".sml"], out)
			   ; Out.output (out, concat ["val _ = Main.doit ",
						      benchCount bench,
						      "\n"])))
		      val foundOne = ref false
		      val res =
			 List.fold
			 (compilers, ac, fn ({name, abbrv = _, test},
					     ac as {compiles: real data,
						    runs: real data,
						    sizes: Position.int data,
						    outs: string data,
						    errs: string data}) =>
			  if true
			     then
				let
				   val (outTmpFile, _) =
				      File.temp {prefix = "tmp", suffix = "out"}
				   val (errTmpFile, _) =
				      File.temp {prefix = "tmp", suffix = "err"}
				   val {compile, run, size} =
				     ignoreOutput
				     (fn () => test {bench = bench})
				   val _ =
				      if name = base
					 andalso Option.isNone run
					 then List.push (failures, bench)
				      else ()
				   val out = 
				      case !outData of 
					 NONE => NONE
				       | SOME (_, doit) => 
					    File.foldLines
					    (outTmpFile, NONE, fn (s, v) =>
					     let val s = String.removeTrailing
					                 (s, fn c => 
							  Char.equals (c, Char.newline))
					     in
						case doit s of
						   NONE => v
						 | v => v
					     end)
				   val err = 
				      case !errData of 
					 NONE => NONE
				       | SOME (_, doit) =>
					    File.foldLines
					    (errTmpFile, NONE, fn (s, v) =>
					     let val s = String.removeTrailing
						         (s, fn c => 
							  Char.equals (c, Char.newline))
					     in
						case doit s of
						   NONE => v
						 | v => v
					     end)
				   val _ = File.remove outTmpFile
				   val _ = File.remove errTmpFile
				   fun add (v, ac) =
				      case v of
					 NONE => ac
				       | SOME v =>
					    (foundOne := true
					     ; {bench = bench,
						compiler = name,
						value = v} :: ac)
				   val ac =
				      {compiles = add (compile, compiles),
				       runs = add (run, runs),
				       sizes = add (size, sizes),
				       outs = add (out, outs),
				       errs = add (err, errs)}
				   val _ = show (ac, {showAll = false})
				   val _ = Out.flush Out.standard
				in
				   ac
				end
			  else ac)
		      val _ =
			 if !foundOne
			    then ()
			 else List.push (totalFailures, bench)
		   in
		      res
		   end)
	       val _ = show (data, {showAll = true})
	       val totalFailures = !totalFailures
	       val _ =
		  if List.isEmpty totalFailures
		     then ()
		  else (print ("The following benchmarks failed completely.\n")
			; List.foreach (totalFailures, fn s =>
					print (concat [s, "\n"])))
	    in ()
	    end
   end

val main = Process.makeMain main

end