Un-size optimize

[snipe.git] / codegen / solidify.sml

(* L3 Compiler
 * Takes a list of mappings of temporaries to colors and a pseudoasm listing,
 * then produces x86 code.
 * Author: Chris Lu <czl@andrew.cmu.edu>
 * Author: Joshua Wise <jwise@andrew.cmu.edu>
 *)

signature SOLIDIFY =
sig
  type colorings = (Temp.temp * int) list
  type asm = x86.insn list
  
  val solidify : colorings -> asm -> asm
end

structure Solidify :> SOLIDIFY =
struct
  structure X = x86
  structure T = Temp
  
  type colorings = (Temp.temp * int) list
  type asm = x86.insn list
  
  exception Spilled

  structure TempMap = SplayMapFn(struct
                                   type ord_key = Temp.temp
                                   val compare = Temp.compare
                                 end)
  structure Tm = Temp

  fun solidify (regmap : colorings) (instrs : asm) : asm =
    let
      (* r14d and r15d is reserved for spilling *)
      val maxreg = X.regtonum X.R13D
      fun numtoreg n =
          if (n > maxreg)
          then raise Spilled
          else X.numtoreg n

      val tempnums = List.foldr (fn ((t,n),b) => TempMap.insert(b,t,n)) (TempMap.empty) regmap
      fun temptonum (t: T.temp) : int = valOf (TempMap.find (tempnums, t))

      fun temptoreg (t: T.temp) : x86.reg =
        numtoreg (temptonum t)          
        handle Empty => raise ErrorMsg.InternalError ("Uncolored temp "^(Temp.name t)^", agh!")

      val spillreg1 = X.R15D
      val spillreg2 = X.R14D

      (* Determine which need to be saved. *)
      val opsused = (map (fn (_, n) => X.REG (numtoreg n handle Spilled => X.R15D)) regmap) @ [X.REG X.R14D]
      val saveregs = X.OperSet.intersection (
        X.OperSet.addList (X.OperSet.empty, opsused),
        X.OperSet.addList (
          X.OperSet.empty,
          [X.REG X.EBX,
           X.REG X.EBP,
           X.REG X.R12D,
           X.REG X.R13D,
           X.REG X.R14D,
           X.REG X.R15D]))
      val savelist = X.OperSet.listItems saveregs
      val nsave = length savelist

      val numreg = foldr (Int.max) 0 (map (fn (_, n) => n) regmap)    (* Number of registers used. *)
      val nspilled = Int.max (numreg - maxreg, 0)    (* Number of spilled registers. *)
      fun isspilled (X.TEMP temp, _) = (temptonum temp) > maxreg    (* Whether a register is spilled *)
        | isspilled (X.STACKARG _, _) = true
        | isspilled (X.REL _, _) = true
        | isspilled _ = false

      val stacksz = (nspilled + nsave) * 8
      fun stackpos (reg: int) = stacksz - (reg - maxreg + nsave) * 8    (* Stack position of some register number *)

      val prologue =
        (X.SUB ((X.REG X.RSP, Tm.Quad), (X.CONST (Word32.fromInt stacksz), Tm.Quad))) ::
        (ListPair.map
          (fn (num, reg) =>
            X.MOV ((X.REL ((X.REG X.RSP, Tm.Quad), (X.CONST (Word32.fromInt (stacksz - 8*(num+1))), Tm.Quad), 0w1), Tm.Quad), (reg, Tm.Quad)))
          (List.tabulate (nsave, fn x => x), savelist))
      val epilogue =
        (ListPair.map
          (fn (num, reg) =>
            X.MOV ((reg, Tm.Quad), (X.REL ((X.REG X.RSP, Tm.Quad), (X.CONST (Word32.fromInt (stacksz - 8*(num+1))), Tm.Quad), 0w1), Tm.Quad)))
          (List.tabulate (nsave, fn x => x), savelist)) @
        [X.ADD ((X.REG X.RSP, Tm.Quad), (X.CONST (Word32.fromInt stacksz), Tm.Quad))]
      val endlbl = Label.new()

      fun spill ((X.TEMP temp, s), xreg: x86.reg) =    (* Spill a register if need be. *)
            let
              val base = (X.REG X.RSP, Tm.Quad)
              val offs = (X.CONST (Word32.fromInt (stackpos (temptonum temp))), Tm.Quad)
            in
              if (isspilled (X.TEMP temp, s))
                then [X.MOV ((X.REL (base, offs, 0w1), Tm.Quad), (X.REG xreg, Tm.Quad))]
                else nil
            end
        | spill ((X.STACKARG _, s), _) = raise ErrorMsg.InternalError "Cannot spill to a stack arg"
        | spill ((a as X.REL _, s), xreg) = [X.MOV ((a,s), (X.REG xreg,s))]
        | spill _ = nil        (* Nothing else can be spilled. *)
      fun unspill ((X.TEMP temp, s), xreg: x86.reg) =    (* Unspill a register if need be. *)
            let
              val base = (X.REG X.RSP, Tm.Quad)
              val offs = (X.CONST (Word32.fromInt (stackpos (temptonum temp))), Tm.Quad)
            in
              if (isspilled (X.TEMP temp, s))
                then [X.MOV ((X.REG xreg, Tm.Quad), (X.REL (base, offs, 0w1), Tm.Quad))]
                else nil
            end
        | unspill ((X.STACKARG arg, s), xreg) =
            let
              val base = (X.REG X.RSP, Tm.Quad)
              val offs = (X.CONST (Word32.fromInt (stacksz + 8 + (arg * 8))), Tm.Quad)
            in
              [X.MOV ((X.REG xreg, s), (X.REL (base, offs, 0w1), s))]
            end
        | unspill ((a as X.REL _, s), xreg) = [X.MOV ((X.REG xreg, s), (a,s))]
        | unspill _ = nil

      fun realoper (X.TEMP temp, s) = (X.REG (temptoreg temp), s) (* makes an operand 'real' *)
        | realoper (X.STACKARG arg, _) = raise Spilled
        | realoper (X.REL _, _) = raise Spilled
        | realoper r = r

      fun stackoper (X.TEMP temp, s) =
            let
              val base = (X.REG X.RSP, Tm.Quad)
              val offs = (X.CONST (Word32.fromInt (stackpos (temptonum temp))), Tm.Quad)
            in
              if (isspilled (X.TEMP temp, s)) 
                then (X.REL (base, offs, 0w1), s)
                else raise ErrorMsg.InternalError "stackoper on unspilled temp?"
            end
        | stackoper (X.STACKARG arg, s) =
            let
              val base = (X.REG X.RSP, Tm.Quad)
              val offs = (X.CONST (Word32.fromInt (stacksz + 8 + (arg * 8))), Tm.Quad)
            in
              (X.REL (base, offs, 0w1), s)
            end
        | stackoper (a as (X.REL _, s)) = a
        | stackoper (a as (X.CONST _, s)) = a
        | stackoper anous = raise ErrorMsg.InternalError ("stackoper on not temp " ^ X.pp_oper anous)

      fun ophit (X.REL(op1, op2, m), s) =
            if (isspilled op1 andalso isspilled op2) then
              ([X.MOV ((X.REG spillreg1, Tm.Long), stackoper op2),
                X.IMUL((X.REG spillreg1, Tm.Quad), (X.CONST m, Tm.Quad)), 
                X.ADD ((X.REG spillreg1, Tm.Quad), stackoper op1)],
                (X.REL ((X.REG spillreg1, Tm.Quad), (X.CONST 0w0, Tm.Quad), 0w1), s))
            else if(isspilled op1) then
              ([X.MOV ((X.REG spillreg1, Tm.Quad), stackoper op1)],
                (X.REL ((X.REG spillreg1, Tm.Quad), realoper op2, m), s))
            else if(isspilled op2) then
              ([X.MOV ((X.REG spillreg1, Tm.Long), stackoper op2)],
                (X.REL (realoper op1, (X.REG spillreg1, Tm.Quad), m), s))
            else
              ([],
                (X.REL (realoper op1, realoper op2, m), s))
        | ophit a = (nil, realoper a handle Spilled => stackoper a)

      fun transform (X.DIRECTIVE s) = [X.DIRECTIVE s]
        | transform (X.COMMENT s) = [X.COMMENT s]
        | transform (X.LIVEIGN a) = transform a
        | transform (X.MOV (dest, src)) =
            let
              val (insns1, realop1 as (_,s1)) = ophit dest
              val (insns2, realop2 as (_,s2)) = ophit src
            in
              if(isspilled dest andalso isspilled src) then
                insns2 @ [X.MOV ((X.REG spillreg2, s2), realop2)] @ insns1 @ [X.MOV (realop1, (X.REG spillreg2, s1))]
              else
                insns1 @ insns2 @ [X.MOV (realop1, realop2)]
            end
        | transform (X.LEA (dest, src)) =
            let
              val (insns1, realop1 as (_,s1)) = ophit dest
              val (insns2, realop2 as (_,s2)) = ophit src
            in
              if(isspilled dest andalso isspilled src) then
                insns2 @ [X.LEA ((X.REG spillreg2, s2), realop2)] @ insns1 @ [X.MOV (realop1, (X.REG spillreg2, s1))]
              else
                insns1 @ insns2 @ [X.LEA (realop1, realop2)]
            end
        | transform (X.SUB (dest, src)) =
            let
              val (insns, realop) = ophit dest
            in
              unspill (src, spillreg2) @ insns @
              [ X.SUB(realop,
                  realoper src handle Spilled => (X.REG spillreg2, X.osize realop))]
            end
        | transform (X.IMUL (dest, src)) =
            unspill (dest, spillreg1) @
            [ X.IMUL(
                realoper dest handle Spilled => (X.REG spillreg1, X.osize dest),
                realoper src handle Spilled => stackoper src)] @
            spill (dest, spillreg1)
        | transform (X.IMUL3 (dest, src, const)) =
            unspill (src, spillreg2) @
            [ X.IMUL3(
                realoper dest handle Spilled => (X.REG spillreg1, X.osize dest),
                realoper src handle Spilled => (X.REG spillreg2, X.osize src),
                const)] @
            spill (dest, spillreg1)
        | transform (X.ADD (dest, src)) =
            let
              val (insns, realop) = ophit dest
            in
              unspill (src, spillreg2) @ insns @
              [ X.ADD(realop,
                  realoper src handle Spilled => (X.REG spillreg2, X.osize realop))]
            end
        | transform (X.IDIV (src)) = [ X.IDIV(realoper src handle Spilled => stackoper src)]
        | transform (X.NEG (src)) = [ X.NEG(realoper src handle Spilled => stackoper src)]
        | transform (X.NOT (src)) = [ X.NOT(realoper src handle Spilled => stackoper src)]
        | transform (X.SAL (dest, shft)) =
            [ X.SAL (
                realoper dest handle Spilled => stackoper dest,
                shft)]
        | transform (X.SAR (dest, shft)) =
            [ X.SAR (
                realoper dest handle Spilled => stackoper dest,
                shft)]
        | transform (X.CLTD) = [ X.CLTD ]
        | transform (X.AND (dest, src)) =
            unspill (src, spillreg1) @
            [ X.AND(
                realoper dest handle Spilled => stackoper dest,
                realoper src handle Spilled => (X.REG spillreg1, X.osize src))]
        | transform (X.OR (dest, src)) =
            unspill (src, spillreg1) @
            [ X.OR(
                realoper dest handle Spilled => stackoper dest,
                realoper src handle Spilled => (X.REG spillreg1, X.osize src))]
        | transform (X.XOR (dest, src)) =
            unspill (src, spillreg1) @
            [ X.XOR(
                realoper dest handle Spilled => stackoper dest,
                realoper src handle Spilled => (X.REG spillreg1, X.osize src))]
        | transform (X.CMP (op1, op2)) =
            let
              val (insns1, realop1) = ophit op1
            in
              if(isspilled realop1 andalso isspilled op2) then
                unspill (op2, spillreg2) @ insns1 @ [X.CMP (realop1, (X.REG spillreg2, X.osize realop1))]
              else
                insns1 @ [X.CMP (realop1, realoper op2 handle Spilled => stackoper op2)]
            end
        | transform (X.TEST (op1, op2)) =
            unspill (op2, spillreg1) @
            [ X.TEST(
                realoper op1 handle Spilled => stackoper op1,
                realoper op2 handle Spilled => (X.REG spillreg1, X.osize op2))]
        | transform (X.SETcc (c,src)) = [ X.SETcc(c, realoper src handle Spilled => stackoper src)]
        | transform (X.CMOVcc (c, dest, src)) =
            let
              val (insns1, realop1) = ophit dest
              val (insns2, realop2) = ophit src
            in
              if(isspilled dest andalso isspilled src) then
                insns2 @ [X.MOV ((X.REG spillreg2, X.osize src), realop2)] @ insns1 @ [X.CMOVcc (c, realop1, (X.REG spillreg2, X.osize src))]
              else
                insns1 @ insns2 @ [X.CMOVcc (c, realop1, realop2)]
            end
        | transform (X.CALL l) = [ X.CALL l ]
        | transform (X.MOVZB (dest, src)) =
            [ X.MOVZB(
                realoper dest handle Spilled => (X.REG spillreg1, X.osize dest),
                realoper src handle Spilled => stackoper src)]
            @ spill (dest, spillreg1)
        | transform (X.RET) = if nsave < 2 then (epilogue @ [X.RET]) else [X.JMP endlbl]
        | transform (X.LABEL l) = [ X.LABEL l ]
        | transform (X.JMP l) = [ X.JMP l ]
        | transform (X.Jcc (c,l)) = [X.Jcc (c,l)]
        | transform _ = raise ErrorMsg.InternalError "probably movsc: unimplemented"
    in
      if (nsave < 2) then
        List.concat (prologue :: (map transform instrs))
      else
        List.concat (prologue :: ((map transform instrs) @ [[X.LABEL endlbl], epilogue, [X.RET]]))
    end
end