open Tree

let commute = function
  | NAME _
  | CONST _ -> true
  | _ -> false

let rec do_stm = function
  | SEQ (s1, s2) -> do_stm s1 @ do_stm s2
  | EXP e -> fst (do_exp e)
  | (LABEL _) as s -> [s]
  | JUMP (e, l) ->
      let s, e = do_exp e in
      s @ [JUMP (e, l)]
  | CJUMP (o, e1, e2, t, f) ->
      let e1s, e1e = do_exp e1 in
      let e2s, e2e = do_exp e2 in
      if commute e1e then
        e1s @ e2s @ [CJUMP (o, e1e, e2e, t, f)]
      else
        let t1 = Temp.new_temp () in
        e1s @ [MOVE(TEMP t1, e1e)] @
        e2s @ [CJUMP (o, TEMP t1, e2e, t, f)]
  | MOVE (TEMP t, e) ->
      let es, ee = do_exp e in
      es @ [MOVE(TEMP t, ee)]
  | MOVE (MEM(e1), e2) ->
      let e1s, e1e = do_exp e1 in
      let e2s, e2e = do_exp e2 in
      let t1 = Temp.new_temp () in
      e1s @ [MOVE(TEMP t1, e1e)] @
      e2s @ [MOVE(MEM(TEMP t1), e2e)]
  | _ -> raise (Failure "do_stm")

and do_exp = function
  | MEM e ->
      let es, ee = do_exp e in
      es, MEM ee
  | BINOP (o, e1, e2) ->
      let e1s, e1e = do_exp e1 in
      let e2s, e2e = do_exp e2 in
      if commute e1e then
        (e1s @ e2s), BINOP(o, e1e, e2e)
      else
        let t = Temp.new_temp() in
        (e1s @ [MOVE(TEMP t, e1e)] @ e2s), BINOP(o, TEMP t, e2e)
  | ESEQ (s, e) ->
      let ss = do_stm s in
      let es, ee = do_exp e in
      (ss @ es), ee
  | CALL (e, es) ->
      let ess, ts = List.fold_right (fun e (ess, ts) ->
        if commute e then ess, e::ts
        else
          let es, ee = do_exp e in
          let t = Temp.new_temp () in
          es @ [MOVE (TEMP t, ee)] @ ess, (TEMP t)::ts
      ) es ([], []) in
      let t = Temp.new_temp () in
      ess @ [MOVE(TEMP t, CALL(e, ts))], TEMP t
  | e -> [], e

let linearize = do_stm

let basic_blocks stms =
  let donelab = Temp.new_label () in
  let rec blocks stms blist =
    match stms with
      LABEL _ as h :: l ->
        let rec next stms thisblock =
          match stms with
            JUMP _ as h :: l ->
              end_block l (h :: thisblock)
          | CJUMP _ as h :: l ->
              end_block l (h :: thisblock)
          | LABEL lab :: _ ->
              next ((JUMP (NAME lab, [lab])) :: stms) thisblock
          | h :: l ->
              next l (h :: thisblock)
          | [] ->
              next [JUMP (NAME donelab, [donelab])] thisblock 
        and end_block stms thisblock =
          blocks stms (List.rev thisblock :: blist)
        in next l [h]
    | [] ->
        List.rev blist
    | _ ->
        blocks (LABEL (Temp.new_label ()) :: stms) blist
  in (blocks stms [], donelab)

  
  let enter_block block table =
    match block with
      LABEL s :: _ ->
        Symbol.enter s block table
    | _ -> table
  
  let rec split_last = function
      [] -> raise (Failure "split_last")
    | [x] -> ([], x)
    | h :: t -> let (t', last) = split_last t in (h::t', last) 
  
  let rec trace table (LABEL lab :: _ as b) rest =
    let table = Symbol.enter lab [] table in
      match split_last b with
        most, JUMP (NAME lab, _) ->
          begin
            match Symbol.look lab table with
              Some (_ :: _ as b') ->
                most @ trace table b' rest
            | _ ->
                b @ get_next table rest
          end
      | most, CJUMP (op, x, y, t, f) ->
          begin
            match Symbol.look t table, Symbol.look f table with
              _, Some (_ :: _ as b') ->
                b @ trace table b' rest
            | Some (_ :: _ as b'), _ ->
                most @ [CJUMP (not_relop op, x, y, f, t)] @ trace table b' rest
            | _ ->
                let f' = Temp.new_label () in
                  most @ [CJUMP (op, x, y, t, f');
                          LABEL f';
                          JUMP (NAME f, [f])] @ get_next table rest
          end
      | most, JUMP _ ->
          b @ get_next table rest
      | _ -> failwith "trace"
  
  and get_next table = function
      [] -> []
    | (LABEL lab :: _ as b) :: rest ->
        begin
          match Symbol.look lab table with
            Some (_ :: _) -> trace table b rest
          | _ -> get_next table rest
        end
    | _ -> failwith "get_next"
  
  let trace_schedule (blocks, donelab) =
    let table = List.fold_right enter_block blocks Symbol.empty in
      get_next table blocks @ [LABEL donelab]