structure Optimize : sig
  val optimizeIR : PSysIR.block list -> PSysIR.block list
 end = struct
 
  fun printErr s = TextIO.output(TextIO.stdErr, s ^ "\n")

  structure IR = PSysIR
 
  fun countVarReferences(blocks) = let

	fun updateVar(IR.V{useCount, ...}) = useCount := (!useCount) + 1
    fun updateStmt(IR.PRIM(var, prim, vl, stmt)) = (
        List.app updateVar vl;
        updateStmt(stmt))
      | updateStmt(IR.IF(var, s1, s2)) = (
        updateVar(var);
        updateStmt(s1);
        updateStmt(s2))
      | updateStmt(IR.GOTO(block, vl)) = List.app updateVar vl
      | updateStmt(_) = ()
    
    fun countForBlock(IR.BLK{id, params, body}) = (updateStmt(body); ())
   in
    List.app countForBlock blocks
   end (*countVarReferences*)
  
  fun prv(var as IR.V{name, useCount, ...}) = printErr (String.concat([name, ": ", Int.toString (!useCount)]))
  fun prRefCt(IR.PRIM(v, p, vl, stmt)) = (prv(v); prRefCt(stmt))
    | prRefCt(IR.IF(v, s1, s2)) = (prRefCt(s1); prRefCt(s2))
    | prRefCt(_) = ()
  fun prRefCtForBlk(IR.BLK{body, ...}) = prRefCt(body)
  
  fun removeUnused(blocks) = let
    fun isUnused(v as IR.V{useCount, name, ...}) = if !useCount > 0 then false else true
    fun removeUnusedStmt(IR.PRIM(v, p, vl, s)) =
        if isUnused(v) then 
          removeUnusedStmt(s) 
        else 
          IR.PRIM(v, p, vl, removeUnusedStmt(s))
      | removeUnusedStmt(IR.IF(var, s1, s2)) = 
        IR.IF(var, removeUnusedStmt(s1), removeUnusedStmt(s2))
      | removeUnusedStmt(x) = x
    fun removeUnusedInBlock(IR.BLK{id, params, body}) = 
      IR.BLK{id=id, params=params, body=removeUnusedStmt(body)}
   in
    List.map removeUnusedInBlock blocks
   end (* removeUnused *)
  
  fun foldConstants(blocks) = let
    
    fun allConsts([]) = true
      | allConsts(v as IR.V{scope, ...} :: vl) = 
        case scope 
          of IR.S_CONST(c) => allConsts(vl) 
           | _ => false
    
    fun extractConstVec(IR.V{scope, ...}) = (case scope 
      of IR.S_CONST c => (case c 
        of IR.C_VEC v => v
         | _ => raise Fail ("Expected const vec.")
        (* end case *))
       | _ => raise Fail ("Expected const vec.")
      (* end case *))
      
    fun extractConstInt(IR.V{scope, ...}) = (case scope 
      of IR.S_CONST c => (case c 
        of IR.C_INT i => i
         | _ => raise Fail ("Expected const int.")
        (* end case *))
       | _ => raise Fail ("Expected const int.")
      (* end case *))
       
    fun extractConstFloat(IR.V{scope, ...}) = (case scope 
      of IR.S_CONST c => (case c 
        of IR.C_FLOAT f => f
         | _ => raise Fail ("Expected const float.")
        (* end case *))
       | _ => raise Fail ("Expected const float.")
      (* end case *))

    fun extractConstBool(IR.V{scope, ...}) = (case scope 
      of IR.S_CONST c => (case c 
        of IR.C_BOOL b => b
         | _ => raise Fail ("Expected const bool.")
        (* end case *))
       | _ => raise Fail ("Expected const bool.")
      (* end case *))
    
    fun createConst(v as IR.V{scope, varType, name, ...}) = (case scope
      of IR.S_LOCAL(p, vl) => (case p 
          of IR.ADD_VEC => v
          | IR.SUB_VEC => v
          | IR.LEN_SQ => v
          | IR.LEN => v
          | IR.NORM => v
          | IR.SCALE => 
            if allConsts(vl) then 
              let 
               val scale = extractConstFloat(List.nth(vl, 0))
               val vec = extractConstVec(List.nth(vl, 1))
              in
               IR.newConst(name, IR.C_VEC(Vec3f.scale(scale, vec)))
              end (*let*)
            else v
          | IR.DOT => v
          | IR.CROSS => v
          | IR.ADD => v
          | IR.SUB => v
          | IR.MULT => v
          | IR.DIV => v
          | IR.SQRT => v
          | IR.COS => v
          | IR.SIN => v
          | IR.GT => v
          | IR.EQUALS => v
          | IR.AND => v
          | IR.OR => v
          | IR.NOT => v
          | IR.RAND => v
          | IR.ITOF => 
            if allConsts(vl) then
              let 
               val intVal = extractConstInt(List.nth(vl, 0))
              in
               IR.newConst(name, IR.C_FLOAT(Float.fromInt (intVal)))
              end
            else v
          | IR.COPY => if allConsts(vl) then List.nth(vl, 0) else v
         (* end case *))
       | _ => v
      (* end case *))
    
    fun foldStmt(IR.PRIM(v, p, vl, s)) =
        IR.PRIM(v, p, List.map createConst vl, foldStmt(s))
      | foldStmt(IR.IF(v, s1, s2)) =
        IR.IF(v, foldStmt(s1), foldStmt(s2))
      | foldStmt(IR.GOTO(blk, vl)) =
        IR.GOTO(blk, List.map createConst vl)
      | foldStmt(IR.RETURN(vl)) =
        IR.RETURN(List.map createConst vl)
      | foldStmt(IR.DISCARD) = 
        IR.DISCARD
    
    fun foldConstantsInBlock(IR.BLK{id, params, body}) =
      IR.BLK{id=id, params=params, body=foldStmt(body)}
      
   in
    List.map foldConstantsInBlock blocks
   end
  
  fun optimizeIR(blocks) = 
   let
    val foldedBlks = foldConstants(blocks)
    val counted = countVarReferences(foldedBlks)
    val optBlks = removeUnused(foldedBlks)
   in
    (*List.app prRefCtForBlk optBlks; *)
    optBlks    
   end
 
 end (* structure *)